refactor: pull out database code

Co-authored-by: mxttc <mxttc@users.noreply.github.com>

Author: c-git

Cargo.lock

@@ -14,7 +14,6 @@ serde_json = "1.0.143"
 shuttle-runtime = { version = "0.57.0", default-features = false }
 shuttle-serenity = "0.57.0" # Since poise is a serenity command framework, it can run on Shuttle with shuttle-serenity
 shuttle-shared-db = { version = "0.57.0", default-features = false, features = ["sqlx", "postgres"] }
-sqlx = { version = "0.8.6", default-features = false, features = ["postgres", "runtime-tokio-rustls"] }
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.20", features = ["env-filter"] }
 version = "3.0.0"

Cargo.toml

@@ -20,7 +20,6 @@ pub struct SharedConfig {
     pub start_instant: Instant,
     pub auth_role_id: RoleId,
     pub channel_unranked: ChannelId,
-    pub db_pool: sqlx::PgPool,
     pub channel_bot_status: Option<ChannelId>,
 }
 
@@ -49,30 +48,26 @@ impl StartupConfig {
 }
 
 impl SharedConfig {
-    pub fn try_new(
-        secret_store: &SecretStore,
-        db_pool: sqlx::PgPool,
-    ) -> anyhow::Result<&'static Self> {
+    const KEY_VALUE_STORE_FOLDER: &str = "KV";
+    pub fn try_new(secret_store: &SecretStore) -> anyhow::Result<&'static Self> {
         let auth_role_id = KeyName::AuthRoleId.get_non_secret_parse(secret_store)?;
         let channel_unranked = KeyName::ChannelUnrankedId.get_non_secret_parse(secret_store)?;
         let channel_bot_status = KeyName::ChannelBotStatus.get_non_secret_parse_opt(secret_store);
         let result = Box::new(Self {
             start_instant: Instant::now(),
             auth_role_id,
             channel_unranked,
-            db_pool,
             channel_bot_status,
         });
         Ok(Box::leak(result))
     }
 
     /// Doesn't actually perform the save but spawns a task to do it in the background
     pub fn save_kv<T: serde::Serialize>(&self, key: &str, value: &T) -> anyhow::Result<()> {
-        let pool = self.db_pool.clone();
         let key = key.to_string();
         let value = serde_json::to_string(value).context("failed to convert to json")?;
         tokio::spawn(async move {
-            crate::db::save_kv(&pool, &key, value).await;
+            crate::db::save_kv(&key, value).await;
         });
         Ok(())
     }
