ss-tools/backend/src/models/mapping.py

# #region MappingModels [C:5] [TYPE Module] [SEMANTICS sqlalchemy, mapping, model, schema, resource-type]
# @defgroup Models Module group.
#
# @BRIEF Defines the database schema for environment metadata and database mappings using SQLAlchemy.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [EXT:Library:sqlalchemy]

#
# @INVARIANT  All primary keys are UUID strings.
# CONSTRAINT: source_env_id and target_env_id must be valid environment IDs.
# @PRE  Database engine initialized
# @POST  Mapping ORM models registered with UUID primary keys
# @SIDE_EFFECT  Defines environment/database/resource mapping tables
# @DATA_CONTRACT  MappingData -> MappingRecord

import enum
import uuid

from sqlalchemy import Boolean, Column, DateTime, Enum as SQLEnum, ForeignKey, String
from sqlalchemy.orm import declarative_base
from sqlalchemy.sql import func

# #region Base [C:1] [TYPE Class]
# @BRIEF SQLAlchemy declarative base for all domain models.
# #endregion Base
Base = declarative_base()

# #region ResourceType [C:1] [TYPE Class]
# @BRIEF Enumeration of possible Superset resource types for ID mapping.
class ResourceType(str, enum.Enum):
    CHART = "chart"
    DATASET = "dataset"
    DASHBOARD = "dashboard"
# #endregion ResourceType


# #region MigrationStatus [C:1] [TYPE Class]
# @BRIEF Enumeration of possible migration job statuses.
class MigrationStatus(enum.Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
    AWAITING_MAPPING = "AWAITING_MAPPING"
# #endregion MigrationStatus

# #region Environment [C:3] [TYPE Class]
# @defgroup Models Module group.
# @BRIEF Represents a Superset instance environment.
# @RELATION DEPENDS_ON -> MappingModels
class Environment(Base):
    __tablename__ = "environments"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    name = Column(String, nullable=False)
    url = Column(String, nullable=False)
    credentials_id = Column(String, nullable=False)
# #endregion Environment

# #region DatabaseMapping [C:3] [TYPE Class]
# @defgroup Models Module group.
# @BRIEF Represents a mapping between source and target databases.
class DatabaseMapping(Base):
    __tablename__ = "database_mappings"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    source_env_id = Column(String, ForeignKey("environments.id", ondelete="CASCADE"), nullable=False)
    target_env_id = Column(String, ForeignKey("environments.id", ondelete="CASCADE"), nullable=False)
    source_db_uuid = Column(String, nullable=False)
    target_db_uuid = Column(String, nullable=False)
    source_db_name = Column(String, nullable=False)
    target_db_name = Column(String, nullable=False)
    engine = Column(String, nullable=True)
# #endregion DatabaseMapping

# #region MigrationJob [C:2] [TYPE Class]
# @defgroup Models Module group.
# @BRIEF Represents a single migration execution job.
class MigrationJob(Base):
    __tablename__ = "migration_jobs"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    source_env_id = Column(String, ForeignKey("environments.id", ondelete="CASCADE"), nullable=False)
    target_env_id = Column(String, ForeignKey("environments.id", ondelete="CASCADE"), nullable=False)
    status = Column(SQLEnum(MigrationStatus), default=MigrationStatus.PENDING)
    replace_db = Column(Boolean, default=False)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
# #endregion MigrationJob

# #region ResourceMapping [C:3] [TYPE Class]
# @defgroup Models Module group.
# @BRIEF Maps a universal UUID for a resource to its actual ID on a specific environment.
# @TEST_DATA: resource_mapping_record -> {'environment_id': 'prod-env-1', 'resource_type': 'chart', 'uuid': '123e4567-e89b-12d3-a456-426614174000', 'remote_integer_id': '42'}
# @RELATION DEPENDS_ON -> MappingModels
class ResourceMapping(Base):
    __tablename__ = "resource_mappings"

    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
    environment_id = Column(String, ForeignKey("environments.id", ondelete="CASCADE"), nullable=False)
    resource_type = Column(SQLEnum(ResourceType), nullable=False)
    uuid = Column(String, nullable=False)
    remote_integer_id = Column(String, nullable=False)  # Stored as string to handle potentially large or composite IDs safely, though Superset usually uses integers.
    resource_name = Column(String, nullable=True) # Used for UI display
    last_synced_at = Column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())
# #endregion ResourceMapping

# #endregion MappingModels