metaxy 0.0.0__py3-none-any.whl

metaxy/ext/sqlmodel.py

@@ -0,0 +1,172 @@
+"""SQLModel integration for Metaxy.
+
+This module provides a combined metaclass that allows Metaxy Feature classes
+to also be SQLModel table classes, enabling seamless integration with SQLAlchemy/SQLModel ORMs.
+"""
+
+from typing import Any
+
+from sqlalchemy.types import JSON
+from sqlmodel import Field, SQLModel
+from sqlmodel.main import SQLModelMetaclass
+
+from metaxy.config import MetaxyConfig
+from metaxy.models.feature import Feature, MetaxyMeta
+from metaxy.models.feature_spec import FeatureSpec
+
+
+class SQLModelFeatureMeta(MetaxyMeta, SQLModelMetaclass):  # pyright: ignore[reportUnsafeMultipleInheritance]
+    """Combined metaclass for SQLModel + Metaxy Feature.
+
+    This metaclass inherits from both MetaxyMeta and SQLModelMetaclass,
+    allowing classes to be both SQLModel tables and Metaxy features.
+
+    The MRO (Method Resolution Order) ensures that:
+    1. MetaxyMeta handles Metaxy feature registration
+    2. SQLModelMetaclass handles SQLAlchemy table creation
+
+    Note: MetaxyMeta comes first in the inheritance to ensure spec parameter is handled.
+
+    Automatic table naming:
+        __tablename__ is automatically generated from the feature key if not explicitly provided.
+        This ensures consistency with Metaxy's metadata store table naming conventions.
+
+    Example:
+        >>> from metaxy.integrations.sqlmodel import SQLModelFeature
+        >>> from metaxy import FeatureSpec, FeatureKey, FieldSpec, FieldKey
+        >>> from sqlmodel import Field
+        >>>
+        >>> class MyFeature(
+        ...     SQLModelFeature,
+        ...     table=True,
+        ...     spec=FeatureSpec(
+        ...         key=FeatureKey(["my", "feature"]),
+        ...         deps=None,
+        ...         fields=[
+        ...             FieldSpec(
+        ...                 key=FieldKey(["data"]),
+        ...                 code_version=1,
+        ...             ),
+        ...         ],
+        ...     ),
+        ... ):
+        ...     __tablename__ = "my_feature"
+        ...
+        ...     uid: str = Field(primary_key=True)
+        ...     data: str
+    """
+
+    def __new__(
+        cls,
+        cls_name: str,
+        bases: tuple[type[Any], ...],
+        namespace: dict[str, Any],
+        *,
+        spec: FeatureSpec | None = None,
+        **kwargs: Any,
+    ) -> type[Any]:
+        """Create a new SQLModel + Feature class.
+
+        Args:
+            cls_name: Name of the class being created
+            bases: Base classes
+            namespace: Class namespace (attributes and methods)
+            spec: Metaxy FeatureSpec (required for concrete features)
+            **kwargs: Additional keyword arguments (e.g., table=True for SQLModel)
+
+        Returns:
+            New class that is both a SQLModel table and a Metaxy feature
+        """
+        # If this is a concrete table (table=True) with a spec
+        if kwargs.get("table") and spec is not None:
+            # Automatically set __tablename__ from the feature key if not provided
+            if (
+                "__tablename__" not in namespace
+                and MetaxyConfig.get().ext.sqlmodel.infer_db_table_names
+            ):
+                namespace["__tablename__"] = spec.key.table_name
+
+        # Call super().__new__ which follows MRO: MetaxyMeta -> SQLModelMetaclass -> ...
+        # MetaxyMeta will consume the spec parameter and pass remaining kwargs to SQLModelMetaclass
+        return super().__new__(cls, cls_name, bases, namespace, spec=spec, **kwargs)
+
+
+# pyright: reportIncompatibleMethodOverride=false, reportIncompatibleVariableOverride=false
+class SQLModelFeature(SQLModel, Feature, metaclass=SQLModelFeatureMeta, spec=None):  # type: ignore[misc]
+    """Base class for features that are also SQLModel tables.
+
+    Use this as a base class when you want to create Metaxy features
+    that are also SQLAlchemy/SQLModel ORM models.
+
+    This class combines:
+    - Metaxy's Feature functionality (versioning, dependency tracking)
+    - SQLModel's ORM functionality (database mapping, queries)
+
+    Note: Unlike regular Feature classes, SQLModelFeature instances are mutable
+    to support SQLModel ORM operations. Only the spec and graph class attributes
+    are used from Feature, not the instance behavior.
+
+    System-managed fields are defined as optional here and will be populated
+    by the metadata store when reading/writing data.
+
+    Example:
+        >>> from metaxy.integrations.sqlmodel import SQLModelFeature
+        >>> from metaxy import FeatureSpec, FeatureKey, FieldSpec, FieldKey
+        >>> from sqlmodel import Field
+        >>>
+        >>> class VideoFeature(
+        ...     SQLModelFeature,
+        ...     table=True,
+        ...     spec=FeatureSpec(
+        ...         key=FeatureKey(["video"]),
+        ...         deps=None,  # Root feature
+        ...         fields=[
+        ...             FieldSpec(
+        ...                 key=FieldKey(["video_file"]),
+        ...                 code_version=1,
+        ...             ),
+        ...         ],
+        ...     ),
+        ... ):
+        ...     __tablename__ = "video"
+        ...
+        ...     uid: str = Field(primary_key=True)
+        ...     path: str
+        ...     duration: float
+        ...
+        ...     # Now you can use both Metaxy and SQLModel features:
+        ...     # - VideoFeature.feature_version() -> Metaxy versioning
+        ...     # - session.exec(select(VideoFeature)) -> SQLModel queries
+    """
+
+    # Override the frozen config from Feature's FrozenBaseModel
+    # SQLModel instances need to be mutable for ORM operations
+    model_config = {"frozen": False}  # pyright: ignore[reportAssignmentType]
+
+    # System-managed metadata fields - these are optional and populated by Metaxy
+    # The metadata store will populate these when reading/writing data
+    # Users can override these definitions in subclasses if they need different constraints
+    sample_uid: str | None = Field(default=None, nullable=True)
+
+    # Using sa_column_kwargs to map to the actual column names used by Metaxy
+    metaxy_data_version: str | None = Field(
+        default=None,
+        sa_type=JSON,
+        sa_column_kwargs={"name": "data_version", "nullable": True},
+    )
+
+    metaxy_feature_version: str | None = Field(
+        default=None, sa_column_kwargs={"name": "feature_version", "nullable": True}
+    )
+
+    metaxy_snapshot_version: str | None = Field(
+        default=None, sa_column_kwargs={"name": "snapshot_version", "nullable": True}
+    )
+
+    # All Feature class methods and attributes are inherited from Feature base class:
+    # - spec: ClassVar[FeatureSpec]
+    # - graph: ClassVar[FeatureGraph]
+    # - feature_version() -> str
+    # - data_version() -> dict[str, str]
+    # - load_input(...)
+    # - resolve_data_version_diff(...)

