Use connection instead of session.get_bind()

apache · May 26, 2023 · edbbb7e · edbbb7e
1 parent 4ea473a
commit edbbb7e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
@@ -1643,7 +1643,7 @@ def resetdb(session: Session = NEW_SESSION, skip_init: bool = False):
     with create_global_lock(session=session, lock=DBLocks.MIGRATIONS):
         with connection.begin():
             drop_airflow_models(connection)
-            drop_airflow_moved_tables(session)
+            drop_airflow_moved_tables(connection)
 
     if not skip_init:
         initdb(session=session)
@@ -1725,11 +1725,11 @@ def drop_airflow_models(connection):
         version.drop(connection)
 
 
-def drop_airflow_moved_tables(session):
+def drop_airflow_moved_tables(connection):
     from airflow.models.base import Base
     from airflow.settings import AIRFLOW_MOVED_TABLE_PREFIX
 
-    tables = set(inspect(session.get_bind()).get_table_names())
+    tables = set(inspect(connection).get_table_names())
     to_delete = [Table(x, Base.metadata) for x in tables if x.startswith(AIRFLOW_MOVED_TABLE_PREFIX)]
     for tbl in to_delete:
         tbl.drop(settings.engine, checkfirst=False)

diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py
@@ -230,7 +230,7 @@ def test_resetdb(
         session_mock = MagicMock()
         resetdb(session_mock, skip_init=skip_init)
         mock_drop_airflow.assert_called_once_with(mock_connect.return_value)
-        mock_drop_moved.assert_called_once_with(session_mock)
+        mock_drop_moved.assert_called_once_with(mock_connect.return_value)
         if skip_init:
             mock_init.assert_not_called()
         else: