metaxy 0.0.0__py3-none-any.whl

metaxy/entrypoints.py

@@ -0,0 +1,309 @@
+"""Entrypoint discovery and loading for Metaxy features.
+
+This module provides functionality to automatically discover and load Feature
+classes from modules, supporting both:
+- Config-based entrypoints (list of module paths)
+- Environment-based entrypoints (environment variables starting with METAXY_ENTRYPOINT)
+- Package-based entrypoints (via importlib.metadata)
+
+Features are automatically registered to the active FeatureGraph when their
+containing modules are imported (via the Feature metaclass).
+"""
+
+import importlib
+import logging
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from metaxy.models.feature import FeatureGraph
+
+# Conditional import for Python 3.10+ compatibility
+from importlib.metadata import entry_points  # type: ignore[import-not-found]
+
+logger = logging.getLogger(__name__)
+
+# Default entry point group name for package-based discovery
+DEFAULT_ENTRY_POINT_GROUP = "metaxy.features"
+
+
+class EntrypointLoadError(Exception):
+    """Raised when an entrypoint fails to load."""
+
+    pass
+
+
+def load_module_entrypoint(
+    module_path: str,
+    *,
+    graph: "FeatureGraph | None" = None,
+) -> None:
+    """Load a single module entrypoint.
+
+    Imports the specified module, which should contain Feature class definitions.
+    Features will be automatically registered to the active graph via the
+    Feature metaclass.
+
+    Args:
+        module_path: Fully qualified module path (e.g., "myapp.features.video")
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If module import fails
+
+    Example:
+        >>> from metaxy.entrypoints import load_module_entrypoint
+        >>> load_module_entrypoint("myapp.features.core")
+        >>> # Features from myapp.features.core are now registered
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = graph or FeatureGraph.get_active()
+
+    try:
+        # Set graph as active during import so Features register to it
+        with target_graph.use():
+            logger.debug(f"Loading entrypoint module: {module_path}")
+            importlib.import_module(module_path)
+            logger.info(f"Successfully loaded entrypoint: {module_path}")
+    except ImportError as e:
+        raise EntrypointLoadError(
+            f"Failed to import entrypoint module '{module_path}': {e}"
+        ) from e
+    except Exception as e:
+        raise EntrypointLoadError(
+            f"Error loading entrypoint module '{module_path}': {e}"
+        ) from e
+
+
+def load_entrypoints(
+    entrypoints: list[str],
+    *,
+    graph: "FeatureGraph | None" = None,
+) -> None:
+    """Load multiple module entrypoints from a list.
+
+    Args:
+        entrypoints: List of module paths to import
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If any module import fails
+
+    Example:
+        >>> from metaxy.entrypoints import load_config_entrypoints
+        >>> load_config_entrypoints([
+        ...     "myapp.features.video",
+        ...     "myapp.features.audio",
+        ...     "myapp.features.text"
+        ... ])
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = graph or FeatureGraph.get_active()
+
+    logger.info(f"Loading {len(entrypoints)} entrypoints")
+
+    for module_path in entrypoints:
+        load_module_entrypoint(module_path, graph=target_graph)
+
+
+def load_package_entrypoints(
+    group: str = DEFAULT_ENTRY_POINT_GROUP,
+    *,
+    graph: "FeatureGraph | None" = None,
+) -> None:
+    """Load entrypoints from installed packages using importlib.metadata.
+
+    Discovers and loads all entry points registered in the specified group.
+    This is the package-based entrypoint mechanism using standard Python
+    packaging infrastructure.
+
+    Packages declare entrypoints in their pyproject.toml:
+        [project.entry-points."metaxy.features"]
+        myfeature = "mypackage.features.module"
+
+    Args:
+        group: Entry point group name (default: "metaxy.features")
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If any entrypoint fails to load
+
+    Example:
+        >>> from metaxy.entrypoints import load_package_entrypoints
+        >>> # Discover and load all installed plugins
+        >>> load_package_entrypoints()
+    """
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = graph or FeatureGraph.get_active()
+
+    logger.info(f"Discovering package entrypoints in group: {group}")
+
+    # Discover entry points
+    # Note: Python 3.10+ returns SelectableGroups, 3.9 returns dict
+    discovered = entry_points()
+
+    # Handle different return types across Python versions
+    if hasattr(discovered, "select"):
+        # Python 3.10+: SelectableGroups with select() method
+        eps = discovered.select(group=group)
+    else:
+        # Python 3.9: dict-like interface
+        eps = discovered.get(group, [])
+
+    eps_list = list(eps)
+
+    if not eps_list:
+        logger.debug(f"No package entrypoints found in group: {group}")
+        return
+
+    logger.info(f"Found {len(eps_list)} package entrypoints in group: {group}")
+
+    for ep in eps_list:
+        try:
+            logger.debug(f"Loading package entrypoint: {ep.name} = {ep.value}")
+            # Load the entry point (imports the module)
+            with target_graph.use():
+                ep.load()
+            logger.info(f"Successfully loaded package entrypoint: {ep.name}")
+        except Exception as e:
+            raise EntrypointLoadError(
+                f"Failed to load package entrypoint '{ep.name}' ({ep.value}): {e}"
+            ) from e
+
+
+def load_env_entrypoints() -> None:
+    """Load entrypoints from environment variables.
+
+    Discovers and loads all entry points from environment variables matching
+    the pattern METAXY_ENTRYPOINT*. Each variable should contain a
+    comma-separated list of module paths.
+
+    Environment variables:
+        METAXY_ENTRYPOINT="myapp.features.core,myapp.features.extra"
+        METAXY_ENTRYPOINT_PLUGINS="plugin1.features,plugin2.features"
+
+    Args:
+        graph: Target graph. If None, uses FeatureGraph.get_active()
+
+    Raises:
+        EntrypointLoadError: If any entrypoint fails to load
+
+    Example:
+        >>> import os
+        >>> os.environ["METAXY_ENTRYPOINT"] = "myapp.features.core"
+        >>> from metaxy.entrypoints import load_env_entrypoints
+        >>> load_env_entrypoints()
+    """
+
+    logger.info("Discovering environment-based entrypoints (METAXY_ENTRYPOINT*)")
+
+    # Find all environment variables matching METAXY_ENTRYPOINT*
+    env_vars = {
+        key: value
+        for key, value in os.environ.items()
+        if key.startswith("METAXY_ENTRYPOINT")
+    }
+
+    if not env_vars:
+        logger.debug("No environment entrypoints found (METAXY_ENTRYPOINT* not set)")
+        return
+
+    logger.info(f"Found {len(env_vars)} METAXY_ENTRYPOINT* environment variable(s)")
+
+    # Collect all module paths from all matching env vars
+    all_module_paths = []
+    for env_var, value in sorted(env_vars.items()):
+        logger.debug(f"Processing {env_var}={value}")
+        # Split by comma and strip whitespace
+        module_paths = [path.strip() for path in value.split(",") if path.strip()]
+        all_module_paths.extend(module_paths)
+
+    if not all_module_paths:
+        logger.debug("No module paths found in METAXY_ENTRYPOINT* variables")
+        return
+
+    logger.info(
+        f"Loading {len(all_module_paths)} module(s) from environment entrypoints"
+    )
+
+    # Load each module path
+    for module_path in all_module_paths:
+        load_module_entrypoint(module_path)
+
+
+def load_features(
+    entrypoints: list[str] | None = None,
+    package_entrypoint_group: str = DEFAULT_ENTRY_POINT_GROUP,
+    *,
+    load_config: bool = True,
+    load_packages: bool = True,
+    load_env: bool = True,
+) -> "FeatureGraph":
+    """Discover and load all entrypoints from config, packages, and environment.
+
+    This is the main entry point for loading features. It combines config-based,
+    package-based, and environment-based entrypoint discovery.
+
+    Args:
+        entrypoints: List of module paths (optional)
+        package_entrypoint_group: Entry point group for package discovery
+        load_config: Whether to load config-based entrypoints (default: True)
+        load_packages: Whether to load package-based entrypoints (default: True)
+        load_env: Whether to load environment-based entrypoints (default: True)
+
+    Returns:
+        The graph that was populated (useful for testing/inspection)
+
+    Raises:
+        EntrypointLoadError: If any entrypoint fails to load
+
+    Example:
+        >>> from metaxy.entrypoints import load_features
+        >>>
+        >>> # Load from all sources
+        >>> graph = load_features(
+        ...     entrypoints=["myapp.features.core"],
+        ...     load_packages=True,
+        ...     load_env=True
+        ... )
+        >>>
+        >>> # Load only from config
+        >>> graph = load_features(
+        ...     entrypoints=["myapp.features.core"],
+        ...     load_packages=False,
+        ...     load_env=False
+        ... )
+    """
+    from metaxy.config import MetaxyConfig
+    from metaxy.models.feature import FeatureGraph
+
+    target_graph = FeatureGraph.get_active()
+
+    logger.info("Starting entrypoint discovery and loading")
+
+    # Load explicit entrypoints
+    if entrypoints:
+        load_entrypoints(entrypoints)
+
+    # Load config-based entrypoints
+    if load_config:
+        config = MetaxyConfig.load(search_parents=True)
+        load_entrypoints(config.entrypoints)
+
+    # Load package-based entrypoints
+    if load_packages:
+        load_package_entrypoints(package_entrypoint_group)
+
+    # Load environment-based entrypoints
+    if load_env:
+        load_env_entrypoints()
+
+    num_features = len(target_graph.features_by_key)
+    logger.info(
+        f"Entrypoint loading complete. Registry now contains {num_features} features."
+    )
+
+    return target_graph

