Source code for sqlspec.migrations.schema

"""Additive schema ensure and diff helpers."""

from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any

from sqlglot import exp, parse

from sqlspec.builder import AlterTable, CreateTable, sql

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

__all__ = ("SchemaEnsureResult", "SchemaTarget", "ensure_schema_async", "ensure_schema_sync")




class SchemaTarget:
    """Describe one table's target schema.

    Args:
        table_name: Unqualified table name used for introspection and DDL.
        create_table: Builder containing the target columns and create DDL.
        schema: Optional schema containing the table.
    """

    __slots__ = ("create_statement", "create_table", "schema", "table_name")



    def __init__(
        self,
        table_name: str,
        create_table: CreateTable,
        schema: str | None = None,
        create_statement: "str | CreateTable | None" = None,
    ) -> None:
        self.table_name = table_name
        self.create_table = create_table
        self.schema = schema
        self.create_statement = create_statement or create_table

@property def identity(self) -> str: """Return the schema-qualified identity used in results.""" if self.schema: return f"{self.schema}.{self.table_name}" return self.table_name

@classmethod def from_ddl( cls, table_name: str, create_statement: str, *, schema: str | None = None, dialect: Any = None ) -> "SchemaTarget": """Build a target descriptor from an adapter's canonical CREATE TABLE DDL. Args: table_name: Unqualified table name used for introspection. create_statement: Canonical adapter DDL, including any companion statements. schema: Optional schema containing the table. dialect: SQLGlot dialect used to parse column definitions. Returns: Target descriptor that executes the original DDL while deriving additive column statements from the parsed table definition. Raises: ValueError: If no CREATE TABLE definition can be parsed. """ create_expression = _find_create_table_expression(create_statement, table_name, dialect) target = sql.create_table(table_name, dialect=dialect) if schema: target.in_schema(schema) for column in create_expression.find_all(exp.ColumnDef): kind = column.args.get("kind") if not isinstance(kind, exp.DataType): continue constraints = [constraint.args.get("kind") for constraint in column.args.get("constraints", [])] default = next( ( constraint.this.sql(dialect=dialect) for constraint in constraints if isinstance(constraint, exp.DefaultColumnConstraint) and constraint.this is not None ), None, ) target.column( column.name, kind.sql(dialect=dialect), default=default, not_null=any(isinstance(constraint, exp.NotNullColumnConstraint) for constraint in constraints), primary_key=any(isinstance(constraint, exp.PrimaryKeyColumnConstraint) for constraint in constraints), unique=any(isinstance(constraint, exp.UniqueColumnConstraint) for constraint in constraints), ) if not target.columns: msg = f"CREATE TABLE DDL for {table_name!r} has no parseable columns" raise ValueError(msg) return cls(table_name, target, schema, create_statement)