cocoindex 0.2.3__cp311-abi3-macosx_10_12_x86_64.whl

cocoindex/setting.py

@@ -0,0 +1,181 @@
+"""
+Data types for settings of the cocoindex library.
+"""
+
+import os
+
+from typing import Callable, Self, Any, overload
+from dataclasses import dataclass
+from .validation import validate_app_namespace_name
+
+_app_namespace: str = ""
+
+
+def get_app_namespace(*, trailing_delimiter: str | None = None) -> str:
+    """Get the application namespace. Append the `trailing_delimiter` if not empty."""
+    if _app_namespace == "" or trailing_delimiter is None:
+        return _app_namespace
+    return f"{_app_namespace}{trailing_delimiter}"
+
+
+def split_app_namespace(full_name: str, delimiter: str) -> tuple[str, str]:
+    """Split the full name into the application namespace and the rest."""
+    parts = full_name.split(delimiter, 1)
+    if len(parts) == 1:
+        return "", parts[0]
+    return (parts[0], parts[1])
+
+
+def set_app_namespace(app_namespace: str) -> None:
+    """Set the application namespace."""
+    if app_namespace:
+        validate_app_namespace_name(app_namespace)
+    global _app_namespace  # pylint: disable=global-statement
+    _app_namespace = app_namespace
+
+
+@dataclass
+class DatabaseConnectionSpec:
+    """
+    Connection spec for relational database.
+    Used by both internal and target storage.
+    """
+
+    url: str
+    user: str | None = None
+    password: str | None = None
+    max_connections: int = 25
+    min_connections: int = 5
+
+
+@dataclass
+class GlobalExecutionOptions:
+    """Global execution options."""
+
+    # The maximum number of concurrent inflight requests, shared among all sources from all flows.
+    source_max_inflight_rows: int | None = 1024
+    source_max_inflight_bytes: int | None = None
+
+
+def _load_field(
+    target: dict[str, Any],
+    name: str,
+    env_name: str,
+    required: bool = False,
+    parse: Callable[[str], Any] | None = None,
+) -> None:
+    value = os.getenv(env_name)
+    if value is None:
+        if required:
+            raise ValueError(f"{env_name} is not set")
+    else:
+        if parse is None:
+            target[name] = value
+        else:
+            try:
+                target[name] = parse(value)
+            except Exception as e:
+                raise ValueError(
+                    f"failed to parse environment variable {env_name}: {value}"
+                ) from e
+
+
+@dataclass
+class Settings:
+    """Settings for the cocoindex library."""
+
+    database: DatabaseConnectionSpec | None = None
+    app_namespace: str = ""
+    global_execution_options: GlobalExecutionOptions | None = None
+
+    @classmethod
+    def from_env(cls) -> Self:
+        """Load settings from environment variables."""
+
+        database_url = os.getenv("COCOINDEX_DATABASE_URL")
+        if database_url is not None:
+            db_kwargs: dict[str, Any] = dict()
+            _load_field(db_kwargs, "url", "COCOINDEX_DATABASE_URL", required=True)
+            _load_field(db_kwargs, "user", "COCOINDEX_DATABASE_USER")
+            _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
+            _load_field(
+                db_kwargs,
+                "max_connections",
+                "COCOINDEX_DATABASE_MAX_CONNECTIONS",
+                parse=int,
+            )
+            _load_field(
+                db_kwargs,
+                "min_connections",
+                "COCOINDEX_DATABASE_MIN_CONNECTIONS",
+                parse=int,
+            )
+            database = DatabaseConnectionSpec(**db_kwargs)
+        else:
+            database = None
+
+        exec_kwargs: dict[str, Any] = dict()
+        _load_field(
+            exec_kwargs,
+            "source_max_inflight_rows",
+            "COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS",
+            parse=int,
+        )
+        _load_field(
+            exec_kwargs,
+            "source_max_inflight_bytes",
+            "COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES",
+            parse=int,
+        )
+        global_execution_options = GlobalExecutionOptions(**exec_kwargs)
+
+        app_namespace = os.getenv("COCOINDEX_APP_NAMESPACE", "")
+
+        return cls(
+            database=database,
+            app_namespace=app_namespace,
+            global_execution_options=global_execution_options,
+        )
+
+
+@dataclass
+class ServerSettings:
+    """Settings for the cocoindex server."""
+
+    # The address to bind the server to.
+    address: str = "127.0.0.1:49344"
+
+    # The origins of the clients (e.g. CocoInsight UI) to allow CORS from.
+    cors_origins: list[str] | None = None
+
+    @classmethod
+    def from_env(cls) -> Self:
+        """Load settings from environment variables."""
+        kwargs: dict[str, Any] = dict()
+        _load_field(kwargs, "address", "COCOINDEX_SERVER_ADDRESS")
+        _load_field(
+            kwargs,
+            "cors_origins",
+            "COCOINDEX_SERVER_CORS_ORIGINS",
+            parse=ServerSettings.parse_cors_origins,
+        )
+        return cls(**kwargs)
+
+    @overload
+    @staticmethod
+    def parse_cors_origins(s: str) -> list[str]: ...
+
+    @overload
+    @staticmethod
+    def parse_cors_origins(s: str | None) -> list[str] | None: ...
+
+    @staticmethod
+    def parse_cors_origins(s: str | None) -> list[str] | None:
+        """
+        Parse the CORS origins from a string.
+        """
+        return (
+            [o for e in s.split(",") if (o := e.strip()) != ""]
+            if s is not None
+            else None
+        )

