metaxy 0.0.0__py3-none-any.whl

metaxy/metadata_store/clickhouse.py

@@ -0,0 +1,161 @@
+"""ClickHouse metadata store - thin wrapper around IbisMetadataStore."""
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from metaxy.data_versioning.calculators.ibis import HashSQLGenerator
+    from metaxy.metadata_store.base import MetadataStore
+
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+from metaxy.metadata_store.ibis import IbisMetadataStore
+
+
+class ClickHouseMetadataStore(IbisMetadataStore):
+    """
+    ClickHouse metadata store using Ibis backend.
+
+    Convenience wrapper that configures IbisMetadataStore for ClickHouse.
+
+    Hash algorithm support:
+    - MD5: Always available (built-in)
+    - XXHASH32, XXHASH64: Available via ClickHouse's xxHash32/xxHash64 functions
+
+    Components:
+        - joiner: NarwhalsJoiner (works with any backend)
+        - calculator: IbisDataVersionCalculator (native SQL hash computation with xxHash64/xxHash32/MD5)
+        - diff_resolver: NarwhalsDiffResolver
+
+    Examples:
+        >>> # Local ClickHouse instance
+        >>> with ClickHouseMetadataStore("clickhouse://localhost:9000/default") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # With authentication
+        >>> with ClickHouseMetadataStore("clickhouse://user:pass@host:9000/db") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # Using connection params
+        >>> store = ClickHouseMetadataStore(
+        ...     backend="clickhouse",
+        ...     connection_params={
+        ...         "host": "localhost",
+        ...         "port": 9000,
+        ...         "database": "default",
+        ...         "user": "default",
+        ...         "password": ""
+        ...     },
+        ...     hash_algorithm=HashAlgorithm.XXHASH64
+        ... )
+        >>> with store:
+        ...     store.write_metadata(MyFeature, df)
+    """
+
+    def __init__(
+        self,
+        connection_string: str | None = None,
+        *,
+        connection_params: dict[str, Any] | None = None,
+        fallback_stores: list["MetadataStore"] | None = None,
+        **kwargs,
+    ):
+        """
+        Initialize ClickHouse metadata store.
+
+        Args:
+            connection_string: ClickHouse connection string.
+                Format: "clickhouse://[user[:password]@]host[:port]/database[?param=value]"
+                Examples:
+                    - "clickhouse://localhost:9000/default"
+                    - "clickhouse://user:pass@host:9000/db"
+                    - "clickhouse://host:9000/db?secure=true"
+            connection_params: Alternative to connection_string, specify params as dict:
+                - host: Server host (default: "localhost")
+                - port: Server port (default: 9000)
+                - database: Database name (default: "default")
+                - user: Username (default: "default")
+                - password: Password (default: "")
+                - secure: Use secure connection (default: False)
+            fallback_stores: Ordered list of read-only fallback stores.
+            **kwargs: Passed to IbisMetadataStore (e.g., hash_algorithm, graph)
+
+        Raises:
+            ImportError: If ibis-clickhouse not installed
+            ValueError: If neither connection_string nor connection_params provided
+        """
+        if connection_string is None and connection_params is None:
+            raise ValueError(
+                "Must provide either connection_string or connection_params. "
+                "Example: connection_string='clickhouse://localhost:9000/default'"
+            )
+
+        # Initialize Ibis store with ClickHouse backend
+        super().__init__(
+            connection_string=connection_string,
+            backend="clickhouse" if connection_string is None else None,
+            connection_params=connection_params,
+            fallback_stores=fallback_stores,
+            **kwargs,
+        )
+
+    def _get_default_hash_algorithm(self) -> HashAlgorithm:
+        """Get default hash algorithm for ClickHouse stores.
+
+        Uses XXHASH64 which is built-in to ClickHouse.
+        """
+        return HashAlgorithm.XXHASH64
+
+    def _supports_native_components(self) -> bool:
+        """ClickHouse stores support native data version calculations when connection is open."""
+        return self._conn is not None
+
+    def _get_hash_sql_generators(self) -> dict[HashAlgorithm, "HashSQLGenerator"]:
+        """Get hash SQL generators for ClickHouse.
+
+        ClickHouse supports:
+        - MD5: Always available (built-in)
+        - XXHASH32, XXHASH64: Always available (built-in xxHash32/xxHash64 functions)
+
+        Returns:
+            Dictionary mapping HashAlgorithm to SQL generator functions
+        """
+
+        def md5_generator(table, concat_columns: dict[str, str]) -> str:
+            hash_selects: list[str] = []
+            for field_key, concat_col in concat_columns.items():
+                hash_col = f"__hash_{field_key}"
+                # MD5() in ClickHouse returns FixedString(16) binary, convert to lowercase hex string
+                # Use lower(hex(MD5(...))) to match DuckDB's md5() lowercase hex output
+                hash_expr = f"lower(hex(MD5({concat_col})))"
+                hash_selects.append(f"{hash_expr} as {hash_col}")
+
+            hash_clause = ", ".join(hash_selects)
+            table_sql = table.compile()
+            return f"SELECT *, {hash_clause} FROM ({table_sql}) AS __metaxy_temp"
+
+        def xxhash32_generator(table, concat_columns: dict[str, str]) -> str:
+            hash_selects: list[str] = []
+            for field_key, concat_col in concat_columns.items():
+                hash_col = f"__hash_{field_key}"
+                hash_expr = f"CAST(xxHash32({concat_col}) AS String)"
+                hash_selects.append(f"{hash_expr} as {hash_col}")
+
+            hash_clause = ", ".join(hash_selects)
+            table_sql = table.compile()
+            return f"SELECT *, {hash_clause} FROM ({table_sql}) AS __metaxy_temp"
+
+        def xxhash64_generator(table, concat_columns: dict[str, str]) -> str:
+            hash_selects: list[str] = []
+            for field_key, concat_col in concat_columns.items():
+                hash_col = f"__hash_{field_key}"
+                hash_expr = f"CAST(xxHash64({concat_col}) AS String)"
+                hash_selects.append(f"{hash_expr} as {hash_col}")
+
+            hash_clause = ", ".join(hash_selects)
+            table_sql = table.compile()
+            return f"SELECT *, {hash_clause} FROM ({table_sql}) AS __metaxy_temp"
+
+        return {
+            HashAlgorithm.MD5: md5_generator,
+            HashAlgorithm.XXHASH32: xxhash32_generator,
+            HashAlgorithm.XXHASH64: xxhash64_generator,
+        }

