fix(setup): restore check for setup up-to-date before evaluation (#672)

Author: badmonster0

examples/code_embedding/main.py

@@ -67,6 +67,7 @@ def code_embedding_flow(
                 location=chunk["location"],
                 code=chunk["text"],
                 embedding=chunk["embedding"],
+                start=chunk["start"],
             )
 
     code_embeddings.export(

src/execution/live_updater.rs

@@ -60,16 +60,16 @@ impl SharedAckFn {
 }
 
 async fn update_source(
-    flow_ctx: Arc<FlowContext>,
+    flow: Arc<builder::AnalyzedFlow>,
     plan: Arc<plan::ExecutionPlan>,
+    execution_ctx: Arc<tokio::sync::OwnedRwLockReadGuard<crate::lib_context::ExecutionContext>>,
     source_update_stats: Arc<stats::UpdateStats>,
     source_idx: usize,
     pool: PgPool,
     options: FlowLiveUpdaterOptions,
 ) -> Result<()> {
-    let execution_ctx = flow_ctx.execution_ctx.read().await;
     let source_context = execution_ctx
-        .get_source_indexing_context(&flow_ctx.flow, source_idx, &pool)
+        .get_source_indexing_context(&flow, source_idx, &pool)
         .await?;
 
     let import_op = &plan.import_ops[source_idx];
@@ -97,15 +97,9 @@ async fn update_source(
             delta
         };
         if options.print_stats {
-            println!(
-                "{}.{}: {}",
-                flow_ctx.flow.flow_instance.name, import_op.name, delta
-            );
+            println!("{}.{}: {}", flow.flow_instance.name, import_op.name, delta);
         } else {
-            trace!(
-                "{}.{}: {}",
-                flow_ctx.flow.flow_instance.name, import_op.name, delta
-            );
+            trace!("{}.{}: {}", flow.flow_instance.name, import_op.name, delta);
         }
     };
 
@@ -219,14 +213,16 @@ impl FlowLiveUpdater {
         options: FlowLiveUpdaterOptions,
     ) -> Result<Self> {
         let plan = flow_ctx.flow.get_execution_plan().await?;
+        let execution_ctx = Arc::new(flow_ctx.use_owned_execution_ctx().await?);
 
         let mut tasks = JoinSet::new();
         let sources_update_stats = (0..plan.import_ops.len())
             .map(|source_idx| {
                 let source_update_stats = Arc::new(stats::UpdateStats::default());
                 tasks.spawn(update_source(
-                    flow_ctx.clone(),
+                    flow_ctx.flow.clone(),
                     plan.clone(),
+                    execution_ctx.clone(),
                     source_update_stats.clone(),
                     source_idx,
                     pool.clone(),

src/lib_context.rs

@@ -4,6 +4,7 @@ use crate::builder::AnalyzedFlow;
 use crate::execution::source_indexer::SourceIndexingContext;
 use crate::service::error::ApiError;
 use crate::settings;
+use crate::setup::ObjectSetupStatus;
 use axum::http::StatusCode;
 use sqlx::PgPool;
 use sqlx::postgres::PgConnectOptions;
@@ -67,21 +68,52 @@ impl ExecutionContext {
 }
 pub struct FlowContext {
     pub flow: Arc<AnalyzedFlow>,
-    pub execution_ctx: tokio::sync::RwLock<ExecutionContext>,
+    execution_ctx: Arc<tokio::sync::RwLock<ExecutionContext>>,
 }
 
 impl FlowContext {
     pub async fn new(
         flow: Arc<AnalyzedFlow>,
         existing_flow_ss: Option<&setup::FlowSetupState<setup::ExistingMode>>,
     ) -> Result<Self> {
-        let execution_ctx =
-            tokio::sync::RwLock::new(ExecutionContext::new(&flow, existing_flow_ss).await?);
+        let execution_ctx = Arc::new(tokio::sync::RwLock::new(
+            ExecutionContext::new(&flow, existing_flow_ss).await?,
+        ));
         Ok(Self {
             flow,
             execution_ctx,
         })
     }
+
+    pub async fn use_execution_ctx(
+        &self,
+    ) -> Result<tokio::sync::RwLockReadGuard<ExecutionContext>> {
+        let execution_ctx = self.execution_ctx.read().await;
+        if !execution_ctx.flow_setup_status.is_up_to_date() {
+            api_bail!(
+                "Flow setup is not up-to-date. Please run `cocoindex setup` to update the setup."
+            );
+        }
+        Ok(execution_ctx)
+    }
+
+    pub async fn use_owned_execution_ctx(
+        &self,
+    ) -> Result<tokio::sync::OwnedRwLockReadGuard<ExecutionContext>> {
+        let execution_ctx = self.execution_ctx.clone().read_owned().await;
+        if !execution_ctx.flow_setup_status.is_up_to_date() {
+            api_bail!(
+                "Flow setup is not up-to-date. Please run `cocoindex setup` to update the setup."
+            );
+        }
+        Ok(execution_ctx)
+    }
+
+    pub async fn get_execution_ctx_for_setup(
+        &self,
+    ) -> tokio::sync::RwLockReadGuard<ExecutionContext> {
+        self.execution_ctx.read().await
+    }
 }
 
 static TOKIO_RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());

src/py/mod.rs

@@ -202,7 +202,7 @@ impl Flow {
                 .block_on(async {
                     let exec_plan = self.0.flow.get_execution_plan().await?;
                     let lib_context = get_lib_context()?;
-                    let execution_ctx = self.0.execution_ctx.read().await;
+                    let execution_ctx = self.0.use_execution_ctx().await?;
                     execution::dumper::evaluate_and_dump(
                         &exec_plan,
                         &execution_ctx.setup_execution_context,

src/service/flows.rs

@@ -189,7 +189,7 @@ pub async fn evaluate_data(
 ) -> Result<Json<EvaluateDataResponse>, ApiError> {
     let flow_ctx = lib_context.get_flow_context(&flow_name)?;
     let source_row_key_ctx = SourceRowKeyContextHolder::create(&flow_ctx, query).await?;
-    let execution_ctx = flow_ctx.execution_ctx.read().await;
+    let execution_ctx = flow_ctx.use_execution_ctx().await?;
     let evaluate_output = row_indexer::evaluate_source_entry_with_memory(
         &source_row_key_ctx.as_context(),
         &execution_ctx.setup_execution_context,
@@ -239,7 +239,7 @@ pub async fn get_row_indexing_status(
     let flow_ctx = lib_context.get_flow_context(&flow_name)?;
     let source_row_key_ctx = SourceRowKeyContextHolder::create(&flow_ctx, query).await?;
 
-    let execution_ctx = flow_ctx.execution_ctx.read().await;
+    let execution_ctx = flow_ctx.use_execution_ctx().await?;
     let indexing_status = indexing_status::get_source_row_indexing_status(
         &source_row_key_ctx.as_context(),
         &execution_ctx.setup_execution_context,

src/setup/driver.rs

@@ -366,7 +366,7 @@ pub async fn sync_setup(
     let mut flow_setup_status = BTreeMap::new();
     for (flow_name, flow_context) in flows {
         let existing_state = all_setup_state.flows.get(flow_name);
-        let execution_ctx = flow_context.execution_ctx.read().await;
+        let execution_ctx = flow_context.get_execution_ctx_for_setup().await;
         flow_setup_status.insert(
             flow_name.clone(),
             check_flow_setup_status(