feat: plumbing FlowInstanceContext to apply_setup_changes() (#805)

Author: georgeh0

src/builder/analyzed_flow.rs

@@ -8,6 +8,8 @@ pub struct AnalyzedFlow {
     pub data_schema: schema::FlowSchema,
     pub setup_state: exec_ctx::AnalyzedSetupState,
 
+    pub flow_instance_ctx: Arc<FlowInstanceContext>,
+
     /// It's None if the flow is not up to date
     pub execution_plan: Shared<BoxFuture<'static, Result<Arc<plan::ExecutionPlan>, SharedError>>>,
 }
@@ -18,7 +20,7 @@ impl AnalyzedFlow {
         flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<Self> {
         let (data_schema, setup_state, execution_plan_fut) =
-            analyzer::analyze_flow(&flow_instance, flow_instance_ctx).await?;
+            analyzer::analyze_flow(&flow_instance, flow_instance_ctx.clone()).await?;
         let execution_plan = async move {
             shared_ok(Arc::new(
                 execution_plan_fut.await.map_err(SharedError::new)?,
@@ -30,6 +32,7 @@ impl AnalyzedFlow {
             flow_instance,
             data_schema,
             setup_state,
+            flow_instance_ctx,
             execution_plan,
         };
         Ok(result)

src/lib_context.rs

@@ -33,6 +33,7 @@ async fn build_setup_context(
     let setup_status = setup::check_flow_setup_status(
         Some(&setup_execution_context.setup_state),
         existing_flow_ss,
+        &analyzed_flow.flow_instance_ctx,
     )
     .await?;
 

src/ops/factory_bases.rs

@@ -320,7 +320,7 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
         key: Self::Key,
         desired_state: Option<Self::SetupState>,
         existing_states: setup::CombinedState<Self::SetupState>,
-        auth_registry: &Arc<AuthRegistry>,
+        flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<Self::SetupStatus>;
 
     fn check_state_compatibility(
@@ -420,7 +420,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
         key: &serde_json::Value,
         desired_state: Option<serde_json::Value>,
         existing_states: setup::CombinedState<serde_json::Value>,
-        auth_registry: &Arc<AuthRegistry>,
+        flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<Box<dyn setup::ResourceSetupStatus>> {
         let key: T::Key = Self::deserialize_setup_key(key.clone())?;
         let desired_state: Option<T::SetupState> = desired_state
@@ -432,7 +432,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
             key,
             desired_state,
             existing_states,
-            auth_registry,
+            flow_instance_ctx,
         )
         .await?;
         Ok(Box::new(setup_status))

src/ops/interface.rs

@@ -272,7 +272,7 @@ pub trait ExportTargetFactory: Send + Sync {
         key: &serde_json::Value,
         desired_state: Option<serde_json::Value>,
         existing_states: setup::CombinedState<serde_json::Value>,
-        auth_registry: &Arc<AuthRegistry>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<Box<dyn setup::ResourceSetupStatus>>;
 
     /// Normalize the key. e.g. the JSON format may change (after code change, e.g. new optional field or field ordering), even if the underlying value is not changed.

src/ops/targets/kuzu.rs

@@ -847,7 +847,7 @@ impl StorageFactoryBase for Factory {
         _key: KuzuGraphElement,
         desired: Option<SetupState>,
         existing: CombinedState<SetupState>,
-        _auth_registry: &Arc<AuthRegistry>,
+        _flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<Self::SetupStatus> {
         let existing_invalidated = desired.as_ref().is_some_and(|desired| {
             existing

src/ops/targets/neo4j.rs

@@ -1009,9 +1009,11 @@ impl StorageFactoryBase for Factory {
         key: Neo4jGraphElement,
         desired: Option<SetupState>,
         existing: CombinedState<SetupState>,
-        auth_registry: &Arc<AuthRegistry>,
+        flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<Self::SetupStatus> {
-        let conn_spec = auth_registry.get::<ConnectionSpec>(&key.connection)?;
+        let conn_spec = flow_instance_ctx
+            .auth_registry
+            .get::<ConnectionSpec>(&key.connection)?;
         let data_status = GraphElementDataSetupStatus::new(desired.as_ref(), &existing);
         let components = components::SetupStatus::create(
             SetupComponentOperator {

src/ops/targets/postgres.rs

@@ -719,7 +719,7 @@ impl StorageFactoryBase for Factory {
         _key: TableId,
         desired: Option<SetupState>,
         existing: setup::CombinedState<SetupState>,
-        _auth_registry: &Arc<AuthRegistry>,
+        _flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<SetupStatus> {
         Ok(SetupStatus::new(desired, existing))
     }

src/ops/targets/qdrant.rs

@@ -489,7 +489,7 @@ impl StorageFactoryBase for Factory {
         _key: CollectionKey,
         desired: Option<SetupState>,
         existing: setup::CombinedState<SetupState>,
-        _auth_registry: &Arc<AuthRegistry>,
+        _flow_instance_ctx: Arc<FlowInstanceContext>,
     ) -> Result<Self::SetupStatus> {
         let desired_exists = desired.is_some();
         let add_collection = desired.filter(|state| {

src/setup/driver.rs

@@ -1,6 +1,9 @@
 use crate::{
     lib_context::{FlowExecutionContext, LibSetupContext, get_auth_registry},
-    ops::{get_optional_executor_factory, interface::ExportTargetFactory},
+    ops::{
+        get_optional_executor_factory,
+        interface::{ExportTargetFactory, FlowInstanceContext},
+    },
     prelude::*,
 };
 
@@ -252,6 +255,7 @@ fn group_resource_states<'a>(
 pub async fn check_flow_setup_status(
     desired_state: Option<&FlowSetupState<DesiredMode>>,
     existing_state: Option<&FlowSetupState<ExistingMode>>,
+    flow_instance_ctx: &Arc<FlowInstanceContext>,
 ) -> Result<FlowSetupStatus> {
     let metadata_change = diff_state(
         existing_state.map(|e| &e.metadata),
@@ -331,7 +335,7 @@ pub async fn check_flow_setup_status(
                         &resource_id.key,
                         target_state,
                         existing_without_setup_by_user,
-                        get_auth_registry(),
+                        flow_instance_ctx.clone(),
                     )
                     .await?,
             )
@@ -618,16 +622,19 @@ pub struct SetupChangeBundle {
 impl SetupChangeBundle {
     async fn get_flow_setup_status<'a>(
         setup_ctx: &LibSetupContext,
-        flow_name: &str,
+        flow_ctx: &'a FlowContext,
         flow_exec_ctx: &'a FlowExecutionContext,
         action: &FlowSetupChangeAction,
         buffer: &'a mut Option<FlowSetupStatus>,
     ) -> Result<&'a FlowSetupStatus> {
         let result = match action {
             FlowSetupChangeAction::Setup => &flow_exec_ctx.setup_status,
             FlowSetupChangeAction::Drop => {
-                let existing_state = setup_ctx.all_setup_states.flows.get(flow_name);
-                buffer.insert(check_flow_setup_status(None, existing_state).await?)
+                let existing_state = setup_ctx.all_setup_states.flows.get(flow_ctx.flow_name());
+                buffer.insert(
+                    check_flow_setup_status(None, existing_state, &flow_ctx.flow.flow_instance_ctx)
+                        .await?,
+                )
             }
         };
         Ok(result)
@@ -662,7 +669,7 @@ impl SetupChangeBundle {
             let mut setup_status_buffer = None;
             let setup_status = Self::get_flow_setup_status(
                 setup_ctx,
-                flow_name,
+                &flow_ctx,
                 &flow_exec_ctx,
                 &self.action,
                 &mut setup_status_buffer,
@@ -714,7 +721,7 @@ impl SetupChangeBundle {
             let mut setup_status_buffer = None;
             let setup_status = Self::get_flow_setup_status(
                 setup_ctx,
-                flow_name,
+                &flow_ctx,
                 &flow_exec_ctx,
                 &self.action,
                 &mut setup_status_buffer,