PGExtension (olirice#53)

* PGExtension entity with schema alter conflict issue

* PGExtension create/drop

* PGExtension docs

Author: olirice

docs/api.md

@@ -63,6 +63,19 @@ trigger = PGTrigger(
 )
 ```
 
+::: alembic_utils.pg_extension.PGExtension
+    :docstring:
+
+
+```python
+from alembic_utils.pg_extension import PGExtension
+
+extension = PGExtension(
+    schema="public",
+    signature="uuid-ossp",
+)
+```
+
 
 ::: alembic_utils.pg_policy.PGPolicy
     :docstring:

mkdocs.yml

@@ -9,6 +9,7 @@ extra_css:
 
 repo_name: olirice/alembic_utils
 repo_url: https://github.com/olirice/alembic_utils
+site_url: https://olirice.github.io/alembic_utils
 
 nav:
     - Introduction: 'index.md'

src/alembic_utils/pg_extension.py

@@ -0,0 +1,84 @@
+# pylint: disable=unused-argument,invalid-name,line-too-long
+
+
+from sqlalchemy import text as sql_text
+from sqlalchemy.sql.elements import TextClause
+
+from alembic_utils.replaceable_entity import ReplaceableEntity
+from alembic_utils.statement import (
+    coerce_to_quoted,
+    coerce_to_unquoted,
+    escape_colon_for_sql,
+    normalize_whitespace,
+    strip_terminating_semicolon,
+)
+
+
+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) -> TextClause:
+        """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 = [PGExtension(x[0], x[1]) for x in rows]
+        return db_exts

src/alembic_utils/replaceable_entity.py

@@ -143,7 +143,7 @@ def identity(self) -> str:
     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()
+        object_name = self.signature.split("(")[0].strip().lower().replace("-", "_")
         return f"{schema_name}_{object_name}"
 
     def get_required_migration_op(

src/test/test_pg_extension.py

@@ -0,0 +1,126 @@
+import pytest
+
+from alembic_utils.pg_extension import PGExtension
+from alembic_utils.replaceable_entity import register_entities
+from alembic_utils.testbase import TEST_VERSIONS_ROOT, run_alembic_command
+
+TEST_EXT = PGExtension(schema="public", signature="uuid-ossp")
+
+
+def test_create_revision(engine) -> None:
+    register_entities([TEST_EXT], entity_types=[PGExtension])
+
+    output = run_alembic_command(
+        engine=engine,
+        command="revision",
+        command_kwargs={"autogenerate": True, "rev_id": "1", "message": "create"},
+    )
+
+    migration_create_path = TEST_VERSIONS_ROOT / "1_create.py"
+
+    with migration_create_path.open() as migration_file:
+        migration_contents = migration_file.read()
+
+    assert "op.create_entity" in migration_contents
+    assert "op.drop_entity" in migration_contents
+    assert "op.replace_entity" not in migration_contents
+    assert "from alembic_utils.pg_extension import PGExtension" in migration_contents
+
+    # Execute upgrade
+    run_alembic_command(engine=engine, command="upgrade", command_kwargs={"revision": "head"})
+    # Execute Downgrade
+    run_alembic_command(engine=engine, command="downgrade", command_kwargs={"revision": "base"})
+
+
+def test_create_or_replace_raises():
+    with pytest.raises(NotImplementedError):
+        TEST_EXT.to_sql_statement_create_or_replace()
+
+
+def test_update_is_unreachable(engine) -> None:
+    # Updates are not possible. The only parameter that may change is
+    # schema, and that will result in a drop and create due to schema
+    # scoping assumptions made for all other entities
+
+    # Create the view outside of a revision
+    engine.execute(TEST_EXT.to_sql_statement_create())
+
+    UPDATED_TEST_EXT = PGExtension("DEV", TEST_EXT.signature)
+
+    register_entities([UPDATED_TEST_EXT], schemas=["public", "DEV"], entity_types=[PGExtension])
+
+    # Autogenerate a new migration
+    # It should detect the change we made and produce a "replace_function" statement
+    output = run_alembic_command(
+        engine=engine,
+        command="revision",
+        command_kwargs={"autogenerate": True, "rev_id": "2", "message": "replace"},
+    )
+
+    migration_replace_path = TEST_VERSIONS_ROOT / "2_replace.py"
+
+    with migration_replace_path.open() as migration_file:
+        migration_contents = migration_file.read()
+
+    assert "op.replace_entity" not in migration_contents
+    assert "from alembic_utils.pg_extension import PGExtension" in migration_contents
+
+
+def test_noop_revision(engine) -> None:
+    # Create the view outside of a revision
+    engine.execute(TEST_EXT.to_sql_statement_create())
+
+    register_entities([TEST_EXT], entity_types=[PGExtension])
+
+    # Create a third migration without making changes.
+    # This should result in no create, drop or replace statements
+    run_alembic_command(engine=engine, command="upgrade", command_kwargs={"revision": "head"})
+
+    output = run_alembic_command(
+        engine=engine,
+        command="revision",
+        command_kwargs={"autogenerate": True, "rev_id": "3", "message": "do_nothing"},
+    )
+    migration_do_nothing_path = TEST_VERSIONS_ROOT / "3_do_nothing.py"
+
+    with migration_do_nothing_path.open() as migration_file:
+        migration_contents = migration_file.read()
+
+    assert "op.create_entity" not in migration_contents
+    assert "op.drop_entity" not in migration_contents
+    assert "op.replace_entity" not in migration_contents
+    assert "from alembic_utils" not in migration_contents
+
+    # Execute upgrade
+    run_alembic_command(engine=engine, command="upgrade", command_kwargs={"revision": "head"})
+    # Execute Downgrade
+    run_alembic_command(engine=engine, command="downgrade", command_kwargs={"revision": "base"})
+
+
+def test_drop_revision(engine) -> None:
+    # Register no functions locally
+    register_entities([], entity_types=[PGExtension])
+
+    # Manually create a SQL function
+    engine.execute(TEST_EXT.to_sql_statement_create())
+
+    output = run_alembic_command(
+        engine=engine,
+        command="revision",
+        command_kwargs={"autogenerate": True, "rev_id": "1", "message": "drop"},
+    )
+
+    migration_create_path = TEST_VERSIONS_ROOT / "1_drop.py"
+
+    with migration_create_path.open() as migration_file:
+        migration_contents = migration_file.read()
+
+    assert "op.drop_entity" in migration_contents
+    assert "op.create_entity" in migration_contents
+    assert "from alembic_utils" in migration_contents
+    assert migration_contents.index("op.drop_entity") < migration_contents.index("op.create_entity")
+
+    # Execute upgrade
+    run_alembic_command(engine=engine, command="upgrade", command_kwargs={"revision": "head"})
+    # Execute Downgrade
+    run_alembic_command(engine=engine, command="downgrade", command_kwargs={"revision": "base"})