Pre/Post SQL Statements in Python models during creating phase?

[philieeas]

I am not sure if this is already possible but I think it would be useful to be able pre_statements/post_statements during the table creation phase in python.

My use case is that I want to setup and index for a postgres table but only do that after the underlying snapshot table was created.
I don't want to do this every time during model evaluation.

I could achieve the desired behavior in SQL models by defining a macro in python like this and using it in the model.

@macro()
def create_index(
    evaluator: MacroEvaluator,
    model_name: Literal,
    column: Literal,
):
    if evaluator.runtime_stage == "creating":
        snapshot: Snapshot | None = evaluator.get_snapshot(model_name.text("this"))
        if snapshot is None:
            raise ValueError("Can't resolve snapshot table")
        table_name = snapshot.table_name()
        ix_fingerprint = snapshot.fingerprint.to_version()
        index_name = f"{model_name}_{column}_{ix_fingerprint}"
        create_idx_statement = f"CREATE INDEX {index_name} ON {table_name} ({column})"
        stmt = parse(create_idx_statement, dialect="postgres")[0]
        assert stmt is not None
        return stmt

The documentation says that you can use fetchdf but I think that is only done during the evaluation phase?

I looked around the code and found that you can set pre_statements in @model decorator, but those give me an error in python models too:
Error: Invalid extra fields {'post_statements'} in the model definition at '/home/python/sqlmesh/models/pymodel.py'

I wonder though if its also possible to do this in a python model?
Ideally I would like to execute pre/post statements like this, reusing the macros too:

@model(
    "model",
    columns={
        "time": "datetime",
        "value": "float",
    },
    kind=IncrementalByTimeRangeKind(
        time_column="time", # type: ignore
        batch_size=12,
    ),
    start=datetime(2024, 4, 1),
    cron="@hourly",
    post_statements=["@CREATE_INDEX('example.model', time);"]
)
def execute(
    context: ExecutionContext,
    start: datetime,
...