metaxy 0.0.1.dev3__py3-none-any.whl

metaxy/metadata_store/system/__init__.py

@@ -0,0 +1,57 @@
+"""System table components for metadata store.
+
+This package provides system table functionality for Metaxy:
+- events: Migration event types with builder pattern
+- keys: System table keys and constants
+- models: Pydantic models and schemas for system tables
+- storage: Storage layer for system tables
+"""
+
+from metaxy.metadata_store.system.events import (
+    COL_EVENT_TYPE,
+    COL_EXECUTION_ID,
+    COL_FEATURE_KEY,
+    COL_PAYLOAD,
+    COL_PROJECT,
+    COL_TIMESTAMP,
+    Event,
+    EventType,
+    MigrationStatus,
+    PayloadType,
+)
+from metaxy.metadata_store.system.keys import (
+    EVENTS_KEY,
+    FEATURE_VERSIONS_KEY,
+    METAXY_SYSTEM_KEY_PREFIX,
+)
+from metaxy.metadata_store.system.models import (
+    FEATURE_VERSIONS_SCHEMA,
+    FeatureVersionsModel,
+)
+from metaxy.metadata_store.system.storage import (
+    SystemTableStorage,
+)
+
+__all__ = [
+    # Events
+    "Event",
+    "EventType",
+    "MigrationStatus",
+    "PayloadType",
+    # Column names
+    "COL_PROJECT",
+    "COL_EXECUTION_ID",
+    "COL_EVENT_TYPE",
+    "COL_TIMESTAMP",
+    "COL_FEATURE_KEY",
+    "COL_PAYLOAD",
+    # Keys
+    "METAXY_SYSTEM_KEY_PREFIX",
+    "FEATURE_VERSIONS_KEY",
+    "EVENTS_KEY",
+    # Models
+    "FEATURE_VERSIONS_SCHEMA",
+    "FeatureVersionsModel",
+    # Storage
+    "SystemTableStorage",
+]

