from __future__ import annotations from dataclasses import dataclass from enum import Enum from typing import Literal, Protocol Dialect = Literal["sqlite", "postgres", "mysql", "unknown"] class SchemaOpRisk(str, Enum): SAFE = "safe" GUARDED = "guarded" MANUAL = "manual" class SchemaOp(Protocol): risk: SchemaOpRisk def to_sql(self, dialect: Dialect) -> list[str]: ... def quote_identifier(identifier: str, dialect: Dialect) -> str: if dialect == "mysql": return f"`{identifier.replace('`', '``')}`" return f'"{identifier.replace(chr(34), chr(34) + chr(34))}"' def _column_type(column_type: str | dict[str, str], dialect: Dialect) -> str: if isinstance(column_type, dict): return column_type.get(dialect) or column_type.get("default") or "TEXT" return column_type def _default_sql(default: str | float | bool | None, dialect: Dialect) -> str: if default is None: return "" if isinstance(default, bool): if dialect == "postgres": return " DEFAULT TRUE" if default else " DEFAULT FALSE" return " DEFAULT 1" if default else " DEFAULT 0" if isinstance(default, int | float): return f" DEFAULT {default}" escaped = default.replace("'", "''") return f" DEFAULT '{escaped}'" @dataclass(frozen=True, slots=True) class AddColumn: table: str column: str column_type: str | dict[str, str] nullable: bool = True default: str | float | bool | None = None risk: SchemaOpRisk = SchemaOpRisk.SAFE def to_sql(self, dialect: Dialect) -> list[str]: table = quote_identifier(self.table, dialect) column = quote_identifier(self.column, dialect) column_type = _column_type(self.column_type, dialect) null_sql = "" if self.nullable else " NOT NULL" default_sql = _default_sql(self.default, dialect) if not self.nullable and not default_sql: return [] if dialect == "postgres": return [ "ALTER TABLE " f"{table} ADD COLUMN IF NOT EXISTS {column} " f"{column_type}{default_sql}{null_sql}" ] return [ "ALTER TABLE " f"{table} ADD COLUMN {column} {column_type}{default_sql}{null_sql}" ] @dataclass(frozen=True, slots=True) class CreateIndex: table: str columns: tuple[str, ...] name: str | None = None if_not_exists: bool = True unique: bool = False where: str | None = None risk: SchemaOpRisk = SchemaOpRisk.SAFE def __init__( self, table: str, columns: tuple[str, ...] | list[str], name: str | None = None, if_not_exists: bool = True, unique: bool = False, where: str | None = None, risk: SchemaOpRisk | None = None, ) -> None: object.__setattr__(self, "table", table) object.__setattr__(self, "columns", tuple(columns)) object.__setattr__(self, "name", name) object.__setattr__(self, "if_not_exists", if_not_exists) object.__setattr__(self, "unique", unique) object.__setattr__(self, "where", where) resolved_risk = risk or (SchemaOpRisk.MANUAL if unique else SchemaOpRisk.SAFE) object.__setattr__(self, "risk", resolved_risk) def to_sql(self, dialect: Dialect) -> list[str]: if self.unique or not self.columns: return [] name = self.name or f"idx_{self.table}_{'_'.join(self.columns)}"[:62] if dialect == "mysql" and self.where: return [] exists_sql = ( "IF NOT EXISTS " if self.if_not_exists and dialect != "mysql" else "" ) table = quote_identifier(self.table, dialect) index = quote_identifier(name, dialect) columns = ", ".join( quote_identifier(column, dialect) for column in self.columns ) where = f" WHERE {self.where}" if self.where and dialect != "mysql" else "" return [f"CREATE INDEX {exists_sql}{index} ON {table}({columns}){where}"] @dataclass(frozen=True, slots=True) class RenameColumn: table: str old: str new: str risk: SchemaOpRisk = SchemaOpRisk.GUARDED def to_sql(self, dialect: Dialect) -> list[str]: table = quote_identifier(self.table, dialect) old = quote_identifier(self.old, dialect) new = quote_identifier(self.new, dialect) return [f"ALTER TABLE {table} RENAME COLUMN {old} TO {new}"] @dataclass(frozen=True, slots=True) class DropColumn: table: str column: str risk: SchemaOpRisk = SchemaOpRisk.GUARDED def to_sql(self, dialect: Dialect) -> list[str]: table = quote_identifier(self.table, dialect) column = quote_identifier(self.column, dialect) return [f"ALTER TABLE {table} DROP COLUMN {column}"] @dataclass(frozen=True, slots=True) class AlterColumnType: table: str column: str column_type: str | dict[str, str] nullable: bool | None = None risk: SchemaOpRisk = SchemaOpRisk.GUARDED def to_sql(self, dialect: Dialect) -> list[str]: column_type = _column_type(self.column_type, dialect) table = quote_identifier(self.table, dialect) column = quote_identifier(self.column, dialect) if dialect == "postgres": return [f"ALTER TABLE {table} ALTER COLUMN {column} TYPE {column_type}"] if dialect == "mysql": null_sql = " NULL" if self.nullable else " NOT NULL" if self.nullable is None: null_sql = "" return [ f"ALTER TABLE {table} MODIFY COLUMN {column} {column_type}{null_sql}" ] return [] def normalize_schema_ops(items: list[str | SchemaOp], dialect: Dialect) -> list[str]: sql_list: list[str] = [] for item in items: if isinstance(item, str): sql_list.append(item) else: sql_list.extend(item.to_sql(dialect)) return sql_list