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 44c0edb34e Add option `--skip-init` to db reset command (#22989)
44c0edb34e is described below
commit 44c0edb34ec5a91b1e0021eba2b37707c378ce11
Author: Daniel Standish <[email protected]>
AuthorDate: Wed Apr 13 13:50:01 2022 -0700
Add option `--skip-init` to db reset command (#22989)
This is useful when doing testing, if we want to clear out the tables but
no re-initialize the database, e.g. because we plan to only initialize it to a
certain revision with `db upgrade --to-version`.
---
airflow/cli/cli_parser.py | 8 +++++++-
airflow/cli/commands/db_command.py | 7 +++----
airflow/utils/db.py | 5 +++--
tests/cli/commands/test_db_command.py | 7 ++++++-
tests/utils/test_db.py | 12 ++++++++++--
5 files changed, 29 insertions(+), 10 deletions(-)
diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py
index 8aeb759e81..0789b4ee88 100644
--- a/airflow/cli/cli_parser.py
+++ b/airflow/cli/cli_parser.py
@@ -559,6 +559,12 @@ ARG_DB_SQL_ONLY = Arg(
action="store_true",
default=False,
)
+ARG_DB_SKIP_INIT = Arg(
+ ("-s", "--skip-init"),
+ help="Only remove tables; do not perform db init.",
+ action="store_true",
+ default=False,
+)
# webserver
ARG_PORT = Arg(
@@ -1384,7 +1390,7 @@ DB_COMMANDS = (
name='reset',
help="Burn down and rebuild the metadata database",
func=lazy_load_command('airflow.cli.commands.db_command.resetdb'),
- args=(ARG_YES,),
+ args=(ARG_YES, ARG_DB_SKIP_INIT),
),
ActionCommand(
name='upgrade',
diff --git a/airflow/cli/commands/db_command.py
b/airflow/cli/commands/db_command.py
index 4d96d2a7d3..c9201ad59b 100644
--- a/airflow/cli/commands/db_command.py
+++ b/airflow/cli/commands/db_command.py
@@ -39,10 +39,9 @@ def initdb(args):
def resetdb(args):
"""Resets the metadata database"""
print("DB: " + repr(settings.engine.url))
- if args.yes or input("This will drop existing tables if they exist.
Proceed? (y/n)").upper() == "Y":
- db.resetdb()
- else:
- print("Cancelled")
+ if not (args.yes or input("This will drop existing tables if they exist.
Proceed? (y/n)").upper() == "Y"):
+ raise SystemExit("Cancelled")
+ db.resetdb(skip_init=args.skip_init)
@cli_utils.action_cli(check_db=False)
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index f898cfefc7..c9b7ad09e5 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -1265,7 +1265,7 @@ def upgradedb(
@provide_session
-def resetdb(session: Session = NEW_SESSION):
+def resetdb(session: Session = NEW_SESSION, skip_init: bool = False):
"""Clear out the database"""
if not settings.engine:
raise RuntimeError("The settings.engine must be set. This is a
critical assertion")
@@ -1278,7 +1278,8 @@ def resetdb(session: Session = NEW_SESSION):
drop_flask_models(connection)
drop_airflow_moved_tables(session)
- initdb(session=session)
+ if not skip_init:
+ initdb(session=session)
@provide_session
diff --git a/tests/cli/commands/test_db_command.py
b/tests/cli/commands/test_db_command.py
index 62b9079051..125e5d7c3e 100644
--- a/tests/cli/commands/test_db_command.py
+++ b/tests/cli/commands/test_db_command.py
@@ -43,7 +43,12 @@ class TestCliDb:
def test_cli_resetdb(self, mock_resetdb):
db_command.resetdb(self.parser.parse_args(['db', 'reset', '--yes']))
- mock_resetdb.assert_called_once_with()
+ mock_resetdb.assert_called_once_with(skip_init=False)
+
+ @mock.patch("airflow.cli.commands.db_command.db.resetdb")
+ def test_cli_resetdb_skip_init(self, mock_resetdb):
+ db_command.resetdb(self.parser.parse_args(['db', 'reset', '--yes',
'--skip-init']))
+ mock_resetdb.assert_called_once_with(skip_init=True)
@mock.patch("airflow.cli.commands.db_command.db.check_migrations")
def test_cli_check_migrations(self, mock_wait_for_migrations):
diff --git a/tests/utils/test_db.py b/tests/utils/test_db.py
index 680955d8a5..187253cc76 100644
--- a/tests/utils/test_db.py
+++ b/tests/utils/test_db.py
@@ -178,20 +178,28 @@ class TestDb:
actual = mock_om.call_args[1]['revision']
assert actual == 'abc'
+ @pytest.mark.parametrize('skip_init', [False, True])
@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.drop_airflow_moved_tables')
@mock.patch('airflow.utils.db.initdb')
@mock.patch('airflow.settings.engine.connect')
def test_resetdb(
self,
mock_connect,
mock_init,
+ mock_drop_moved,
mock_drop_flask,
mock_drop_airflow,
+ skip_init,
):
session_mock = MagicMock()
- resetdb(session_mock)
+ resetdb(session_mock, skip_init=skip_init)
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)
+ mock_drop_moved.assert_called_once_with(session_mock)
+ if skip_init:
+ mock_init.assert_not_called()
+ else:
+ mock_init.assert_called_once_with(session=session_mock)