metaxy/metadata_store/system/events.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Annotated, Literal
+
+import polars as pl
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class EventType(str, Enum):
+    """Metaxy event types."""
+
+    MIGRATION_STARTED = "migration_started"
+    MIGRATION_COMPLETED = "migration_completed"
+    MIGRATION_FAILED = "migration_failed"
+    FEATURE_MIGRATION_STARTED = "feature_migration_started"
+    FEATURE_MIGRATION_COMPLETED = "feature_migration_completed"
+    FEATURE_MIGRATION_FAILED = "feature_migration_failed"
+
+
+class PayloadType(str, Enum):
+    """Payload types for event payloads."""
+
+    EMPTY = "empty"
+    ERROR = "error"
+    ROWS_AFFECTED = "rows_affected"
+
+
+class MigrationStatus(str, Enum):
+    """Migration execution status."""
+
+    NOT_STARTED = "not_started"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+# Column name constants (to avoid drift between Event model and storage)
+COL_PROJECT = "project"
+COL_EXECUTION_ID = "execution_id"
+COL_EVENT_TYPE = "event_type"
+COL_TIMESTAMP = "timestamp"
+COL_FEATURE_KEY = "feature_key"
+COL_PAYLOAD = "payload"
+
+# Events schema (for Polars storage)
+EVENTS_SCHEMA = {
+    COL_PROJECT: pl.String,
+    COL_EXECUTION_ID: pl.String,
+    COL_EVENT_TYPE: pl.Enum(EventType),
+    COL_TIMESTAMP: pl.Datetime("us"),
+    COL_FEATURE_KEY: pl.String,
+    COL_PAYLOAD: pl.String,  # JSON string with arbitrary event data
+}
+
+
+class EmptyPayload(BaseModel):
+    """Empty payload for events with no additional data."""
+
+    model_config = ConfigDict(frozen=True)
+    type: Literal[PayloadType.EMPTY] = PayloadType.EMPTY
+
+
+class ErrorPayload(BaseModel):
+    """Payload for events with error information."""
+
+    model_config = ConfigDict(frozen=True)
+    type: Literal[PayloadType.ERROR] = PayloadType.ERROR
+    error_message: str
+    rows_affected: int | None = None  # Optional: rows processed before failure
+
+
+class RowsAffectedPayload(BaseModel):
+    """Payload for events tracking rows affected."""
+
+    model_config = ConfigDict(frozen=True)
+    type: Literal[PayloadType.ROWS_AFFECTED] = PayloadType.ROWS_AFFECTED
+    rows_affected: int
+
+
+# Discriminated union for payloads
+Payload = EmptyPayload | ErrorPayload | RowsAffectedPayload
+
+
+class Event(BaseModel):
+    """Migration event with typed payload.
+
+    All event types use this single class and are distinguished by event_type and payload.type fields.
+    """
+
+    model_config = ConfigDict(frozen=True)
+
+    event_type: EventType
+    project: str
+    execution_id: str  # Generic ID for the execution (migration, job, etc.)
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    feature_key: str | None = (
+        None  # Feature key for feature-level events, empty for execution-level events
+    )
+    payload: Annotated[
+        Payload, Field(default_factory=EmptyPayload, discriminator="type")
+    ]
+
+    def to_polars(self) -> pl.DataFrame:
+        """Convert this model instance to a single-row Polars DataFrame.
+
+        Returns:
+            Polars DataFrame with one row matching EVENTS_SCHEMA
+        """
+        data = {
+            COL_PROJECT: self.project,
+            COL_EXECUTION_ID: self.execution_id,
+            COL_EVENT_TYPE: self.event_type,
+            COL_TIMESTAMP: self.timestamp,
+            COL_FEATURE_KEY: self.feature_key,
+            COL_PAYLOAD: self.payload.model_dump_json(),
+        }
+        return pl.DataFrame([data], schema=EVENTS_SCHEMA)
+
+    @classmethod
+    def migration_started(cls, project: str, migration_id: str) -> Event:
+        """Create a migration started event.
+
+        Args:
+            project: Project name
+            migration_id: Migration ID (maps to execution_id internally)
+
+        Returns:
+            Event with started type and empty payload
+        """
+        return cls(
+            project=project,
+            execution_id=migration_id,
+            event_type=EventType.MIGRATION_STARTED,
+            payload=EmptyPayload(),
+        )
+
+    @classmethod
+    def migration_completed(cls, project: str, migration_id: str) -> Event:
+        """Create a migration completed event.
+
+        Args:
+            project: Project name
+            migration_id: Migration ID (maps to execution_id internally)
+
+        Returns:
+            Event with migration_completed type and empty payload
+        """
+        return cls(
+            project=project,
+            execution_id=migration_id,
+            event_type=EventType.MIGRATION_COMPLETED,
+            payload=EmptyPayload(),
+        )
+
+    @classmethod
+    def migration_failed(
+        cls,
+        project: str,
+        migration_id: str,
+        error_message: str,
+        rows_affected: int | None = None,
+    ) -> Event:
+        """Create a migration failed event.
+
+        Args:
+            project: Project name
+            migration_id: Migration ID (maps to execution_id internally)
+            error_message: Error message describing the failure
+            rows_affected: Optional number of rows processed before failure
+
+        Returns:
+            Event with migration_failed type and error payload
+        """
+        return cls(
+            project=project,
+            execution_id=migration_id,
+            event_type=EventType.MIGRATION_FAILED,
+            payload=ErrorPayload(
+                error_message=error_message, rows_affected=rows_affected
+            ),
+        )
+
+    @classmethod
+    def feature_migration_started(
+        cls, project: str, migration_id: str, feature_key: str
+    ) -> Event:
+        """Create a feature started event.
+
+        Args:
+            project: Project name
+            migration_id: Migration ID (maps to execution_id internally)
+            feature_key: Feature key being processed
+
+        Returns:
+            Event with feature_started type and empty payload
+        """
+        return cls(
+            project=project,
+            execution_id=migration_id,
+            event_type=EventType.FEATURE_MIGRATION_STARTED,
+            payload=EmptyPayload(),
+            feature_key=feature_key,
+        )
+
+    @classmethod
+    def feature_migration_completed(
+        cls, project: str, migration_id: str, feature_key: str, rows_affected: int
+    ) -> Event:
+        """Create a feature completed event (successful).
+
+        Args:
+            project: Project name
+            migration_id: Migration ID (maps to execution_id internally)
+            feature_key: Feature key that was processed
+            rows_affected: Number of rows affected
+
+        Returns:
+            Event with feature_completed type and rows_affected payload
+        """
+        return cls(
+            project=project,
+            execution_id=migration_id,
+            event_type=EventType.FEATURE_MIGRATION_COMPLETED,
+            feature_key=feature_key,
+            payload=RowsAffectedPayload(rows_affected=rows_affected),
+        )
+
+    @classmethod
+    def feature_migration_failed(
+        cls,
+        project: str,
+        migration_id: str,
+        feature_key: str,
+        error_message: str,
+        rows_affected: int | None = None,
+    ) -> Event:
+        """Create a feature failed event.
+
+        Args:
+            project: Project name
+            migration_id: Migration ID (maps to execution_id internally)
+            feature_key: Feature key that failed
+            error_message: Error message describing the failure
+            rows_affected: Optional number of rows processed before failure
+
+        Returns:
+            Event with feature_failed type and error payload
+        """
+        return cls(
+            project=project,
+            execution_id=migration_id,
+            event_type=EventType.FEATURE_MIGRATION_FAILED,
+            feature_key=feature_key,
+            payload=ErrorPayload(
+                error_message=error_message, rows_affected=rows_affected
+            ),
+        )
+
+    # Shorter aliases for convenience
+    feature_started = feature_migration_started
+    feature_completed = feature_migration_completed
+    feature_failed = feature_migration_failed

