feat: plumbing FlowInstanceContext to more setup methods

Author: georgeh0

src/ops/factory_bases.rs

@@ -358,7 +358,7 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
     async fn apply_setup_changes(
         &self,
         setup_status: Vec<TypedResourceSetupChangeItem<'async_trait, Self>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()>;
 }
 
@@ -499,7 +499,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
     async fn apply_setup_changes(
         &self,
         setup_status: Vec<ResourceSetupChangeItem<'async_trait>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()> {
         StorageFactoryBase::apply_setup_changes(
             self,
@@ -514,7 +514,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
                     })
                 })
                 .collect::<Result<Vec<_>>>()?,
-            auth_registry,
+            context,
         )
         .await
     }

src/ops/interface.rs

@@ -302,7 +302,7 @@ pub trait ExportTargetFactory: Send + Sync {
     async fn apply_setup_changes(
         &self,
         setup_status: Vec<ResourceSetupChangeItem<'async_trait>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()>;
 }
 

src/ops/targets/kuzu.rs

@@ -1040,7 +1040,7 @@ impl StorageFactoryBase for Factory {
     async fn apply_setup_changes(
         &self,
         changes: Vec<TypedResourceSetupChangeItem<'async_trait, Self>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()> {
         let mut changes_by_conn = IndexMap::new();
         for change in changes.into_iter() {
@@ -1050,7 +1050,7 @@ impl StorageFactoryBase for Factory {
                 .push(change);
         }
         for (conn, changes) in changes_by_conn.into_iter() {
-            let conn_spec = auth_registry.get::<ConnectionSpec>(&conn)?;
+            let conn_spec = context.auth_registry.get::<ConnectionSpec>(&conn)?;
             let kuzu_client = KuzuThinClient::new(&conn_spec, self.reqwest_client.clone());
 
             let (node_changes, rel_changes): (Vec<_>, Vec<_>) =

src/ops/targets/neo4j.rs

@@ -1084,7 +1084,7 @@ impl StorageFactoryBase for Factory {
     async fn apply_setup_changes(
         &self,
         changes: Vec<TypedResourceSetupChangeItem<'async_trait, Self>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()> {
         // Relationships first, then nodes, as relationships need to be deleted before nodes they referenced.
         let mut relationship_types = IndexSet::<&Neo4jGraphElement>::new();
@@ -1116,15 +1116,15 @@ impl StorageFactoryBase for Factory {
         for rel_type in relationship_types.into_iter() {
             let graph = self
                 .graph_pool
-                .get_graph_for_key(rel_type, auth_registry)
+                .get_graph_for_key(rel_type, &context.auth_registry)
                 .await?;
             clear_graph_element_data(&graph, rel_type, true).await?;
         }
         // Clear standalone nodes, which is simpler than dependent nodes.
         for node_label in node_labels.iter() {
             let graph = self
                 .graph_pool
-                .get_graph_for_key(node_label, auth_registry)
+                .get_graph_for_key(node_label, &context.auth_registry)
                 .await?;
             clear_graph_element_data(&graph, node_label, true).await?;
         }
@@ -1133,7 +1133,7 @@ impl StorageFactoryBase for Factory {
             if !node_labels.contains(node_label) {
                 let graph = self
                     .graph_pool
-                    .get_graph_for_key(node_label, auth_registry)
+                    .get_graph_for_key(node_label, &context.auth_registry)
                     .await?;
                 clear_graph_element_data(&graph, node_label, false).await?;
             }

src/ops/targets/postgres.rs

@@ -777,10 +777,10 @@ impl StorageFactoryBase for Factory {
     async fn apply_setup_changes(
         &self,
         changes: Vec<TypedResourceSetupChangeItem<'async_trait, Self>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()> {
         for change in changes.iter() {
-            let db_pool = get_db_pool(change.key.database.as_ref(), auth_registry).await?;
+            let db_pool = get_db_pool(change.key.database.as_ref(), &context.auth_registry).await?;
             change
                 .setup_status
                 .apply_change(&db_pool, &change.key.table_name)

src/ops/targets/qdrant.rs

@@ -544,19 +544,19 @@ impl StorageFactoryBase for Factory {
     async fn apply_setup_changes(
         &self,
         setup_status: Vec<TypedResourceSetupChangeItem<'async_trait, Self>>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<()> {
         for setup_change in setup_status.iter() {
             let qdrant_client =
-                self.get_qdrant_client(&setup_change.key.connection, auth_registry)?;
+                self.get_qdrant_client(&setup_change.key.connection, &context.auth_registry)?;
             setup_change
                 .setup_status
                 .apply_delete(&setup_change.key.collection_name, &qdrant_client)
                 .await?;
         }
         for setup_change in setup_status.iter() {
             let qdrant_client =
-                self.get_qdrant_client(&setup_change.key.connection, auth_registry)?;
+                self.get_qdrant_client(&setup_change.key.connection, &context.auth_registry)?;
             setup_change
                 .setup_status
                 .apply_create(&setup_change.key.collection_name, &qdrant_client)

src/setup/driver.rs

@@ -1,5 +1,5 @@
 use crate::{
-    lib_context::{FlowExecutionContext, LibSetupContext, get_auth_registry},
+    lib_context::{FlowContext, FlowExecutionContext, LibSetupContext},
     ops::{
         get_optional_executor_factory,
         interface::{ExportTargetFactory, FlowInstanceContext},
@@ -411,7 +411,7 @@ async fn maybe_update_resource_setup<
 
 async fn apply_changes_for_flow(
     write: &mut (dyn std::io::Write + Send),
-    flow_name: &str,
+    flow_ctx: &FlowContext,
     flow_status: &FlowSetupStatus,
     existing_setup_state: &mut Option<setup::FlowSetupState<setup::ExistingMode>>,
     pool: &PgPool,
@@ -425,7 +425,7 @@ async fn apply_changes_for_flow(
         ObjectStatus::Existing => "Updating resources for ",
         _ => bail!("invalid flow status"),
     };
-    write!(write, "\n{verb} flow {flow_name}:\n")?;
+    write!(write, "\n{verb} flow {}:\n", flow_ctx.flow_name())?;
 
     let mut update_info =
         HashMap::<db_metadata::ResourceTypeKey, db_metadata::StateUpdateInfo>::new();
@@ -475,7 +475,7 @@ async fn apply_changes_for_flow(
     }
 
     let new_version_id = db_metadata::stage_changes_for_flow(
-        flow_name,
+        flow_ctx.flow_name(),
         flow_status.seen_flow_metadata_version,
         &update_info,
         pool,
@@ -517,7 +517,7 @@ async fn apply_changes_for_flow(
                                 setup_status: s.setup_status.as_ref(),
                             })
                             .collect(),
-                        get_auth_registry(),
+                        flow_ctx.flow.flow_instance_ctx.clone(),
                     )
                     .await?;
                 Ok(())
@@ -528,7 +528,7 @@ async fn apply_changes_for_flow(
 
     let is_deletion = status == ObjectStatus::Deleted;
     db_metadata::commit_changes_for_flow(
-        flow_name,
+        flow_ctx.flow_name(),
         new_version_id,
         &update_info,
         is_deletion,
@@ -580,7 +580,7 @@ async fn apply_changes_for_flow(
         });
     }
 
-    writeln!(write, "Done for flow {flow_name}")?;
+    writeln!(write, "Done for flow {}", flow_ctx.flow_name())?;
     Ok(())
 }
 
@@ -734,7 +734,7 @@ impl SetupChangeBundle {
             let mut flow_states = setup_ctx.all_setup_states.flows.remove(flow_name);
             apply_changes_for_flow(
                 write,
-                flow_name,
+                &flow_ctx,
                 setup_status,
                 &mut flow_states,
                 &persistence_ctx.builtin_db_pool,