Enable to configure DynamoDbHandle

Author: coord-e

src/main.rs

@@ -9,7 +9,7 @@ use serenity::{
 
 use kaisantantoudaijin::{
     context::{ChannelContext, ContextBuilder},
-    database::RedisHandle,
+    database::{AnyDatabaseHandle, DynamoDbHandle, RedisHandle},
     model::message::Message,
 };
 
@@ -19,10 +19,40 @@ fn strip_affix<'a>(content: &'a str, affix: &str) -> Option<&'a str> {
         .or_else(|| content.strip_suffix(affix))
 }
 
+#[derive(Debug, Clone)]
+enum DatabaseClient {
+    Redis {
+        prefix: String,
+        pool: deadpool_redis::Pool,
+    },
+    DynamoDb {
+        table_name: String,
+        client: aws_sdk_dynamodb::Client,
+    },
+}
+
+impl DatabaseClient {
+    pub async fn obtain_handle(
+        &self,
+        guild_id: serenity::model::id::GuildId,
+    ) -> Result<AnyDatabaseHandle> {
+        match self {
+            DatabaseClient::Redis { prefix, pool } => {
+                let conn = pool.get().await?;
+                let db = RedisHandle::new(prefix.clone(), guild_id, conn);
+                Ok(db.into())
+            }
+            DatabaseClient::DynamoDb { table_name, client } => {
+                let db = DynamoDbHandle::new(client.clone(), guild_id, table_name.clone());
+                Ok(db.into())
+            }
+        }
+    }
+}
+
 struct Handler {
     command_prefix: String,
-    redis_prefix: String,
-    redis: deadpool_redis::Pool,
+    database_client: DatabaseClient,
 }
 
 #[async_trait::async_trait]
@@ -54,16 +84,15 @@ impl EventHandler for Handler {
             return;
         };
 
-        let redis_conn = match self.redis.get().await {
+        let db = match self.database_client.obtain_handle(guild_id).await {
             Ok(x) => x,
             Err(e) => {
-                tracing::error!("error in getting redis connection: {:#}", e);
+                tracing::error!("error in getting DB connection: {:#}", e);
                 let _ = msg.channel_id.say(&ctx.http, "エラーが発生しました").await;
                 return;
             }
         };
 
-        let db = RedisHandle::new(self.redis_prefix.clone(), guild_id, redis_conn);
         let ctx = ContextBuilder::with_serenity(&ctx)
             .db(db)
             .guild_id(guild_id)
@@ -88,6 +117,7 @@ impl EventHandler for Handler {
 
 #[derive(Parser)]
 #[command(group(clap::ArgGroup::new("tokens").required(true).multiple(false).args(["token", "token_file"])))]
+#[command(group(clap::ArgGroup::new("database_config").required(true).multiple(false).args(["redis_uri", "dynamodb_table_name"])))]
 struct Args {
     #[arg(long, default_value = "!kaisan", env = "KAISANDAIJIN_COMMAND_PREFIX")]
     command_prefix: String,
@@ -96,14 +126,16 @@ struct Args {
     #[arg(long, env = "KAISANDAIJIN_DISCORD_TOKEN_FILE")]
     token_file: Option<PathBuf>,
     #[arg(short, long, env = "KAISANDAIJIN_REDIS_URI")]
-    redis_uri: String,
+    redis_uri: Option<String>,
     #[arg(
         short = 'p',
         long,
         default_value = "kaisandaijin",
         env = "KAISANDAIJIN_REDIS_PREFIX"
     )]
     redis_prefix: String,
+    #[arg(short, long, env = "KAISANDAIJIN_DYNAMODB_TABLE_NAME")]
+    dynamodb_table_name: Option<String>,
     /// Specify log level filter, configured in conjunction with KAISANDAIJIN_LOG environment variable
     #[arg(short, long)]
     log_level: Option<tracing_subscriber::filter::LevelFilter>,
@@ -113,9 +145,6 @@ struct Args {
 async fn main() -> Result<()> {
     let args = Args::parse();
 
-    let redis = deadpool_redis::Config::from_url(args.redis_uri)
-        .create_pool(Some(deadpool_redis::Runtime::Tokio1))?;
-
     let token = if let Some(token) = args.token {
         token
     } else {
@@ -135,6 +164,22 @@ async fn main() -> Result<()> {
         .with_writer(std::io::stderr)
         .init();
 
+    let database_client = match (&args.redis_uri, &args.dynamodb_table_name) {
+        (Some(redis_uri), None) => DatabaseClient::Redis {
+            prefix: args.redis_prefix.clone(),
+            pool: deadpool_redis::Config::from_url(redis_uri)
+                .create_pool(Some(deadpool_redis::Runtime::Tokio1))?,
+        },
+        (None, Some(table_name)) => DatabaseClient::DynamoDb {
+            table_name: table_name.clone(),
+            client: {
+                let config = aws_config::load_from_env().await;
+                aws_sdk_dynamodb::Client::new(&config)
+            },
+        },
+        _ => anyhow::bail!("either Redis URI or DynamoDB table name must be specified"),
+    };
+
     let intents = [
         GatewayIntents::GUILDS,
         GatewayIntents::GUILD_MESSAGES,
@@ -146,8 +191,7 @@ async fn main() -> Result<()> {
     let mut client = Client::builder(token, intents)
         .event_handler(Handler {
             command_prefix: args.command_prefix,
-            redis_prefix: args.redis_prefix,
-            redis,
+            database_client,
         })
         .await
         .context("Failed to create client")?;