alembic-utils-extended 0.0.0__py3-none-any.whl

alembic_utils_extended/depends.py

@@ -0,0 +1,145 @@
+import logging
+from contextlib import contextmanager
+from typing import Generator, List
+
+from sqlalchemy import exc as sqla_exc
+from sqlalchemy.orm import Session
+
+from alembic_utils_extended.simulate import simulate_entity
+
+logger = logging.getLogger(__name__)
+
+
+def solve_resolution_order(sess: Session, entities):
+    """Solve for an entity resolution order that increases the probability that
+    a migration will suceed if, for example, two new views are created and one
+    refers to the other
+
+    This strategy will only solve for simple cases
+    """
+
+    resolved = []
+
+    # Resolve the entities with 0 dependencies first (faster)
+    logger.info("Resolving entities with no dependencies")
+    for entity in entities:
+        try:
+            with simulate_entity(sess, entity):
+                resolved.append(entity)
+        except (sqla_exc.ProgrammingError, sqla_exc.InternalError) as exc:
+            continue
+
+    # Resolve entities with possible dependencies
+    logger.info("Resolving entities with dependencies. This may take a minute")
+    for _ in range(len(entities)):
+        n_resolved = len(resolved)
+
+        for entity in entities:
+            if entity in resolved:
+                continue
+
+            try:
+                with simulate_entity(sess, entity, dependencies=resolved):
+                    resolved.append(entity)
+            except (sqla_exc.ProgrammingError, sqla_exc.InternalError):
+                continue
+
+        if len(resolved) == n_resolved:
+            # No new entities resolved in the last iteration. Exit
+            break
+
+    for entity in entities:
+        if entity not in resolved:
+            resolved.append(entity)
+
+    return resolved
+
+
+@contextmanager
+def recreate_dropped(connection) -> Generator[Session, None, None]:
+    """Recreate any dropped all ReplaceableEntities that were dropped within block
+
+    This is useful for making cascading updates. For example, updating a table's column type when it has dependent views.
+
+    def upgrade() -> None:
+
+        my_view = PGView(...)
+
+        with recreate_dropped(op.get_bind()) as conn:
+
+            op.drop_entity(my_view)
+
+            # change an integer column to a bigint
+            op.alter_column(
+                table_name="account",
+                column_name="id",
+                schema="public"
+                type_=sa.BIGINT()
+                existing_type=sa.Integer(),
+            )
+    """
+    from alembic_utils_extended.pg_function import PGFunction
+    from alembic_utils_extended.pg_materialized_view import PGMaterializedView
+    from alembic_utils_extended.pg_trigger import PGTrigger
+    from alembic_utils_extended.pg_view import PGView
+    from alembic_utils_extended.replaceable_entity import ReplaceableEntity
+
+    # Do not include permissions here e.g. PGGrantTable. If columns granted to users are dropped, it will cause an error
+
+    def collect_all_db_entities(sess: Session) -> List[ReplaceableEntity]:
+        """Collect all entities from the database"""
+
+        return [
+            *PGFunction.from_database(sess, "%"),
+            *PGTrigger.from_database(sess, "%"),
+            *PGView.from_database(sess, "%"),
+            *PGMaterializedView.from_database(sess, "%"),
+        ]
+
+    sess = Session(bind=connection)
+
+    # All existing entities, before the upgrade
+    before = collect_all_db_entities(sess)
+
+    # In the yield, do a
+    #     op.drop_entity(my_mat_view, cascade=True)
+    #     op.create_entity(my_mat_view)
+    try:
+        yield sess
+    except:
+        sess.rollback()
+        raise
+
+    # All existing entities, after the upgrade
+    after = collect_all_db_entities(sess)
+    after_identities = {x.identity for x in after}
+
+    # Entities that were not impacted, or that we have "recovered"
+    resolved = []
+    unresolved = []
+
+    # First, ignore the ones that were not impacted by the upgrade
+    for ent in before:
+        if ent.identity in after_identities:
+            resolved.append(ent)
+        else:
+            unresolved.append(ent)
+
+    # Attempt to find an acceptable order of creation for the unresolved entities
+    ordered_unresolved = solve_resolution_order(sess, unresolved)
+
+    # Attempt to recreate the missing entities in the specified order
+    for ent in ordered_unresolved:
+        sess.execute(ent.to_sql_statement_create())
+
+    # Sanity check that everything is now fine
+    sanity_check = collect_all_db_entities(sess)
+    # Fail and rollback if the sanity check is wrong
+    try:
+        assert len(before) == len(sanity_check)
+    except:
+        sess.rollback()
+        raise
+
+    # Close out the session
+    sess.commit()