@@ -81,7 +76,7 @@ impl SharedConfig {
         &self,
         key: &str,
     ) -> T {
-        let Some(content) = crate::db::load_kv(&self.db_pool, key).await else {
+        let Some(content) = crate::db::load_kv(key).await else {
             return T::default();
         };
         match serde_json::from_str(&content) {

src/config.rs

@@ -1,48 +1,48 @@
-use tracing::{debug, error, info};
-
-pub async fn save_kv(pool: &sqlx::PgPool, key: &str, value: String) {
-    let query = sqlx::query!(
-        "\
-INSERT INTO kv_store (id, content)
-    VALUES ($1, $2)
-    ON CONFLICT(id)
-    DO UPDATE SET
-    content = EXCLUDED.content;",
-        key,
-        value
-    );
-    match query.execute(pool).await {
-        Ok(query_result) => {
-            if query_result.rows_affected() == 1 {
-                debug!("Save completed for key: {key}");
-            } else {
-                error!(
-                    ?key,
-                    "Expected 1 row to be affected by save but got: {}",
-                    query_result.rows_affected()
-                )
-            }
-        }
-        Err(err_msg) => error!(
-            ?err_msg,
-            "Failed to save content for key: {key} to kv store"
-        ),
-    }
+pub async fn save_kv(key: &str, value: String) {
+    todo!("save to disk")
+    //     let query = sqlx::query!(
+    //         "\
+    // INSERT INTO kv_store (id, content)
+    //     VALUES ($1, $2)
+    //     ON CONFLICT(id)
+    //     DO UPDATE SET
+    //     content = EXCLUDED.content;",
+    //         key,
+    //         value
+    //     );
+    //     match query.execute(pool).await {
+    //         Ok(query_result) => {
+    //             if query_result.rows_affected() == 1 {
+    //                 debug!("Save completed for key: {key}");
+    //             } else {
+    //                 error!(
+    //                     ?key,
+    //                     "Expected 1 row to be affected by save but got: {}",
+    //                     query_result.rows_affected()
+    //                 )
+    //             }
+    //         }
+    //         Err(err_msg) => error!(
+    //             ?err_msg,
+    //             "Failed to save content for key: {key} to kv store"
+    //         ),
+    //     }
 }
 
-pub async fn load_kv(pool: &sqlx::PgPool, key: &str) -> Option<String> {
-    match sqlx::query!("SELECT content FROM kv_store where id = $1", key)
-        .fetch_optional(pool)
-        .await
-    {
-        Ok(Some(record)) => Some(record.content),
-        Ok(None) => {
-            info!("No content found in DB for key: {key}");
-            None
-        }
-        Err(err_msg) => {
-            error!(?err_msg, "Failed to get content for key: {key}");
-            None
-        }
-    }
+pub async fn load_kv(key: &str) -> Option<String> {
+    todo!("load from disk")
+    // match sqlx::query!("SELECT content FROM kv_store where id = $1", key)
+    //     .fetch_optional(pool)
+    //     .await
+    // {
+    //     Ok(Some(record)) => Some(record.content),
+    //     Ok(None) => {
+    //         info!("No content found in DB for key: {key}");
+    //         None
+    //     }
+    //     Err(err_msg) => {
+    //         error!(?err_msg, "Failed to get content for key: {key}");
+    //         None
+    //     }
+    // }
 }

src/db.rs

@@ -9,7 +9,7 @@ use tracing::{error, info};
 
 const KEY: &str = "HEARTBEAT";
 
-pub fn start_heartbeat(db_pool: sqlx::PgPool) {
+pub fn start_heartbeat() {
     shuttle_runtime::tokio::spawn(async move {
         info!("Heartbeat started");
         loop {
@@ -20,14 +20,14 @@ pub fn start_heartbeat(db_pool: sqlx::PgPool) {
                     break;
                 }
             };
-            save_kv(&db_pool, KEY, timestamp.to_db_fmt()).await;
+            save_kv(KEY, timestamp.to_db_fmt()).await;
             shuttle_runtime::tokio::time::sleep(std::time::Duration::from_secs(600)).await;
         }
     });
 }
 
-pub async fn last_heartbeat_info(db_pool: sqlx::PgPool) -> String {
-    match load_kv(&db_pool, KEY).await {
+pub async fn last_heartbeat_info() -> String {
+    match load_kv(KEY).await {
         Some(db_value) => match UnixTimestamp::from_db_fmt(&db_value) {
             Ok(last_heartbeat) => {
                 let Ok(now) = UnixTimestamp::now() else {

src/heartbeat.rs

@@ -16,13 +16,7 @@ use tracing_subscriber::{
 use version::version;
 
 #[shuttle_runtime::main]
-async fn main(
-    #[shuttle_runtime::Secrets] secret_store: SecretStore,
-    #[shuttle_shared_db::Postgres(
-        local_uri = "postgres://db_user:password@localhost:5432/bazooka_bot"
-    )]
-    db_pool: sqlx::PgPool,
-) -> ShuttleSerenity {
+async fn main(#[shuttle_runtime::Secrets] secret_store: SecretStore) -> ShuttleSerenity {
     tracing_subscriber::registry()
         .with(fmt::layer().with_span_events(FmtSpan::NEW | FmtSpan::CLOSE))
         .with(
@@ -33,17 +27,12 @@ async fn main(
 
     info!("Bot version is {}", version::version!());
 
-    sqlx::migrate!("./migrations")
-        .run(&db_pool)
-        .await
-        .expect("Migrations failed");
-
     // Load setup values
     let discord_token = get_secret_discord_token(&secret_store)?;
     let startup_config =
         StartupConfig::try_new(&secret_store).context("failed to create setup config")?;
-    let shared_config = SharedConfig::try_new(&secret_store, db_pool.clone())
-        .context("failed to created shared_config")?;
+    let shared_config =
+        SharedConfig::try_new(&secret_store).context("failed to created shared_config")?;
 
     let framework = poise::Framework::builder()
         .options(poise::FrameworkOptions {
@@ -86,7 +75,7 @@ async fn main(
                 let connect_msg = format!(
                     "{} is connected! Version: {}\n{}", 
                     ready.user.name, version!(),
-                    heartbeat::last_heartbeat_info(db_pool.clone()).await,
+                    heartbeat::last_heartbeat_info().await,
                 );
                 info!("{connect_msg}");
                 if let Some(channel) = shared_config.channel_bot_status{
@@ -95,7 +84,7 @@ async fn main(
                     warn!("Not sending connection notification because channel_bot_status not set");
                 }
                 let data = Data::new(shared_config, ctx.clone()).await;
-                heartbeat::start_heartbeat(db_pool);
+                heartbeat::start_heartbeat();
                 info!("END OF SETUP CLOSURE");
                 Ok(data)
             })