sqlalchemy / alembic

A database migrations tool for SQLAlchemy.
MIT License
2.86k stars 247 forks source link

`server_default` with custom function element #641

Closed sdg32 closed 4 years ago

sdg32 commented 4 years ago

I use custom function element to support multiple database:

from sqlalchemy import DateTime
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.functions import FunctionElement

class GetNow(FunctionElement):
    type = DateTime()

@compiles(GetNow, 'mssql')
def mssql_get_now(element, compiler, **kwargs):
    return 'GETDATE()'

@compiles(GetNow, 'oracle')
def oracle_get_now(element, compiler, **kwargs):
    return 'SYSDATE'

The model class:

class User(db.Model):
    __bind_key__ = 'core'  # MSSQL dialect
    __tablename__ = 'user'

    id = db.Column(db.Integer(), primary_key=True)
    created_at = db.Column(db.DateTime(), nullable=False, server_default=GetNow())

When use flask-migrate generate migration scripts, GetNow alaways render SYSDATE. So I custom the render_item function and found the render_item always receive the last engine dialect:

def render_item(type_, obj, autogen_context):
    if type_ == 'server_default':
        if isinstance(obj.arg, GetNow):
            return "db.text('%s')" % obj.arg.compile(dialect=autogen_context.dialect)

    return False

def run_migrations_online():
    ...
    for name, rec in engines.items():
        context.configure(
            ...
            render_item=render_item,
        )
        context.run_migrations(engine_name=name)

Is this a bug or I missing some config?

Versions:

flask-sqlalchemy config:

SQLALCHEMY_DATABASE_URI = 'sqlite:///:memory:'  # not in use
SQLALCHEMY_BINDS = {
    'core': 'mssql+pyodbc://***',
    'spring': 'oracle+cx_oracle://***',
}
zzzeek commented 4 years ago

hi there -

can you clarify what "the last engine" means? are you using multiple engines on one run at the same time and if so can you your env.py please?

sdg32 commented 4 years ago

@zzzeek Yes, I run at the same time with multiple engines.

The env.py file is generated by flask-migrate:

import logging
from logging.config import fileConfig

from alembic import context
from flask import current_app
from sqlalchemy import engine_from_config
from sqlalchemy import MetaData
from sqlalchemy import pool

from flask_sqla import sqltypes as custom_types

USE_TWOPHASE = False
config = context.config
fileConfig(config.config_file_name)
logger = logging.getLogger('alembic.env')

config.set_main_option(
    'sqlalchemy.url',
    current_app.config.get('SQLALCHEMY_DATABASE_URI').replace('%', '%%'),
)
bind_names = []
for name, url in current_app.config.get('SQLALCHEMY_BINDS').items():
    context.config.set_section_option(name, 'sqlalchemy.url', url.replace('%', '%%'))
    bind_names.append(name)
target_metadata = current_app.extensions['migrate'].db.metadata

def get_metadata(bind):
    """Return the metadata for a bind."""
    if bind == '':
        bind = None
    m = MetaData()
    for t in target_metadata.tables.values():
        if t.info.get('bind_key') == bind:
            t.tometadata(m)
    return m

def render_custom_item(type_, obj, autogen_context):
    if type_ in ('server_default', 'server_onupdate'):
        print(obj, f'dialect={autogen_context.dialect.name}')
    return False

def run_migrations_offline():
    engines = {
        '': {
            'url': context.config.get_main_option('sqlalchemy.url')
        }
    }
    for name in bind_names:
        engines[name] = rec = {}
        rec['url'] = context.config.get_section_option(name, "sqlalchemy.url")

    for name, rec in engines.items():
        logger.info("Migrating database %s" % (name or '<default>'))
        file_ = "%s.sql" % name
        logger.info("Writing output to %s" % file_)
        with open(file_, 'w') as buffer:
            context.configure(
                url=rec['url'],
                output_buffer=buffer,
                target_metadata=get_metadata(name),
                literal_binds=True,
            )
            with context.begin_transaction():
                context.run_migrations(engine_name=name)

