Store data in mongodb

Cargo.toml

@@ -24,3 +24,4 @@ reqwest = "0.12.12"
 
 tracing = "0.1.37"
 tracing-subscriber = "0.3.16"
+mongodb = "3.2.1"

src/config.rs

@@ -13,6 +13,8 @@ pub struct Config {
 
     pub twitch_webhook_url: String,
     pub twitch_webhook_port: u16,
+
+    pub mongodb_connection_string: String,
 }
 
 impl Config {
@@ -41,6 +43,9 @@ impl Config {
                 .expect("TWITCH_WEBHOOK_PORT is not set")
                 .parse()
                 .expect("TWITCH_WEBHOOK_PORT is not a valid u16"),
+
+            mongodb_connection_string: std::env::var("MONGODB_CONNECTION_STRING")
+                .expect("MONGODB_CONNECTION_STRING is not set"),
         }
     }
 }

src/main.rs

@@ -17,7 +17,7 @@ async fn main() {
 
     let subscription_manager = Arc::new(SubscriptionManager::new());
 
-    subscription_manager.init().await;
+    subscription_manager.load().await;
 
     tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
 

src/subscription_manager.rs

@@ -1,7 +1,14 @@
 use std::collections::{HashMap, HashSet};
 
+use futures::StreamExt;
+use mongodb::{
+    Client, Collection,
+    bson::{Document, doc},
+};
 use tokio::sync::RwLock;
 
+use crate::config::CONFIG;
+
 pub struct SubscriptionManager {
     pub subscriptions: RwLock<HashMap<String, HashSet<u64>>>,
 }
@@ -13,19 +20,55 @@ impl SubscriptionManager {
         }
     }
 
-    pub async fn init(&self) {
-        tracing::debug!("Initializing subscription manager");
+    async fn get_collection() -> mongodb::error::Result<Collection<Document>> {
+        let client = Client::with_uri_str(CONFIG.mongodb_connection_string.clone()).await?;
+
+        let database = client.database("telegram-twitch-notifier");
+
+        Ok(database.collection("subscriptions"))
+    }
+
+    pub async fn load(&self) -> mongodb::error::Result<()> {
+        let collection = Self::get_collection().await?;
+
+        let mut subs = collection.find(doc! {}).await?;
+
+        while let Some(sub) = subs.next().await {
+            let sub = sub?;
+
+            let username = sub.get_str("streamer").unwrap();
+            let telegram_user_id = sub.get_i64("telegram_user_id").unwrap() as u64;
+
+            self.subscribe(telegram_user_id, username.to_string()).await;
+        }
+
+        Ok(())
     }
 
     pub async fn subscribe(&self, telegram_user_id: u64, username: String) {
         tracing::debug!("Subscribing {} to {}", telegram_user_id, username);
 
-        self.subscriptions
+        let inserted = self
+            .subscriptions
             .write()
             .await
-            .entry(username)
+            .entry(username.clone())
             .or_insert(HashSet::new())
             .insert(telegram_user_id);
+
+        if !inserted {
+            return;
+        }
+
+        Self::get_collection()
+            .await
+            .unwrap()
+            .insert_one(doc! {
+                "streamer": username,
+                "telegram_user_id": telegram_user_id as i64,
+            })
+            .await
+            .unwrap();
     }
 
     pub async fn unsubscribe(&self, telegram_user_id: u64, username: String) {
@@ -34,9 +77,19 @@ impl SubscriptionManager {
         self.subscriptions
             .write()
             .await
-            .entry(username)
+            .entry(username.clone())
             .and_modify(|set| {
                 set.remove(&telegram_user_id);
             });
+
+        Self::get_collection()
+            .await
+            .unwrap()
+            .delete_one(doc! {
+                "streamer": username,
+                "telegram_user_id": telegram_user_id as i64,
+            })
+            .await
+            .unwrap();
     }
 }

src/twitch_webhook.rs

@@ -24,9 +24,7 @@ use twitch_api::{
 use twitch_oauth2::AppAccessToken;
 
 use crate::{
-    config::CONFIG,
-    subscription_manager::{self, SubscriptionManager},
-    telegram_bot::get_telegram_bot,
+    config::CONFIG, subscription_manager::SubscriptionManager, telegram_bot::get_telegram_bot,
 };
 
 pub async fn eventsub_register(