cocoindex/setup.py

@@ -0,0 +1,92 @@
+"""
+This module provides APIs to manage the setup of flows.
+"""
+
+from . import setting
+from . import _engine  # type: ignore
+from .runtime import execution_context
+
+
+class SetupChangeBundle:
+    """
+    This class represents a bundle of setup changes.
+    """
+
+    _engine_bundle: _engine.SetupChangeBundle
+
+    def __init__(self, _engine_bundle: _engine.SetupChangeBundle):
+        self._engine_bundle = _engine_bundle
+
+    def __str__(self) -> str:
+        desc, _ = execution_context.run(self._engine_bundle.describe_async())
+        return desc  # type: ignore
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def apply(self, report_to_stdout: bool = False) -> None:
+        """
+        Apply the setup changes.
+        """
+        execution_context.run(self.apply_async(report_to_stdout=report_to_stdout))
+
+    async def apply_async(self, report_to_stdout: bool = False) -> None:
+        """
+        Apply the setup changes. Async version of `apply`.
+        """
+        await self._engine_bundle.apply_async(report_to_stdout=report_to_stdout)
+
+    def describe(self) -> tuple[str, bool]:
+        """
+        Describe the setup changes.
+        """
+        return execution_context.run(self.describe_async())  # type: ignore
+
+    async def describe_async(self) -> tuple[str, bool]:
+        """
+        Describe the setup changes. Async version of `describe`.
+        """
+        return await self._engine_bundle.describe_async()  # type: ignore
+
+    def describe_and_apply(self, report_to_stdout: bool = False) -> None:
+        """
+        Describe the setup changes and apply them if `report_to_stdout` is True.
+        Silently apply setup changes otherwise.
+        """
+        execution_context.run(
+            self.describe_and_apply_async(report_to_stdout=report_to_stdout)
+        )
+
+    async def describe_and_apply_async(self, *, report_to_stdout: bool = False) -> None:
+        """
+        Describe the setup changes and apply them if `report_to_stdout` is True.
+        Silently apply setup changes otherwise. Async version of `describe_and_apply`.
+        """
+        if report_to_stdout:
+            desc, is_up_to_date = await self.describe_async()
+            print("Setup status:\n")
+            print(desc)
+            if is_up_to_date:
+                print("No setup changes to apply.")
+                return
+        await self.apply_async(report_to_stdout=report_to_stdout)
+
+
+def flow_names_with_setup() -> list[str]:
+    """
+    Get the names of all flows that have been setup.
+    """
+    return execution_context.run(flow_names_with_setup_async())  # type: ignore
+
+
+async def flow_names_with_setup_async() -> list[str]:
+    """
+    Get the names of all flows that have been setup. Async version of `flow_names_with_setup`.
+    """
+    result = []
+    all_flow_names = await _engine.flow_names_with_setup_async()
+    for name in all_flow_names:
+        app_namespace, name = setting.split_app_namespace(name, ".")
+        if app_namespace == setting.get_app_namespace():
+            result.append(name)
+    return result

