cocoindex 0.1.55__tar.gz → 0.1.56__tar.gz

{cocoindex-0.1.55 → cocoindex-0.1.56}/Cargo.lock

@@ -1040,7 +1040,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.1.55"
+version = "0.1.56"
 dependencies = [
  "anyhow",
  "async-openai",

{cocoindex-0.1.55 → cocoindex-0.1.56}/Cargo.toml

@@ -2,7 +2,7 @@
 name = "cocoindex"
 # Version used for local development is always higher than others to take precedence.
 # Will be overridden for specific release versions.
-version = "0.1.55"
+version = "0.1.56"
 edition = "2024"
 rust-version = "1.86"
 

{cocoindex-0.1.55 → cocoindex-0.1.56}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.55
+Version: 0.1.56
 Requires-Dist: sentence-transformers>=3.3.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0

{cocoindex-0.1.55 → cocoindex-0.1.56}/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(

{cocoindex-0.1.55 → cocoindex-0.1.56}/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(),

{cocoindex-0.1.55 → cocoindex-0.1.56}/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,7 +68,7 @@ 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 {
@@ -75,13 +76,44 @@ impl FlowContext {
         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());

{cocoindex-0.1.55 → cocoindex-0.1.56}/src/ops/targets/postgres.rs

@@ -194,17 +194,17 @@ impl ExportContext {
     ) -> Result<Self> {
         let key_fields = key_fields_schema
             .iter()
-            .map(|f| f.name.as_str())
+            .map(|f| format!("\"{}\"", f.name))
             .collect::<Vec<_>>()
             .join(", ");
         let value_fields = value_fields_schema
             .iter()
-            .map(|f| f.name.as_str())
+            .map(|f| format!("\"{}\"", f.name))
             .collect::<Vec<_>>()
             .join(", ");
         let set_value_fields = value_fields_schema
             .iter()
-            .map(|f| format!("{} = EXCLUDED.{}", f.name, f.name))
+            .map(|f| format!("\"{}\" = EXCLUDED.\"{}\"", f.name, f.name))
             .collect::<Vec<_>>()
             .join(", ");
 
@@ -288,7 +288,9 @@ impl ExportContext {
                 if i > 0 {
                     query_builder.push(" AND ");
                 }
+                query_builder.push("\"");
                 query_builder.push(schema.name.as_str());
+                query_builder.push("\"");
                 query_builder.push("=");
                 bind_key_field(&mut query_builder, value)?;
             }
@@ -547,14 +549,16 @@ impl SetupStatus {
                 .await?;
         }
         for index_name in self.actions.indexes_to_delete.iter() {
-            let sql = format!("DROP INDEX IF EXISTS {}", index_name);
+            let sql = format!("DROP INDEX IF EXISTS {index_name}");
             sqlx::query(&sql).execute(db_pool).await?;
         }
         if let Some(table_upsertion) = &self.actions.table_action.table_upsertion {
             match table_upsertion {
                 TableUpsertionAction::Create { keys, values } => {
-                    let mut fields = (keys.iter().map(|(k, v)| format!("{k} {v} NOT NULL")))
-                        .chain(values.iter().map(|(k, v)| format!("{k} {v}")));
+                    let mut fields = (keys
+                        .iter()
+                        .map(|(name, typ)| format!("\"{name}\" {typ} NOT NULL")))
+                    .chain(values.iter().map(|(name, typ)| format!("\"{name}\" {typ}")));
                     let sql = format!(
                         "CREATE TABLE IF NOT EXISTS {table_name} ({}, PRIMARY KEY ({}))",
                         fields.join(", "),
@@ -568,13 +572,13 @@ impl SetupStatus {
                 } => {
                     for column_name in columns_to_delete.iter() {
                         let sql = format!(
-                            "ALTER TABLE {table_name} DROP COLUMN IF EXISTS {column_name}",
+                            "ALTER TABLE {table_name} DROP COLUMN IF EXISTS \"{column_name}\"",
                         );
                         sqlx::query(&sql).execute(db_pool).await?;
                     }
                     for (column_name, column_type) in columns_to_upsert.iter() {
                         let sql = format!(
-                            "ALTER TABLE {table_name} DROP COLUMN IF EXISTS {column_name}, ADD COLUMN {column_name} {column_type}"
+                            "ALTER TABLE {table_name} DROP COLUMN IF EXISTS \"{column_name}\", ADD COLUMN \"{column_name}\" {column_type}"
                         );
                         sqlx::query(&sql).execute(db_pool).await?;
                     }

{cocoindex-0.1.55 → cocoindex-0.1.56}/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,

{cocoindex-0.1.55 → cocoindex-0.1.56}/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,

{cocoindex-0.1.55 → cocoindex-0.1.56}/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(