cocoindex 0.1.69__tar.gz → 0.1.70__tar.gz

{cocoindex-0.1.69 → cocoindex-0.1.70}/Cargo.lock

@@ -1297,7 +1297,7 @@ dependencies = [
 
 [[package]]
 name = "cocoindex"
-version = "0.1.69"
+version = "0.1.70"
 dependencies = [
  "anyhow",
  "async-openai",

{cocoindex-0.1.69 → cocoindex-0.1.70}/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.69"
+version = "0.1.70"
 edition = "2024"
 rust-version = "1.88"
 

{cocoindex-0.1.69 → cocoindex-0.1.70}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cocoindex
-Version: 0.1.69
+Version: 0.1.70
 Requires-Dist: click>=8.1.8
 Requires-Dist: rich>=14.0.0
 Requires-Dist: python-dotenv>=1.1.0

{cocoindex-0.1.69 → cocoindex-0.1.70}/docs/docs/core/flow_def.mdx

@@ -416,22 +416,28 @@ flow_builder.declare(
 ### Auth Registry
 
 CocoIndex manages an auth registry. It's an in-memory key-value store, mainly to store authentication information for a backend.
+It's usually used for targets, where key stability is important for backend cleanup.
 
-Operation spec is the default way to configure a persistent backend. But it has the following limitations:
+Operation spec is the default way to configure sources, functions and targets. But it has the following limitations:
 
 *   The spec isn't supposed to contain secret information, and it's frequently shown in various places, e.g. `cocoindex show`.
-*   Once an operation is removed after flow definition code change, the spec is also gone.
-    But we still need to be able to drop the backend (e.g. a table) when [setup / drop flow](/docs/core/flow_methods#setup--drop-flow).
+*   For targets, once an operation is removed after flow definition code change, the spec is also gone.
+    But we still need to be able to drop the persistent backend (e.g. a table) when [setup / drop flow](/docs/core/flow_methods#setup--drop-flow).
 
-Auth registry is introduced to solve the problems above. It works as follows:
+Auth registry is introduced to solve the problems above.
 
-*   You can create new **auth entry** by a key and a value.
-*   You can references the entry by the key, and pass it as part of spec for certain operations. e.g. `Neo4j` takes `connection` field in the form of auth entry reference.
+
+#### Auth Entry
+
+An auth entry is an entry in the auth registry with an explicit key.
+
+*   You can create new *auth entry* by a key and a value.
+*   You can reference the entry by the key, and pass it as part of spec for certain operations. e.g. `Neo4j` takes `connection` field in the form of auth entry reference.
 
 <Tabs>
 <TabItem value="python" label="Python" default>
 
-You can add an auth entry by `cocoindex.add_auth_entry()` function, which returns a `cocoindex.AuthEntryReference`:
+You can add an auth entry by `cocoindex.add_auth_entry()` function, which returns a `cocoindex.AuthEntryReference[T]`:
 
 ```python
 my_graph_conn = cocoindex.add_auth_entry(
@@ -445,7 +451,7 @@ my_graph_conn = cocoindex.add_auth_entry(
 
 Then reference it when building a spec that takes an auth entry:
 
-*   You can either reference by the `AuthEntryReference` object directly:
+*   You can either reference by the `AuthEntryReference[T]` object directly:
 
     ```python
     demo_collector.export(
@@ -472,3 +478,28 @@ Note that CocoIndex backends use the key of an auth entry to identify the backen
 
 *   If a key is no longer referenced in any operation spec, keep it until the next flow setup / drop action,
     so that CocoIndex will be able to clean up the backends.
+
+#### Transient Auth Entry
+
+A transient auth entry is an entry in the auth registry with an automatically generated key.
+It's usually used for sources and functions, where key stability is not important.
+
+<Tabs>
+<TabItem value="python" label="Python" default>
+
+You can create a new *transient auth entry* by `cocoindex.add_transient_auth_entry()` function, which returns a `cocoindex.TransientAuthEntryReference[T]`, and pass it to a source or function spec that takes it, e.g.
+
+```python
+flow_builder.add_source(
+    cocoindex.sources.AzureBlob(
+        ...
+        sas_token=cocoindex.add_transient_auth_entry("...")
+    )
+)
+```
+
+
+</TabItem>
+</Tabs>
+
+Whenever a `TransientAuthEntryReference[T]` is expected, you can also pass a `AuthEntryReference[T]` instead, as `AuthEntryReference[T]` is a subtype of `TransientAuthEntryReference[T]`.

{cocoindex-0.1.69 → cocoindex-0.1.70}/docs/docs/ops/sources.md

@@ -170,22 +170,33 @@ These are actions you need to take:
 
 #### Authentication
 
-We use Azure’s **Default Credential** system (DefaultAzureCredential) for secure and flexible authentication.
-This allows you to connect to Azure services without putting any secrets in the code or flow spec.
-It automatically chooses the best authentication method based on your environment:
-
-*   On your local machine: uses your Azure CLI login (`az login`) or environment variables.
-
-    ```sh
-    az login
-    # Optional: Set a default subscription if you have more than one
-    az account set --subscription "<YOUR_SUBSCRIPTION_NAME_OR_ID>"
-    ```
-*   In Azure (VM, App Service, AKS, etc.): uses the resource’s Managed Identity.
-*   In automated environments: supports Service Principals via environment variables
-    *   `AZURE_CLIENT_ID`
-    *   `AZURE_TENANT_ID`
-    *   `AZURE_CLIENT_SECRET`
+We support the following authentication methods:
+
+*   Shared access signature (SAS) tokens.
+    You can generate it from the Azure Portal in the settings for a specific container.
+    You need to provide at least *List* and *Read* permissions when generating the SAS token.
+    It's a query string in the form of
+    `sp=rl&st=2025-07-20T09:33:00Z&se=2025-07-19T09:48:53Z&sv=2024-11-04&sr=c&sig=i3FDjsadfklj3%23adsfkk`.
+
+*   Storage account access key. You can find it in the Azure Portal in the settings for a specific storage account.
+
+*   Default credential. When none of the above is provided, it will use the default credential.
+
+    This allows you to connect to Azure services without putting any secrets in the code or flow spec.
+    It automatically chooses the best authentication method based on your environment:
+
+    *   On your local machine: uses your Azure CLI login (`az login`) or environment variables.
+
+        ```sh
+        az login
+        # Optional: Set a default subscription if you have more than one
+        az account set --subscription "<YOUR_SUBSCRIPTION_NAME_OR_ID>"
+        ```
+    *   In Azure (VM, App Service, AKS, etc.): uses the resource’s Managed Identity.
+    *   In automated environments: supports Service Principals via environment variables
+        *   `AZURE_CLIENT_ID`
+        *   `AZURE_TENANT_ID`
+        *   `AZURE_CLIENT_SECRET`
 
 You can refer to [this doc](https://learn.microsoft.com/en-us/azure/developer/python/sdk/authentication/overview) for more details.
 
@@ -202,6 +213,8 @@ The spec takes the following fields:
 *   `excluded_patterns` (`list[str]`, optional): a list of glob patterns to exclude files, e.g. `["*.tmp", "**/*.log"]`.
     Any file or directory matching these patterns will be excluded even if they match `included_patterns`.
     If not specified, no files will be excluded.
+*   `sas_token` (`cocoindex.TransientAuthEntryReference[str]`, optional): a SAS token for authentication.
+*   `account_access_key` (`cocoindex.TransientAuthEntryReference[str]`, optional): an account access key for authentication.
 
     :::info
 

{cocoindex-0.1.69 → cocoindex-0.1.70}/examples/amazon_s3_embedding/main.py

@@ -102,12 +102,6 @@ def _main() -> None:
 
     amazon_s3_text_embedding_flow.setup()
     with cocoindex.FlowLiveUpdater(amazon_s3_text_embedding_flow) as updater:
-        while True:
-            updates = updater.next_status_updates()
-            print(f"Updates: {updates}")
-            if not updates.active_sources:
-                break
-
         # Run queries in a loop to demonstrate the query capabilities.
         while True:
             query = input("Enter search query (or Enter to quit): ")

{cocoindex-0.1.69 → cocoindex-0.1.70}/python/cocoindex/__init__.py

@@ -6,7 +6,7 @@ from . import functions, sources, targets, cli, utils
 
 from . import targets as storages  # Deprecated: Use targets instead
 
-from .auth_registry import AuthEntryReference, add_auth_entry, ref_auth_entry
+from .auth_registry import AuthEntryReference, add_auth_entry, add_transient_auth_entry
 from .flow import FlowBuilder, DataScope, DataSlice, Flow, transform_flow
 from .flow import flow_def
 from .flow import EvaluateAndDumpOptions, GeneratedField
@@ -42,6 +42,7 @@ __all__ = [
     # Auth registry
     "AuthEntryReference",
     "add_auth_entry",
+    "add_transient_auth_entry",
     "ref_auth_entry",
     # Flow
     "FlowBuilder",

cocoindex-0.1.70/python/cocoindex/auth_registry.py

@@ -0,0 +1,51 @@
+"""
+Auth registry is used to register and reference auth entries.
+"""
+
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+import threading
+
+from . import _engine  # type: ignore
+from .convert import dump_engine_object
+
+T = TypeVar("T")
+
+# Global atomic counter for generating unique auth entry keys
+_counter_lock = threading.Lock()
+_auth_key_counter = 0
+
+
+def _generate_auth_key() -> str:
+    """Generate a unique auth entry key using a global atomic counter."""
+    global _auth_key_counter  # pylint: disable=global-statement
+    with _counter_lock:
+        _auth_key_counter += 1
+        return f"__auth_{_auth_key_counter}"
+
+
+@dataclass
+class TransientAuthEntryReference(Generic[T]):
+    """Reference an auth entry, may or may not have a stable key."""
+
+    key: str
+
+
+class AuthEntryReference(TransientAuthEntryReference[T]):
+    """Reference an auth entry, with a key stable across ."""
+
+
+def add_transient_auth_entry(value: T) -> TransientAuthEntryReference[T]:
+    """Add an auth entry to the registry. Returns its reference."""
+    return add_auth_entry(_generate_auth_key(), value)
+
+
+def add_auth_entry(key: str, value: T) -> AuthEntryReference[T]:
+    """Add an auth entry to the registry. Returns its reference."""
+    _engine.add_auth_entry(key, dump_engine_object(value))
+    return AuthEntryReference(key)
+
+
+def ref_auth_entry(key: str) -> AuthEntryReference[T]:
+    """Reference an auth entry by its key."""
+    return AuthEntryReference(key)

{cocoindex-0.1.69 → cocoindex-0.1.70}/python/cocoindex/sources.py

@@ -1,6 +1,7 @@
 """All builtin sources."""
 
 from . import op
+from .auth_registry import TransientAuthEntryReference
 import datetime
 
 
@@ -48,6 +49,11 @@ class AmazonS3(op.SourceSpec):
 class AzureBlob(op.SourceSpec):
     """
     Import data from an Azure Blob Storage container. Supports optional prefix and file filtering by glob patterns.
+
+    Authentication mechanisms taken in the following order:
+    - SAS token (if provided)
+    - Account access key (if provided)
+    - Default Azure credential
     """
 
     _op_category = op.OpCategory.SOURCE
@@ -58,3 +64,6 @@ class AzureBlob(op.SourceSpec):
     binary: bool = False
     included_patterns: list[str] | None = None
     excluded_patterns: list[str] | None = None
+
+    sas_token: TransientAuthEntryReference[str] | None = None
+    account_access_key: TransientAuthEntryReference[str] | None = None

{cocoindex-0.1.69 → cocoindex-0.1.70}/src/execution/live_updater.rs

@@ -3,7 +3,7 @@ use crate::{execution::stats::UpdateStats, prelude::*};
 use super::stats;
 use futures::future::try_join_all;
 use sqlx::PgPool;
-use tokio::{sync::watch, time::MissedTickBehavior};
+use tokio::{sync::watch, task::JoinSet, time::MissedTickBehavior};
 
 pub struct FlowLiveUpdaterUpdates {
     pub active_sources: Vec<String>,
@@ -22,7 +22,8 @@ struct UpdateReceiveState {
 
 pub struct FlowLiveUpdater {
     flow_ctx: Arc<FlowContext>,
-    tasks: Vec<(tokio::task::JoinHandle<Result<()>>, Arc<stats::UpdateStats>)>,
+    join_set: Mutex<Option<JoinSet<Result<()>>>>,
+    stats_per_task: Vec<Arc<stats::UpdateStats>>,
     recv_state: tokio::sync::Mutex<UpdateReceiveState>,
     num_remaining_tasks_rx: watch::Receiver<usize>,
 
@@ -267,7 +268,11 @@ impl SourceUpdateTask {
             .boxed()
         });
 
-        try_join_all(futs).await?;
+        let join_result = try_join_all(futs).await;
+        if let Err(err) = join_result {
+            error!("Error in source `{}`: {:?}", import_op.name, err);
+            return Err(err);
+        }
         Ok(())
     }
 }
@@ -288,27 +293,30 @@ impl FlowLiveUpdater {
 
         let (num_remaining_tasks_tx, num_remaining_tasks_rx) =
             watch::channel(plan.import_ops.len());
-        let tasks = (0..plan.import_ops.len())
-            .map(|source_idx| {
-                let source_update_stats = Arc::new(stats::UpdateStats::default());
-                let source_update_task = SourceUpdateTask {
-                    source_idx,
-                    flow: flow_ctx.flow.clone(),
-                    plan: plan.clone(),
-                    execution_ctx: execution_ctx.clone(),
-                    source_update_stats: source_update_stats.clone(),
-                    pool: pool.clone(),
-                    options: options.clone(),
-                    status_tx: status_tx.clone(),
-                    num_remaining_tasks_tx: num_remaining_tasks_tx.clone(),
-                };
-                let task = tokio::spawn(source_update_task.run());
-                (task, source_update_stats)
-            })
-            .collect();
+
+        let mut join_set = JoinSet::new();
+        let mut stats_per_task = Vec::new();
+
+        for source_idx in 0..plan.import_ops.len() {
+            let source_update_stats = Arc::new(stats::UpdateStats::default());
+            let source_update_task = SourceUpdateTask {
+                source_idx,
+                flow: flow_ctx.flow.clone(),
+                plan: plan.clone(),
+                execution_ctx: execution_ctx.clone(),
+                source_update_stats: source_update_stats.clone(),
+                pool: pool.clone(),
+                options: options.clone(),
+                status_tx: status_tx.clone(),
+                num_remaining_tasks_tx: num_remaining_tasks_tx.clone(),
+            };
+            join_set.spawn(source_update_task.run());
+            stats_per_task.push(source_update_stats);
+        }
         Ok(Self {
             flow_ctx,
-            tasks,
+            join_set: Mutex::new(Some(join_set)),
+            stats_per_task,
             recv_state: tokio::sync::Mutex::new(UpdateReceiveState {
                 status_rx,
                 last_num_source_updates: vec![0; plan.import_ops.len()],
@@ -322,17 +330,33 @@ impl FlowLiveUpdater {
     }
 
     pub async fn wait(&self) -> Result<()> {
-        let mut rx = self.num_remaining_tasks_rx.clone();
-        if *rx.borrow() == 0 {
+        {
+            let mut rx = self.num_remaining_tasks_rx.clone();
+            rx.wait_for(|v| *v == 0).await?;
+        }
+
+        let Some(mut join_set) = self.join_set.lock().unwrap().take() else {
             return Ok(());
+        };
+        while let Some(task_result) = join_set.join_next().await {
+            match task_result {
+                Ok(Ok(_)) => {}
+                Ok(Err(err)) => {
+                    return Err(err);
+                }
+                Err(err) if err.is_cancelled() => {}
+                Err(err) => {
+                    return Err(err.into());
+                }
+            }
         }
-        rx.wait_for(|v| *v == 0).await?;
         Ok(())
     }
 
     pub fn abort(&self) {
-        for (task, _) in &self.tasks {
-            task.abort();
+        let mut join_set = self.join_set.lock().unwrap();
+        if let Some(join_set) = &mut *join_set {
+            join_set.abort_all();
         }
     }
 
@@ -340,9 +364,9 @@ impl FlowLiveUpdater {
         stats::IndexUpdateInfo {
             sources: std::iter::zip(
                 self.flow_ctx.flow.flow_instance.import_ops.iter(),
-                self.tasks.iter(),
+                self.stats_per_task.iter(),
             )
-            .map(|(import_op, (_, stats))| stats::SourceUpdateInfo {
+            .map(|(import_op, stats)| stats::SourceUpdateInfo {
                 source_name: import_op.name.clone(),
                 stats: stats.as_ref().clone(),
             })

{cocoindex-0.1.69 → cocoindex-0.1.70}/src/ops/sources/azure_blob.rs

@@ -19,6 +19,11 @@ pub struct Spec {
     binary: bool,
     included_patterns: Option<Vec<String>>,
     excluded_patterns: Option<Vec<String>>,
+
+    /// SAS token for authentication. Takes precedence over account_access_key.
+    sas_token: Option<AuthEntryReference<String>>,
+    /// Account access key for authentication. If not provided, will use default Azure credential.
+    account_access_key: Option<AuthEntryReference<String>>,
 }
 
 struct Executor {
@@ -209,15 +214,22 @@ impl SourceFactoryBase for Factory {
     async fn build_executor(
         self: Arc<Self>,
         spec: Spec,
-        _context: Arc<FlowInstanceContext>,
+        context: Arc<FlowInstanceContext>,
     ) -> Result<Box<dyn SourceExecutor>> {
-        let default_credential = Arc::new(DefaultAzureCredential::create(
-            TokenCredentialOptions::default(),
-        )?);
-        let client = BlobServiceClient::new(
-            &spec.account_name,
-            StorageCredentials::token_credential(default_credential),
-        );
+        let credential = if let Some(sas_token) = spec.sas_token {
+            let sas_token = context.auth_registry.get(&sas_token)?;
+            StorageCredentials::sas_token(sas_token)?
+        } else if let Some(account_access_key) = spec.account_access_key {
+            let account_access_key = context.auth_registry.get(&account_access_key)?;
+            StorageCredentials::access_key(spec.account_name.clone(), account_access_key)
+        } else {
+            let default_credential = Arc::new(DefaultAzureCredential::create(
+                TokenCredentialOptions::default(),
+            )?);
+            StorageCredentials::token_credential(default_credential)
+        };
+
+        let client = BlobServiceClient::new(&spec.account_name, credential);
         Ok(Box::new(Executor {
             client,
             container_name: spec.container_name,

cocoindex-0.1.69/python/cocoindex/auth_registry.py

@@ -1,29 +0,0 @@
-"""
-Auth registry is used to register and reference auth entries.
-"""
-
-from dataclasses import dataclass
-from typing import Generic, TypeVar
-
-from . import _engine  # type: ignore
-from .convert import dump_engine_object
-
-T = TypeVar("T")
-
-
-@dataclass
-class AuthEntryReference(Generic[T]):
-    """Reference an auth entry by its key."""
-
-    key: str
-
-
-def add_auth_entry(key: str, value: T) -> AuthEntryReference[T]:
-    """Add an auth entry to the registry. Returns its reference."""
-    _engine.add_auth_entry(key, dump_engine_object(value))
-    return AuthEntryReference(key)
-
-
-def ref_auth_entry(key: str) -> AuthEntryReference[T]:
-    """Reference an auth entry by its key."""
-    return AuthEntryReference(key)