-
-
Save h4/fc9b6d350544ff66491308b535762fee to your computer and use it in GitHub Desktop.
from __future__ import with_statement | |
from alembic import context | |
from sqlalchemy import engine_from_config, pool | |
from logging.config import fileConfig | |
from models import Base | |
config = context.config | |
fileConfig(config.config_file_name) | |
""" | |
Load models metadata. We should define schema in this class firstly, | |
or set schema implicit with `__table_args__ = {'schema' : 'test'}` in model class | |
""" | |
target_metadata = Base.metadata | |
def run_migrations_offline(): | |
url = config.get_main_option("sqlalchemy.url") | |
context.configure( | |
url=url, target_metadata=target_metadata, literal_binds=True) | |
with context.begin_transaction(): | |
context.run_migrations() | |
def run_migrations_online(): | |
connectable = engine_from_config( | |
config.get_section(config.config_ini_section), | |
prefix='sqlalchemy.', | |
poolclass=pool.NullPool) | |
with connectable.connect() as connection: | |
""" | |
Configure migration context | |
1. Pass our models metadata | |
2. Set schema for alembic_version table | |
3. Load all available schemas | |
""" | |
context.configure( | |
connection=connection, | |
target_metadata=target_metadata, | |
version_table_schema=target_metadata.schema, | |
include_schemas=True | |
) | |
with context.begin_transaction(): | |
""" | |
By default search_path is setted to "$user",public | |
that why alembic can't create foreign keys correctly | |
""" | |
context.execute('SET search_path TO public') | |
context.run_migrations() | |
if context.is_offline_mode(): | |
run_migrations_offline() | |
else: | |
run_migrations_online() |
After spending a lot of time struggling to get my alembic to work, with a schema_name other than "public", I've come up with the following solution:
[The only solution that worked for me on February 16, 2024]
def run_migrations_online() -> None:
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
connectable = engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
with connectable.connect() as connection:
context.configure(connection=connection, target_metadata=target_metadata)
connection.execute(text('set search_path to "%s"' % settings.postgres_db_schema)) # <-- The magic line
with context.begin_transaction():
context.run_migrations()
Hope this helps
After spending a lot of time struggling to get my alembic to work, with a schema_name other than "public", I've come up with the following solution: [The only solution that worked for me on February 16, 2024]
def run_migrations_online() -> None: """Run migrations in 'online' mode. In this scenario we need to create an Engine and associate a connection with the context. """ connectable = engine_from_config( config.get_section(config.config_ini_section, {}), prefix="sqlalchemy.", poolclass=pool.NullPool, ) with connectable.connect() as connection: context.configure(connection=connection, target_metadata=target_metadata) connection.execute(text('set search_path to "%s"' % settings.postgres_db_schema)) # <-- The magic line with context.begin_transaction(): context.run_migrations()Hope this helps
This solution works for me. Thanks to share it!
connection.execute(text('set search_path to "%s"' % settings.postgres_db_schema)) # <-- The magic line with context.begin_transaction(): context.run_migrations()Hope this helps
Thank you, @Ilyes-git! This definitely helps a lot! 🙇🏼
This solution, in case your app uses a single custom schema, allows you to avoid adding schema configuration to each model in the application.
However, there's a minor thing that can be improved here. Although it works perfectly fine with psycopg
, when you use asyncpg
the whole migration rolls back since, I assume, connection.execute
implicitly starts a transaction. So it's worth moving the magic line
under context.run_migrations()
just like this:
with context.begin_transaction():
context.execute(text(f'SET search_path TO {settings.db_schema}'))
context.run_migrations()
And it fixes the issue for asyncpg
.
I'm having a concern with the odd behavior of the alembic revision --autogenerate
that doesn't respect the version_table_schema
argument of context.configure()
,that sets to db_schema so detect the version_table
as an external table and add a drop_table
command for that in the initial migration script (the genesis)!
I came across some discussion threads, but no resolution for that (like its kind of a technical debt). anyone has any idea of the fix?
A working solution for me with async pg:
My base.py
"""This module contains the base class for all models in the application."""
from sqlalchemy.ext.declarative import declarative_base
from app.core.config import settings
class Base:
__table_args__ = {"schema": settings.SCHEMA}
Base = declarative_base(cls=Base)
My env.py
import asyncio
from logging.config import fileConfig
from sqlalchemy import pool, text
from sqlalchemy.engine import Connection
from sqlalchemy.ext.asyncio import async_engine_from_config
from alembic import context
from app.core.config import settings
from app.db.models import *
from app.db.base import Base
# Database Schmea for SecurityBot
SCHEMA = settings.SCHEMA # <<<<<<<< THIS COMES FORM MY CONFIG SYSTEM config.py
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config
# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
fileConfig(config.config_file_name)
# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.
# update database URL
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL)
def run_migrations_offline() -> None:
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
url = config.get_main_option("sqlalchemy.url")
context.configure(
url=url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)
with context.begin_transaction():
context.run_migrations()
def do_run_migrations(connection: Connection) -> None:
context.configure(connection=connection, target_metadata=target_metadata)
with context.begin_transaction():
# Set search path to your schema
context.execute(text(f'CREATE SCHEMA IF NOT EXISTS {SCHEMA}'))
# Create schema if it doesn't exist
context.execute(text(f'SET search_path TO {SCHEMA}, public'))
context.run_migrations()
async def run_async_migrations() -> None:
"""In this scenario we need to create an Engine
and associate a connection with the context.
"""
connectable = async_engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
async with connectable.connect() as connection:
await connection.run_sync(do_run_migrations)
await connectable.dispose()
def run_migrations_online() -> None:
"""Run migrations in 'online' mode."""
asyncio.run(run_async_migrations())
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()
Cheers hope it helps
here is my solution for work properly with PostgreSQL schema using asyncpg and asyncio in sqlalchemy.
what counts is like the official document and above said
configure filter object from specified schema