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])