Adding a pseudo MigrationOp to add autocommit context wrapper to generated revision

[lachaib]

Describe the use case

We're using a PostgreSQL database and Squawk as linter in CI to minimise the risks of breaking production when applying migrations.
Among the checks, Squawk recommends creating or dropping indexes with the CONCURRENTLY parameter, however this operation cannot be performed in transactionnal DDL, which means we need to amend the generated migration script to add with op.get_context().autocommit_block(): around the operation
Manually updating the generated script is error prone and often detected after a couple round trips of failing CI.

I have written a pseudo MigrationOp called AutoCommitBlockOp, and I would like to share it with the community.
See example use below for more code

Databases / Backends / Drivers targeted

PostgreSQL (probably others, autocommit block is not restricted to a dialect)

Example Use

class AutoCommitBlockOp(ops.MigrateOperation):
    """Wraps a list of ops inside a with op.get_context().autocommit_block(): block.

    Use this for any DDL that PostgreSQL requires to run outside a transaction
    (e.g. CREATE INDEX CONCURRENTLY).
    """

    def __init__(self, nested_ops: list[ops.MigrateOperation]) -> None:
        self.nested_ops = nested_ops

@writer.rewrites(ops.CreateIndexOp)
def create_index_concurrently_if_not_exists(_autogen_context, _revision, op: ops.CreateIndexOp):
    op.if_not_exists = True
    op.kw["postgresql_concurrently"] = True
    return AutoCommitBlockOp([op])

would properly generated some code like this:

def upgrade() -> None:
    
    with op.get_context().autocommit_block():
        op.create_index(
            "uq_my_index",
            "my_table",
            ["col1", "col2", "col3"],
            unique=True,
            if_not_exists=True,
            postgresql_concurrently=True,
        )

Additional context
Happy to share a pull request if this feature can serve the greater good, but I could as well understand if it should not be streamlined.

Have a nice day!

[zzzeek]

right you'd like the rewriter to be able to render an autocommit block.

I think this is a valid use case and a good use of rewriter.

not sure about rendering though, perhaps autocommit_block() should ideally be part of the op.* namespace directly? not sure. this would be the first case where we make an "op directive" out of something that's not strictly an "op". Need to be careful about reinventing all of Python AST in "ops" :) (a slippery slope but that's a very long way down the slope so far).

it seems like maybe we could start with a recipe of some kind? after all you can define ops and renderers, although this one needs to indent what's contained within it, i'd have to look at how things work (it's been awhile) to see how feasible that is.