remdb 0.3.0__py3-none-any.whl

rem/services/postgres/migration_service.py

@@ -0,0 +1,427 @@
+"""
+Migration service for managing Alembic-based database migrations.
+
+This service provides:
+1. SQL diff generation between target database and Pydantic models
+2. Migration planning (dry-run)
+3. Migration application
+4. SQL file execution
+5. Safety validation for migration operations
+"""
+
+import io
+import re
+import subprocess
+import sys
+import tempfile
+from contextlib import redirect_stdout
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+from alembic import command
+from alembic.config import Config
+from alembic.script import ScriptDirectory
+from loguru import logger
+
+from ...settings import settings
+
+
+@dataclass
+class MigrationSafetyResult:
+    """Result of migration safety validation."""
+
+    is_safe: bool
+    errors: list[str]
+    warnings: list[str]
+    sql: str
+
+
+class MigrationService:
+    """
+    Service for managing database migrations using Alembic.
+
+    Integrates Alembic with REM's Pydantic model-first approach.
+    """
+
+    def __init__(self, alembic_cfg_path: Optional[Path] = None):
+        """
+        Initialize migration service.
+
+        Args:
+            alembic_cfg_path: Path to alembic.ini (defaults to package alembic.ini)
+        """
+        if alembic_cfg_path is None:
+            # Find alembic.ini in package root
+            import rem
+
+            package_root = Path(rem.__file__).parent.parent.parent
+            alembic_cfg_path = package_root / "alembic.ini"
+
+        if not alembic_cfg_path.exists():
+            raise FileNotFoundError(f"Alembic config not found: {alembic_cfg_path}")
+
+        self.alembic_cfg = Config(str(alembic_cfg_path))
+        self.alembic_cfg.set_main_option(
+            "script_location", str(alembic_cfg_path.parent / "alembic")
+        )
+
+    def get_connection_string(self, target_db: Optional[str] = None) -> str:
+        """
+        Build PostgreSQL connection string.
+
+        Args:
+            target_db: Override database name (useful for comparing against different DB)
+
+        Returns:
+            PostgreSQL connection string
+        """
+        pg = settings.postgres
+
+        url = f"postgresql://{pg.user}"
+        if pg.password:
+            url += f":{pg.password}"
+
+        db_name = target_db or pg.database
+        url += f"@{pg.host}:{pg.port}/{db_name}"
+
+        return url
+
+    def generate_migration_sql(
+        self,
+        output_file: Optional[Path] = None,
+        target_db: Optional[str] = None,
+        message: str = "Auto-generated migration",
+    ) -> str:
+        """
+        Generate SQL diff between current models and target database using Alembic autogenerate.
+
+        This uses Alembic's autogenerate feature to compare the Pydantic models
+        (via SQLAlchemy metadata) with the target database schema.
+
+        Args:
+            output_file: Path to write SQL file (if None, returns SQL string)
+            target_db: Target database name (overrides settings)
+            message: Migration message/description
+
+        Returns:
+            Generated SQL as string
+        """
+        # Override database URL if target_db provided
+        url = self.get_connection_string(target_db)
+        self.alembic_cfg.set_main_option("sqlalchemy.url", url)
+
+        # Generate migration SQL using autogenerate
+        # We'll use upgrade --sql to generate SQL without applying
+        stdout_buffer = io.StringIO()
+
+        try:
+            # First, create an autogenerate revision to detect changes
+            with redirect_stdout(stdout_buffer):
+                command.revision(
+                    self.alembic_cfg,
+                    message=message,
+                    autogenerate=True,
+                )
+
+            # Get the revision that was just created
+            script_dir = ScriptDirectory.from_config(self.alembic_cfg)
+            revision = script_dir.get_current_head()
+
+            if revision is None:
+                logger.warning("No changes detected - models match database")
+                return "-- No changes detected"
+
+            # Now generate SQL from that revision
+            sql_buffer = io.StringIO()
+            with redirect_stdout(sql_buffer):
+                command.upgrade(
+                    self.alembic_cfg,
+                    f"{revision}:head",
+                    sql=True,
+                )
+
+            sql = sql_buffer.getvalue()
+
+            # If no SQL was generated, check the revision file
+            if not sql.strip() or "-- Running upgrade" not in sql:
+                # Read the revision file directly
+                version_path = script_dir.get_revision(revision).path
+                sql = f"-- Generated from Alembic revision: {revision}\n"
+                sql += f"-- Migration file: {version_path}\n"
+                sql += "-- Run: alembic upgrade head\n"
+                sql += "-- Or use this migration file to review/edit before applying\n\n"
+
+                # Try to extract upgrade operations from the revision file
+                if version_path:
+                    with open(version_path, "r") as f:
+                        content = f.read()
+                        # Extract the upgrade function
+                        import re
+
+                        upgrade_match = re.search(
+                            r"def upgrade\(\).*?:\s*(.*?)(?=def downgrade|$)",
+                            content,
+                            re.DOTALL,
+                        )
+                        if upgrade_match:
+                            upgrade_code = upgrade_match.group(1).strip()
+                            if upgrade_code and upgrade_code != "pass":
+                                sql += f"-- Upgrade operations:\n{upgrade_code}\n"
+
+            # Write to output file if specified
+            if output_file and sql.strip():
+                output_file.write_text(sql)
+                logger.info(f"Migration SQL written to {output_file}")
+
+            return sql if sql.strip() else "-- No changes detected"
+
+        except Exception as e:
+            logger.error(f"Failed to generate migration: {e}")
+            return f"-- Error generating migration: {e}"
+
+    def plan_migration(self, target_db: Optional[str] = None) -> str:
+        """
+        Plan a migration (dry-run) showing what would change.
+
+        Args:
+            target_db: Target database to compare against
+
+        Returns:
+            Human-readable migration plan
+        """
+        sql = self.generate_migration_sql(target_db=target_db)
+
+        if "No changes detected" in sql:
+            return "No changes detected between models and database."
+
+        return f"Migration Plan:\n\n{sql}"
+
+    def apply_sql_file(
+        self, sql_file: Path, connection_string: Optional[str] = None
+    ) -> bool:
+        """
+        Apply a SQL file to the database using psql.
+
+        Args:
+            sql_file: Path to SQL file
+            connection_string: PostgreSQL connection string (uses settings if not provided)
+
+        Returns:
+            True if successful, False otherwise
+        """
+        if not sql_file.exists():
+            raise FileNotFoundError(f"SQL file not found: {sql_file}")
+
+        conn_str = connection_string or self.get_connection_string()
+
+        try:
+            result = subprocess.run(
+                ["psql", conn_str, "-f", str(sql_file), "-v", "ON_ERROR_STOP=1"],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+
+            logger.info(f"Successfully applied {sql_file}")
+            if result.stdout:
+                logger.debug(result.stdout)
+
+            return True
+
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Failed to apply {sql_file}: {e.stderr or e.stdout}")
+            return False
+
+        except FileNotFoundError:
+            logger.error(
+                "psql command not found. Ensure PostgreSQL client is installed."
+            )
+            return False
+
+    def apply_migration(self, target_db: Optional[str] = None) -> bool:
+        """
+        Generate and apply migration in one step.
+
+        Args:
+            target_db: Target database to migrate
+
+        Returns:
+            True if successful, False otherwise
+        """
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".sql", delete=False
+        ) as f:
+            sql_file = Path(f.name)
+
+        try:
+            sql = self.generate_migration_sql(output_file=sql_file, target_db=target_db)
+
+            if "No changes detected" in sql:
+                logger.info("No migration needed")
+                return True
+
+            return self.apply_sql_file(sql_file)
+
+        finally:
+            # Clean up temp file
+            if sql_file.exists():
+                sql_file.unlink()
+
+    def get_current_revision(self) -> Optional[str]:
+        """
+        Get current database revision.
+
+        Returns:
+            Current revision ID or None if no migrations applied
+        """
+        try:
+            stdout_buffer = io.StringIO()
+            with redirect_stdout(stdout_buffer):
+                command.current(self.alembic_cfg)
+            output = stdout_buffer.getvalue()
+            # Parse the output to get revision ID
+            import re
+
+            match = re.search(r"([a-f0-9]+)\s+\(head\)", output)
+            if match:
+                return match.group(1)
+            return None
+        except Exception as e:
+            logger.error(f"Could not get current revision: {e}")
+            return None
+
+    def validate_migration_safety(
+        self, sql: str, safe_mode: Optional[str] = None
+    ) -> MigrationSafetyResult:
+        """
+        Validate migration SQL against safety rules.
+
+        Args:
+            sql: Migration SQL to validate
+            safe_mode: Override safety mode (uses settings if not provided)
+
+        Returns:
+            MigrationSafetyResult with validation results
+        """
+        mode = safe_mode or settings.migration.mode
+        errors: list[str] = []
+        warnings: list[str] = []
+
+        # Normalize SQL for pattern matching
+        sql_upper = sql.upper()
+
+        # Check for DROP COLUMN
+        if re.search(r"\bDROP\s+COLUMN\b", sql_upper, re.IGNORECASE):
+            if mode in ("additive", "strict"):
+                errors.append("DROP COLUMN not allowed in safe mode")
+            elif not settings.migration.allow_drop_columns:
+                errors.append(
+                    "DROP COLUMN not allowed (MIGRATION__ALLOW_DROP_COLUMNS=false)"
+                )
+            elif settings.migration.unsafe_alter_warning:
+                warnings.append("Migration contains DROP COLUMN operations")
+
+        # Check for DROP TABLE
+        if re.search(r"\bDROP\s+TABLE\b", sql_upper, re.IGNORECASE):
+            if mode in ("additive", "strict"):
+                errors.append("DROP TABLE not allowed in safe mode")
+            elif not settings.migration.allow_drop_tables:
+                errors.append(
+                    "DROP TABLE not allowed (MIGRATION__ALLOW_DROP_TABLES=false)"
+                )
+            elif settings.migration.unsafe_alter_warning:
+                warnings.append("Migration contains DROP TABLE operations")
+
+        # Check for ALTER COLUMN (type changes)
+        if re.search(r"\bALTER\s+COLUMN\s+\w+\s+TYPE\b", sql_upper, re.IGNORECASE):
+            if mode == "strict":
+                errors.append("ALTER COLUMN TYPE not allowed in strict mode")
+            elif not settings.migration.allow_alter_columns:
+                errors.append(
+                    "ALTER COLUMN TYPE not allowed (MIGRATION__ALLOW_ALTER_COLUMNS=false)"
+                )
+            elif settings.migration.unsafe_alter_warning:
+                warnings.append("Migration contains ALTER COLUMN TYPE operations")
+
+        # Check for RENAME COLUMN
+        if re.search(r"\bRENAME\s+COLUMN\b", sql_upper, re.IGNORECASE):
+            if mode == "strict":
+                errors.append("RENAME COLUMN not allowed in strict mode")
+            elif not settings.migration.allow_rename_columns:
+                errors.append(
+                    "RENAME COLUMN not allowed (MIGRATION__ALLOW_RENAME_COLUMNS=false)"
+                )
+            elif settings.migration.unsafe_alter_warning:
+                warnings.append("Migration contains RENAME COLUMN operations")
+
+        # Check for RENAME TABLE / ALTER TABLE RENAME
+        if re.search(
+            r"\b(RENAME\s+TABLE|ALTER\s+TABLE\s+\w+\s+RENAME\s+TO)\b",
+            sql_upper,
+            re.IGNORECASE,
+        ):
+            if mode == "strict":
+                errors.append("RENAME TABLE not allowed in strict mode")
+            elif not settings.migration.allow_rename_tables:
+                errors.append(
+                    "RENAME TABLE not allowed (MIGRATION__ALLOW_RENAME_TABLES=false)"
+                )
+            elif settings.migration.unsafe_alter_warning:
+                warnings.append("Migration contains RENAME TABLE operations")
+
+        # Check for other ALTER operations
+        if (
+            re.search(r"\bALTER\s+TABLE\b", sql_upper, re.IGNORECASE)
+            and settings.migration.unsafe_alter_warning
+        ):
+            # Only warn if not already warned above
+            if not any("ALTER" in w for w in warnings):
+                warnings.append("Migration contains ALTER TABLE operations")
+
+        is_safe = len(errors) == 0
+
+        return MigrationSafetyResult(
+            is_safe=is_safe,
+            errors=errors,
+            warnings=warnings,
+            sql=sql,
+        )
+
+    def generate_migration_sql_safe(
+        self,
+        output_file: Optional[Path] = None,
+        target_db: Optional[str] = None,
+        message: str = "Auto-generated migration",
+        safe_mode: Optional[str] = None,
+    ) -> MigrationSafetyResult:
+        """
+        Generate migration SQL with safety validation.
+
+        Args:
+            output_file: Path to write SQL file (if None, returns SQL string)
+            target_db: Target database name (overrides settings)
+            message: Migration message/description
+            safe_mode: Override safety mode (permissive, additive, strict)
+
+        Returns:
+            MigrationSafetyResult with validation results
+        """
+        # Generate SQL
+        sql = self.generate_migration_sql(
+            output_file=None,  # Don't write yet, validate first
+            target_db=target_db,
+            message=message,
+        )
+
+        # Validate safety
+        result = self.validate_migration_safety(sql, safe_mode=safe_mode)
+
+        # Write to file only if safe and output_file provided
+        if result.is_safe and output_file:
+            output_file.write_text(sql)
+            logger.info(f"Migration SQL written to {output_file}")
+        elif not result.is_safe:
+            logger.warning("Migration contains unsafe operations, not writing to file")
+
+        return result