cocoindex/sources.py

@@ -0,0 +1,102 @@
+"""All builtin sources."""
+
+from . import op
+from .auth_registry import TransientAuthEntryReference
+from .setting import DatabaseConnectionSpec
+from dataclasses import dataclass
+import datetime
+
+
+class LocalFile(op.SourceSpec):
+    """Import data from local file system."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    path: str
+    binary: bool = False
+
+    # If provided, only files matching these patterns will be included.
+    # See https://docs.rs/globset/latest/globset/index.html#syntax for the syntax of the patterns.
+    included_patterns: list[str] | None = None
+
+    # If provided, files matching these patterns will be excluded.
+    # See https://docs.rs/globset/latest/globset/index.html#syntax for the syntax of the patterns.
+    excluded_patterns: list[str] | None = None
+
+
+class GoogleDrive(op.SourceSpec):
+    """Import data from Google Drive."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    service_account_credential_path: str
+    root_folder_ids: list[str]
+    binary: bool = False
+    recent_changes_poll_interval: datetime.timedelta | None = None
+
+
+class AmazonS3(op.SourceSpec):
+    """Import data from an Amazon S3 bucket. Supports optional prefix and file filtering by glob patterns."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    bucket_name: str
+    prefix: str | None = None
+    binary: bool = False
+    included_patterns: list[str] | None = None
+    excluded_patterns: list[str] | None = None
+    sqs_queue_url: str | None = None
+
+
+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
+
+    account_name: str
+    container_name: str
+    prefix: str | None = None
+    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
+
+
+@dataclass
+class PostgresNotification:
+    """Notification for a PostgreSQL table."""
+
+    # Optional: name of the PostgreSQL channel to use.
+    # If not provided, will generate a default channel name.
+    channel_name: str | None = None
+
+
+class Postgres(op.SourceSpec):
+    """Import data from a PostgreSQL table."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    # Table name to read from (required)
+    table_name: str
+
+    # Database connection reference (optional - uses default if not provided)
+    database: TransientAuthEntryReference[DatabaseConnectionSpec] | None = None
+
+    # Optional: specific columns to include (if None, includes all columns)
+    included_columns: list[str] | None = None
+
+    # Optional: column name to use for ordinal tracking (for incremental updates)
+    # Should be a timestamp, serial, or other incrementing column
+    ordinal_column: str | None = None
+
+    # Optional: when set, supports change capture from PostgreSQL notification.
+    notification: PostgresNotification | None = None

cocoindex/subprocess_exec.py