alembic_utils_extended/exceptions.py

@@ -0,0 +1,18 @@
+class AlembicUtilsException(Exception):
+    """Base exception for AlembicUtils package"""
+
+
+class SQLParseFailure(AlembicUtilsException):
+    """An entity could not be parsed"""
+
+
+class FailedToGenerateComparable(AlembicUtilsException):
+    """Failed to generate a comparable entity"""
+
+
+class UnreachableException(AlembicUtilsException):
+    """An exception no one should ever see"""
+
+
+class BadInputException(AlembicUtilsException):
+    """Invalid user input"""

alembic_utils_extended/experimental/__init__.py

@@ -0,0 +1,6 @@
+from alembic_utils_extended.experimental._collect_instances import (
+    collect_instances,
+    collect_subclasses,
+)
+
+__all__ = ["collect_instances", "collect_subclasses"]

alembic_utils_extended/experimental/_collect_instances.py

@@ -0,0 +1,90 @@
+import importlib
+import os
+from pathlib import Path
+from types import ModuleType
+from typing import Generator, List, Type, TypeVar
+
+from flupy import walk_files
+
+T = TypeVar("T")
+
+
+def walk_modules(module: ModuleType) -> Generator[ModuleType, None, None]:
+    """Recursively yield python import paths to submodules in *module*
+
+    Example:
+        module_iter = walk_modules(alembic_utils_extended)
+
+        for module_path in module_iter:
+            print(module_path)
+
+        # alembic_utils_extended.exceptions
+        # alembic_utils_extended.on_entity_mixin
+        # ...
+    """
+    top_module = module
+    top_path = Path(top_module.__path__[0])
+    top_path_absolute = top_path.resolve()
+
+    directories = (
+        walk_files(str(top_path_absolute))
+        .filter(lambda x: x.endswith(".py"))
+        .map(Path)
+        .group_by(lambda x: x.parent)
+        .map(lambda x: (x[0], x[1].collect()))
+    )
+
+    for base_path, files in directories:
+        if str(base_path / "__init__.py") in [str(x) for x in files]:
+            for module_path in files:
+                if "__init__.py" not in str(module_path):
+
+                    # Example: elt.settings
+                    module_import_path = str(module_path)[
+                        len(str(top_path_absolute)) - len(top_module.__name__) :
+                    ].replace(os.path.sep, ".")[:-3]
+
+                    module = importlib.import_module(module_import_path)
+                    yield module
+
+
+def collect_instances(module: ModuleType, class_: Type[T]) -> List[T]:
+    """Collect all instances of *class_* defined in *module*
+
+    Note: Will import all submodules in *module*. Beware of import side effects
+    """
+
+    found: List[T] = []
+
+    for module_ in walk_modules(module):
+
+        for _, variable in module_.__dict__.items():
+
+            if isinstance(variable, class_):
+                # Ensure variable is not a subclass
+                if variable.__class__ == class_:
+                    found.append(variable)
+    return found
+
+
+def collect_subclasses(module: ModuleType, class_: Type[T]) -> List[Type[T]]:
+    """Collect all subclasses of *class_* currently imported or defined in *module*
+
+    Note: Will import all submodules in *module*. Beware of import side effects
+    """
+
+    found: List[Type[T]] = []
+
+    for module_ in walk_modules(module):
+
+        for _, variable in module_.__dict__.items():
+            try:
+                if issubclass(variable, class_) and not class_ == variable:
+                    found.append(variable)
+            except TypeError:
+                # argument 2 to issubclass must be a class ....
+                pass
+
+    imported: List[Type[T]] = list(class_.__subclasses__())  # type: ignore
+
+    return list(set(found + imported))

