feat(store): Create a trait for the taskbroker store

Cargo.toml

@@ -10,6 +10,7 @@ debug = 1
 
 [dependencies]
 anyhow = "1.0.92"
+async-trait = "0.1"
 bytes = "1.10.0"
 chrono = { version = "0.4.26" }
 clap = { version = "4.5.20", features = ["derive"] }

benches/store_bench.rs

@@ -6,6 +6,7 @@ use rand::Rng;
 use taskbroker::{
     store::inflight_activation::{
         InflightActivationStatus, InflightActivationStore, InflightActivationStoreConfig,
+        SqliteActivationStore,
     },
     test_utils::{
         generate_temp_filename, generate_unique_namespace, make_activations_with_namespace,
@@ -25,7 +26,7 @@ async fn get_pending_activations(num_activations: u32, num_workers: u32) {
         generate_temp_filename()
     };
     let store = Arc::new(
-        InflightActivationStore::new(
+        SqliteActivationStore::new(
             &url,
             InflightActivationStoreConfig {
                 max_processing_attempts: 1,
@@ -89,7 +90,7 @@ async fn set_status(num_activations: u32, num_workers: u32) {
     };
 
     let store = Arc::new(
-        InflightActivationStore::new(
+        SqliteActivationStore::new(
             &url,
             InflightActivationStoreConfig {
                 max_processing_attempts: 1,

src/grpc/server.rs

@@ -13,7 +13,7 @@ use crate::store::inflight_activation::{InflightActivationStatus, InflightActiva
 use tracing::{error, instrument};
 
 pub struct TaskbrokerServer {
-    pub store: Arc<InflightActivationStore>,
+    pub store: Arc<dyn InflightActivationStore>,
 }
 
 #[tonic::async_trait]

src/grpc/server_tests.rs

@@ -1,12 +1,12 @@
+use crate::grpc::server::TaskbrokerServer;
+use crate::store::inflight_activation::InflightActivationStore;
 use prost::Message;
 use sentry_protos::taskbroker::v1::consumer_service_server::ConsumerService;
 use sentry_protos::taskbroker::v1::{
     FetchNextTask, GetTaskRequest, SetTaskStatusRequest, TaskActivation,
 };
 use tonic::{Code, Request};
 
-use crate::grpc::server::TaskbrokerServer;
-
 use crate::test_utils::{create_test_store, make_activations};
 
 #[tokio::test]

src/kafka/inflight_activation_writer.rs

@@ -44,12 +44,12 @@ impl ActivationWriterConfig {
 
 pub struct InflightActivationWriter {
     config: ActivationWriterConfig,
-    store: Arc<InflightActivationStore>,
+    store: Arc<dyn InflightActivationStore>,
     batch: Option<Vec<InflightActivation>>,
 }
 
 impl InflightActivationWriter {
-    pub fn new(store: Arc<InflightActivationStore>, config: ActivationWriterConfig) -> Self {
+    pub fn new(store: Arc<dyn InflightActivationStore>, config: ActivationWriterConfig) -> Self {
         Self {
             config,
             store,
@@ -214,6 +214,7 @@ mod tests {
 
     use crate::store::inflight_activation::{
         InflightActivationStatus, InflightActivationStore, InflightActivationStoreConfig,
+        SqliteActivationStore,
     };
     use crate::test_utils::generate_unique_namespace;
     use crate::test_utils::make_activations;
@@ -231,7 +232,7 @@ mod tests {
         };
         let mut writer = InflightActivationWriter::new(
             Arc::new(
-                InflightActivationStore::new(
+                SqliteActivationStore::new(
                     &generate_temp_filename(),
                     InflightActivationStoreConfig::from_config(&create_integration_config()),
                 )
@@ -346,7 +347,7 @@ mod tests {
         };
         let mut writer = InflightActivationWriter::new(
             Arc::new(
-                InflightActivationStore::new(
+                SqliteActivationStore::new(
                     &generate_temp_filename(),
                     InflightActivationStoreConfig::from_config(&create_integration_config()),
                 )
@@ -415,7 +416,7 @@ mod tests {
         };
         let mut writer = InflightActivationWriter::new(
             Arc::new(
-                InflightActivationStore::new(
+                SqliteActivationStore::new(
                     &generate_temp_filename(),
                     InflightActivationStoreConfig::from_config(&create_integration_config()),
                 )
@@ -488,7 +489,7 @@ mod tests {
         };
         let mut writer = InflightActivationWriter::new(
             Arc::new(
-                InflightActivationStore::new(
+                SqliteActivationStore::new(
                     &generate_temp_filename(),
                     InflightActivationStoreConfig::from_config(&create_integration_config()),
                 )
@@ -604,7 +605,7 @@ mod tests {
         };
         let mut writer = InflightActivationWriter::new(
             Arc::new(
-                InflightActivationStore::new(
+                SqliteActivationStore::new(
                     &generate_temp_filename(),
                     InflightActivationStoreConfig::from_config(&create_integration_config()),
                 )
@@ -719,7 +720,7 @@ mod tests {
             write_failure_backoff_ms: 4000,
         };
         let store = Arc::new(
-            InflightActivationStore::new(
+            SqliteActivationStore::new(
                 &generate_temp_filename(),
                 InflightActivationStoreConfig::from_config(&create_integration_config()),
             )
@@ -879,7 +880,7 @@ mod tests {
             write_failure_backoff_ms: 4000,
         };
         let store = Arc::new(
-            InflightActivationStore::new(
+            SqliteActivationStore::new(
                 &generate_temp_filename(),
                 InflightActivationStoreConfig::from_config(&create_integration_config()),
             )
@@ -913,7 +914,7 @@ mod tests {
             write_failure_backoff_ms: 4000,
         };
         let store = Arc::new(
-            InflightActivationStore::new(
+            SqliteActivationStore::new(
                 &generate_temp_filename(),
                 InflightActivationStoreConfig::from_config(&create_integration_config()),
             )

src/main.rs

@@ -31,7 +31,7 @@ use taskbroker::metrics;
 use taskbroker::processing_strategy;
 use taskbroker::runtime_config::RuntimeConfigManager;
 use taskbroker::store::inflight_activation::{
-    InflightActivationStore, InflightActivationStoreConfig,
+    InflightActivationStore, InflightActivationStoreConfig, SqliteActivationStore,
 };
 use taskbroker::{Args, get_version};
 use tonic_health::ServingStatus;
@@ -62,8 +62,8 @@ async fn main() -> Result<(), Error> {
 
     logging::init(logging::LoggingConfig::from_config(&config));
     metrics::init(metrics::MetricsConfig::from_config(&config));
-    let store = Arc::new(
-        InflightActivationStore::new(
+    let store: Arc<dyn InflightActivationStore> = Arc::new(
+        SqliteActivationStore::new(
             &config.db_path,
             InflightActivationStoreConfig::from_config(&config),
         )