feat(setup): batch applying setup changes for the same target type (#461)

Author: badmonster0

src/execution/db_tracking_setup.rs

@@ -102,7 +102,6 @@ impl TrackingTableSetupStatus {
     }
 }
 
-#[async_trait]
 impl ResourceSetupStatus for TrackingTableSetupStatus {
     fn describe_changes(&self) -> Vec<String> {
         let mut changes: Vec<String> = vec![];
@@ -157,7 +156,13 @@ impl ResourceSetupStatus for TrackingTableSetupStatus {
         }
     }
 
-    async fn apply_change(&self) -> Result<()> {
+    fn as_any(&self) -> &dyn Any {
+        self as &dyn Any
+    }
+}
+
+impl TrackingTableSetupStatus {
+    pub async fn apply_change(&self) -> Result<()> {
         let pool = &get_lib_context()?.builtin_db_pool;
         if let Some(desired) = &self.desired_state {
             for lagacy_name in self.legacy_table_names.iter() {

src/ops/factory_bases.rs

@@ -1,4 +1,5 @@
 use crate::prelude::*;
+use crate::setup::ResourceSetupStatus;
 use std::fmt::Debug;
 use std::hash::Hash;
 
@@ -288,6 +289,7 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
     type DeclarationSpec: DeserializeOwned + Send + Sync;
     type Key: Debug + Clone + Serialize + DeserializeOwned + Eq + Hash + Send + Sync;
     type SetupState: Debug + Clone + Serialize + DeserializeOwned + Send + Sync;
+    type SetupStatus: ResourceSetupStatus;
     type ExportContext: Send + Sync + 'static;
 
     fn name(&self) -> &str;
@@ -310,7 +312,7 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
         desired_state: Option<Self::SetupState>,
         existing_states: setup::CombinedState<Self::SetupState>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl setup::ResourceSetupStatus + 'static>;
+    ) -> Result<Self::SetupStatus>;
 
     fn check_state_compatibility(
         &self,
@@ -334,6 +336,11 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
         &self,
         mutations: Vec<ExportTargetMutationWithContext<'async_trait, Self::ExportContext>>,
     ) -> Result<()>;
+
+    async fn apply_setup_changes(
+        &self,
+        setup_status: Vec<&'async_trait Self::SetupStatus>,
+    ) -> Result<()>;
 }
 
 #[async_trait]
@@ -456,8 +463,25 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
             .collect::<Result<_>>()?;
         StorageFactoryBase::apply_mutation(self, mutations).await
     }
-}
 
+    async fn apply_setup_changes(
+        &self,
+        setup_status: Vec<&'async_trait dyn ResourceSetupStatus>,
+    ) -> Result<()> {
+        StorageFactoryBase::apply_setup_changes(
+            self,
+            setup_status
+                .into_iter()
+                .map(|s| -> anyhow::Result<_> {
+                    Ok(s.as_any()
+                        .downcast_ref::<T::SetupStatus>()
+                        .ok_or_else(|| anyhow!("Unexpected setup status type"))?)
+                })
+                .collect::<Result<Vec<_>>>()?,
+        )
+        .await
+    }
+}
 fn from_json_combined_state<T: Debug + Clone + Serialize + DeserializeOwned>(
     existing_states: setup::CombinedState<serde_json::Value>,
 ) -> Result<setup::CombinedState<T>> {

src/ops/interface.rs

@@ -214,6 +214,11 @@ pub trait ExportTargetFactory: Send + Sync {
         &self,
         mutations: Vec<ExportTargetMutationWithContext<'async_trait, dyn Any + Send + Sync>>,
     ) -> Result<()>;
+
+    async fn apply_setup_changes(
+        &self,
+        setup_status: Vec<&'async_trait dyn setup::ResourceSetupStatus>,
+    ) -> Result<()>;
 }
 
 #[derive(Clone)]

src/ops/storages/neo4j.rs

@@ -1,7 +1,7 @@
 use crate::prelude::*;
 
 use super::spec::{GraphDeclaration, GraphElementMapping, NodeFromFieldsSpec, TargetFieldMapping};
-use crate::setup::components::{self, State};
+use crate::setup::components::{self, apply_component_changes, State};
 use crate::setup::{ResourceSetupStatus, SetupChangeType};
 use crate::{ops::sdk::*, setup::CombinedState};
 
@@ -688,7 +688,7 @@ impl ComponentKind {
     }
 }
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
-struct ComponentKey {
+pub struct ComponentKey {
     kind: ComponentKind,
     name: String,
 }
@@ -754,13 +754,13 @@ impl components::State<ComponentKey> for ComponentState {
     }
 }
 