@@ -0,0 +1,279 @@
+"""
+Lightweight subprocess-backed executor stub.
+
+- Uses a single global ProcessPoolExecutor (max_workers=1), created lazily.
+- In the subprocess, maintains a registry of executor instances keyed by
+  (executor_factory, pickled spec) to enable reuse.
+- Caches analyze() and prepare() results per key to avoid repeated calls
+  even if key collision happens.
+"""
+
+from __future__ import annotations
+
+from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures.process import BrokenProcessPool
+from dataclasses import dataclass, field
+from typing import Any, Callable
+import pickle
+import threading
+import asyncio
+import os
+import time
+import atexit
+from .user_app_loader import load_user_app
+from .runtime import execution_context
+import logging
+import multiprocessing as mp
+
+WATCHDOG_INTERVAL_SECONDS = 10.0
+
+# ---------------------------------------------
+# Main process: single, lazily-created pool
+# ---------------------------------------------
+_pool_lock = threading.Lock()
+_pool: ProcessPoolExecutor | None = None
+_pool_cleanup_registered = False
+_user_apps: list[str] = []
+_logger = logging.getLogger(__name__)
+
+
+def shutdown_pool_at_exit() -> None:
+    """Best-effort shutdown of the global ProcessPoolExecutor on interpreter exit."""
+    global _pool, _pool_cleanup_registered  # pylint: disable=global-statement
+    with _pool_lock:
+        if _pool is not None:
+            try:
+                _pool.shutdown(wait=True, cancel_futures=True)
+            except Exception as e:
+                _logger.error(
+                    "Error during ProcessPoolExecutor shutdown at exit: %s",
+                    e,
+                    exc_info=True,
+                )
+            finally:
+                _pool = None
+                _pool_cleanup_registered = False
+
+
+def _get_pool() -> ProcessPoolExecutor:
+    global _pool, _pool_cleanup_registered  # pylint: disable=global-statement
+    with _pool_lock:
+        if _pool is None:
+            if not _pool_cleanup_registered:
+                # Register the shutdown at exit at creation time (rather than at import time)
+                # to make sure it's executed earlier in the shutdown sequence.
+                atexit.register(shutdown_pool_at_exit)
+                _pool_cleanup_registered = True
+
+            # Single worker process as requested
+            _pool = ProcessPoolExecutor(
+                max_workers=1,
+                initializer=_subprocess_init,
+                initargs=(_user_apps, os.getpid()),
+                mp_context=mp.get_context("spawn"),
+            )
+        return _pool
+
+
+def add_user_app(app_target: str) -> None:
+    with _pool_lock:
+        _user_apps.append(app_target)
+
+
+def _restart_pool(old_pool: ProcessPoolExecutor | None = None) -> None:
+    """Safely restart the global ProcessPoolExecutor.
+
+    Thread-safe via `_pool_lock`. Shuts down the old pool and re-creates a new
+    one with the same initializer/args.
+    """
+    global _pool
+    with _pool_lock:
+        # If another thread already swapped the pool, skip restart
+        if old_pool is not None and _pool is not old_pool:
+            return
+        _logger.error("Detected dead subprocess pool; restarting and retrying.")
+        prev_pool = _pool
+        _pool = ProcessPoolExecutor(
+            max_workers=1,
+            initializer=_subprocess_init,
+            initargs=(_user_apps, os.getpid()),
+            mp_context=mp.get_context("spawn"),
+        )
+        if prev_pool is not None:
+            # Best-effort shutdown of previous pool; letting exceptions bubble up
+            # is acceptable here and signals irrecoverable executor state.
+            prev_pool.shutdown(cancel_futures=True)
+
+
+async def _submit_with_restart(fn: Callable[..., Any], *args: Any) -> Any:
+    """Submit and await work, restarting the subprocess until it succeeds.
+
+    Retries on BrokenProcessPool or pool-shutdown RuntimeError; re-raises other
+    exceptions.
+    """
+    while True:
+        pool = _get_pool()
+        try:
+            fut = pool.submit(fn, *args)
+            return await asyncio.wrap_future(fut)
+        except BrokenProcessPool:
+            _restart_pool(old_pool=pool)
+            # loop and retry
+
+
+# ---------------------------------------------
+# Subprocess: executor registry and helpers
+# ---------------------------------------------
+
+
+def _start_parent_watchdog(
+    parent_pid: int, interval_seconds: float = WATCHDOG_INTERVAL_SECONDS
+) -> None:
+    """Terminate this process if the parent process exits or PPID changes.
+
+    This runs in a background daemon thread so it never blocks pool work.
+    """
+
+    def _watch() -> None:
+        while True:
+            # If PPID changed (parent died and we were reparented), exit.
+            if os.getppid() != parent_pid:
+                os._exit(1)
+
+            # Best-effort liveness probe in case PPID was reused.
+            try:
+                os.kill(parent_pid, 0)
+            except OSError:
+                os._exit(1)
+
+            time.sleep(interval_seconds)
+
+    threading.Thread(target=_watch, name="parent-watchdog", daemon=True).start()
+
+
+def _subprocess_init(user_apps: list[str], parent_pid: int) -> None:
+    _start_parent_watchdog(parent_pid)
+
+    # In case any user app is already in this subprocess, e.g. the subprocess is forked, we need to avoid loading it again.
+    with _pool_lock:
+        already_loaded_apps = set(_user_apps)
+
+    loaded_apps = []
+    for app_target in user_apps:
+        if app_target not in already_loaded_apps:
+            load_user_app(app_target)
+            loaded_apps.append(app_target)
+
+    with _pool_lock:
+        _user_apps.extend(loaded_apps)
+
+
+class _OnceResult:
+    _result: Any = None
+    _done: bool = False
+
+    def run_once(self, method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        if self._done:
+            return self._result
+        self._result = _call_method(method, *args, **kwargs)
+        self._done = True
+        return self._result
+
+
+@dataclass
+class _ExecutorEntry:
+    executor: Any
+    prepare: _OnceResult = field(default_factory=_OnceResult)
+    analyze: _OnceResult = field(default_factory=_OnceResult)
+    ready_to_call: bool = False
+
+
+_SUBPROC_EXECUTORS: dict[bytes, _ExecutorEntry] = {}
+
+
+def _call_method(method: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    """Run an awaitable/coroutine to completion synchronously, otherwise return as-is."""
+    if asyncio.iscoroutinefunction(method):
+        return asyncio.run(method(*args, **kwargs))
+    else:
+        return method(*args, **kwargs)
+
+
+def _get_or_create_entry(key_bytes: bytes) -> _ExecutorEntry:
+    entry = _SUBPROC_EXECUTORS.get(key_bytes)
+    if entry is None:
+        executor_factory, spec = pickle.loads(key_bytes)
+        inst = executor_factory()
+        inst.spec = spec
+        entry = _ExecutorEntry(executor=inst)
+        _SUBPROC_EXECUTORS[key_bytes] = entry
+    return entry
+
+
+def _sp_analyze(key_bytes: bytes) -> Any:
+    entry = _get_or_create_entry(key_bytes)
+    return entry.analyze.run_once(entry.executor.analyze)
+
+
+def _sp_prepare(key_bytes: bytes) -> Any:
+    entry = _get_or_create_entry(key_bytes)
+    return entry.prepare.run_once(entry.executor.prepare)
+
+
+def _sp_call(key_bytes: bytes, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+    entry = _get_or_create_entry(key_bytes)
+    # There's a chance that the subprocess crashes and restarts in the middle.
+    # So we want to always make sure the executor is ready before each call.
+    if not entry.ready_to_call:
+        if analyze_fn := getattr(entry.executor, "analyze", None):
+            entry.analyze.run_once(analyze_fn)
+        if prepare_fn := getattr(entry.executor, "prepare", None):
+            entry.prepare.run_once(prepare_fn)
+        entry.ready_to_call = True
+    return _call_method(entry.executor.__call__, *args, **kwargs)
+
+
+# ---------------------------------------------
+# Public stub
+# ---------------------------------------------
+
+
+class _ExecutorStub:
+    _key_bytes: bytes
+
+    def __init__(self, executor_factory: type[Any], spec: Any) -> None:
+        self._key_bytes = pickle.dumps(
+            (executor_factory, spec), protocol=pickle.HIGHEST_PROTOCOL
+        )
+
+        # Conditionally expose analyze if underlying class has it
+        if hasattr(executor_factory, "analyze"):
+            # Bind as attribute so getattr(..., "analyze", None) works upstream
+            def analyze() -> Any:
+                return execution_context.run(
+                    _submit_with_restart(_sp_analyze, self._key_bytes)
+                )
+
+            # Attach method
+            setattr(self, "analyze", analyze)
+
+        if hasattr(executor_factory, "prepare"):
+
+            async def prepare() -> Any:
+                return await _submit_with_restart(_sp_prepare, self._key_bytes)
+
+            setattr(self, "prepare", prepare)
+
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return await _submit_with_restart(_sp_call, self._key_bytes, args, kwargs)
+
+
+def executor_stub(executor_factory: type[Any], spec: Any) -> Any:
+    """
+    Create a subprocess-backed stub for the given executor class/spec.
+
+    - Lazily initializes a singleton ProcessPoolExecutor (max_workers=1).
+    - Returns a stub object exposing async __call__ and async prepare; analyze is
+      exposed if present on the original class.
+    """
+    return _ExecutorStub(executor_factory, spec)