alembic_utils_extended/on_entity_mixin.py

@@ -0,0 +1,51 @@
+from typing import TYPE_CHECKING
+
+from alembic_utils_extended.statement import coerce_to_unquoted
+
+if TYPE_CHECKING:
+    from alembic_utils_extended.replaceable_entity import ReplaceableEntity
+
+    _Base = ReplaceableEntity
+else:
+    _Base = object
+
+
+class OnEntityMixin(_Base):
+    """Mixin to ReplaceableEntity providing setup for entity types requiring an "ON" clause"""
+
+    def __init__(self, schema: str, signature: str, definition: str, on_entity: str):
+        super().__init__(schema=schema, signature=signature, definition=definition)
+
+        if "." not in on_entity:
+            on_entity = "public." + on_entity
+
+        # Guarenteed to have a schema
+        self.on_entity = coerce_to_unquoted(on_entity)
+
+    @property
+    def identity(self) -> str:
+        """A string that consistently and globally identifies a function
+
+        Overriding default to add the "on table" clause
+        """
+        return f"{self.__class__.__name__}: {self.schema}.{self.signature} {self.on_entity}"
+
+    def render_self_for_migration(self, omit_definition=False) -> str:
+        """Render a string that is valid python code to reconstruct self in a migration"""
+        var_name = self.to_variable_name()
+        class_name = self.__class__.__name__
+        escaped_definition = self.definition if not omit_definition else "# not required for op"
+
+        return f"""{var_name} = {class_name}(
+    schema="{self.schema}",
+    signature="{self.signature}",
+    on_entity="{self.on_entity}",
+    definition={repr(escaped_definition)}
+)\n"""
+
+    def to_variable_name(self) -> str:
+        """A deterministic variable name based on PGFunction's contents"""
+        schema_name = self.schema.lower()
+        object_name = self.signature.split("(")[0].strip().lower()
+        _, _, unqualified_entity_name = self.on_entity.lower().partition(".")
+        return f"{schema_name}_{unqualified_entity_name}_{object_name}"

alembic_utils_extended/pg_extension.py