rem/services/postgres/pydantic_to_sqlalchemy.py

@@ -0,0 +1,232 @@
+"""
+Convert Pydantic models to SQLAlchemy metadata for Alembic autogenerate.
+
+This module bridges REM's Pydantic-first approach with Alembic's SQLAlchemy requirement
+by dynamically building SQLAlchemy Table objects from Pydantic model definitions.
+"""
+
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+from pydantic import BaseModel
+from sqlalchemy import (
+    JSON,
+    Boolean,
+    Column,
+    DateTime,
+    Float,
+    Integer,
+    MetaData,
+    String,
+    Table,
+    Text,
+)
+from sqlalchemy.dialects.postgresql import ARRAY, JSONB, UUID
+
+from .schema_generator import SchemaGenerator
+
+
+def pydantic_type_to_sqlalchemy(
+    field_type: Any, field_info: Any
+) -> Any:
+    """
+    Map Pydantic field type to SQLAlchemy column type.
+
+    Args:
+        field_type: Pydantic field type annotation
+        field_info: Pydantic FieldInfo object
+
+    Returns:
+        SQLAlchemy column type
+    """
+    # Get the origin type (handles Optional, List, etc.)
+    import typing
+
+    origin = typing.get_origin(field_type)
+    args = typing.get_args(field_type)
+
+    # Handle Optional types
+    if origin is typing.Union:
+        # Optional[X] is Union[X, None]
+        non_none_types = [t for t in args if t is not type(None)]
+        if non_none_types:
+            field_type = non_none_types[0]
+            origin = typing.get_origin(field_type)
+            args = typing.get_args(field_type)
+
+    # Handle list types -> PostgreSQL ARRAY
+    if origin is list:
+        if args:
+            inner_type = args[0]
+            if inner_type is str:
+                return ARRAY(Text)
+            elif inner_type is int:
+                return ARRAY(Integer)
+            elif inner_type is float:
+                return ARRAY(Float)
+        return ARRAY(Text)  # Default to text array
+
+    # Handle dict types -> JSONB
+    if origin is dict or field_type is dict:
+        return JSONB
+
+    # Handle basic types
+    if field_type is str:
+        # Check if there's a max_length constraint
+        max_length = getattr(field_info, "max_length", None)
+        if max_length:
+            return String(max_length)
+        return Text
+
+    if field_type is int:
+        return Integer
+
+    if field_type is float:
+        return Float
+
+    if field_type is bool:
+        return Boolean
+
+    # Handle datetime
+    from datetime import datetime
+
+    if field_type is datetime:
+        return DateTime
+
+    # Handle UUID
+    from uuid import UUID as UUIDType
+
+    if field_type is UUIDType:
+        return UUID(as_uuid=True)
+
+    # Handle enums
+    import enum
+
+    if isinstance(field_type, type) and issubclass(field_type, enum.Enum):
+        return String(50)
+
+    # Default to Text for unknown types
+    logger.warning(f"Unknown field type {field_type}, defaulting to Text")
+    return Text
+
+
+def build_sqlalchemy_metadata_from_pydantic(models_dir: Path) -> MetaData:
+    """
+    Build SQLAlchemy MetaData from Pydantic models.
+
+    This function:
+    1. Discovers Pydantic models in the given directory
+    2. Infers table names and column definitions
+    3. Creates SQLAlchemy Table objects
+    4. Returns a MetaData object for Alembic
+
+    Args:
+        models_dir: Directory containing Pydantic models
+
+    Returns:
+        SQLAlchemy MetaData object
+    """
+    metadata = MetaData()
+    generator = SchemaGenerator()
+
+    # Discover models
+    models = generator.discover_models(models_dir)
+    logger.info(f"Discovered {len(models)} models for metadata generation")
+
+    for model_name, model_class in models.items():
+        # Infer table name
+        table_name = generator.infer_table_name(model_class)
+        logger.debug(f"Building table {table_name} from model {model_name}")
+
+        # Build columns
+        columns = []
+
+        for field_name, field_info in model_class.model_fields.items():
+            # Get field type
+            field_type = field_info.annotation
+
+            # Map to SQLAlchemy type
+            sa_type = pydantic_type_to_sqlalchemy(field_type, field_info)
+
+            # Determine nullable
+            nullable = not field_info.is_required()
+
+            # Get default value
+            from pydantic_core import PydanticUndefined
+
+            default = None
+            if field_info.default is not PydanticUndefined and field_info.default is not None:
+                default = field_info.default
+            elif field_info.default_factory is not None:
+                # For default_factory, we'll use the server default if possible
+                factory = field_info.default_factory
+                # Handle common default factories
+                if factory.__name__ == "list":
+                    default = "ARRAY[]::TEXT[]"  # PostgreSQL empty array
+                elif factory.__name__ == "dict":
+                    default = "'{}'::jsonb"  # PostgreSQL empty JSON
+                else:
+                    default = None
+
+            # Handle special fields
+            server_default = None
+            primary_key = False
+
+            if field_name == "id":
+                primary_key = True
+                if sa_type == UUID(as_uuid=True):
+                    server_default = "uuid_generate_v4()"
+            elif field_name in ("created_at", "updated_at"):
+                server_default = "CURRENT_TIMESTAMP"
+            elif isinstance(default, str) and default.startswith("ARRAY["):
+                server_default = default
+                default = None
+            elif isinstance(default, str) and "::jsonb" in default:
+                server_default = default
+                default = None
+
+            # Create column - only pass server_default if it's a string SQL expression
+            column_kwargs = {
+                "type_": sa_type,
+                "primary_key": primary_key,
+                "nullable": nullable,
+            }
+
+            if server_default is not None:
+                from sqlalchemy import text
+                column_kwargs["server_default"] = text(server_default)
+
+            column = Column(field_name, **column_kwargs)
+
+            columns.append(column)
+
+        # Create table
+        if columns:
+            Table(table_name, metadata, *columns)
+            logger.debug(f"Created table {table_name} with {len(columns)} columns")
+
+    logger.info(f"Built metadata with {len(metadata.tables)} tables")
+    return metadata
+
+
+def get_target_metadata() -> MetaData:
+    """
+    Get SQLAlchemy metadata for Alembic autogenerate.
+
+    This is the main entry point used by alembic/env.py.
+
+    Returns:
+        SQLAlchemy MetaData object representing current Pydantic models
+    """
+    # Find models directory
+    import rem
+
+    package_root = Path(rem.__file__).parent.parent.parent
+    models_dir = package_root / "src" / "rem" / "models" / "entities"
+
+    if not models_dir.exists():
+        logger.error(f"Models directory not found: {models_dir}")
+        return MetaData()
+
+    return build_sqlalchemy_metadata_from_pydantic(models_dir)