PostgreSQL 是一款天然支持异步操作的高性能开源关系型数据库。本章将讨论如何在 axum 中使用 PostgreSQL。包括：数据的增加、修改、删除、查找以及开始事务保证业务的原子性。如果你对 PostgreSQL 不是很了解，可以通过[PostgreSQL 轻松学](https://pg.sjk66.com)网站进行学习。 [ElephantSQL](https://elephantsql.com)提供了免费的 PostgreSQL 数据库，对于学习而言，它提供的免费实例就足够了。你可以创建多个免费实例，并且它的实例可以选择包括香港在内的多个节点。 ## 依赖本章代码涉及多个依赖，下面对各依赖项进行简单说明： - `deadpool-postgres`：连接池 - `tokio-postgres`：异步操作 PostgreSQL 的库 - `tokio-pg-mapper`：自动将数据库的`SELECT`语句与结构体进行映射 - `tokio-pg-mapper-derive`：使用 derive 简化数据库与结构体的映射 ## 准备工作开始之前，你需要做好以下准备工作 - 有一个 PostgreSQL 实例，如果你不想在本地安装，可以使用 ElephantSQL 提供的免费实例。 - 导入[SQL 文件](https://raw.githubusercontent.com/axumrs/roaming-axum/main/postgres/axum.rs.sql) 该 SQL 文件内容如下： ```sql CREATE TABLE account ( id SERIAL PRIMARY KEY, username VARCHAR(50) NOT NULL, balance INTEGER NOT NULL DEFAULT 0, UNIQUE(username) ); INSERT INTO account(username, balance) VALUES ('axum.rs', 999999), ('foo', 0), ('bar', 0); ``` 下面可以开始学习了。 ## 创建连接池为了提升并发效率，我们不直接连接数据库，而是通过连接也来维护数据库连接。rust 有很多连接池的库，我们选择的是 `deadpool-postgres`。创建连接池的方法很多，最方便的莫过于从配置中创建，所以我们从数据库配置开始： ### 数据库连接池配置 ```rust let mut cfg = deadpool_postgres::Config::new(); cfg.user = Some("axum_rs".to_string()); //数据库用户名 cfg.password = Some("axum.rs".to_string()); //数据库密码 cfg.dbname = Some("axum_rs".to_string()); //数据库名称 cfg.host = Some("pg.axum.rs".to_string()); // 数据库主机 cfg.port = Some(5432); //数据库端口 cfg ``` > 本示例的配置仅供演示说明，请根据你自身的环境对数据库配置进行修改。 ### 从配置中创建连接池有了配置之后，可以方便的地从中创建连接池，比如： ```rust let pool = cfg.create_pool(tokio_postgres::NoTls).unwrap(); ``` ## 从连接池中获取数据库连接接下来，我们调用连接池的`get()`方法，从中获取一个数据库连接： ```rust let client = pool.get().await.unwrap(); ``` 现在，可以开始数据库操作了。 ## 获取数据库连接的完整代码为了方便你对照学习，现将获取数据库连接的完整代码提供给你： ```rust /// 数据库配置 fn get_cfg() -> deadpool_postgres::Config { let mut cfg = deadpool_postgres::Config::new(); cfg.user = Some("axum_rs".to_string()); //数据库用户名 cfg.password = Some("axum.rs".to_string()); //数据库密码 cfg.dbname = Some("axum_rs".to_string()); //数据库名称 cfg.host = Some("pg.axum.rs".to_string()); // 数据库主机 cfg.port = Some(5432); //数据库端口 cfg } /// 从连接池中获取数据库连接 async fn get_client() -> Result { // 通过配置文件创建连接池 let pool = get_cfg() .create_pool(tokio_postgres::NoTls) .map_err(|err| err.to_string())?; // 从连接池中获取数据库连接 pool.get().await.map_err(|err| err.to_string()) } ``` ## 定义模型为了与数据表进行映射，我们需要定义结构体。该结构体和数据表字段需要对应好。 ```rust #[derive(PostgresMapper, Debug)] #[pg_mapper(table = "account")] pub struct Account { pub id: i32, pub username: String, pub balance: i32, } ``` 眼尖的你可能发现了，除了`Debug`之外，我们还使用了`PostgresMapper`和`pg_mapper`。 - `PostgresMapper`：让你的结构体能使用快速映射方法，将数据库查询结果自动填充到结构体 - `pg_mapper`：指定这个结构体关联的数据表 ## 插入数据首先，我们来看一下如何向数据库中添加一条新数据： ```rust async fn insert(Path(username): Path) -> Result<&'static str, String> { let create_user = CreateAccount { username, balance: 0, }; let client = get_client().await?; let stmt = client .prepare("INSERT INTO account (username, balance) VALUES ($1, $2)") .await .map_err(|err| err.to_string())?; let rows = client .execute(&stmt, &[&create_user.username, &create_user.balance]) .await .map_err(|err| err.to_string())?; if rows < 1 { return Err("Insert account failed".to_string()); } Ok("Successfully insert account") } ``` 让我们对这个函数进行分解说明：首先，创建了一个 `CreateAccount`的实例，它是一个结构体，该结构体的定义如下： ```rust pub struct CreateAccount { pub username: String, pub balance: i32, } ``` 然后，通过 `get_clinet()` 函数从连接池中获取连接。之后，通过`client.prepare()` 方法获取一个“语句”对象。`prepare`方法会对 SQL 语句进行预编译，这是防止 SQL 注入最有效的方法。再然后，调用`client.execute()`执行 SQL 语句，它将返回受影响的行数。通过判断受影响的行数就能知道这个 SQL 语句是否执行成功。 ## `execute()` 对于 `INSERT/UPDATE/DELETE` 等语句，通常使用`execute()`方法来执行 SQL 语句。它接收两个参数，一个是预编译 SQL 语句之后返回的对象，另一个是 SQL 语句所需要的参数。通常我们是通过它的返回值来判断 SQL 语句是否成功。 > 注意，由于 PostgreSQL 不提供类似其它数据库的 `LastInsertedID`功能，所以使用`execute()`执行`INSERT`语句并不能获得最后插入的 ID。 ## 修改数据我们以修改指定用户余额作为演示： ```rust async fn update(Path((id, balance)): Path<(i32, i32)>) -> Result<&'static str, String> { let client = get_client().await?; let stmt = client .prepare("UPDATE account SET balance=$1 WHERE id=$2") .await .map_err(|err| err.to_string())?; let rows = client .execute(&stmt, &[&balance, &id]) .await .map_err(|err| err.to_string())?; if rows < 1 { return Err("Update account failed".to_string()); } Ok("Successfully update account") } ``` 和插入数据类似，我们使用的是`execute()`，不同之处在于，这里的 SQL 语句是`UPDATE`。 ## 删除数据和修改数据类型，不同之处在于，这里的 SQL 语句是`DELETE`，如下所示： ```rust async fn delete(Path(id): Path) -> Result<&'static str, String> { let client = get_client().await?; let stmt = client .prepare("DELETE FROM account WHERE id=$1") .await .map_err(|err| err.to_string())?; let rows = client .execute(&stmt, &[&id]) .await .map_err(|err| err.to_string())?; if rows < 1 { return Err("Delete account failed".to_string()); } Ok("Successfully delete account") } ``` ## 查找多条数据现在我们来看一下怎么查找数据。 ```rust async fn list() -> Result { let client = get_client().await?; let stmt = client .prepare("SELECT id,username,balance FROM account ORDER BY id DESC") .await .map_err(|err| err.to_string())?; let account_list = client .query(&stmt, &[]) .await .map_err(|err| err.to_string())? .iter() .map(|row| Account::from_row_ref(&row).unwrap()) .collect::>(); let mut output = Vec::with_capacity(account_list.len()); for account in account_list.iter() { output.push(format!("{:?}", account)); } Ok(output.join("\n")) } ``` 前两步都一样，获取数据库连接，然后对 SQL 语句进行预编译。之后，我们通过`client.query()`方法来执行数据的查询。它是一个迭代器，我们对所有查询到的数据用`map()`进行处理。 `map()`里的闭包只有一行： ```rust |row| Account::from_row_ref(&row).unwrap() ``` 它的含义是，将查询到的数据，逐行映射到 `Account` 结构体。我们可以把它理解为： ```rust |row| { Account: { id: row.id, username: row.username.clone(), // ... } } ``` 这里之所以能写成 `Account::from_row_ref(&row)` 是因为我们使用了 `tokio-pg-mapper`。之后，使用`collect()` 将`map()` 处理过后的数据变成`Vec`集合。 ## `query()` 该方法通常用于查询多条数据。配合迭代器和`map()`、`collect()`，可以很方便、直观、高效地获取到包含满足条件的记录的集合。实际上，对于单条记录的查询，也可以使用该方法。另外，如果要获取 PostgreSQL 最后插入的 ID，我们也使用该方法。 ## 查询单条记录和查询多条记录一样，我们也用 `query()`。 ```rust async fn find(Path(id): Path) -> Result { let client = get_client().await?; let stmt = client .prepare("SELECT id,username,balance FROM account WHERE id=$1 ORDER BY id DESC LIMIT 1") .await .map_err(|err| err.to_string())?; let account = client .query(&stmt, &[&id]) .await .map_err(|err| err.to_string())? .iter() .map(|row| Account::from_row_ref(&row).unwrap()) .collect::>() .pop() .ok_or(format!("Couldn't find account #{}", id))?; Ok(format!("{:?}", account)) } ``` 之前的操作和查询多条一毛一样。在 `collect()` 获取集合后，我们调用了 `pop()` 方法，如果有记录就能得到一条记录，如果没有记录，就会用后面的`ok_or()`来返回一个`Err`，就是我们的报错信息。 ## 事务我们通过模拟账户之间转账来演示事务的处理。 ```rust async fn transfer( Path((from_id, to_id, balance)): Path<(i32, i32, i32)>, ) -> Result<&'static str, String> { let mut client = get_client().await?; let tx = client.transaction().await.map_err(|err| err.to_string())?; // 修改出账记录 let stmt = tx .prepare("UPDATE account SET balance=balance-$1 WHERE id=$2 AND balance>=$1") .await .map_err(|err| err.to_string())?; match tx.execute(&stmt, &[&balance, &from_id]).await { Ok(_rows) if _rows > 0 => { // 检查受影响的行数 // 如果大于零表示账户存在并且余额足够 // 不必做其余操作，等待最终的事务提交 } _ => { // 回滚事务 tx.rollback().await.map_err(|err| err.to_string())?; // 提前结束函数，将错误信息返回 return Err("Step 1 failed".to_string()); } }; // 修改入账记录 let stmt = tx .prepare("UPDATE account SET balance=balance+$1 WHERE id=$2") .await .map_err(|err| err.to_string())?; match tx.execute(&stmt, &[&balance, &to_id]).await { Ok(_rows) if _rows > 0 => { // 检查受影响的行数 // 如果大于零表示入账记录修改成功 // 不必做其余操作，等待最终的事务提交 } _ => { // 回滚事务 tx.rollback().await.map_err(|err| err.to_string())?; // 提前结束函数，将错误信息返回 return Err("Step 2 failed".to_string()); } }; // 提交事务 tx.commit().await.map_err(|err| err.to_string())?; Ok("Successfully transfer") } ``` 获取到连接之后，我们并没有像之前那样调用它的 `prepare()`，而是调用了它的 `transaction()` 来开启事务。注意，此时要求连接是`mut`的。之后就是`prepare`、`execute`了。不同的是，我们用的是事务的相关方法，而不是数据库连接的相关方法。值得注意的是，我们判断`execute`如果没有成功执行，调用了`tx.rollback()` 对事务进行回滚，并返回错误信息。最后，通过`tx.commit()` 提交事务。本章讨论了在 axum 集成 PostgreSQL 数据库的方法。完整代码可以在[代码仓库](https://github.com/axumrs/roaming-axum/tree/main/postgres)找到。 ## 思考题如何获取 PostgreSQL 最后插入的 ID？提示： - 在 `INSERT语句使用RETURNING`，你可以查看[这里](https://pg.sjk66.com/postgresql/insert.html) - 使用`query()`而不是`execute()`（类似查找单条记录）

标签

事务

axum 操作 Postgres 数据库