@@ -0,0 +1,83 @@
+# pylint: disable=unused-argument,invalid-name,line-too-long
+
+
+from typing import Generator
+
+from sqlalchemy import text as sql_text
+from sqlalchemy.sql.elements import TextClause
+
+from alembic_utils_extended.replaceable_entity import ReplaceableEntity
+from alembic_utils_extended.statement import (
+    coerce_to_unquoted,
+    normalize_whitespace,
+)
+
+
+class PGExtension(ReplaceableEntity):
+    """A PostgreSQL Extension compatible with `alembic revision --autogenerate`
+
+    **Parameters:**
+
+    * **schema** - *str*: A SQL schema name
+    * **signature** - *str*: A PostgreSQL extension's name
+    """
+
+    type_ = "extension"
+
+    def __init__(self, schema: str, signature: str):
+        self.schema: str = coerce_to_unquoted(normalize_whitespace(schema))
+        self.signature: str = coerce_to_unquoted(normalize_whitespace(signature))
+        # Include schema in definition since extensions can only exist once per
+        # database and we want to detect schema changes and emit alter schema
+        self.definition: str = f"{self.__class__.__name__}: {self.schema} {self.signature}"
+
+    def to_sql_statement_create(self) -> TextClause:
+        """Generates a SQL "create extension" statement"""
+        return sql_text(f'CREATE EXTENSION "{self.signature}" WITH SCHEMA {self.literal_schema};')
+
+    def to_sql_statement_drop(self, cascade=False) -> TextClause:
+        """Generates a SQL "drop extension" statement"""
+        cascade = "CASCADE" if cascade else ""
+        return sql_text(f'DROP EXTENSION "{self.signature}" {cascade}')
+
+    def to_sql_statement_create_or_replace(self) -> Generator[TextClause, None, None]:
+        """Generates SQL equivalent to "create or replace" statement"""
+        raise NotImplementedError()
+
+    @property
+    def identity(self) -> str:
+        """A string that consistently and globally identifies an extension"""
+        # Extensions may only be installed once per db, schema is not a
+        # component of identity
+        return f"{self.__class__.__name__}: {self.signature}"
+
+    def render_self_for_migration(self, omit_definition=False) -> str:
+        """Render a string that is valid python code to reconstruct self in a migration"""
+        var_name = self.to_variable_name()
+        class_name = self.__class__.__name__
+
+        return f"""{var_name} = {class_name}(
+    schema="{self.schema}",
+    signature="{self.signature}"
+)\n"""
+
+    @classmethod
+    def from_database(cls, sess, schema):
+        """Get a list of all extensions defined in the db"""
+        sql = sql_text(
+            f"""
+        select
+            np.nspname schema_name,
+            ext.extname extension_name
+        from
+            pg_extension ext
+            join pg_namespace np
+                on ext.extnamespace = np.oid
+        where
+            np.nspname not in ('pg_catalog')
+            and np.nspname like :schema;
+        """
+        )
+        rows = sess.execute(sql, {"schema": schema}).fetchall()
+        db_exts = [cls(x[0], x[1]) for x in rows]
+        return db_exts

alembic_utils_extended/pg_function.py

