Ingestor refactor complete baby~

ingestor/src/collections.rs

@@ -21,7 +21,9 @@ macro_rules! create_commit_collection {
     struct Ingestor;
 
     impl $crate::collections::Collection for $collection {
-      fn get_nsid() -> String {
+      fn new() -> Self { Self }
+
+      fn get_nsid(&self) -> String {
         $nsid
       }
 
@@ -77,7 +79,8 @@ pub use my_spoor_activity::MySpoorActivity;
 struct CommitIngestorState;
 
 pub trait Collection {
-  fn get_nsid() -> String;
+  fn new() -> Self;
+  fn get_nsid(&self) -> String;
   fn get_ingestor(&self) -> Box<dyn LexiconIngestor + Send + Sync>;
 }
 

ingestor/src/ingestor.rs

@@ -1,46 +1,43 @@
-use atproto::{
-  Collection,
-  lexicons::my::spoor::log::{Activity, Session},
-};
-use async_trait::async_trait;
+use crate::collections::Collection;
 use rocketman::{
   options::JetstreamOptions,
   ingestion::LexiconIngestor,
-  // types::event::{Event, Operation},
-  types::event::Event,
   connection::JetstreamConnection,
   handler,
 };
-use serde_json::value::Value;
 use std::{
   collections::HashMap,
   sync::{Arc, Mutex},
 };
-use tracing::{error, info};
+use tracing::{info, error};
 
-pub async fn start_ingestor() {
-  info!("Starting ingestor");
-  let nsids = vec![
-    Activity::NSID.to_string(),
-    Session::NSID.to_string(),
-  ];
-  let opts = JetstreamOptions::builder().wanted_collections(nsids.clone()).build();
-  let jetstream = JetstreamConnection::new(opts);
-
-  let mut ingesters: HashMap<String, Box<dyn LexiconIngestor + Send + Sync>> = HashMap::new();
-  for nsid in nsids {
-    ingesters.insert(nsid, Box::new(SpoorJetstream));
+pub struct Ingestor {
+  ingesters: HashMap<String, Box<dyn LexiconIngestor + Send + Sync>>,
 }
 
-  let cursor: Arc<Mutex<Option<u64>>> = Arc::new(Mutex::new(None));
+impl Ingestor {
+  pub fn new() -> Self {
+    Self { ingesters: HashMap::new() }
+  }
 
+  pub fn add_collection<C: Collection>(&mut self, collection: C) {
+    self.ingesters.insert(collection.get_nsid(), collection.get_ingestor());
+  }
+
+  pub async fn start(self) {
+    info!("Starting ingestor");
+    let opts = JetstreamOptions::builder()
+      .wanted_collections(self.ingesters.keys().cloned().collect())
+      .build();
+    let jetstream = JetstreamConnection::new(opts);
+    let cursor: Arc<Mutex<Option<u64>>> = Arc::new(Mutex::new(None));
     let msg_rx = jetstream.get_msg_rx();
     let reconnect_tx = jetstream.get_reconnect_tx();
 
     let cursor_clone = cursor.clone();
     tokio::spawn(async move {
       while let Ok(message) = msg_rx.recv_async().await {
-      if let Err(e) = handler::handle_message(message, &ingesters, 
+        if let Err(e) = handler::handle_message(message, &self.ingesters,
             reconnect_tx.clone(), cursor_clone.clone()).await {
           error!("Error processing message: {}", e);
         }
@@ -52,3 +49,4 @@ pub async fn start_ingestor() {
       std::process::exit(1);
     }
   }
+}

ingestor/src/main.rs

@@ -1,3 +1,10 @@
+use crate::{
+  collections::{
+    Collection,
+    MySpoorActivity, MySpoorSession,
+  },
+  ingestor::Ingestor,
+};
 
 mod ingestor;
 mod collections;
@@ -7,5 +14,8 @@ async fn main() {
   let subscriber = tracing_subscriber::FmtSubscriber::new();
   let _ = tracing::subscriber::set_global_default(subscriber);
 
-  
+  let mut ingestor = Ingestor::new();
+  ingestor.add_collection(MySpoorActivity::new());
+  ingestor.add_collection(MySpoorSession::new());
+  ingestor.start().await;
 }