def run_migrations_online():
    def process_revision_directives(context, revision, directives):
        if getattr(config.cmd_opts, 'autogenerate', False):
            script = directives[0]
            if len(script.upgrade_ops_list) >= len(bind_names) + 1:
                empty = True
                for upgrade_ops in script.upgrade_ops_list:
                    if not upgrade_ops.is_empty():
                        empty = False
                if empty:
                    directives[:] = []
                    logger.info('No changes in schema detected.')

    # for the direct-to-DB use case, start a transaction on all
    # engines, then run all migrations, then commit all transactions.
    engines = {
        '': {
            'engine': engine_from_config(
                config.get_section(config.config_ini_section),
                prefix='sqlalchemy.',
                poolclass=pool.NullPool,
            )
        }
    }
    for name in bind_names:
        engines[name] = rec = {}
        rec['engine'] = engine_from_config(
            context.config.get_section(name),
            prefix='sqlalchemy.',
            poolclass=pool.NullPool,
        )

    for name, rec in engines.items():
        engine = rec['engine']
        rec['connection'] = conn = engine.connect()

        if USE_TWOPHASE:
            rec['transaction'] = conn.begin_twophase()
        else:
            rec['transaction'] = conn.begin()

    try:
        for name, rec in engines.items():
            logger.info("Migrating database %s" % (name or '<default>'))
            context.configure(
                connection=rec['connection'],
                upgrade_token="%s_upgrades" % name,
                downgrade_token="%s_downgrades" % name,
                target_metadata=get_metadata(name),
                process_revision_directives=process_revision_directives,
                render_item=render_custom_item,
                **current_app.extensions['migrate'].configure_args,
            )
            context.run_migrations(engine_name=name)

        if USE_TWOPHASE:
            for rec in engines.values():
                rec['transaction'].prepare()

        for rec in engines.values():
            rec['transaction'].commit()
    except:
        for rec in engines.values():
            rec['transaction'].rollback()
        raise
    finally:
        for rec in engines.values():
            rec['connection'].close()

if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()

And the flask-sqlalchemy config:

SQLALCHEMY_DATABASE_URI = 'sqlite:///:memory:'  # not in use
SQLALCHEMY_BINDS = {
    'core': 'mssql+pyodbc://***',
    'spring': 'oracle+cx_oracle://***',
}
SQLALCHEMY_TRACK_MODIFICATIONS = False

When I run command: flask db migrate, the output:

INFO  [alembic.env] Migrating database <default>
INFO  [alembic.runtime.migration] Context impl SQLiteImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.env] Migrating database core
INFO  [alembic.runtime.migration] Context impl MSSQLImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.autogenerate.compare] Detected added table 'user'
INFO  [alembic.env] Migrating database spring
INFO  [alembic.runtime.migration] Context impl OracleImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
DefaultClause(<GetNow at 0x286c944fac0; GetNow object>, for_update=False) dialect=oracle
Generating migrations\versions\b4e6b461df58_.py ...  done

When I change the flask-sqlalchemy binds order:

SQLALCHEMY_BINDS = {
    'spring': 'oracle+cx_oracle://***',
    'core': 'mssql+pyodbc://***',
}

The migration output will be:

INFO  [alembic.env] Migrating database <default>
INFO  [alembic.runtime.migration] Context impl SQLiteImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.env] Migrating database spring
INFO  [alembic.runtime.migration] Context impl OracleImpl.
INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
INFO  [alembic.env] Migrating database core
INFO  [alembic.runtime.migration] Context impl MSSQLImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.autogenerate.compare] Detected added table 'user'
DefaultClause(<GetNow at 0x1a4b3d8da90; GetNow object>, for_update=False) dialect=mssql
Generating migrations\versions\f7c9ce3870c9_.py ...  done

Seems that render_item only called once after compare.

zzzeek commented 4 years ago

OK so yes, the rendering step for the migration structure that's generated is run at the end, and it makes use of the last autogenerate context for what it needs, which does not actually include the dialect at all, however your code is digging in and finding it, which is not an assumption made by the current autogenerate architecture. There is an expectation that there is no dialect-specific text that would be needed for the render. if you wanted it to work this way (which it does not have to), you would need to carry along the dialect information that you need inside the MigrateOps structure and I'd have to think about what hooks you would need to do that, but let's look at how render_item was intended to be used:

def render_item(type_, obj, autogen_context):
    if type_ == 'server_default':
        if isinstance(obj.arg, GetNow):
            return "GetNow()"

    return False

that is, "render_item" does not generate SQL, it generates Python code. So here, you have your own GetNow() structure that is the Python code you want; just generate that. Your @compile hook is not intended to be invoked within the autogenerate step, it is invoked in the "alembic upgrade" step, which is where all the other SQL is compiled. hope this helps.

sdg32 commented 4 years ago

Thanks, that help me a lot.

NChandan commented 1 year ago