@@ -0,0 +1,169 @@
+# pylint: disable=unused-argument,invalid-name,line-too-long
+from typing import List
+
+from parse import parse
+from sqlalchemy import text as sql_text
+
+from alembic_utils_extended.exceptions import SQLParseFailure
+from alembic_utils_extended.replaceable_entity import ReplaceableEntity
+from alembic_utils_extended.statement import (
+    escape_colon_for_plpgsql,
+    escape_colon_for_sql,
+    normalize_whitespace,
+    strip_terminating_semicolon,
+)
+
+
+class PGFunction(ReplaceableEntity):
+    """A PostgreSQL Function compatible with `alembic revision --autogenerate`
+
+    **Parameters:**
+
+    * **schema** - *str*: A SQL schema name
+    * **signature** - *str*: A SQL function's call signature
+    * **definition** - *str*:  The remainig function body and identifiers
+    """
+
+    type_ = "function"
+
+    def __init__(self, schema: str, signature: str, definition: str):
+        super().__init__(schema, signature, definition)
+        # Detect if function uses plpgsql and update escaping rules to not escape ":="
+        is_plpgsql: bool = "language plpgsql" in normalize_whitespace(definition).lower().replace(
+            "'", ""
+        )
+        escaping_callable = escape_colon_for_plpgsql if is_plpgsql else escape_colon_for_sql
+        # Override definition with correct escaping rules
+        self.definition: str = escaping_callable(strip_terminating_semicolon(definition))
+
+    @classmethod
+    def from_sql(cls, sql: str) -> "PGFunction":
+        """Create an instance instance from a SQL string"""
+        template = "create{}function{:s}{schema}.{signature}{:s}returns{:s}{definition}"
+        result = parse(template, sql.strip(), case_sensitive=False)
+        if result is not None:
+            # remove possible quotes from signature
+            raw_signature = result["signature"]
+            signature = (
+                "".join(raw_signature.split('"', 2))
+                if raw_signature.startswith('"')
+                else raw_signature
+            )
+            return cls(
+                schema=result["schema"],
+                signature=signature,
+                definition="returns " + result["definition"],
+            )
+        raise SQLParseFailure(f'Failed to parse SQL into PGFunction """{sql}"""')
+
+    @property
+    def literal_signature(self) -> str:
+        """Adds quoting around the functions name when emitting SQL statements
+
+        e.g.
+        'toUpper(text) returns text' -> '"toUpper"(text) returns text'
+        """
+        # May already be quoted if loading from database or SQL file
+        name, remainder = self.signature.split("(", 1)
+        return '"' + name.strip() + '"(' + remainder
+
+    def to_sql_statement_create(self):
+        """Generates a SQL "create function" statement for PGFunction"""
+        return sql_text(
+            f"CREATE FUNCTION {self.literal_schema}.{self.literal_signature} {self.definition}"
+        )
+
+    def to_sql_statement_drop(self, cascade=False):
+        """Generates a SQL "drop function" statement for PGFunction"""
+        cascade = "cascade" if cascade else ""
+        template = "{function_name}({parameters})"
+        result = parse(template, self.signature, case_sensitive=False)
+        try:
+            function_name = result["function_name"].strip()
+            parameters_str = result["parameters"].strip()
+        except TypeError:
+            # Did not match, NoneType is not scriptable
+            result = parse("{function_name}()", self.signature, case_sensitive=False)
+            function_name = result["function_name"].strip()
+            parameters_str = ""
+
+        # NOTE: Will fail if a text field has a default and that deafult contains a comma...
+        parameters: List[str] = parameters_str.split(",")
+        parameters = [x[: len(x.lower().split("default")[0])] for x in parameters]
+        parameters = [x.strip() for x in parameters]
+        drop_params = ", ".join(parameters)
+        return sql_text(
+            f'DROP FUNCTION {self.literal_schema}."{function_name}"({drop_params}) {cascade}'
+        )
+
+    def to_sql_statement_create_or_replace(self):
+        """Generates a SQL "create or replace function" statement for PGFunction"""
+        yield sql_text(
+            f"CREATE OR REPLACE FUNCTION {self.literal_schema}.{self.literal_signature} {self.definition}"
+        )
+
+    @classmethod
+    def from_database(cls, sess, schema):
+        """Get a list of all functions defined in the db"""
+
+        # Prior to postgres 11, pg_proc had different columns
+        # https://github.com/candidhealth/alembic-utils-extended/issues/12
+        PG_GTE_11 = """
+            and p.prokind = 'f'
+        """
+
+        PG_LT_11 = """
+            and not p.proisagg
+            and not p.proiswindow
+        """
+
+        # Retrieve the postgres server version e.g. 90603 for 9.6.3 or 120003 for 12.3
+        pg_version_str = sess.execute(sql_text("show server_version_num")).fetchone()[0]
+        pg_version = int(pg_version_str)
+
+        sql = sql_text(
+            f"""
+        with extension_functions as (
+            select
+                objid as extension_function_oid
+            from
+                pg_depend
+            where
+                -- depends on an extension
+                deptype='e'
+                -- is a proc/function
+                and classid = 'pg_proc'::regclass
+        )
+
+        select
+            n.nspname as function_schema,
+            p.proname as function_name,
+            pg_get_function_arguments(p.oid) as function_arguments,
+            case
+                when l.lanname = 'internal' then p.prosrc
+                else pg_get_functiondef(p.oid)
+            end as create_statement,
+            t.typname as return_type,
+            l.lanname as function_language
+        from
+            pg_proc p
+            left join pg_namespace n on p.pronamespace = n.oid
+            left join pg_language l on p.prolang = l.oid
+            left join pg_type t on t.oid = p.prorettype
+            left join extension_functions ef on p.oid = ef.extension_function_oid
+        where
+            n.nspname not in ('pg_catalog', 'information_schema')
+            -- Filter out functions from extensions
+            and ef.extension_function_oid is null
+            and n.nspname = :schema
+        """
+            + (PG_GTE_11 if pg_version >= 110000 else PG_LT_11)
+        )
+
+        rows = sess.execute(sql, {"schema": schema}).fetchall()
+        db_functions = [cls.from_sql(x[3]) for x in rows]
+
+        for func in db_functions:
+            assert func is not None
+
+        return db_functions