v0019_add_env_suffix_target.py

"""Add support for environment suffix target."""

from sqlglot import exp


def migrate_ddl(state_sync, **kwargs):  # type: ignore
    engine_adapter = state_sync.engine_adapter
    environments_table = "_environments"
    if state_sync.schema:
        environments_table = f"{state_sync.schema}.{environments_table}"

    alter_table_exp = exp.Alter(
        this=exp.to_table(environments_table),
        kind="TABLE",
        actions=[
            exp.ColumnDef(
                this=exp.to_column("suffix_target"),
                kind=exp.DataType.build("text"),
            )
        ],
    )
    engine_adapter.execute(alter_table_exp)


def migrate_dml(state_sync, **kwargs):  # type: ignore
    engine_adapter = state_sync.engine_adapter
    environments_table = "_environments"
    if state_sync.schema:
        environments_table = f"{state_sync.schema}.{environments_table}"

    state_sync.engine_adapter.update_table(
        environments_table,
        {"suffix_target": "schema"},
        where="1=1",
    )