[airflow] branch main updated: Bind engine before attempting to drop archive tables (#33622)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk 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 911cf46621 Bind engine before attempting to drop archive tables (#33622)
911cf46621 is described below

commit 911cf466218bcd548519a50c9a32c9df58ec8b2e
Author: Sam Wheating <sa...@gmail.com>
AuthorDate: Wed Aug 23 02:38:52 2023 -0700

    Bind engine before attempting to drop archive tables (#33622)
---
 airflow/utils/db_cleanup.py    |  1 +
 tests/utils/test_db_cleanup.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/airflow/utils/db_cleanup.py b/airflow/utils/db_cleanup.py
index 79a804a9e5..4c82f90ab3 100644
--- a/airflow/utils/db_cleanup.py
+++ b/airflow/utils/db_cleanup.py
@@ -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")
diff --git a/tests/utils/test_db_cleanup.py b/tests/utils/test_db_cleanup.py
index f9cbc25feb..b6819342e7 100644
--- a/tests/utils/test_db_cleanup.py
+++ b/tests/utils/test_db_cleanup.py
@@ -41,6 +41,7 @@ from airflow.utils.db_cleanup import (
     _cleanup_table,
     _confirm_drop_archives,
     _dump_table_to_file,
+    _get_archived_table_names,
     config_dict,
     drop_archived_tables,
     export_archived_records,
@@ -266,6 +267,34 @@ class TestDBCleanup:
             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,