Rust、最近使い出しました。
テストを書くにあたって、テスト用のデータベース作成と削除をうまく実装できないか考えていて、調べるとDropというトレイトが使えるとわかりました。
ただ、Dropはasyncに対応していないので、どういう代替案があるかな〜と気になったので調べてみました。
そもそもDropトレイトとは？
dropメソッドのみが定義されており、ここに実装を入れるとオブジェクトがメモリから開放される際に勝手に呼び出されるという仕組みです。
async-dropper
調べてみたら一応メジャーリリースしているasyncに対応したDropトレイトがありました。
ちょっと使ってみましょう。
以下のようにパッケージを追加してあげます。
[dependencies]
async-dropper = { version = "0.2.3", features = [ "tokio" ] }
async-trait = "0.1.73"
sqlx = { version = "0.7", features = [ "runtime-tokio", "postgres" ] }
普段からtokioを使っているので、そのfeaturesも入れてあげます。
tokioの対応バージョンは^1.29.1からとなっているので、比較的新しいバージョンじゃないと動かないようになってますね。
async-traitはトレイトにasyncメソッドを定義できるようになるマクロで、async-dropperを使用する際に役立つので入れておきます。
データベースの作成・削除に使用するためSQLXも入れておきます。
ついでにdocker-composeで使えるpostgresサーバーも作っておきましょう。
version: '3'
services:
db:
image: postgres:15
container_name: postgres
ports:
- 5432:5432
environment:
- POSTGRES_PASSWORD=example
次に実装です。
とりあえずデータベース作成をランダム文字列で行うような実装をしました。
use sqlx::postgres::{PgPool, PgPoolOptions};
use rand::{Rng};
const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyz";
const DBNAME_LEN: usize = 30;
struct MockDBServer {
pool: PgPool,
}
impl MockDBServer {
async fn new() -> Self {
let mut rng = rand::thread_rng();
let dbname: String = (0..DBNAME_LEN)
.map(|_| {
let idx = rng.gen_range(0..CHARSET.len());
char::from(unsafe { *CHARSET.get_unchecked(idx) })
})
.collect();
let pool = PgPoolOptions::new()
.max_connections(5)
.connect("postgres://postgres:example@localhost/postgres").await.unwrap();
sqlx::query(format!("CREATE DATABASE {};", dbname).as_str())
.bind(dbname)
.execute(&pool).await.unwrap();
Self {
pool
}
}
}
#[tokio::main]
async fn main() {
MockDBServer::new().await;
}
これを実行するとデータベースが作成されます。ただ作成されっぱなしになってしまいます。
ここにasync-dropperを使ってデストラクタを追加します。
use sqlx::postgres::{PgPool, PgPoolOptions};
use async_dropper::simple::{AsyncDrop, AsyncDropper};
use async_trait::async_trait;
use rand::Rng;
#[derive(Default)] // デフォルト実装を追加
struct MockDBServer {
pool: Option<PgPool>, // デフォルトが必要なため、Optionにしました
temp_dbname: String,
}
// ...中略...
#[async_trait]
impl AsyncDrop for MockDBServer {
async fn async_drop(&mut self) {
if self.pool.is_some() {
let pool = self.pool.clone().unwrap();
sqlx::query(format!("DROP DATABASE {}", &self.temp_dbname).as_str())
.execute(&pool).await.unwrap();
}
}
}
この実装を追加すると、MockDBServerのドロップとともにDBも削除されていることが確認できます！
まとめ
Asyncにドロップ処理を定義することができました。
ただasync-traitに関してはメジャーリリースしていないっぽいので、あまり本番では使わないほうがいいかも？