Merge pull request olirice#105 from jdimmerman/feature/sqlalchemy-2-support

sqlalchemy 1.4 & 2.0 support (1.3 dropped)

Author: olirice

README.md

@@ -116,3 +116,23 @@ def downgrade():
 Visit the [quickstart guide](https://olirice.github.io/alembic_utils/quickstart/) for usage instructions.
 
 <p align="center">&mdash;&mdash;  &mdash;&mdash;</p>
+
+### Contributing
+
+To run the tests
+```
+# install pip dependencies
+pip install wheel && pip install -e ".[dev]"
+
+# run the tests
+pytest src/test
+```
+
+To invoke the linter automated formatting and generally make use of precommit checks:
+```
+pip install pre-commit
+pre-commit install
+
+# manually run
+pre-commit run --all
+```

setup.py

@@ -37,7 +37,7 @@ def get_version(package):
         "alembic>=1.5.7",
         "flupy",
         "parse>=1.8.4",
-        "sqlalchemy>=1.3.0,<2",
+        "sqlalchemy>=1.4",
         "typing_extensions",
     ],
     extras_require={

src/alembic_utils/__init__.py

@@ -1 +1 @@
-__version__ = "0.8.0"
+__version__ = "0.8.1"

src/alembic_utils/simulate.py

@@ -32,9 +32,8 @@ def simulate_entity(
 
     deps: List["ReplaceableEntity"] = dependencies or []
 
+    outer_transaction = sess.begin_nested()
     try:
-        sess.begin_nested()
-
         dependency_managers = [simulate_entity(sess, x) for x in deps]
 
         with ExitStack() as stack:
@@ -43,8 +42,8 @@ def simulate_entity(
                 stack.enter_context(mgr)
 
             did_drop = False
+            inner_transaction = sess.begin_nested()
             try:
-                sess.begin_nested()
                 sess.execute(entity.to_sql_statement_drop(cascade=True))
                 did_drop = True
                 sess.execute(entity.to_sql_statement_create())
@@ -58,11 +57,11 @@ def simulate_entity(
 
                 # Try again without the drop in case the drop raised
                 # a does not exist error
-                sess.rollback()
-                sess.begin_nested()
+                inner_transaction.rollback()
+                inner_transaction = sess.begin_nested()
                 sess.execute(entity.to_sql_statement_create())
                 yield sess
             finally:
-                sess.rollback()
+                inner_transaction.rollback()
     finally:
-        sess.rollback()
+        outer_transaction.rollback()

src/alembic_utils/testbase.py

@@ -27,8 +27,10 @@ def build_alembic_config(engine: Engine) -> Config:
     path_to_alembic_ini = REPO_ROOT / "alembic.ini"
 
     alembic_cfg = Config(path_to_alembic_ini)
+
     # Make double sure alembic references the test database
-    alembic_cfg.set_main_option("sqlalchemy.url", str(engine.url))
+    alembic_cfg.set_main_option("sqlalchemy.url", engine.url.render_as_string(hide_password=False))
+
     alembic_cfg.set_main_option("script_location", str((Path("src") / "test" / "alembic_config")))
     return alembic_cfg
 

src/test/alembic_config/env.py

@@ -57,7 +57,6 @@ def run_migrations_online():
         config.get_section(config.config_ini_section),
         prefix="sqlalchemy.",
         poolclass=pool.NullPool,
-        future=True,
     )
 
     with connectable.connect() as connection:

src/test/conftest.py

@@ -9,7 +9,7 @@
 
 import pytest
 from parse import parse
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, text
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session, sessionmaker
 
@@ -105,14 +105,15 @@ def raw_engine(maybe_start_pg: None) -> Generator[Engine, None, None]:
 
 
 @pytest.fixture(scope="function")
-def engine(raw_engine) -> Generator[Engine, None, None]:
+def engine(raw_engine: Engine) -> Generator[Engine, None, None]:
     """Engine that has been reset between tests"""
 
     def run_cleaners():
         registry.clear()
-        raw_engine.execute("drop schema public cascade; create schema public;")
-        raw_engine.execute('drop schema if exists "DEV" cascade; create schema "DEV";')
-        raw_engine.execute('drop role if exists "anon_user"')
+        with raw_engine.begin() as connection:
+            connection.execute(text("drop schema public cascade; create schema public;"))
+            connection.execute(text('drop schema if exists "DEV" cascade; create schema "DEV";'))
+            connection.execute(text('drop role if exists "anon_user"'))
         # Remove any migrations that were left behind
         TEST_VERSIONS_ROOT.mkdir(exist_ok=True, parents=True)
         shutil.rmtree(TEST_VERSIONS_ROOT)

src/test/test_create_function.py

@@ -1,3 +1,5 @@
+from sqlalchemy import text
+
 from alembic_utils.pg_function import PGFunction
 
 to_upper = PGFunction(
@@ -18,8 +20,9 @@ def test_create_and_drop(engine) -> None:
     down_sql = to_upper.to_sql_statement_drop()
 
     # Testing that the following two lines don't raise
-    engine.execute(up_sql)
-    result = engine.execute("select public.to_upper('hello');").fetchone()
-    assert result[0] == "HELLO"
-    engine.execute(down_sql)
-    assert True
+    with engine.begin() as connection:
+        connection.execute(up_sql)
+        result = connection.execute(text("select public.to_upper('hello');")).fetchone()
+        assert result[0] == "HELLO"
+        connection.execute(down_sql)
+        assert True

src/test/test_include_filters.py

@@ -75,8 +75,9 @@
 
 
 def test_create_revision_with_filters(engine) -> None:
-    for entity in reflected_entities:
-        engine.execute(entity.to_sql_statement_create())
+    with engine.begin() as connection:
+        for entity in reflected_entities:
+            connection.execute(entity.to_sql_statement_create())
     register_entities(registered_entities)
 
     run_alembic_command(

src/test/test_op_drop_cascade.py

@@ -20,8 +20,9 @@
 
 def test_drop_fails_without_cascade(engine) -> None:
 
-    engine.execute(A.to_sql_statement_create())
-    engine.execute(B_A.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(A.to_sql_statement_create())
+        connection.execute(B_A.to_sql_statement_create())
 
     register_entities([B_A], schemas=["DEV"], entity_types=[PGView])
 
@@ -48,8 +49,9 @@ def test_drop_fails_without_cascade(engine) -> None:
 
 def test_drop_fails_with_cascade(engine, sess) -> None:
 
-    engine.execute(A.to_sql_statement_create())
-    engine.execute(B_A.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(A.to_sql_statement_create())
+        connection.execute(B_A.to_sql_statement_create())
 
     register_entities([B_A], schemas=["DEV"], entity_types=[PGView])
 

src/test/test_pg_constraint_trigger.py

@@ -1,4 +1,5 @@
 import pytest
+from sqlalchemy import text
 
 from alembic_utils.exceptions import SQLParseFailure
 from alembic_utils.pg_function import PGFunction
@@ -9,18 +10,21 @@
 
 @pytest.fixture(scope="function")
 def sql_setup(engine):
-    conn = engine
-    conn.execute(
+    with engine.begin() as connection:
+        connection.execute(
+            text(
+                """
+        create table public.account (
+            id serial primary key,
+            email text not null
+        );
         """
-    create table public.account (
-        id serial primary key,
-        email text not null
-    );
-    """
-    )
+            )
+        )
 
     yield
-    conn.execute("drop table public.account cascade")
+    with engine.begin() as connection:
+        connection.execute(text("drop table public.account cascade"))
 
 
 FUNC = PGFunction.from_sql(
@@ -45,7 +49,8 @@ def sql_setup(engine):
 
 
 def test_create_revision(sql_setup, engine) -> None:
-    engine.execute(FUNC.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(FUNC.to_sql_statement_create())
 
     register_entities([FUNC, TRIG], entity_types=[PGTrigger])
     run_alembic_command(
@@ -71,8 +76,9 @@ def test_create_revision(sql_setup, engine) -> None:
 
 
 def test_trig_update_revision(sql_setup, engine) -> None:
-    engine.execute(FUNC.to_sql_statement_create())
-    engine.execute(TRIG.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(FUNC.to_sql_statement_create())
+        connection.execute(TRIG.to_sql_statement_create())
 
     UPDATED_TRIG = PGTrigger(
         schema="public",
@@ -113,8 +119,9 @@ def test_trig_update_revision(sql_setup, engine) -> None:
 
 
 def test_noop_revision(sql_setup, engine) -> None:
-    engine.execute(FUNC.to_sql_statement_create())
-    engine.execute(TRIG.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(FUNC.to_sql_statement_create())
+        connection.execute(TRIG.to_sql_statement_create())
 
     register_entities([FUNC, TRIG], entity_types=[PGTrigger])
 
@@ -141,8 +148,9 @@ def test_noop_revision(sql_setup, engine) -> None:
 
 def test_drop(sql_setup, engine) -> None:
     # Manually create a SQL function
-    engine.execute(FUNC.to_sql_statement_create())
-    engine.execute(TRIG.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(FUNC.to_sql_statement_create())
+        connection.execute(TRIG.to_sql_statement_create())
 
     # Register no functions locally
     register_entities([], schemas=["public"], entity_types=[PGTrigger])

src/test/test_pg_extension.py

@@ -43,7 +43,8 @@ def test_update_is_unreachable(engine) -> None:
     # scoping assumptions made for all other entities
 
     # Create the view outside of a revision
-    engine.execute(TEST_EXT.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TEST_EXT.to_sql_statement_create())
 
     UPDATED_TEST_EXT = PGExtension("DEV", TEST_EXT.signature)
 
@@ -68,7 +69,8 @@ def test_update_is_unreachable(engine) -> None:
 
 def test_noop_revision(engine) -> None:
     # Create the view outside of a revision
-    engine.execute(TEST_EXT.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TEST_EXT.to_sql_statement_create())
 
     register_entities([TEST_EXT], entity_types=[PGExtension])
 
@@ -102,7 +104,8 @@ def test_drop_revision(engine) -> None:
     register_entities([], entity_types=[PGExtension])
 
     # Manually create a SQL function
-    engine.execute(TEST_EXT.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TEST_EXT.to_sql_statement_create())
 
     output = run_alembic_command(
         engine=engine,

src/test/test_pg_function.py

@@ -1,3 +1,5 @@
+from sqlalchemy import text
+
 from alembic_utils.pg_function import PGFunction
 from alembic_utils.replaceable_entity import register_entities
 from alembic_utils.testbase import TEST_VERSIONS_ROOT, run_alembic_command
@@ -39,7 +41,8 @@ def test_create_revision(engine) -> None:
 
 
 def test_update_revision(engine) -> None:
-    engine.execute(TO_UPPER.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TO_UPPER.to_sql_statement_create())
 
     # Update definition of TO_UPPER
     UPDATED_TO_UPPER = PGFunction(
@@ -79,7 +82,8 @@ def test_update_revision(engine) -> None:
 
 
 def test_noop_revision(engine) -> None:
-    engine.execute(TO_UPPER.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TO_UPPER.to_sql_statement_create())
 
     register_entities([TO_UPPER], entity_types=[PGFunction])
 
@@ -106,7 +110,8 @@ def test_noop_revision(engine) -> None:
 
 def test_drop(engine) -> None:
     # Manually create a SQL function
-    engine.execute(TO_UPPER.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TO_UPPER.to_sql_statement_create())
 
     # Register no functions locally
     register_entities([], schemas=["public"], entity_types=[PGFunction])
@@ -174,7 +179,8 @@ def test_ignores_extension_functions(engine) -> None:
     # Unless they are excluded, every autogenerate migration will produce
     # drop statements for those functions
     try:
-        engine.execute("create extension if not exists unaccent;")
+        with engine.begin() as connection:
+            connection.execute(text("create extension if not exists unaccent;"))
         register_entities([], schemas=["public"], entity_types=[PGFunction])
         run_alembic_command(
             engine=engine,
@@ -189,7 +195,8 @@ def test_ignores_extension_functions(engine) -> None:
 
         assert "op.drop_entity" not in migration_contents
     finally:
-        engine.execute("drop extension if exists unaccent;")
+        with engine.begin() as connection:
+            connection.execute(text("drop extension if exists unaccent;"))
 
 
 def test_plpgsql_colon_esacpe(engine) -> None:

src/test/test_pg_function_overloading.py

@@ -49,8 +49,9 @@ def test_create_revision(engine) -> None:
 
 
 def test_update_revision(engine) -> None:
-    engine.execute(TO_FLOAT_FROM_INT.to_sql_statement_create())
-    engine.execute(TO_FLOAT_FROM_TEXT.to_sql_statement_create())
+    with engine.begin() as connection:
+        connection.execute(TO_FLOAT_FROM_INT.to_sql_statement_create())
+        connection.execute(TO_FLOAT_FROM_TEXT.to_sql_statement_create())
 
     UPDATE = PGFunction(
         schema="public",