Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
- Added a unit test
  • Loading branch information
harjeevanmaan authored Oct 9, 2024
1 parent 04ba8ee commit f166467
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 10 deletions.
23 changes: 13 additions & 10 deletions airflow/utils/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -1203,19 +1203,22 @@ def resetdb(session: Session = NEW_SESSION, skip_init: bool = False):
if not settings.engine:
raise RuntimeError("The settings.engine must be set. This is a critical assertion")
log.info("Dropping tables that exist")
original_logging_level = logging.root.level
try:
import_all_models()

import_all_models()

connection = settings.engine.connect()
connection = settings.engine.connect()

with create_global_lock(session=session, lock=DBLocks.MIGRATIONS), connection.begin():
drop_airflow_models(connection)
drop_airflow_moved_tables(connection)
external_db_manager = RunDBManager()
external_db_manager.drop_tables(session, connection)
with create_global_lock(session=session, lock=DBLocks.MIGRATIONS), connection.begin():
drop_airflow_models(connection)
drop_airflow_moved_tables(connection)
external_db_manager = RunDBManager()
external_db_manager.drop_tables(session, connection)

if not skip_init:
initdb(session=session)
if not skip_init:
initdb(session=session)
finally:
logging.root.setLevel(original_logging_level)


@provide_session
Expand Down
9 changes: 9 additions & 0 deletions tests/utils/test_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from __future__ import annotations

import inspect
import logging
import os
import re
from contextlib import redirect_stdout
Expand Down Expand Up @@ -229,6 +230,14 @@ def test_resetdb(
else:
mock_init.assert_called_once_with(session=session_mock)

def test_resetdb_logging_level(self):
unset_logging_level = logging.root.level
logging.root.setLevel(logging.DEBUG)
set_logging_level = logging.root.level
resetdb()
assert logging.root.level == set_logging_level
assert logging.root.level != unset_logging_level

def test_alembic_configuration(self):
with mock.patch.dict(
os.environ, {"AIRFLOW__DATABASE__ALEMBIC_INI_FILE_PATH": "/tmp/alembic.ini"}, clear=True
Expand Down

0 comments on commit f166467

Please sign in to comment.