Skip to main content

api_handler/
db.rs

1//! # データベース接続モジュール
2//!
3//! このモジュールは、Amazon Aurora DSQL への接続管理を担当します。
4//! [`aurora_dsql_sqlx_connector`] クレートを使用して IAM 認証付きの接続プールを構築し、
5//! [`sea_orm`] の [`DatabaseConnection`] として提供します。
6//!
7//! ## Aurora DSQL について
8//!
9//! Amazon Aurora DSQL は、IAM ロールベースの認証を使用するサーバーレス分散 SQL データベースです。
10//! 通常の PostgreSQL と異なり、ユーザー名はロール名、パスワードは IAM 認証トークンで自動生成されます。
11//! このモジュールでは [`aurora_dsql_sqlx_connector`] がトークン生成と更新を自動的に処理します。
12//!
13//! ## 接続文字列の形式
14//!
15//! ```text
16//! postgres://<role>@<endpoint>/postgres?region=<region>
17//! ```
18//!
19//! - `<role>`: データベースロール名(例: `lambda`, `selectview`)
20//! - `<endpoint>`: Aurora DSQL クラスターのエンドポイント(例: `abc123.dsql.ap-northeast-1.on.aws`)
21//! - `<region>`: AWSリージョン(例: `ap-northeast-3`)
22//!
23//! ## 使用するロール
24//!
25//! Lambda 関数では以下のロールを使用します:
26//! - `lambda`: `messages` テーブルの SELECT / INSERT に使用
27
28use aurora_dsql_sqlx_connector::pool;
29use lambda_runtime::Error;
30use sea_orm::{DatabaseConnection, SqlxPostgresConnector};
31
32/// Aurora DSQL への接続文字列を構築する
33///
34/// 指定されたロール名、エンドポイント、リージョンから PostgreSQL 形式の接続文字列を生成します。
35/// [`aurora_dsql_sqlx_connector`] がこの文字列を解析して IAM 認証トークンを取得し、
36/// 実際のデータベース接続を確立します。
37///
38/// # Arguments
39///
40/// * `role` - データベースロール名(例: `"lambda"`, `"selectview"`)
41/// * `endpoint` - Aurora DSQL クラスターのエンドポイントホスト名
42///   (例: `"abc123.dsql.ap-northeast-1.on.aws"`)
43/// * `region` - Aurora DSQL クラスターが存在する AWS リージョン
44///   (例: `"ap-northeast-3"`)
45///
46/// # Returns
47///
48/// PostgreSQL URI 形式の接続文字列。
49/// `aurora_dsql_sqlx_connector` に渡すことで、IAM 認証が自動的に処理されます。
50///
51/// # Examples
52///
53/// ```
54/// let conn_str = build_connection_string(
55///     "lambda",
56///     "abc123.dsql.ap-northeast-1.on.aws",
57///     "ap-northeast-1"
58/// );
59/// assert_eq!(
60///     conn_str,
61///     "postgres://lambda@abc123.dsql.ap-northeast-1.on.aws/postgres?region=ap-northeast-1"
62/// );
63/// ```
64fn build_connection_string(role: &str, endpoint: &str, region: &str) -> String {
65    format!("postgres://{role}@{endpoint}/postgres?region={region}")
66}
67
68/// Aurora DSQL への SeaORM データベース接続を作成する
69///
70/// 指定されたロール・エンドポイント・リージョンを使用して Aurora DSQL への接続プールを構築し、
71/// [`sea_orm::DatabaseConnection`] として返します。
72///
73/// 接続確立には IAM 認証が使用されます。[`aurora_dsql_sqlx_connector`] が AWS STS と通信して
74/// 認証トークンを自動取得します。そのため、Lambda 関数の実行ロールに
75/// `dsql:DbConnectAdmin` または `dsql:DbConnect` 権限が必要です。
76///
77/// # Arguments
78///
79/// * `role` - Aurora DSQL のデータベースロール名(例: `"lambda"`)
80/// * `endpoint` - Aurora DSQL クラスターのエンドポイントホスト名
81/// * `region` - Aurora DSQL クラスターが存在する AWS リージョン
82///
83/// # Returns
84///
85/// * `Ok(DatabaseConnection)` - 正常に接続が確立された場合
86/// * `Err(Error)` - 接続に失敗した場合(IAM権限不足、ネットワークエラー等)
87///
88/// # Errors
89///
90/// - `aurora_dsql_sqlx_connector::pool::connect` が失敗した場合(IAM認証エラー、接続拒否等)
91///   `"Failed to connect to database: ..."` メッセージを含む [`anyhow::Error`] を返します。
92///
93/// # Panics
94///
95/// この関数はパニックしません。
96pub(crate) async fn create_db(
97    role: &str,
98    endpoint: &str,
99    region: &str,
100) -> Result<DatabaseConnection, Error> {
101    tracing::info!("Creating database connection with Aurora DSQL SQLx connector...");
102    let connection_string = build_connection_string(role, endpoint, region);
103    let pool = pool::connect(&connection_string)
104        .await
105        .map_err(|e| anyhow::anyhow!("Failed to connect to database: {}", e))?;
106
107    Ok(SqlxPostgresConnector::from_sqlx_postgres_pool(pool))
108}
109
110#[cfg(test)]
111mod tests {
112    use super::build_connection_string;
113
114    #[test]
115    fn test_build_connection_string() {
116        let endpoint = "foo0bar1baz2quux3quuux4.dsql.ap-northeast-1.on.aws";
117        let region = "ap-northeast-1";
118        let role = "selectview";
119
120        let result = build_connection_string(role, endpoint, region);
121
122        assert_eq!(
123            result,
124            "postgres://selectview@foo0bar1baz2quux3quuux4.dsql.ap-northeast-1.on.aws/postgres?region=ap-northeast-1"
125        );
126    }
127}
128
129#[cfg(test)]
130mod prop_tests {
131    use super::build_connection_string;
132    use proptest::prelude::*;
133
134    proptest! {
135        /// 任意のエンドポイントとリージョンで接続文字列が正しい形式になる
136        #[test]
137        fn prop_connection_string_contains_endpoint_and_region(
138            endpoint in "[a-z0-9][a-z0-9\\-]{0,30}\\.[a-z0-9\\.]{1,30}",
139            region in "[a-z]{2,6}-[a-z]{4,10}-[0-9]",
140        ) {
141            let role = "selectview";
142            let result = build_connection_string(role, &endpoint, &region);
143            let role_prefix = format!("postgres://{role}@");
144            prop_assert!(result.starts_with(&role_prefix));
145            prop_assert!(result.contains(&endpoint));
146            let region_param = format!("region={}", region);
147            prop_assert!(result.contains(&region_param));
148            prop_assert!(result.contains("/postgres?"));
149        }
150
151        /// 接続文字列は常に postgres:// スキームで始まる
152        #[test]
153        fn prop_connection_string_scheme(
154            endpoint in "[a-z0-9\\.]{5,40}",
155            region in "[a-z0-9\\-]{5,20}",
156        ) {
157            let role = "selectview";
158            let result = build_connection_string(role, &endpoint, &region);
159            let role_prefix = format!("postgres://{role}@");
160            prop_assert!(result.starts_with(&role_prefix));
161        }
162    }
163}