OwlCyberSecurity - MANAGER

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):
                yield from self._flatten_diffs(d)
            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()