From ffd246ca4ca76381f157b821f9c794cc0e53814e Mon Sep 17 00:00:00 2001
From: Lee Bradley <bradley@bitwise.io>
Date: Mon, 6 Jun 2022 15:36:07 -0500
Subject: [PATCH] Add batch tracking to the store factory

This change updates Grid so that batch tracking database functionality
can be accessed from the central store factory.

Signed-off-by: Lee Bradley <bradley@bitwise.io>
---
 sdk/src/store/mod.rs      |  4 ++++
 sdk/src/store/postgres.rs | 15 +++++++++++++++
 sdk/src/store/sqlite.rs   | 15 +++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/sdk/src/store/mod.rs b/sdk/src/store/mod.rs
index acaa5051e4..51adfd2487 100644
--- a/sdk/src/store/mod.rs
+++ b/sdk/src/store/mod.rs
@@ -22,6 +22,8 @@ use std::str::FromStr;
 #[cfg(feature = "diesel")]
 use diesel::r2d2::{ConnectionManager, Pool};
 
+#[cfg(feature = "batch-tracking")]
+use crate::batch_tracking::store::BatchTrackingStore;
 #[cfg(feature = "batch-store")]
 use crate::batches::store::BatchStore;
 use crate::commits::store::CommitStore;
@@ -62,6 +64,8 @@ pub trait StoreFactory {
     fn get_batch_store<'a>(&'a self) -> Box<dyn BatchStore + 'a>;
     #[cfg(feature = "purchase-order")]
     fn get_grid_purchase_order_store<'a>(&'a self) -> Box<dyn PurchaseOrderStore + 'a>;
+    #[cfg(feature = "batch-tracking")]
+    fn get_batch_tracking_store<'a>(&'a self) -> Box<dyn BatchTrackingStore + 'a>;
 }
 
 pub trait TransactionalStoreFactory: StoreFactory + Send + Sync {
diff --git a/sdk/src/store/postgres.rs b/sdk/src/store/postgres.rs
index a82bec7bfb..9fb5030969 100644
--- a/sdk/src/store/postgres.rs
+++ b/sdk/src/store/postgres.rs
@@ -19,6 +19,11 @@ use diesel::{
     Connection,
 };
 
+#[cfg(feature = "batch-tracking")]
+use crate::batch_tracking::store::{
+    diesel::{DieselBatchTrackingStore, DieselConnectionBatchTrackingStore},
+    BatchTrackingStore,
+};
 #[cfg(feature = "batch-store")]
 use crate::batches::store::{BatchStore, DieselBatchStore, DieselConnectionBatchStore};
 use crate::commits::store::{CommitStore, DieselCommitStore, DieselConnectionCommitStore};
@@ -93,6 +98,11 @@ impl StoreFactory for PgStoreFactory {
     fn get_grid_purchase_order_store<'a>(&'a self) -> Box<dyn PurchaseOrderStore + 'a> {
         Box::new(DieselPurchaseOrderStore::new(self.pool.clone()))
     }
+
+    #[cfg(feature = "batch-tracking")]
+    fn get_batch_tracking_store<'a>(&'a self) -> Box<dyn BatchTrackingStore + 'a> {
+        Box::new(DieselBatchTrackingStore::new(self.pool.clone()))
+    }
 }
 
 impl TransactionalStoreFactory for PgStoreFactory {
@@ -165,6 +175,11 @@ impl StoreFactory for InContextPgStoreFactory {
     fn get_grid_purchase_order_store<'a>(&'a self) -> Box<dyn PurchaseOrderStore + 'a> {
         Box::new(DieselConnectionPurchaseOrderStore::new(&*self.conn))
     }
+
+    #[cfg(feature = "batch-tracking")]
+    fn get_batch_tracking_store<'a>(&'a self) -> Box<dyn BatchTrackingStore + 'a> {
+        Box::new(DieselConnectionBatchTrackingStore::new(&*self.conn))
+    }
 }
 
 impl<'a> InContextStoreFactory<'a> for InContextPgStoreFactory {
diff --git a/sdk/src/store/sqlite.rs b/sdk/src/store/sqlite.rs
index d46264fdec..72aeb935af 100644
--- a/sdk/src/store/sqlite.rs
+++ b/sdk/src/store/sqlite.rs
@@ -19,6 +19,11 @@ use diesel::{
     Connection,
 };
 
+#[cfg(feature = "batch-tracking")]
+use crate::batch_tracking::store::{
+    diesel::{DieselBatchTrackingStore, DieselConnectionBatchTrackingStore},
+    BatchTrackingStore,
+};
 #[cfg(feature = "batch-store")]
 use crate::batches::store::{BatchStore, DieselBatchStore, DieselConnectionBatchStore};
 use crate::commits::store::{CommitStore, DieselCommitStore, DieselConnectionCommitStore};
@@ -93,6 +98,11 @@ impl StoreFactory for SqliteStoreFactory {
     fn get_grid_purchase_order_store<'a>(&'a self) -> Box<dyn PurchaseOrderStore + 'a> {
         Box::new(DieselPurchaseOrderStore::new(self.pool.clone()))
     }
+
+    #[cfg(feature = "batch-tracking")]
+    fn get_batch_tracking_store<'a>(&'a self) -> Box<dyn BatchTrackingStore + 'a> {
+        Box::new(DieselBatchTrackingStore::new(self.pool.clone()))
+    }
 }
 
 impl TransactionalStoreFactory for SqliteStoreFactory {
@@ -165,6 +175,11 @@ impl StoreFactory for InContextSqliteStoreFactory {
     fn get_grid_purchase_order_store<'a>(&'a self) -> Box<dyn PurchaseOrderStore + 'a> {
         Box::new(DieselConnectionPurchaseOrderStore::new(&*self.conn))
     }
+
+    #[cfg(feature = "batch-tracking")]
+    fn get_batch_tracking_store<'a>(&'a self) -> Box<dyn BatchTrackingStore + 'a> {
+        Box::new(DieselConnectionBatchTrackingStore::new(&*self.conn))
+    }
 }
 
 impl<'a> InContextStoreFactory<'a> for InContextSqliteStoreFactory {