Aurora DSQL Rust コネクタで IAM 認証・OCC リトライを自動化する

aws dsql create-cluster \
  --tags Name=rust-connector-test \
  --region us-east-1

aws dsql get-cluster \
  --identifier <cluster-id> \
  --region us-east-1

cargo init dsql-rust-test
cd dsql-rust-test

[package]
name = "dsql-rust-test"
version = "0.1.0"
edition = "2024"
 
[dependencies]
aurora-dsql-sqlx-connector = { version = "0.1.2", features = ["pool", "occ"] }
sqlx = { version = "0.8", features = ["runtime-tokio", "postgres"] }
tokio = { version = "1", features = ["full"] }
anyhow = "1"

use aurora_dsql_sqlx_connector::DsqlConnectOptions;
use sqlx::{postgres::PgPoolOptions, Executor, Row};
 
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let endpoint = std::env::var("CLUSTER_ENDPOINT")?;
    let conn_str = format!("postgres://admin@{}/postgres", endpoint);
 
    // プール接続（IAM トークン生成・SSL・リフレッシュはすべて自動）
    let config = DsqlConnectOptions::from_connection_string(&conn_str)?;
    let pool = aurora_dsql_sqlx_connector::pool::connect_with(
        &config,
        PgPoolOptions::new().max_connections(10),
    )
    .await?;
    println!("Pool connected successfully");
 
    // CREATE
    pool.execute(
        "CREATE TABLE IF NOT EXISTS test_items(
            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
            name VARCHAR(100) NOT NULL,
            value INT NOT NULL
        )",
    )
    .await?;
 
    // INSERT
    sqlx::query("INSERT INTO test_items(name, value) VALUES($1, $2)")
        .bind("item-a")
        .bind(42)
        .execute(&pool)
        .await?;
 
    // SELECT
    let row = sqlx::query("SELECT name, value FROM test_items WHERE name = $1")
        .bind("item-a")
        .fetch_one(&pool)
        .await?;
    println!("Selected: name={}, value={}",
        row.get::<&str, _>("name"), row.get::<i32, _>("value"));
 
    // UPDATE
    sqlx::query("UPDATE test_items SET value = $1 WHERE name = $2")
        .bind(100).bind("item-a").execute(&pool).await?;
 
    // DELETE
    sqlx::query("DELETE FROM test_items WHERE name = $1")
        .bind("item-a").execute(&pool).await?;
 
    pool.execute("DROP TABLE test_items").await?;
    pool.close().await; // バックグラウンドのトークンリフレッシュタスクを停止
    Ok(())
}

use aurora_dsql_sqlx_connector::{retry_on_occ, DsqlConnectOptions, DsqlError};
use sqlx::{postgres::PgPoolOptions, Executor, Row};
 
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let endpoint = std::env::var("CLUSTER_ENDPOINT")?;
    let conn_str = format!("postgres://admin@{}/postgres", endpoint);
 
    let config = DsqlConnectOptions::from_connection_string(&conn_str)?;
    let pool = aurora_dsql_sqlx_connector::pool::connect_with(
        &config,
        PgPoolOptions::new().max_connections(10),
    ).await?;
 
    // カウンタ行を初期化
    pool.execute(
        "CREATE TABLE IF NOT EXISTS counter(id INT PRIMARY KEY, value INT NOT NULL)",
    ).await?;
    sqlx::query("INSERT INTO counter(id, value) VALUES(1, 0) ON CONFLICT (id) DO UPDATE SET value = 0")
        .execute(&pool).await?;
 
    // デフォルト設定（max_attempts: 3）で並行更新
    let occ_config = aurora_dsql_sqlx_connector::OCCRetryConfig::default();
    let mut handles = Vec::new();
 
    for task_id in 0..5 {
        let pool = pool.clone();
        let occ_config = occ_config.clone();
        handles.push(tokio::spawn(async move {
            let result: Result<(), DsqlError> = retry_on_occ(&occ_config, || async {
                let mut tx = pool.begin().await?;
                let row = sqlx::query("SELECT value FROM counter WHERE id = 1")
                    .fetch_one(&mut *tx).await?;
                let current: i32 = row.get("value");
                sqlx::query("UPDATE counter SET value = $1 WHERE id = 1")
                    .bind(current + 1).execute(&mut *tx).await?;
                tx.commit().await?;
                Ok(())
            }).await;
            match &result {
                Ok(()) => println!("Task {} succeeded", task_id),
                Err(e) => println!("Task {} failed: {}", task_id, e),
            }
            result
        }));
    }
 
    let (mut ok, mut ng) = (0, 0);
    for h in handles {
        match h.await? { Ok(()) => ok += 1, Err(_) => ng += 1 }
    }
    let final_val = sqlx::query_scalar::<_, i32>("SELECT value FROM counter WHERE id = 1")
        .fetch_one(&pool).await?;
    println!("Succeeded: {ok}, Failed: {ng}, Final value: {final_val}");
 
    pool.execute("DROP TABLE counter").await?;
    pool.close().await; // バックグラウンドのトークンリフレッシュタスクを停止
    Ok(())
}