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,
...

[georgesittas]

Hey @philieeas, apologies for the delayed response here, it seems like this thread got lost in the fray.

You can indeed execute pre/post-statements in python models. For python models that generate SQL, the kwarg syntax in your example should work fine, but for models that return dataframes you need to include any post-statements after the dataframe has been yielded.

Feel free to open a ticket in the future for better visibility, or just reach out in Slack!

[philieeas]

Hey, thank you for the reply, but I was already aware of executing the post-statements there. I linked the same guide in my question.
The issue is that this would be executed each time when the model is evaluated, not just when it is created.
This gets expensive with a growing table of time-series in Postgres/TimescaleDB, which is why I wanted to only execute it during the create-phase, which is afaik not possible yet.

[georgesittas]

Ah, whoops -- I missed that link earlier when I read the question. Perhaps there's a workaround using the runtime stage? Need to think about this for a bit and see if we expose it in the execution context. Will get back to you.

[georgesittas]

The issue is that this would be executed each time when the model is evaluated, not just when it is created.

Right, makes sense, I see what you mean. The runtime stage enum won't be of much help here I think because you only reach the model's logic when evaluating it.

[georgesittas]

I've opened a ticket for this, see #2866. Let's continue the discussion there.