Merge branch 'cocoindex-io:main' into feat-anthropic-dataflow

Author: par4m

docs/docs/core/initialization.mdx

@@ -49,12 +49,29 @@ This takes care of the following effects:
 
 The following environment variables are supported:
 
-*   `COCOINDEX_DATABASE_URL`: The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`
+*   `COCOINDEX_DATABASE_URL` (required): The URI of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`
+*   `COCOINDEX_DATABASE_USER` (optional): The username for the Postgres database. If not provided, username will come from `COCOINDEX_DATABASE_URL`.
+*   `COCOINDEX_DATABASE_PASSWORD` (optional): The password for the Postgres database. If not provided, password will come from `COCOINDEX_DATABASE_URL`.
 
 ## Explicit Initialization
 
 Alternatively, for flexibility, you can also explicitly initialize the library by the `init()` function:
 
+### Settings
+
+It takes a `Settings` object as argument, which is a dataclass that contains the following fields:
+
+*   `database` (type: `DatabaseConnectionSpec`, required): The connection to the Postgres database.
+
+#### DatabaseConnectionSpec
+
+`DatabaseConnectionSpec` has the following fields:
+*   `uri` (type: `str`, required): The URI of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`.
+*   `user` (type: `str`, optional): The username for the Postgres database. If not provided, username will come from `uri`.
+*   `password` (type: `str`, optional): The password for the Postgres database. If not provided, password will come from `uri`.
+
+### Example
+
 <Tabs>
 <TabItem value="python" label="Python" default>
 
@@ -63,7 +80,11 @@ import cocoindex
 
 def main():
   ...
-  cocoindex.init(cocoindex.Settings(database_url="postgres://cocoindex:cocoindex@localhost/cocoindex"))
+  cocoindex.init(
+      cocoindex.Settings(
+          database=cocoindex.DatabaseConnectionSpec(
+              uri="postgres://cocoindex:cocoindex@localhost/cocoindex"
+  )))
   ...
 
 ...

docs/docs/ops/storages.md

@@ -36,7 +36,9 @@ It should be a unique table, meaning that no other export target should export t
 
 The spec takes the following fields:
 
-*   `database_url` (type: `str`, optional): The URL of the Postgres database to use as the internal storage, e.g. `postgres://cocoindex:cocoindex@localhost/cocoindex`. If unspecified, will use the same database as the [internal storage](/docs/core/basics#internal-storage).
+*   `database` (type: [auth reference](../core/flow_def#auth-registry) to `DatabaseConnectionSpec`, optional): The connection to the Postgres database.
+    See [DatabaseConnectionSpec](../core/initialization#databaseconnectionspec) for its specific fields.
+    If not provided, will use the same database as the [internal storage](/docs/core/basics#internal-storage).
 
 *   `table_name` (type: `str`, optional): The name of the table to store to. If unspecified, will generate a new automatically. We recommend specifying a name explicitly if you want to directly query the table. It can be omitted if you want to use CocoIndex's query handlers to query the table.
 

python/cocoindex/lib.py

@@ -1,7 +1,6 @@
 """
 Library level functions and states.
 """
-import asyncio
 import os
 import sys
 import functools
@@ -12,6 +11,7 @@
 
 from . import _engine
 from . import flow, query, cli
+from .convert import dump_engine_object
 
 
 def _load_field(target: dict[str, str], name: str, env_name: str, required: bool = False):
@@ -22,24 +22,32 @@ def _load_field(target: dict[str, str], name: str, env_name: str, required: bool
     else:
         target[name] = value
 
+@dataclass
+class DatabaseConnectionSpec:
+    uri: str
+    user: str | None = None
+    password: str | None = None
+
 @dataclass
 class Settings:
     """Settings for the cocoindex library."""
-    database_url: str
+    database: DatabaseConnectionSpec
 
     @classmethod
     def from_env(cls) -> Self:
         """Load settings from environment variables."""
 
-        kwargs: dict[str, str] = dict()
-        _load_field(kwargs, "database_url", "COCOINDEX_DATABASE_URL", required=True)
-
-        return cls(**kwargs)
+        db_kwargs: dict[str, str] = dict()
+        _load_field(db_kwargs, "uri", "COCOINDEX_DATABASE_URL", required=True)
+        _load_field(db_kwargs, "user", "COCOINDEX_DATABASE_USER")
+        _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
+        database = DatabaseConnectionSpec(**db_kwargs)
+        return cls(database=database)
 
 
 def init(settings: Settings):
     """Initialize the cocoindex library."""
-    _engine.init(settings.__dict__)
+    _engine.init(dump_engine_object(settings))
 
 @dataclass
 class ServerSettings:

python/cocoindex/storages.py

@@ -9,7 +9,7 @@
 class Postgres(op.StorageSpec):
     """Storage powered by Postgres and pgvector."""
 
-    database_url: str | None = None
+    database: AuthEntryReference | None = None
     table_name: str | None = None
 
 @dataclass

src/base/spec.rs

@@ -296,7 +296,69 @@ pub struct SimpleSemanticsQueryHandlerSpec {
     pub default_similarity_metric: VectorSimilarityMetric,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
-pub struct AuthEntryReference {
+pub struct AuthEntryReference<T> {
     pub key: String,
+    _phantom: std::marker::PhantomData<T>,
+}
+
+impl<T> std::fmt::Debug for AuthEntryReference<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "AuthEntryReference({})", self.key)
+    }
+}
+
+impl<T> std::fmt::Display for AuthEntryReference<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "AuthEntryReference({})", self.key)
+    }
+}
+
+impl<T> Clone for AuthEntryReference<T> {
+    fn clone(&self) -> Self {
+        Self {
+            key: self.key.clone(),
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize)]
+struct UntypedAuthEntryReference<T> {
+    key: T,
+}
+
+impl<T> Serialize for AuthEntryReference<T> {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        UntypedAuthEntryReference { key: &self.key }.serialize(serializer)
+    }
+}
+
+impl<'de, T> Deserialize<'de> for AuthEntryReference<T> {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let untyped_ref = UntypedAuthEntryReference::<String>::deserialize(deserializer)?;
+        Ok(AuthEntryReference {
+            key: untyped_ref.key,
+            _phantom: std::marker::PhantomData,
+        })
+    }
+}
+
+impl<T> PartialEq for AuthEntryReference<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.key == other.key
+    }
+}
+
+impl<T> Eq for AuthEntryReference<T> {}
+
+impl<T> std::hash::Hash for AuthEntryReference<T> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.key.hash(state);
+    }
 }

