JezK
Edit File: _autogen_fixtures.py
from __future__ import annotations from typing import Any from typing import Dict from typing import Set from sqlalchemy import CHAR from sqlalchemy import CheckConstraint from sqlalchemy import Column from sqlalchemy import event from sqlalchemy import ForeignKey from sqlalchemy import Index from sqlalchemy import inspect from sqlalchemy import Integer from sqlalchemy import MetaData from sqlalchemy import Numeric from sqlalchemy import String from sqlalchemy import Table from sqlalchemy import Text from sqlalchemy import text from sqlalchemy import UniqueConstraint from ... import autogenerate from ... import util from ...autogenerate import api from ...ddl.base import _fk_spec from ...migration import MigrationContext from ...operations import ops from ...testing import config from ...testing import eq_ from ...testing.env import clear_staging_env from ...testing.env import staging_env names_in_this_test: Set[Any] = set() @event.listens_for(Table, "after_parent_attach") def new_table(table, parent): names_in_this_test.add(table.name) def _default_include_object(obj, name, type_, reflected, compare_to): if type_ == "table": return name in names_in_this_test else: return True _default_object_filters: Any = _default_include_object _default_name_filters: Any = None class ModelOne: __requires__ = ("unique_constraint_reflection",) schema: Any = None @classmethod def _get_db_schema(cls): schema = cls.schema m = MetaData(schema=schema) Table( "user", m, Column("id", Integer, primary_key=True), Column("name", String(50)), Column("a1", Text), Column("pw", String(50)), Index("pw_idx", "pw"), ) Table( "address", m, Column("id", Integer, primary_key=True), Column("email_address", String(100), nullable=False), ) Table( "order", m, Column("order_id", Integer, primary_key=True), Column( "amount", Numeric(8, 2), nullable=False, server_default=text("0"), ), CheckConstraint("amount >= 0", name="ck_order_amount"), ) Table( "extra", m, Column("x", CHAR), Column("uid", Integer, ForeignKey("user.id")), ) return m @classmethod def _get_model_schema(cls): schema = cls.schema m = MetaData(schema=schema) Table( "user", m, Column("id", Integer, primary_key=True), Column("name", String(50), nullable=False), Column("a1", Text, server_default="x"), ) Table( "address", m, Column("id", Integer, primary_key=True), Column("email_address", String(100), nullable=False), Column("street", String(50)), UniqueConstraint("email_address", name="uq_email"), ) Table( "order", m, Column("order_id", Integer, primary_key=True), Column( "amount", Numeric(10, 2), nullable=True, server_default=text("0"), ), Column("user_id", Integer, ForeignKey("user.id")), CheckConstraint("amount > -1", name="ck_order_amount"), ) Table( "item", m, Column("id", Integer, primary_key=True), Column("description", String(100)), Column("order_id", Integer, ForeignKey("order.order_id")), CheckConstraint("len(description) > 5"), ) return m class _ComparesFKs: def _assert_fk_diff( self, diff, type_, source_table, source_columns, target_table, target_columns, name=None, conditional_name=None, source_schema=None, onupdate=None, ondelete=None, initially=None, deferrable=None, ): # the public API for ForeignKeyConstraint was not very rich # in 0.7, 0.8, so here we use the well-known but slightly # private API to get at its elements ( fk_source_schema, fk_source_table, fk_source_columns, fk_target_schema, fk_target_table, fk_target_columns, fk_onupdate, fk_ondelete, fk_deferrable, fk_initially, ) = _fk_spec(diff[1]) eq_(diff[0], type_) eq_(fk_source_table, source_table) eq_(fk_source_columns, source_columns) eq_(fk_target_table, target_table) eq_(fk_source_schema, source_schema) eq_(fk_onupdate, onupdate) eq_(fk_ondelete, ondelete) eq_(fk_initially, initially) eq_(fk_deferrable, deferrable) eq_([elem.column.name for elem in diff[1].elements], target_columns) if conditional_name is not None: if conditional_name == "servergenerated": fks = inspect(self.bind).get_foreign_keys(source_table) server_fk_name = fks[0]["name"] eq_(diff[1].name, server_fk_name) else: eq_(diff[1].name, conditional_name) else: eq_(diff[1].name, name) class AutogenTest(_ComparesFKs): def _flatten_diffs(self, diffs): for d in diffs: if isinstance(d, list): for fd in self._flatten_diffs(d): yield fd else: yield d @classmethod def _get_bind(cls): return config.db configure_opts: Dict[Any, Any] = {} @classmethod def setup_class(cls): staging_env() cls.bind = cls._get_bind() cls.m1 = cls._get_db_schema() cls.m1.create_all(cls.bind) cls.m2 = cls._get_model_schema() @classmethod def teardown_class(cls): cls.m1.drop_all(cls.bind) clear_staging_env() def setUp(self): self.conn = conn = self.bind.connect() ctx_opts = { "compare_type": True, "compare_server_default": True, "target_metadata": self.m2, "upgrade_token": "upgrades", "downgrade_token": "downgrades", "alembic_module_prefix": "op.", "sqlalchemy_module_prefix": "sa.", "include_object": _default_object_filters, "include_name": _default_name_filters, } if self.configure_opts: ctx_opts.update(self.configure_opts) self.context = context = MigrationContext.configure( connection=conn, opts=ctx_opts ) self.autogen_context = api.AutogenContext(context, self.m2) def tearDown(self): self.conn.close() def _update_context( self, object_filters=None, name_filters=None, include_schemas=None ): if include_schemas is not None: self.autogen_context.opts["include_schemas"] = include_schemas if object_filters is not None: self.autogen_context._object_filters = [object_filters] if name_filters is not None: self.autogen_context._name_filters = [name_filters] return self.autogen_context class AutogenFixtureTest(_ComparesFKs): def _fixture( self, m1, m2, include_schemas=False, opts=None, object_filters=_default_object_filters, name_filters=_default_name_filters, return_ops=False, max_identifier_length=None, ): if max_identifier_length: dialect = self.bind.dialect existing_length = dialect.max_identifier_length dialect.max_identifier_length = ( dialect._user_defined_max_identifier_length ) = max_identifier_length try: self._alembic_metadata, model_metadata = m1, m2 for m in util.to_list(self._alembic_metadata): m.create_all(self.bind) with self.bind.connect() as conn: ctx_opts = { "compare_type": True, "compare_server_default": True, "target_metadata": model_metadata, "upgrade_token": "upgrades", "downgrade_token": "downgrades", "alembic_module_prefix": "op.", "sqlalchemy_module_prefix": "sa.", "include_object": object_filters, "include_name": name_filters, "include_schemas": include_schemas, } if opts: ctx_opts.update(opts) self.context = context = MigrationContext.configure( connection=conn, opts=ctx_opts ) autogen_context = api.AutogenContext(context, model_metadata) uo = ops.UpgradeOps(ops=[]) autogenerate._produce_net_changes(autogen_context, uo) if return_ops: return uo else: return uo.as_diffs() finally: if max_identifier_length: dialect = self.bind.dialect dialect.max_identifier_length = ( dialect._user_defined_max_identifier_length ) = existing_length def setUp(self): staging_env() self.bind = config.db def tearDown(self): if hasattr(self, "_alembic_metadata"): for m in util.to_list(self._alembic_metadata): m.drop_all(self.bind) clear_staging_env()