metaxy/ext/__init__.py

@@ -0,0 +1 @@
+"""Integrations with third-party software."""

metaxy/ext/alembic.py

@@ -0,0 +1,326 @@
+"""Alembic integration helpers for metaxy.
+
+This module provides utilities for including metaxy system tables
+in Alembic migrations when using SQLModel integration.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sqlalchemy import MetaData
+
+
+def get_metaxy_metadata() -> MetaData:
+    """Get SQLAlchemy metadata containing metaxy system tables.
+
+    This function returns the metadata object that should be included
+    in your Alembic configuration to manage metaxy system tables.
+
+    Returns:
+        SQLAlchemy MetaData containing metaxy system table definitions
+
+    Raises:
+        ImportError: If SQLModel is not installed
+
+    Example:
+        >>> # In your alembic/env.py:
+        >>> from metaxy.ext.alembic import get_metaxy_metadata
+        >>> from sqlalchemy import MetaData
+        >>>
+        >>> # Combine your app metadata with metaxy metadata
+        >>> target_metadata = MetaData()
+        >>> target_metadata.reflect(get_metaxy_metadata())
+        >>> target_metadata.reflect(my_app.metadata)
+    """
+    from metaxy.ext.sqlmodel_system_tables import get_system_metadata
+
+    return get_system_metadata()
+
+
+def include_metaxy_tables(target_metadata: MetaData) -> MetaData:
+    """Include metaxy system tables in existing metadata.
+
+    This is a convenience function that adds metaxy system tables
+    to your existing SQLAlchemy metadata object.
+
+    Args:
+        target_metadata: Your application's metadata object
+
+    Returns:
+        The same metadata object with metaxy tables added
+
+    Example:
+        >>> # In your alembic/env.py:
+        >>> from metaxy.ext.alembic import include_metaxy_tables
+        >>> from myapp.models import metadata
+        >>>
+        >>> # Add metaxy tables to your metadata
+        >>> target_metadata = include_metaxy_tables(metadata)
+    """
+    metaxy_metadata = get_metaxy_metadata()
+
+    # Copy tables from metaxy metadata to target metadata
+    for table_name, table in metaxy_metadata.tables.items():
+        if table_name not in target_metadata.tables:
+            # Create a new table with the same definition in target metadata
+            table.to_metadata(target_metadata)
+
+    return target_metadata
+
+
+def check_sqlmodel_enabled() -> bool:
+    """Check if SQLModel integration is enabled in the configuration.
+
+    Returns:
+        True if SQLModel integration is enabled, False otherwise
+
+    Example:
+        >>> from metaxy.ext.alembic import check_sqlmodel_enabled
+        >>> if check_sqlmodel_enabled():
+        ...     # SQLModel is enabled, include metaxy tables
+        ...     include_metaxy_tables(metadata)
+    """
+    try:
+        from metaxy.config import MetaxyConfig
+
+        config = MetaxyConfig.get()
+        # Check if SQLModel configuration exists and is enabled
+        # The infer_db_table_names setting indicates SQLModel is being used
+        return hasattr(config.ext, "sqlmodel")
+    except Exception:
+        return False
+
+
+def generate_alembic_env_template() -> str:
+    """Generate a template for Alembic env.py with metaxy integration.
+
+    Returns:
+        String containing example Alembic env.py configuration
+
+    Example:
+        >>> from metaxy.ext.alembic import generate_alembic_env_template
+        >>> template = generate_alembic_env_template()
+        >>> print(template)  # Shows example configuration
+    """
+    return '''"""Alembic env.py template with metaxy integration.
+
+This template shows how to include metaxy system tables in your
+Alembic migrations when using SQLModel integration.
+"""
+
+from logging.config import fileConfig
+
+from sqlalchemy import engine_from_config
+from sqlalchemy import pool
+from alembic import context
+
+# Import your models here
+# from myapp import models
+
+# Metaxy integration
+from metaxy.ext.alembic import get_metaxy_metadata, include_metaxy_tables
+from metaxy.config import MetaxyConfig
+
+# this is the Alembic Config object
+config = context.config
+
+# Interpret the config file for Python logging
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# Load metaxy configuration
+metaxy_config = MetaxyConfig.load()
+
+# Add your model's MetaData object here for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+
+# Option 1: Include metaxy tables in your existing metadata
+# from myapp.models import metadata
+# target_metadata = include_metaxy_tables(metadata)
+
+# Option 2: Combine multiple metadata objects
+from sqlalchemy import MetaData
+target_metadata = MetaData()
+
+# Add metaxy system tables (if using SQLModel integration)
+if metaxy_config.ext.sqlmodel:
+    metaxy_metadata = get_metaxy_metadata()
+    for table_name, table in metaxy_metadata.tables.items():
+        table.to_metadata(target_metadata)
+
+# Add your application tables
+# for table_name, table in myapp_metadata.tables.items():
+#     table.to_metadata(target_metadata)
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well. By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+    """
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    with connectable.connect() as connection:
+        context.configure(connection=connection, target_metadata=target_metadata)
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()
+'''
+
+
+def create_initial_migration_script() -> str:
+    """Generate an Alembic migration script for creating metaxy system tables.
+
+    Returns:
+        String containing example migration script
+
+    Example:
+        >>> from metaxy.ext.alembic import create_initial_migration_script
+        >>> script = create_initial_migration_script()
+        >>> # Save to alembic/versions/001_create_metaxy_tables.py
+    """
+    return '''"""Create metaxy system tables.
+
+Revision ID: create_metaxy_system_tables
+Revises:
+Create Date: 2024-01-01 00:00:00.000000
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = 'create_metaxy_system_tables'
+down_revision: Union[str, None] = None
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Create metaxy system tables."""
+
+    # Create feature_versions table
+    op.create_table(
+        'metaxy-system__feature_versions',
+        sa.Column('feature_key', sa.String(), nullable=False),
+        sa.Column('snapshot_version', sa.String(), nullable=False),
+        sa.Column('feature_version', sa.String(), nullable=False),
+        sa.Column('recorded_at', sa.DateTime(), nullable=False),
+        sa.Column('feature_spec', sa.String(), nullable=False),
+        sa.Column('feature_class_path', sa.String(), nullable=False),
+        sa.PrimaryKeyConstraint('feature_key', 'snapshot_version')
+    )
+    op.create_index('idx_feature_versions_recorded', 'metaxy-system__feature_versions', ['recorded_at'])
+    op.create_index('idx_feature_versions_lookup', 'metaxy-system__feature_versions', ['feature_key', 'feature_version'])
+
+    # Create migration_events table
+    op.create_table(
+        'metaxy-system__migration_events',
+        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column('migration_id', sa.String(), nullable=False),
+        sa.Column('event_type', sa.String(), nullable=False),
+        sa.Column('timestamp', sa.DateTime(), nullable=False),
+        sa.Column('feature_key', sa.String(), nullable=False),
+        sa.Column('rows_affected', sa.Integer(), nullable=False),
+        sa.Column('error_message', sa.String(), nullable=False),
+        sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index('idx_migration_events_lookup', 'metaxy-system__migration_events', ['migration_id', 'event_type'])
+    op.create_index('idx_migration_events_feature', 'metaxy-system__migration_events', ['migration_id', 'feature_key'])
+    op.create_index(op.f('ix_metaxy-system__migration_events_migration_id'), 'metaxy-system__migration_events', ['migration_id'])
+    op.create_index(op.f('ix_metaxy-system__migration_events_timestamp'), 'metaxy-system__migration_events', ['timestamp'])
+
+
+def downgrade() -> None:
+    """Drop metaxy system tables."""
+    op.drop_index(op.f('ix_metaxy-system__migration_events_timestamp'), table_name='metaxy-system__migration_events')
+    op.drop_index(op.f('ix_metaxy-system__migration_events_migration_id'), table_name='metaxy-system__migration_events')
+    op.drop_index('idx_migration_events_feature', table_name='metaxy-system__migration_events')
+    op.drop_index('idx_migration_events_lookup', table_name='metaxy-system__migration_events')
+    op.drop_table('metaxy-system__migration_events')
+
+    op.drop_index('idx_feature_versions_lookup', table_name='metaxy-system__feature_versions')
+    op.drop_index('idx_feature_versions_recorded', table_name='metaxy-system__feature_versions')
+    op.drop_table('metaxy-system__feature_versions')
+'''
+
+
+def get_table_creation_sql(dialect: str = "postgresql") -> dict[str, str]:
+    """Get SQL statements for creating metaxy system tables.
+
+    Args:
+        dialect: SQL dialect ('postgresql', 'mysql', 'sqlite', etc.)
+
+    Returns:
+        Dictionary mapping table name to CREATE TABLE SQL statement
+
+    Example:
+        >>> from metaxy.ext.alembic import get_table_creation_sql
+        >>> sql_statements = get_table_creation_sql("postgresql")
+        >>> print(sql_statements["feature_versions"])
+    """
+    from sqlalchemy import create_engine
+    from sqlalchemy.schema import CreateTable
+
+    # Get metaxy metadata
+    metadata = get_metaxy_metadata()
+
+    # Create a mock engine for the specified dialect
+    if dialect == "postgresql":
+        engine_url = "postgresql:///"
+    elif dialect == "mysql":
+        engine_url = "mysql:///"
+    elif dialect == "sqlite":
+        engine_url = "sqlite:///"
+    else:
+        engine_url = f"{dialect}:///"
+
+    engine = create_engine(engine_url)
+
+    # Generate CREATE TABLE statements
+    sql_statements = {}
+    for table_name, table in metadata.tables.items():
+        create_statement = CreateTable(table).compile(engine)
+        sql_statements[table_name] = str(create_statement)
+
+    return sql_statements