Skip to content

Commit

Permalink
Bind engine before attempting to drop archive tables (#33622)
Browse files Browse the repository at this point in the history
  • Loading branch information
SamWheating authored Aug 23, 2023
1 parent 0ca5f70 commit 911cf46
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 0 deletions.
1 change: 1 addition & 0 deletions airflow/utils/db_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,7 @@ def _do_delete(*, query, orm_model, skip_archive, session):
session.execute(delete)
session.commit()
if skip_archive:
metadata.bind = session.get_bind()
target_table.drop()
session.commit()
print("Finished Performing Delete")
Expand Down
29 changes: 29 additions & 0 deletions tests/utils/test_db_cleanup.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
_cleanup_table,
_confirm_drop_archives,
_dump_table_to_file,
_get_archived_table_names,
config_dict,
drop_archived_tables,
export_archived_records,
Expand Down Expand Up @@ -266,6 +267,34 @@ def test__cleanup_table(self, table_name, date_add_kwargs, expected_to_delete, e
else:
raise Exception("unexpected")

@pytest.mark.parametrize(
"skip_archive, expected_archives",
[param(True, 0, id="skip_archive"), param(False, 1, id="do_archive")],
)
def test__skip_archive(self, skip_archive, expected_archives):
"""
Verify that running cleanup_table with drops the archives when requested.
"""
base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC"))
num_tis = 10
create_tis(
base_date=base_date,
num_tis=num_tis,
)
with create_session() as session:
clean_before_date = base_date.add(days=5)
_cleanup_table(
**config_dict["dag_run"].__dict__,
clean_before_timestamp=clean_before_date,
dry_run=False,
session=session,
table_names=["dag_run"],
skip_archive=skip_archive,
)
model = config_dict["dag_run"].orm_model
assert len(session.query(model).all()) == 5
assert len(_get_archived_table_names(["dag_run"], session)) == expected_archives

def test_no_models_missing(self):
"""
1. Verify that for all tables in `airflow.models`, we either have them enabled in db cleanup,
Expand Down

0 comments on commit 911cf46

Please sign in to comment.