Custom sqlalchemy visitor not registered

Question

I have a Postgres DB maintained via SQLAlchemy/Alembic, to which I want to add a new column of type array. I'm trying to follow this page of the SQLAlchemy documentation: https://docs.sqlalchemy.org/en/20/core/compiler.html

The table is currently defined as:

@compiles(CreateColumn, "postgres")
def use_identity(element, compiler, **kw):
    text = compiler.visit_create_column(element, **kw)
    return text.replace("SERIAL", "INT GENERATED BY DEFAULT AS IDENTITY")
    

my_table = Table(
    "my_table",
    my_metadata,
    Column("my_id", Biginteger, primary_key=True),
    # ...
)

The column I'm trying to add is defined as:

Column(
    "authors", 
    ARRAY(String(128)), 
    nullable=False, 
    server_default=text("ARRAY[]::varchar[]"),
),

When I run alembic revision --autogenerate, it correctly creates the revsion file. However, when I then run alembic upgrade head, I see this error message:

psycopg2.errors.SyntaxError: column "my_id" of relation "my_table" is an identity column
HINT:  Use ALTER TABLE ... ALTER COLUMN ... DROP IDENTITY instead.

[SQL: ALTER TABLE my_table ALTER COLUMN my_id DROP DEFAULT]

Which sounds pretty clear to me. So in the same file I tried to add another compilation extension:

@compiles(AlterColumn, "postgres")
def visit_alter_column(element, compiler, **kw):
    text = compiler.visit_alter_column(element, **kw)
    return text.replace("DROP DEFAULT", "DROP IDENTITY IF EXISTS")

However when rerunning Alembic I get the exact same error as before. It's as if the new compilation extension doesn't get registered. I also tried (slightly different, just to see if it got picked up):

@compiles(AlterColumn, "postgres")
def visit_alter_column(element, compiler, **kw):
    return "ALTER TABLE {element.table.name} ALTER COLUMN {element.column.name} DROP IDENTITY IF EXISTS"

But still no cigar. What am I doing wrong?

Answer 1

Answering my own question.

Instead of trying to override the SQL command generated by Alembic, I noticed that the script generated by alembic revision --autogenerate is not immutable--in fact one can find the following comment:

# ### commands auto generated by Alembic - please adjust! ###

With that, and following the stack trace of the error, I jumped to the culprit command:

op.alter_column("my_table", "my_id",
           existing_type=sa.BIGINT(),
           server_default=None,    # BINGO!
           existing_nullable=False,
           autoincrement=True)

The server_default option was causing the generation of the DROP DEFAULT statement. So the solution was to just comment it out:

op.alter_column("my_table", "my_id",
           existing_type=sa.BIGINT(),
           # server_default=None,
           existing_nullable=False,
           autoincrement=True)