-struct SetupComponentOperator {
+pub struct SetupComponentOperator {
     graph_pool: Arc<GraphPool>,
     conn_spec: ConnectionSpec,
 }
 
 #[async_trait]
-impl components::Operator for SetupComponentOperator {
+impl components::SetupOperator for SetupComponentOperator {
     type Key = ComponentKey;
     type State = ComponentState;
     type SetupState = SetupState;
@@ -855,7 +855,7 @@ fn build_composite_field_names(qualifier: &str, field_names: &[String]) -> Strin
 }
 #[derive(Derivative)]
 #[derivative(Debug)]
-struct SetupStatus {
+pub struct GraphElementDataSetupStatus {
     key: GraphElement,
     #[derivative(Debug = "ignore")]
     graph_pool: Arc<GraphPool>,
@@ -864,7 +864,7 @@ struct SetupStatus {
     change_type: SetupChangeType,
 }
 
-impl SetupStatus {
+impl GraphElementDataSetupStatus {
     fn new(
         key: GraphElement,
         graph_pool: Arc<GraphPool>,
@@ -907,8 +907,7 @@ impl SetupStatus {
     }
 }
 
-#[async_trait]
-impl ResourceSetupStatus for SetupStatus {
+impl ResourceSetupStatus for GraphElementDataSetupStatus {
     fn describe_changes(&self) -> Vec<String> {
         let mut result = vec![];
         if let Some(data_clear) = &self.data_clear {
@@ -934,6 +933,12 @@ impl ResourceSetupStatus for SetupStatus {
         self.change_type
     }
 
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
+impl GraphElementDataSetupStatus {
     async fn apply_change(&self) -> Result<()> {
         let graph = self.graph_pool.get_graph(&self.conn_spec).await?;
         if let Some(data_clear) = &self.data_clear {
@@ -1077,6 +1082,10 @@ impl StorageFactoryBase for Factory {
     type Spec = Spec;
     type DeclarationSpec = Declaration;
     type SetupState = SetupState;
+    type SetupStatus = (
+        GraphElementDataSetupStatus,
+        components::SetupStatus<SetupComponentOperator>,
+    );
     type Key = GraphElement;
     type ExportContext = ExportContext;
 
@@ -1255,24 +1264,24 @@ impl StorageFactoryBase for Factory {
         desired: Option<SetupState>,
         existing: CombinedState<SetupState>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl ResourceSetupStatus + 'static> {
+    ) -> Result<Self::SetupStatus> {
         let conn_spec = auth_registry.get::<ConnectionSpec>(&key.connection)?;
-        let base = SetupStatus::new(
+        let data_status = GraphElementDataSetupStatus::new(
             key,
             self.graph_pool.clone(),
             conn_spec.clone(),
             desired.as_ref(),
             &existing,
         );
-        let comp = components::Status::create(
+        let components = components::SetupStatus::create(
             SetupComponentOperator {
                 graph_pool: self.graph_pool.clone(),
                 conn_spec: conn_spec.clone(),
             },
             desired,
             existing,
         )?;
-        Ok(components::combine_setup_statuss(base, comp))
+        Ok((data_status, components))
     }
 
     fn check_state_compatibility(
@@ -1328,4 +1337,15 @@ impl StorageFactoryBase for Factory {
         }
         Ok(())
     }
+
+    async fn apply_setup_changes(
+        &self,
+        changes: Vec<&'async_trait Self::SetupStatus>,
+    ) -> Result<()> {
+        for change in changes.iter() {
+            change.0.apply_change().await?;
+        }
+        apply_component_changes(changes.iter().map(|c| &c.1).collect()).await?;
+        Ok(())
+    }
 }

src/ops/storages/postgres.rs

@@ -738,7 +738,6 @@ fn describe_index_spec(index_name: &str, index_spec: &VectorIndexDef) -> String
     format!("{} {}", index_name, to_index_spec_sql(index_spec))
 }
 
-#[async_trait]
 impl setup::ResourceSetupStatus for SetupStatus {
     fn describe_changes(&self) -> Vec<String> {
         let mut descriptions = vec![];
@@ -821,6 +820,12 @@ impl setup::ResourceSetupStatus for SetupStatus {
         }
     }
 
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}
+
+impl SetupStatus {
     async fn apply_change(&self) -> Result<()> {
         let table_name = &self.table_name;
         if self.drop_existing {
@@ -907,6 +912,7 @@ impl StorageFactoryBase for Factory {
     type Spec = Spec;
     type DeclarationSpec = ();
     type SetupState = SetupState;
+    type SetupStatus = SetupStatus;
     type Key = TableId;
     type ExportContext = ExportContext;
 
@@ -976,7 +982,7 @@ impl StorageFactoryBase for Factory {
         desired: Option<SetupState>,
         existing: setup::CombinedState<SetupState>,
         auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl setup::ResourceSetupStatus + 'static> {
+    ) -> Result<SetupStatus> {
         Ok(SetupStatus::new(
             get_db_pool(key.database.as_ref(), auth_registry).await?,
             key.table_name,
@@ -1049,4 +1055,14 @@ impl StorageFactoryBase for Factory {
         }
         Ok(())
     }
+
+    async fn apply_setup_changes(
+        &self,
+        setup_status: Vec<&'async_trait Self::SetupStatus>,
+    ) -> Result<()> {
+        for setup_status in setup_status.iter() {
+            setup_status.apply_change().await?;
+        }
+        Ok(())
+    }
 }

src/ops/storages/qdrant.rs

@@ -332,6 +332,7 @@ impl StorageFactoryBase for Arc<Factory> {
     type Spec = Spec;
     type DeclarationSpec = ();
     type SetupState = ();
+    type SetupStatus = Infallible;
     type Key = String;
     type ExportContext = ExportContext;
 
@@ -390,7 +391,7 @@ impl StorageFactoryBase for Arc<Factory> {
         _desired: Option<()>,
         _existing: setup::CombinedState<()>,
         _auth_registry: &Arc<AuthRegistry>,
-    ) -> Result<impl setup::ResourceSetupStatus + 'static> {
+    ) -> Result<Self::SetupStatus> {
         Err(anyhow!("Set `setup_by_user` to `true` to export to Qdrant")) as Result<Infallible, _>
     }
 
@@ -418,4 +419,11 @@ impl StorageFactoryBase for Arc<Factory> {
         }
         Ok(())
     }
+
+    async fn apply_setup_changes(
+        &self,
+        _setup_status: Vec<&'async_trait Self::SetupStatus>,
+    ) -> Result<()> {
+        Err(anyhow!("Qdrant does not support setup changes"))
+    }
 }