借助 Rust 语言的强大功能，所构建的程序能够与多个不同类型的数据库建立连接，通过精心设计的逻辑与接口，可依据实际运行时的需求，在这些数据库之间进行无缝且高效的切换，确保数据交互的准确性与流畅性。
笔者以 postgresql 为例，创建两个数据库，分别为 user_db 和 vip_user_db。在 user_db 创建一个 t_user 表，在 vip_user_db 创建一张 vip_user 表。结构如下：
CREATE TABLE public.t_user (
id bigserial NOT NULL,
"name" varchar NOT NULL
);
CREATE TABLE public.vip_user (
id bigserial NOT NULL,
"name" varchar NOT NULL
);
接着分别往两张表里初始化一些数据。
INSERT INTO public.t_user (id, "name") VALUES(1, 'a');
INSERT INTO public.t_user (id, "name") VALUES(2, 'b');
INSERT INTO public.t_user (id, "name") VALUES(3, 'c');
INSERT INTO public.vip_user (id, "name") VALUES(1, 'x');
INSERT INTO public.vip_user (id, "name") VALUES(2, 'y');
INSERT INTO public.vip_user (id, "name") VALUES(3, 'z');
现在去创建一个 rust 项目，同时连接这两个数据库。
cargo new datasource-switch
cd datasource-switch
在 Cargo.toml 添加以下依赖
[dependencies]
tokio = { version = "1.42.0", features = ["full"] }
anyhow = "1.0.95"
axum = "0.7.9"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
once_cell = "1.20.2"
sqlx = { version = "0.8", features = [ "runtime-tokio", "postgres" ] }
创建 src/db.rs，在这里使用 sqlx 来连接数据库，创建数据库连接池。
use std::collections::HashMap;
use anyhow::{anyhow, Result};
use sqlx::{postgres::PgPoolOptions, Pool, Postgres};
type PoolMap = HashMap<String, Pool<Postgres>>;
pub struct DBClient(PoolMap);
pub struct Config {
pub db_name: String,
pub url: String,
}
impl Config {
pub fn new(db_name: &str, url: &str) -> Self {
Config {
db_name: db_name.into(),
url: url.into(),
}
}
}
impl DBClient {
pub async fn new(config: Vec<Config>) -> Result<Self> {
let mut pool_map = HashMap::new();
for c in config {
let pool = PgPoolOptions::new()
.max_connections(10)
.connect(&c.url)
.await?;
pool_map.insert(c.db_name, pool);
}
Ok(DBClient(pool_map))
}
pub fn pool(&self, db: &str) -> Result<Pool<Postgres>> {
Ok(self.0.get(db).ok_or(anyhow!("PoolMap error."))?.clone())
}
}
创建 src/datasource.rs，用来存放数据源，在 src/db.rs 创建的连接会存到这里定义的全局静态变量中，这样的做法主要是为了方便在其他模块中使用数据源，而无需通过参数一层层传递，可以避免大量的 clone，让代码也变得更简洁。
use once_cell::sync::OnceCell;
use crate::db::DBClient;
pub const USER_DB: &'static str = "user_db";
pub const VIP_USER_DB: &'static str = "vip_user_db";
pub static DATA_SOURCE: OnceCell<DBClient> = OnceCell::new();
创建 src/macros.rs，定义一个声明式宏 use_db，当有 struct 使用这个宏时，会自动为这个 struct 实现一个 get_db_connection 函数，这个 struct 就可以通过调用这个函数来获取数据库连接了。
#[macro_export]
macro_rules! use_db {
($struct_name:ident) => {
impl $struct_name {
pub fn get_db_connection(db_name: &str) -> anyhow::Result<sqlx::Pool<sqlx::Postgres>> {
let db_arc =
crate::datasource::DATA_SOURCE
.get()
.ok_or(anyhow::anyhow!(format!(
"DBClient {}.{} not initialized",
db_name,
stringify!($struct_name)
)))?;
let pool = db_arc.pool(db_name)?;
Ok(pool)
}
}
};
}
创建 src/user.rs，在这里有一个与表结构对应的 struct，并且为这个 struct 实现了 get_list 接口，并且指定了这个 struct 对应的数据库为 user_db。
use serde::{Deserialize, Serialize};
use sqlx::prelude::FromRow;
use crate::{datasource::USER_DB, use_db};
#[derive(Serialize, Deserialize, FromRow, Debug)]
pub struct User {
pub id: i64,
pub name: String,
}
use_db!(User);
impl User {
const DB_NAME: &'static str = USER_DB;
pub async fn get_list() -> anyhow::Result<Vec<User>> {
let rows = sqlx::query_as("SELECT * FROM t_user")
.fetch_all(&Self::get_db_connection(Self::DB_NAME)?)
.await?;
Ok(rows)
}
}
创建 src/vip_user.rs，这里的内容与 src/user.rs 差不多是一样的，只是所属数据库不同。在调用这个 struct get_list 函数的时候，会使用 vip_user_db 这个库。
use serde::{Deserialize, Serialize};
use sqlx::prelude::FromRow;
use crate::{datasource::VIP_USER_DB, use_db};
#[derive(Serialize, Deserialize, FromRow, Debug)]
pub struct VipUser {
pub id: i64,
pub name: String,
}
use_db!(VipUser);
impl VipUser {
const DB_NAME: &'static str = VIP_USER_DB;
pub async fn get_list() -> anyhow::Result<Vec<VipUser>> {
let rows = sqlx::query_as("SELECT * FROM vip_user")
.fetch_all(&Self::get_db_connection(Self::DB_NAME)?)
.await?;
Ok(rows)
}
}
创建 src/main.rs，在这里使用 axum 创建了两个 http 接口，在这两个接口中分别调用不同的数据库返回数据。
use anyhow::anyhow;
use axum::{http::StatusCode, routing::get, Json, Router};
use datasource::{DATA_SOURCE, USER_DB, VIP_USER_DB};
use db::{Config, DBClient};
use user::User;
use vip_user::VipUser;
mod datasource;
mod db;
mod macros;
mod user;
mod vip_user;
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let db_config = vec![
Config::new(
USER_DB,
"postgres://postgres:password@127.0.0.1:5432/user_db",
),
Config::new(
VIP_USER_DB,
"postgres://postgres:password@127.0.0.1:5432/vip_user_db",
),
];
let db_client = DBClient::new(db_config).await?;
DATA_SOURCE
.set(db_client)
.map_err(|_err| anyhow!(format!("DATA_SOURCE {} initialization failed.", USER_DB)))?;
// 定义路由
let app = Router::new()
.route("/users", get(get_users))
.route("/vip-users", get(get_vip_users));
// 创建 http 服务
let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?;
axum::serve(listener, app).await?;
Ok(())
}
/// 获取 user_db.user 的列表
async fn get_users() -> (StatusCode, Json<Vec<User>>) {
let users = User::get_list().await.unwrap();
(StatusCode::OK, Json(users))
}
/// 获取 vip_user_db.vip_user 的列表
async fn get_vip_users() -> (StatusCode, Json<Vec<VipUser>>) {
let users = VipUser::get_list().await.unwrap_or_default();
(StatusCode::OK, Json(users))
}
运行
cargo run
测试
curl -X GET http://localhost:3000/users
# /users 结果
[{"id":1,"name":"a"},{"id":2,"name":"b"},{"id":3,"name":"c"}]
curl -X GET http://localhost:3000/vip-users
# /vip-user 结果
[{"id":1,"name":"x"},{"id":2,"name":"y"},{"id":3,"name":"z"}]
两个接口数据按不同数据库中实际对应数据返回，则表示成功。
代码已放到 GitHub，如有需要请自取：https://github.com/maoyutofu/datasource-switch.git