metaxy/metadata_store/system/keys.py

@@ -0,0 +1,9 @@
+"""System table keys and constants."""
+
+from metaxy.models.types import FeatureKey
+
+METAXY_SYSTEM_KEY_PREFIX = "metaxy-system"
+
+# System table keys
+FEATURE_VERSIONS_KEY = FeatureKey([METAXY_SYSTEM_KEY_PREFIX, "feature_versions"])
+EVENTS_KEY = FeatureKey([METAXY_SYSTEM_KEY_PREFIX, "events"])

metaxy/metadata_store/system/models.py

@@ -0,0 +1,129 @@
+"""Pydantic models for system tables."""
+
+from __future__ import annotations
+
+from datetime import datetime
+
+import polars as pl
+from pydantic import BaseModel, Field, field_validator
+
+from metaxy.metadata_store.system import FEATURE_VERSIONS_KEY
+from metaxy.metadata_store.system.events import EVENTS_SCHEMA
+from metaxy.metadata_store.system.keys import EVENTS_KEY
+from metaxy.models.constants import (
+    METAXY_FEATURE_SPEC_VERSION,
+    METAXY_FEATURE_VERSION,
+    METAXY_FULL_DEFINITION_VERSION,
+    METAXY_SNAPSHOT_VERSION,
+)
+
+# Common Polars schemas for system tables
+# Note: tags field schema is omitted - Polars will infer the Struct schema from data
+FEATURE_VERSIONS_SCHEMA = {
+    "project": pl.String,
+    "feature_key": pl.String,
+    METAXY_FEATURE_VERSION: pl.String,
+    METAXY_FEATURE_SPEC_VERSION: pl.String,  # Hash of complete FeatureSpec (all properties)
+    METAXY_FULL_DEFINITION_VERSION: pl.String,  # Hash of feature_spec_version + project (for migration detection)   # TODO: this is probably not needed, we can just use a combination of project and metaxy_feature_version instead
+    "recorded_at": pl.Datetime("us"),
+    "feature_spec": pl.String,  # Full serialized FeatureSpec
+    "feature_schema": pl.String,  # Full Pydantic model schema as JSON
+    "feature_class_path": pl.String,
+    METAXY_SNAPSHOT_VERSION: pl.String,
+    "tags": pl.String,
+}
+
+
+METAXY_TAG = "metaxy"
+METAXY_VERSION_KEY = "version"
+
+
+class FeatureVersionsModel(BaseModel):
+    """Pydantic model for feature_versions system table.
+
+    This table records when feature specifications are pushed to production,
+    tracking the evolution of feature definitions over time.
+    """
+
+    project: str
+    feature_key: str
+    metaxy_feature_version: str = Field(
+        ...,
+        description="Hash of versioned feature topology (combined versions of fields on this feature)",
+    )
+    metaxy_feature_spec_version: str = Field(
+        ..., description="Hash of complete FeatureSpec (all properties)"
+    )
+    metaxy_full_definition_version: str = Field(
+        ..., description="Hash of feature_spec_version + project"
+    )
+    recorded_at: datetime = Field(
+        ..., description="Timestamp when feature version was recorded"
+    )
+    feature_spec: str = Field(
+        ..., description="Full serialized FeatureSpec as JSON string"
+    )
+    feature_schema: str = Field(
+        ..., description="Full Pydantic model schema as JSON string"
+    )
+    feature_class_path: str = Field(
+        ..., description="Python import path to Feature class"
+    )
+    metaxy_snapshot_version: str = Field(
+        ..., description="Deterministic hash of entire Metaxy project"
+    )
+    tags: dict[str, str] | str = Field(
+        default="{}",
+        description="Snapshot tags as JSON string (key-value pairs). The metaxy tag is reserved for internal use.",
+        validate_default=True,
+    )
+
+    @field_validator("tags", mode="before")
+    @classmethod
+    def serialize_tags(cls, v: dict[str, str] | str | None) -> str:
+        """Convert tags dict to JSON string if needed."""
+        import json
+
+        # Parse to dict if string
+        tags_dict: dict[str, str]
+        if isinstance(v, str):
+            try:
+                tags_dict = json.loads(v)
+            except json.JSONDecodeError:
+                tags_dict = {}
+        else:
+            # Handle None or dict
+            tags_dict = v or {}
+
+        # Ensure metaxy.version is set
+        from metaxy._version import __version__
+
+        metaxy_tag_value = tags_dict.get(METAXY_TAG, "{}")
+        metaxy_tag_dict = (
+            json.loads(metaxy_tag_value)
+            if isinstance(metaxy_tag_value, str)
+            else metaxy_tag_value
+        )
+        if not isinstance(metaxy_tag_dict, dict):
+            metaxy_tag_dict = {}
+        metaxy_tag_dict[METAXY_VERSION_KEY] = metaxy_tag_dict.get(
+            METAXY_VERSION_KEY, __version__
+        )
+        tags_dict[METAXY_TAG] = json.dumps(metaxy_tag_dict)
+
+        return json.dumps(tags_dict)
+
+    def to_polars(self) -> pl.DataFrame:
+        """Convert this model instance to a single-row Polars DataFrame.
+
+        Returns:
+            Polars DataFrame with one row matching FEATURE_VERSIONS_SCHEMA
+        """
+        # tags is already a JSON string, no need to serialize
+        return pl.DataFrame([self.model_dump()], schema=FEATURE_VERSIONS_SCHEMA)
+
+
+POLARS_SCHEMAS = {
+    FEATURE_VERSIONS_KEY: FEATURE_VERSIONS_SCHEMA,
+    EVENTS_KEY: EVENTS_SCHEMA,
+}