Alembic Utils

Alembic Utils is an extension for Alembic that adds autogenerate support for various PostgreSQL entity types. While standard Alembic can detect changes to SQLAlchemy models, Alembic Utils extends this capability to include:

• Functions
• Views
• Materialized Views
• Triggers
• Policies

When you use Alembic Utils, the alembic revision --autogenerate command will automatically detect if these entities are new, updated, or removed, and populate the upgrade and downgrade sections of your migration scripts accordingly.

Once your entities are defined and registered in env.py, run the standard Alembic autogenerate command. alembic_utils will detect if your functions or views are new, updated, or removed and automatically populate the upgrade and downgrade sections of the generated revision.

alembic revision --autogenerate -m 'some message'

To enable automatic management of database entities (like functions or views) via Alembic, you must call register_entities within your migrations/env.py file. This allows alembic_utils to track changes to these entities and generate migrations automatically.

Pass a list of entity objects to the entities argument. You can optionally use exclude_schema to prevent tracking entities within specific schemas.

# migrations/env.py

from alembic_utils.replaceable_entity import register_entities
from app.functions import my_function
from app.views import my_view

register_entities(entities=[my_function, my_view], exclude_schema=['audit'])

To allow Alembic to detect changes in your functions and views, you must register them in your <migrations_folder>/env.py file using register_entities from alembic_utils.replaceable_entity.

# <migrations_folder>/env.py
from alembic_utils.replaceable_entity import register_entities

from my_function import to_upper
from my_view import first_view

register_entities([to_upper, first_view])

Install alembic_utils using pip. Requires Python 3.6+.

$ pip install alembic_utils

To enable autogeneration for PostgreSQL entities (like functions or views), you must register them in your Alembic env.py file using register_entities.

1. Define your entity using the appropriate class (e.g., PGFunction).
2. Call register_entities with a list containing your defined entities.

# migrations/env.py
from alembic_utils.pg_function import PGFunction
from alembic_utils.replaceable_entity import register_entities

to_upper = PGFunction(
  schema='public',
  signature='to_upper(some_text text)',
  definition="""
  RETURNS text as
  $$
    SELECT upper(some_text)
  $$
  $$ language SQL;
  """
)

register_entities([to_upper])

By default, Alembic includes all schemas in the target database during an --autogenerate revision. To selectively include or exclude specific objects or schemas, you must implement and pass include_name and include_object callables to context.configure() within your env.py file.

• Use include_name to filter by name and type (e.g., filtering schemas).
• Use include_object to filter by the actual object instance (e.g., excluding specific entity types like PGFunction).

To enable logging for alembic_utils, add alembic_utils to the keys in the [loggers] section of your alembic.ini file, and define a [logger_alembic_utils] section.

[loggers]
keys=root,sqlalchemy,alembic,alembic_utils

[logger_alembic_utils]
level = INFO
handlers =
qualname = alembic_utils

The PGView class allows you to manage PostgreSQL views within Alembic migrations. It is compatible with alembic revision --autogenerate because it inherits from ReplaceableEntity.

You can instantiate a PGView manually, parse it from a raw CREATE VIEW SQL string, or fetch existing views directly from a database session.

from alembic_utils.pg_view import PGView

# Manual instantiation
view = PGView(schema="public", signature="my_view", definition="SELECT * FROM my_table")

# Create from a SQL string
sql = 'CREATE VIEW public.my_view AS SELECT * FROM my_table;'
view_from_sql = PGView.from_sql(sql)

The PGPolicy class allows you to manage PostgreSQL Row Level Security (RLS) policies within Alembic migrations. It is compatible with alembic revision --autogenerate, enabling automatic detection and management of policies. You can instantiate policies from raw SQL strings, fetch existing policies directly from a database connection, and generate SQL statements for creating or dropping policies.

from alembic_utils.pg_policy import PGPolicy

# Create from SQL
policy = PGPolicy.from_sql("CREATE POLICY my_policy ON public.my_table FOR SELECT USING (user_id = current_user_id())")

# Generate SQL statements
print(policy.to_sql_statement_create())
print(policy.to_sql_statement_drop())

# Fetch from database
# policies = PGPolicy.from_database(connection, 'public')

To restrict autogeneration to only specific schemas (for example, only public and bi), implement an include_name function. This function checks if the object type is schema and validates the name against your allowed list. You must also set include_schemas = True in context.configure().

# env.py

def include_name(name, type_, parent_names) -> bool:
    if type_ == "schema":
        return name in ["public", "bi"]
    return True

context.configure(
    # ...
    include_schemas = True,
    include_name = include_name
)