src/builder/analyzed_flow.rs

@@ -30,7 +30,7 @@ impl AnalyzedFlow {
             registry,
         )?;
         let setup_status_check =
-            setup::check_flow_setup_status(Some(&desired_state), existing_flow_ss)?;
+            setup::check_flow_setup_status(Some(&desired_state), existing_flow_ss).await?;
         let execution_plan = if setup_status_check.is_up_to_date() {
             Some(
                 async move {

src/execution/db_tracking_setup.rs

@@ -157,7 +157,7 @@ impl ResourceSetupStatusCheck for TrackingTableSetupStatusCheck {
     }
 
     async fn apply_change(&self) -> Result<()> {
-        let pool = &get_lib_context()?.pool;
+        let pool = &get_lib_context()?.builtin_db_pool;
         if let Some(desired) = &self.desired_state {
             for lagacy_name in self.legacy_table_names.iter() {
                 let query = format!(

src/lib_context.rs

@@ -6,6 +6,7 @@ use crate::settings;
 use crate::setup;
 use crate::{builder::AnalyzedFlow, execution::query::SimpleSemanticsQueryHandler};
 use axum::http::StatusCode;
+use sqlx::postgres::PgConnectOptions;
 use sqlx::PgPool;
 use std::collections::BTreeMap;
 use tokio::runtime::Runtime;
@@ -61,8 +62,40 @@ impl FlowContext {
 static TOKIO_RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
 static AUTH_REGISTRY: LazyLock<Arc<AuthRegistry>> = LazyLock::new(|| Arc::new(AuthRegistry::new()));
 
+#[derive(Default)]
+pub struct DbPools {
+    pub pools: Mutex<HashMap<(String, Option<String>), Arc<tokio::sync::OnceCell<PgPool>>>>,
+}
+
+impl DbPools {
+    pub async fn get_pool(&self, conn_spec: &settings::DatabaseConnectionSpec) -> Result<PgPool> {
+        let db_pool_cell = {
+            let key = (conn_spec.uri.clone(), conn_spec.user.clone());
+            let mut db_pools = self.pools.lock().unwrap();
+            db_pools.entry(key).or_default().clone()
+        };
+        let pool = db_pool_cell
+            .get_or_try_init(|| async move {
+                let mut pg_options: PgConnectOptions = conn_spec.uri.parse()?;
+                if let Some(user) = &conn_spec.user {
+                    pg_options = pg_options.username(user);
+                }
+                if let Some(password) = &conn_spec.password {
+                    pg_options = pg_options.password(password);
+                }
+                let pool = PgPool::connect_with(pg_options)
+                    .await
+                    .context("Failed to connect to database")?;
+                anyhow::Ok(pool)
+            })
+            .await?;
+        Ok(pool.clone())
+    }
+}
+
 pub struct LibContext {
-    pub pool: PgPool,
+    pub db_pools: DbPools,
+    pub builtin_db_pool: PgPool,
     pub flows: Mutex<BTreeMap<String, Arc<FlowContext>>>,
     pub all_setup_states: RwLock<setup::AllSetupState<setup::ExistingMode>>,
 }
@@ -100,13 +133,15 @@ pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
         pyo3_async_runtimes::tokio::init_with_runtime(get_runtime()).unwrap();
     });
 
+    let db_pools = DbPools::default();
     let (pool, all_setup_states) = get_runtime().block_on(async {
-        let pool = PgPool::connect(&settings.database_url).await?;
+        let pool = db_pools.get_pool(&settings.database).await?;
         let existing_ss = setup::get_existing_setup_state(&pool).await?;
         anyhow::Ok((pool, existing_ss))
     })?;
     Ok(LibContext {
-        pool,
+        db_pools,
+        builtin_db_pool: pool,
         all_setup_states: RwLock::new(all_setup_states),
         flows: Mutex::new(BTreeMap::new()),
     })

src/ops/factory_bases.rs

@@ -304,7 +304,7 @@ pub trait StorageFactoryBase: ExportTargetFactory + Send + Sync + 'static {
 
     /// Will not be called if it's setup by user.
     /// It returns an error if the target only supports setup by user.
-    fn check_setup_status(
+    async fn check_setup_status(
         &self,
         key: Self::Key,
         desired_state: Option<Self::SetupState>,
@@ -392,7 +392,7 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
         Ok((data_coll_output, decl_output))
     }
 
-    fn check_setup_status(
+    async fn check_setup_status(
         &self,
         key: &serde_json::Value,
         desired_state: Option<serde_json::Value>,
@@ -410,7 +410,8 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
             desired_state,
             existing_states,
             auth_registry,
-        )?;
+        )
+        .await?;
         Ok(Box::new(status_check))
     }
 
@@ -419,6 +420,11 @@ impl<T: StorageFactoryBase> ExportTargetFactory for T {
         StorageFactoryBase::describe_resource(self, &key)
     }
 
+    fn normalize_setup_key(&self, key: serde_json::Value) -> Result<serde_json::Value> {
+        let key: T::Key = serde_json::from_value(key.clone())?;
+        Ok(serde_json::to_value(key)?)
+    }
+
     fn check_state_compatibility(
         &self,
         desired_state: &serde_json::Value,

src/ops/interface.rs

@@ -190,14 +190,18 @@ pub trait ExportTargetFactory: Send + Sync {
 
     /// Will not be called if it's setup by user.
     /// It returns an error if the target only supports setup by user.
-    fn check_setup_status(
+    async fn check_setup_status(
         &self,
         key: &serde_json::Value,
         desired_state: Option<serde_json::Value>,
         existing_states: setup::CombinedState<serde_json::Value>,
         auth_registry: &Arc<AuthRegistry>,
     ) -> Result<Box<dyn setup::ResourceSetupStatusCheck>>;
 
+    /// 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.
+    /// This should always return the canonical serialized form.
+    fn normalize_setup_key(&self, key: serde_json::Value) -> Result<serde_json::Value>;
+
     fn check_state_compatibility(
         &self,
         desired_state: &serde_json::Value,