1.5 transaction model did not expect branched connections

[gibsondan]

Describe the bug
Hello, after the recent alembic 1.5.0 release our alembic upgrades have started failing with a stack trace in SQLAlchemy code.

The bottom of the stack trace is:

connection = <sqlalchemy.engine.base.Connection object at 0x144efe990>

    def _get_connection_transaction(connection):
        if sqla_14:
            return connection.get_transaction()
        else:
>           return connection._Connection__transaction
E           AttributeError: 'Connection' object has no attribute '_Connection__transaction'

We're currently on SQLAlchemy==1.3.22 (the latest non-beta release) Upgrading to SQLAlchemy=1.4.0b1 fixes the problem.

Expected behavior
Alembic upgrades succeed without needing to upgrade to a beta version of SQLAlchemy

To Reproduce
Check out the repo at https://github.com/dagster-io/dagster and run our automated test suite against master. Can provide more details/repro steps here if needed.

Error

=============================================================================================================== FAILURES ===============================================================================================================
_______________________________________________________________________________________________________ test_asset_key_structure _______________________________________________________________________________________________________

    def test_asset_key_structure():
        src_dir = file_relative_path(__file__, "compat_tests/snapshot_0_9_16_asset_key_structure")
        with copy_directory(src_dir) as test_dir:
            asset_storage = ConsolidatedSqliteEventLogStorage(test_dir)
            asset_keys = asset_storage.get_all_asset_keys()
            assert len(asset_keys) == 5
    
            # get a structured asset key
            asset_key = AssetKey(["dashboards", "cost_dashboard"])
    
            # check that backcompat events are read
            assert asset_storage.has_asset_key(asset_key)
            events = asset_storage.get_asset_events(asset_key)
            assert len(events) == 1
            run_ids = asset_storage.get_asset_run_ids(asset_key)
            assert len(run_ids) == 1
    
>           asset_storage.upgrade()

python_modules/dagster/dagster_tests/core_tests/storage_tests/test_assets.py:273: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
python_modules/dagster/dagster/core/storage/event_log/sqlite/consolidated_sqlite_event_log.py:105: in upgrade
    run_alembic_upgrade(alembic_config, conn)
python_modules/dagster/dagster/core/storage/sql.py:28: in run_alembic_upgrade
    upgrade(alembic_config, rev)
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/command.py:294: in upgrade
    script.run_env()
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/script/base.py:481: in run_env
    util.load_python_file(self.dir, "env.py")
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/util/pyfiles.py:97: in load_python_file
    module = load_module_py(module_id, path)
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/util/compat.py:182: in load_module_py
    spec.loader.exec_module(module)
<frozen importlib._bootstrap_external>:728: in exec_module
    ???
<frozen importlib._bootstrap>:219: in _call_with_frames_removed
    ???
python_modules/dagster/dagster/core/storage/event_log/sqlite/alembic/env.py:15: in <module>
    run_migrations_online(context, config, target_metadata)
python_modules/dagster/dagster/core/storage/sqlite.py:23: in run_migrations_online
    run_migrations_online_(*args, **kwargs)
python_modules/dagster/dagster/core/storage/sql.py:131: in run_migrations_online
    context.run_migrations()
<string>:8: in run_migrations
    ???
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/runtime/environment.py:813: in run_migrations
    self.get_context().run_migrations(**kw)
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/runtime/migration.py:549: in run_migrations
    with self.begin_transaction(_per_migration=True):
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/runtime/migration.py:395: in begin_transaction
    self.connection
../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/util/sqla_compat.py:84: in _safe_begin_connection_transaction
    transaction = _get_connection_transaction(connection)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

connection = <sqlalchemy.engine.base.Connection object at 0x144efe990>

    def _get_connection_transaction(connection):
        if sqla_14:
            return connection.get_transaction()
        else:
>           return connection._Connection__transaction
E           AttributeError: 'Connection' object has no attribute '_Connection__transaction'

../.pyenv/versions/3.7.8/envs/dagster-3.7.8/lib/python3.7/site-packages/alembic/util/sqla_compat.py:105: AttributeError

Versions.

• OS:
• Python: 3.7.8
• Alembic: 1.5.1
• SQLAlchemy: 1.3.22
• Database:
• DBAPI:

Have a nice day!

[zzzeek]

can you share your env.py please? im not able to locate it within the repo given

[zzzeek]

it's ultimately https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/dagster/core/storage/sql.py need to figure out what "Connection" is here

[zzzeek]

no tox file, so yes can you please provide steps to run your test suite.

Is there something special about the SQLAlchemy engine or connecvtion in use? there's an attribute missing and I am not able to reproduce that locally.

[gibsondan]

Here's the env.py:

https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/dagster/core/storage/event_log/sqlite/alembic/env.py

Here's the tox file for the test suite in question:
https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/tox.ini

To reproduce with just one failing test from the root folder (in master, but from before we just landed an alembic pin as a temporary workaround):

cd python_modules/dagster
tox -vv -e py37-core_tests -- -k test_asset_key_structure

Will investigate if there's anything unique about our engine/connection that might be causing this.

[gibsondan]

Updated repro steps (including a specific revision to checkout from before the pin):

git checkout 89b27c7e0d9615356dedf45397a61ec1a7bf05aa
cd python_modules/dagster
tox -vv -e py37-core_tests -- -k test_asset_key_structure

[zzzeek]

ok thanks for the detailed info, this was straightforward.

the bug is that Alembic's test suite is not expecting to have received a "branched" connection, so will fix that now. however "branched" connections are deprecated in 1.4 so you might want to move off of this pattern in any case. the patch below will resolve but this assumes "connection" for you is always a Connection and not an Engine:

diff --git a/python_modules/dagster/dagster/core/storage/sql.py b/python_modules/dagster/dagster/core/storage/sql.py
index cf12888c5..2c3a7d9c5 100644
--- a/python_modules/dagster/dagster/core/storage/sql.py
+++ b/python_modules/dagster/dagster/core/storage/sql.py
@@ -124,8 +124,8 @@ def run_migrations_online(context, config, target_metadata):
             "command line, STOP and read the README."
         )
 
-    with connectable.connect() as connection:
-        context.configure(connection=connection, target_metadata=target_metadata)
+    connection = connectable
+    context.configure(connection=connection, target_metadata=target_metadata)
 
-        with context.begin_transaction():
-            context.run_migrations()
+    with context.begin_transaction():
+        context.run_migrations()

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the master branch:

retrieve 1.3 transaction from branched connection properly https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/2489