This is an automated email from the ASF dual-hosted git repository.
dstandish pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0cbf2c752d Drop "airflow moved" tables in command `db reset` (#22990)
0cbf2c752d is described below
commit 0cbf2c752d7af9ca1c378b013b8f77dd3d858dd9
Author: Daniel Standish <[email protected]>
AuthorDate: Wed Apr 13 11:00:39 2022 -0700
Drop "airflow moved" tables in command `db reset` (#22990)
The `db reset` command does not currently drop the "temporary" tables we
create when purging bad rows as part of an upgrade.
But to truly reset the db, we need to remove them.
---
airflow/utils/db.py | 11 +++++++++++
tests/utils/test_db.py | 21 ++++++++++++++++++++-
2 files changed, 31 insertions(+), 1 deletion(-)
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index 4b0f7e1646..f898cfefc7 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -1276,6 +1276,7 @@ def resetdb(session: Session = NEW_SESSION):
with create_global_lock(session=session, lock=DBLocks.MIGRATIONS):
drop_airflow_models(connection)
drop_flask_models(connection)
+ drop_airflow_moved_tables(session)
initdb(session=session)
@@ -1379,6 +1380,16 @@ def drop_airflow_models(connection):
version.drop(connection)
+def drop_airflow_moved_tables(session):
+ from airflow.models.base import Base
+ from airflow.settings import AIRFLOW_MOVED_TABLE_PREFIX
+
+ tables = set(inspect(session.get_bind()).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=True)
+
+
def drop_flask_models(connection):
"""
Drops all Flask models.
diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py
index c07ae1d750..680955d8a5 100644
--- a/tests/utils/test_db.py
+++ b/tests/utils/test_db.py
@@ -21,6 +21,7 @@ import io
import re
from contextlib import redirect_stdout
from unittest import mock
+from unittest.mock import MagicMock
import pytest
from alembic.autogenerate import compare_metadata
@@ -33,7 +34,7 @@ from sqlalchemy import MetaData
from airflow.exceptions import AirflowException
from airflow.models import Base as airflow_base
from airflow.settings import engine
-from airflow.utils.db import check_migrations, create_default_connections,
downgrade, upgradedb
+from airflow.utils.db import check_migrations, create_default_connections,
downgrade, resetdb, upgradedb
class TestDb:
@@ -176,3 +177,21 @@ class TestDb:
downgrade(to_revision='abc')
actual = mock_om.call_args[1]['revision']
assert actual == 'abc'
+
+ @mock.patch('airflow.utils.db.create_global_lock', new=MagicMock)
+ @mock.patch('airflow.utils.db.drop_airflow_models')
+ @mock.patch('airflow.utils.db.drop_flask_models')
+ @mock.patch('airflow.utils.db.initdb')
+ @mock.patch('airflow.settings.engine.connect')
+ def test_resetdb(
+ self,
+ mock_connect,
+ mock_init,
+ mock_drop_flask,
+ mock_drop_airflow,
+ ):
+ session_mock = MagicMock()
+ resetdb(session_mock)
+ mock_drop_airflow.assert_called_once_with(mock_connect.return_value)
+ mock_drop_flask.assert_called_once_with(mock_connect.return_value)
+ mock_init.assert_called_once_with(session=session_mock)
