How to add custom postgres types (e.g. pgvector) in SQLAlchemy + Alembic?

[MIrinkov]

When I describe a table model like

from pgvector.sqlalchemy import Vector
import numpy as np
from sqlalchemy.orm import Mapped, mapped_column

N_DIM = 1536


class Base(DeclarativeBase):
    pass

class Embedding(Base):
    __tablename__ = "embeddings"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    embedding: Mapped[np.array] = mapped_column(Vector(N_DIM))

I can successfully do alembic revision --autogenerate and then alembic upgrade head. Good and perfect, so far!
However, when I run alembic check afterwards, I am getting a warning like so

PS C:\Users\User\project\migrations> alembic check
INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
C:\Users\User\project\migrations\.\env.py:66: SAWarning: Did not recognize type 'vector' of column 'embedding'
  context.run_migrations()
INFO  [alembic.autogenerate.compare] Couldn't determine database type for column 'embeddings.embedding'
No new upgrade operations detected.

After digging in the code for some time, I found that the reason is this line here:

# venv/Lib/site-packages/sqlalchemy/dialects/postgresql/base.py:3776
coltype = None

Looking into the code in base.py, I figured that it makes sense to force the vector type into the dialect's "list of recognized types". I.e. I want this condition on line venv/Lib/site-packages/sqlalchemy/dialects/postgresql/base.py:3748 to be TRUE

                # attype == 'vector' here
                if attype in self.ischema_names:
                    coltype = self.ischema_names[attype]
                    break

So my workaround in env.py looks like this (I also added the type into a migration_types module as per recommendations at https://alembic.sqlalchemy.org/en/latest/autogenerate.html#controlling-the-module-prefix ):

def do_run_migrations(connection: Connection) -> None:
    # Need to hack the "vector" type into postgres dialect schema types.
    # Otherwise, `alembic check` does not recognize the type
    connection.dialect.ischema_names['vector'] = project.database.migration_types.Vector

    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        user_module_prefix="project.database.migration_types.",
    )

    with context.begin_transaction():
        context.run_migrations()

My question is: Is there a "proper" way to do this?
Is the developer of the library in question supposed to make it "hook into sqlalchemy" somehow?
Is there a way to pass a dictionary of custom types to alembic / sqlalchemy?

It's possible I missed one of the documentation pages, if so please just give me a link, I'll read up.

[zzzeek]

SQLAlchemy itself will not know how to reflect this type, which is the warning you see. There is an event hook called column_reflect, which would be a great place to allow interception of this "vector" name, however, that event hook will give you a NullType and you won't have the information you need (seems like that might be something that can be improved; we can add a field with the raw type reflect information so that third parties can replace the NullType as needed).

We dont otherwise have a hook for custom types so the way you're placing it into ischema_names right now is the only way to do it.

[n-splv]

For anyone wondering, I think I've found a simpler solution based on @MIrinkov one. There's no need to add a custom migration_types file, instead in env.py just do this:

import pgvector
...


def do_run_migrations(connection: Connection) -> None:
    # Need to hack the "vector" type into postgres dialect schema types.
    # Otherwise, `alembic check` does not recognize the type
    connection.dialect.ischema_names['vector'] = pgvector.sqlalchemy.Vector

    context.configure(
        connection=connection,
        target_metadata=target_metadata,
    )

    with context.begin_transaction():
        context.run_migrations()

Then modify your run_migrations_online to call do_run_migrations. Also, don't forget the import in script.py.mako:

import pgvector