metaxy/ext/sqlmodel_system_tables.py

@@ -0,0 +1,139 @@
+"""SQLModel definitions for metaxy system tables.
+
+This module provides SQLModel table classes that mirror the Polars schemas
+from system_tables.py. These are used primarily for Alembic migrations
+when the SQLModel integration is enabled.
+
+The actual data flow remains unchanged (using Polars DataFrames).
+These models are only for schema definition and migrations.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import JSON, Index
+from sqlmodel import Field, SQLModel
+
+# Import the namespace constant from the core module
+from metaxy.metadata_store.system_tables import SYSTEM_NAMESPACE
+
+# System tables that metaxy uses internally
+SYSTEM_TABLES: list[str] = [
+    f"{SYSTEM_NAMESPACE}__feature_versions",
+    f"{SYSTEM_NAMESPACE}__migration_events",
+]
+
+
+class FeatureVersionsTable(SQLModel, table=True):
+    """SQLModel definition for the feature_versions system table.
+
+    This table records when feature specifications are pushed to production,
+    tracking the evolution of feature definitions over time.
+    """
+
+    __tablename__: str = f"{SYSTEM_NAMESPACE}__feature_versions"  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    # Composite primary key: feature_key + snapshot_version
+    feature_key: str = Field(primary_key=True)
+    snapshot_version: str = Field(primary_key=True)
+
+    # Version and timestamp
+    feature_version: str = Field(index=True)
+    recorded_at: datetime = Field(index=True)
+
+    # Serialized feature specification and class path
+    feature_spec: str  # JSON string
+    feature_class_path: str
+
+    # Additional indexes for common queries
+    __table_args__ = (
+        Index("idx_feature_versions_lookup", "feature_key", "feature_version"),
+    )
+
+
+class MigrationEventsTable(SQLModel, table=True):
+    """SQLModel definition for the migration_events system table.
+
+    This table stores append-only events tracking migration execution,
+    enabling recovery from partial failures and progress monitoring.
+    """
+
+    __tablename__: str = f"{SYSTEM_NAMESPACE}__migration_events"  # pyright: ignore[reportIncompatibleVariableOverride]
+
+    # Auto-incrementing ID for append-only events
+    id: int | None = Field(default=None, primary_key=True)
+
+    # Event fields
+    migration_id: str = Field(index=True)
+    event_type: str = (
+        Field()
+    )  # "started", "feature_started", "feature_completed", "completed", "failed"
+    timestamp: datetime = Field(index=True)
+    feature_key: str = Field(default="")  # Empty for migration-level events
+    rows_affected: int = Field(default=0)
+    error_message: str = Field(default="")  # Empty if no error
+
+    # Additional indexes for common queries
+    __table_args__ = (
+        Index("idx_migration_events_lookup", "migration_id", "event_type"),
+    )
+
+
+class MetaxyFeatureMetadataTemplate(SQLModel):
+    """Template showing standard metaxy columns for feature tables.
+
+    This class demonstrates the standard columns that metaxy adds to
+    feature tables when using SQLModel integration.
+
+    Users should NOT inherit from this class directly, but can use it
+    as a reference for the required columns.
+    """
+
+    # Standard metaxy columns
+    sample_uid: str = Field(primary_key=True)
+    data_version: dict[str, Any] | None = Field(
+        default=None, sa_type=JSON
+    )  # Struct column as JSON
+    feature_version: str | None = Field(default=None, index=True)
+    snapshot_version: str | None = Field(default=None, index=True)
+
+    # This is just a template, not meant to be instantiated
+    model_config = {"arbitrary_types_allowed": True}  # pyright: ignore[reportAssignmentType]
+
+
+# Registry of all system tables
+SYSTEM_TABLE_MODELS = {
+    "feature_versions": FeatureVersionsTable,
+    "migration_events": MigrationEventsTable,
+}
+
+
+def get_system_tables() -> list[type[SQLModel]]:
+    """Get all system table SQLModel classes.
+
+    Returns:
+        List of SQLModel table classes
+    """
+    return [FeatureVersionsTable, MigrationEventsTable]
+
+
+def get_system_metadata():
+    """Get SQLAlchemy metadata containing all system tables.
+
+    This function returns the metadata object that contains all
+    system table definitions, which can be used with Alembic
+    for migration generation.
+
+    Returns:
+        SQLModel.metadata containing all system table definitions
+
+    Example:
+        >>> from metaxy.ext.sqlmodel_system_tables import get_system_metadata
+        >>> metadata = get_system_metadata()
+        >>> # Use with Alembic target_metadata
+    """
+    # All tables are automatically registered in SQLModel.metadata
+    # when their classes are defined with table=True
+    return SQLModel.metadata

metaxy/graph/__init__.py

@@ -0,0 +1,21 @@
+"""Graph visualization and rendering utilities."""
+
+from metaxy.graph.diff import GraphData
+from metaxy.graph.diff.rendering import (
+    BaseRenderer,
+    CardsRenderer,
+    GraphvizRenderer,
+    MermaidRenderer,
+    RenderConfig,
+    TerminalRenderer,
+)
+
+__all__ = [
+    "BaseRenderer",
+    "RenderConfig",
+    "GraphData",
+    "TerminalRenderer",
+    "CardsRenderer",
+    "MermaidRenderer",
+    "GraphvizRenderer",
+]

metaxy/graph/diff/__init__.py

@@ -0,0 +1,21 @@
+"""Graph diff infrastructure - models and traversal for graph comparison."""
+
+from metaxy.graph.diff.models import (
+    EdgeData,
+    FieldNode,
+    GraphData,
+    GraphNode,
+    NodeStatus,
+)
+from metaxy.graph.diff.traversal import GraphWalker
+
+__all__ = [
+    # Core models
+    "EdgeData",
+    "FieldNode",
+    "GraphData",
+    "GraphNode",
+    "NodeStatus",
+    # Traversal
+    "GraphWalker",
+]