@zzzeek render_item was working fine for the first time but when I upgrade the table it was trowing error

def render_item(type_, obj, autogen_context):
    if type_ == 'server_default' and obj:
        if isinstance(obj.arg, CustomUtcNow):
            return "CustomUtcNow"
        else:
            return False

    return False
#model.py

class CustomUtcNow(expression.FunctionElement):
    # type = DateTime()
    # inherit_cache = True
    name = "custom_now"

@compiles(CustomUtcNow, "postgresql")
def pg_utcnow(element, compiler, **kw):
    return "timezone('utc'::text, now())"

@compiles(CustomUtcNow, "sqlite")
def ms_utcnow(element, compiler, **kw):
    return "now()"

class ActivityLog(Base):
    __tablename__ = "activity_logs"

    log_id = Column(Integer, primary_key=True)
    project_id = Column(Integer)
    action_name = Column(String(100), nullable=False)
    action_by_module = Column(String(500), nullable=False)
    created_on = Column(DateTime, default=datetime.utcnow, server_default=CustomUtcNow())
    created_by = Column(Integer)

error:

Traceback (most recent call last):
  File "D:\BA\Athena-BA-AdminPortal-API\venv\lib\site-packages\sqlalchemy\engine\base.py",
 line 1900, in _execute_context
    self.dialect.do_execute(
  File "D:\BA\Athena-BA-AdminPortal-API\venv\lib\site-packages\sqlalchemy\engine\default.p
y", line 736, in do_execute
    cursor.execute(statement, parameters)
psycopg2.errors.UndefinedFunction: function customutcnow() does not exist
LINE 1: SELECT timezone('utc'::text, now()) = CustomUtcNow()
                                              ^
HINT:  No function matches the given name and argument types. You might need to add explic
it type casts.

I don't understand what's going wrong?

zzzeek commented 1 year ago

I would need more context to understand the problem including complete stack trace. please start a new discussion at https://github.com/sqlalchemy/alembic/discussions

CaselIT commented 1 year ago

I think it was solved, since a new issue was created and then closed: https://github.com/sqlalchemy/alembic/issues/1141

lukesmurray commented 1 year ago

I think I ran into the same issue. Alembic treats the output of render_item as both python code to be added to tables and compiled code to be rendered in sql statements.

When rendering a custom function type for a table we want the output to be something like app.migration_types.custom_func.

However internally that type of output will cause an error. For example when alembic renders a field for comparison it uses the user_defined_render.

def _render_server_default_for_compare(
    metadata_default: Optional[Any],
    metadata_col: Column,
    autogen_context: AutogenContext,
) -> Optional[str]:
    rendered = _user_defined_render(
        "server_default", metadata_default, autogen_context
    )

Later in the postgres server_default_comparison, and maybe in others, the rendered item is compared in sql.

        return not self.connection.scalar(
            text(
                "SELECT %s = %s"
                % (conn_col_default, rendered_metadata_default)
            )
        )

This comparison is the source of the reported error. The output should be compiled sql but its just the python path to the user's function.

SELECT timezone('utc'::text, now()) = CustomUtcNow()

Currently trying to fix this with an overload for compare_server_default and it seems possible but I think this is a bit of a flaw in either the instructions for rendering custom types or the logic for comparing custom types.

zzzeek commented 1 year ago

yes I did notice this today.

i think the fix is to just remove it

diff --git a/alembic/autogenerate/compare.py b/alembic/autogenerate/compare.py
index 8301e34..828a4cd 100644
--- a/alembic/autogenerate/compare.py
+++ b/alembic/autogenerate/compare.py
@@ -21,7 +21,6 @@ from sqlalchemy import types as sqltypes
 from sqlalchemy.util import OrderedSet

 from alembic.ddl.base import _fk_spec
-from .render import _user_defined_render
 from .. import util
 from ..operations import ops
 from ..util import sqla_compat
@@ -1003,11 +1002,6 @@ def _render_server_default_for_compare(
     metadata_col: Column,
     autogen_context: AutogenContext,
 ) -> Optional[str]:
-    rendered = _user_defined_render(
-        "server_default", metadata_default, autogen_context
-    )
-    if rendered is not False:
-        return rendered

     if isinstance(metadata_default, sa_schema.DefaultClause):
         if isinstance(metadata_default.arg, str):

please open an issue

sqla-tester commented 1 year ago

Mike Bayer referenced this issue:

dont use server_default render_item for SQL compare https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/4453