metaxy/metadata_store/duckdb.py

@@ -0,0 +1,167 @@
+"""DuckDB metadata store - thin wrapper around IbisMetadataStore."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING, TypedDict
+
+if TYPE_CHECKING:
+    from metaxy.metadata_store.base import MetadataStore
+
+from metaxy.data_versioning.hash_algorithms import HashAlgorithm
+from metaxy.metadata_store.ibis import IbisMetadataStore
+
+
+class ExtensionSpec(TypedDict, total=False):
+    """
+    DuckDB extension specification.
+
+    Can be expressed in TOML as:
+        extensions = ["hashfuncs"]  # string form, uses 'community' repo
+        extensions = [{name = "hashfuncs"}]  # dict form, uses 'community' repo
+        extensions = [{name = "spatial", repository = "core_nightly"}]
+        extensions = [{name = "my_ext", repository = "https://my-repo.com"}]
+    """
+
+    name: str
+    repository: str  # defaults to "community" if not specified
+
+
+class DuckDBMetadataStore(IbisMetadataStore):
+    """
+    DuckDB metadata store using Ibis backend.
+
+    Convenience wrapper that configures IbisMetadataStore for DuckDB.
+
+    Hash algorithm support is detected dynamically based on installed extensions:
+    - MD5: Always available (built-in)
+    - XXHASH32, XXHASH64: Available when 'hashfuncs' extension is loaded
+
+    Components:
+        - joiner: NarwhalsJoiner (works with any backend)
+        - calculator: IbisDataVersionCalculator (native SQL hash computation with xxh64/xxh32/md5)
+        - diff_resolver: NarwhalsDiffResolver
+
+    Examples:
+        >>> # Local file database
+        >>> with DuckDBMetadataStore("metadata.db") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # In-memory database
+        >>> with DuckDBMetadataStore(":memory:") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # MotherDuck
+        >>> with DuckDBMetadataStore("md:my_database") as store:
+        ...     store.write_metadata(MyFeature, df)
+
+        >>> # With extensions
+        >>> store = DuckDBMetadataStore(
+        ...     "metadata.db",
+        ...     hash_algorithm=HashAlgorithm.XXHASH64,
+        ...     extensions=["hashfuncs"]
+        ... )
+        >>> with store:
+        ...     store.write_metadata(MyFeature, df)
+    """
+
+    def __init__(
+        self,
+        database: str | Path,
+        *,
+        config: dict[str, str] | None = None,
+        extensions: list[ExtensionSpec | str] | None = None,
+        fallback_stores: list["MetadataStore"] | None = None,
+        **kwargs,
+    ):
+        """
+        Initialize DuckDB metadata store.
+
+        Args:
+            database: Database connection string or path.
+                - File path: "metadata.db" or Path("metadata.db")
+                - In-memory: ":memory:"
+                - MotherDuck: "md:my_database" or "md:my_database?motherduck_token=..."
+                - S3: "s3://bucket/path/database.duckdb" (read-only via ATTACH)
+                - HTTPS: "https://example.com/database.duckdb" (read-only via ATTACH)
+                - Any valid DuckDB connection string
+
+                Note: Parent directories are NOT created automatically. Ensure paths exist
+                before initializing the store.
+            config: Optional DuckDB configuration settings (e.g., {'threads': '4', 'memory_limit': '4GB'})
+            extensions: List of DuckDB extensions to install and load on open.
+                Can be strings (installed from 'community' repository) or dicts
+                specifying both name and repository.
+
+                Examples:
+                    extensions=['hashfuncs']  # Install hashfuncs from community
+                    extensions=[{'name': 'hashfuncs'}]  # Same as above
+                    extensions=[{'name': 'spatial', 'repository': 'core_nightly'}]
+                    extensions=[{'name': 'my_ext', 'repository': 'https://my-repo.com'}]
+            fallback_stores: Ordered list of read-only fallback stores.
+            **kwargs: Passed to IbisMetadataStore (e.g., hash_algorithm, graph)
+        """
+        database_str = str(database)
+
+        # Build connection params for Ibis DuckDB backend
+        # Ibis DuckDB backend accepts config params directly (not nested under 'config')
+        connection_params = {"database": database_str}
+        if config:
+            connection_params.update(config)
+
+        self.database = database_str
+        self.extensions = extensions or []
+
+        # Auto-add hashfuncs extension if not present (needed for default XXHASH64)
+        extension_names = [
+            ext if isinstance(ext, str) else ext.get("name", "")
+            for ext in self.extensions
+        ]
+        if "hashfuncs" not in extension_names:
+            self.extensions.append("hashfuncs")
+
+        # Initialize Ibis store with DuckDB backend
+        super().__init__(
+            backend="duckdb",
+            connection_params=connection_params,
+            fallback_stores=fallback_stores,
+            **kwargs,
+        )
+
+    def _get_default_hash_algorithm(self) -> HashAlgorithm:
+        """Get default hash algorithm for DuckDB stores.
+
+        Uses XXHASH64 which requires the hashfuncs extension (lazily loaded).
+        """
+        return HashAlgorithm.XXHASH64
+
+    def _supports_native_components(self) -> bool:
+        """DuckDB stores support native data version calculations when connection is open."""
+        return self._conn is not None
+
+    def _create_native_components(self):
+        """Create components for native SQL execution with DuckDB.
+
+        Uses DuckDBDataVersionCalculator which handles extension loading lazily.
+        Extensions are loaded when the calculator is created (on-demand), not on store open.
+        """
+        from metaxy.data_versioning.calculators.duckdb import (
+            DuckDBDataVersionCalculator,
+        )
+        from metaxy.data_versioning.diff.narwhals import NarwhalsDiffResolver
+        from metaxy.data_versioning.joiners.narwhals import NarwhalsJoiner
+
+        if self._conn is None:
+            raise RuntimeError(
+                "Cannot create native data version calculations: store is not open. "
+                "Ensure store is used as context manager."
+            )
+
+        # All components accept/return Narwhals LazyFrames
+        # DuckDBDataVersionCalculator loads extensions and generates SQL for hashing
+        joiner = NarwhalsJoiner()
+        calculator = DuckDBDataVersionCalculator(
+            backend=self._conn,
+            extensions=self.extensions,
+        )
+        diff_resolver = NarwhalsDiffResolver()
+
+        return joiner, calculator, diff_resolver

metaxy/metadata_store/exceptions.py

@@ -0,0 +1,43 @@
+"""Exceptions for metadata store operations."""
+
+
+class MetadataStoreError(Exception):
+    """Base exception for metadata store errors."""
+
+    pass
+
+
+class FeatureNotFoundError(MetadataStoreError):
+    """Raised when a feature is not found in the store."""
+
+    pass
+
+
+class FieldNotFoundError(MetadataStoreError):
+    """Raised when a field is not found for a feature."""
+
+    pass
+
+
+class MetadataSchemaError(MetadataStoreError):
+    """Raised when metadata DataFrame has invalid schema."""
+
+    pass
+
+
+class DependencyError(MetadataStoreError):
+    """Raised when upstream dependencies are missing or invalid."""
+
+    pass
+
+
+class StoreNotOpenError(MetadataStoreError):
+    """Raised when attempting to use a store that hasn't been opened."""
+
+    pass
+
+
+class HashAlgorithmNotSupportedError(MetadataStoreError):
+    """Raised when a hash algorithm is not supported by the store or its components."""
+
+    pass