You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "anshuman0123 (via GitHub)" <gi...@apache.org> on 2026/06/11 06:16:22 UTC

[PR] Fix test fixture cleanup for bundles and teams [airflow]

anshuman0123 opened a new pull request, #68376:
URL: https://github.com/apache/airflow/pull/68376

   ## Summary
   
   - Track DAG bundle rows created by the `dag_maker` test fixture and remove them during fixture cleanup.
   - Add teardown for the `testing_dag_bundle` fixture when it creates the shared `testing` bundle row.
   - Add teardown for the `testing_team` fixture when it creates the shared `testing` team row.
   
   ## Details
   
   `dag_maker.cleanup()` already removes DAG-related rows such as `DagModel`, `DagRun`, `TaskInstance`, `DagVersion`, `XCom`, `TaskMap`, and `AssetEvent`, but it did not remove `DagBundleModel` rows that `dag_maker` created. This could leave bundle rows behind and force later tests to defensively call `clear_db_dag_bundles()`.
   
   This change records bundle names created by `dag_maker` and deletes those bundles after DAG model/version cleanup, preserving FK-safe cleanup order. It also converts `testing_dag_bundle` and `testing_team` into fixtures with teardown paths so rows created for a test do not leak into later tests.
   
   ## Testing
   
   - `python -m py_compile devel-common\src\tests_common\pytest_plugin.py`
   - `git diff --check`
   
   Targeted pytest was not run locally because `uv` and `pytest` are not installed in this environment.
   
   Fixes: #68374
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix test fixture cleanup for bundles and teams [airflow]

Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #68376:
URL: https://github.com/apache/airflow/pull/68376#discussion_r3396936025


##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1181,10 +1181,10 @@ def _fetch_remote_logging_conn(conn_id: str, client: Client) -> Connection | Non
         from airflow.sdk.definitions.connection import Connection
 

Review Comment:
   I think changes to this file is not related to the issue attached to this pr. And I think changes to this file can be dropped



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix test fixture cleanup for bundles and teams [airflow]

Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #68376:
URL: https://github.com/apache/airflow/pull/68376#issuecomment-4677740727

   Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our [Contributors' Guide](https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
   Here are some useful points:
   - Pay attention to the quality of your code (ruff, mypy and type annotations). Our [prek-hooks]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-prek-hooks) will help you with that.
   - In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/airflow-core/docs/howto/custom-operator.rst) Consider adding an example Dag that shows how users should use it.
   - Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
   - Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
   - Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
   - Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
   - Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
   Apache Airflow is a community-driven project and together we are making it better 🚀.
   In case of doubts contact the developers at:
   Mailing List: dev@airflow.apache.org
   Slack: https://s.apache.org/airflow-slack
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix test fixture cleanup for bundles and teams [airflow]

Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on PR #68376:
URL: https://github.com/apache/airflow/pull/68376#issuecomment-4681940352

   Thank you so much for the PR. I think the changes to supervisor.py can be dropped, seems unrelated to the issue tagged in the description. And also, I think this is just one part of the issue to be fixed, I would suggest updating the description to related to tagged issue.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix test fixture cleanup for bundles and teams [airflow]

Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #68376:
URL: https://github.com/apache/airflow/pull/68376#discussion_r3397832692


##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -1399,7 +1401,8 @@ def cleanup(self):
             for attempt in run_with_db_retries(logger=self.log):
                 with attempt:
                     dag_ids = list(self.dagbag.dag_ids)
-                    if not dag_ids:
+                    bundle_names = set(self.created_bundle_names)

Review Comment:
   Any reason to not use `self.created_bundle_names` directly?  I think in this case it is overly defensive; I don't see where anything is modifying the set once it reaches this point.



##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -2924,40 +2935,60 @@ def mock_xcom_backend():
 def testing_dag_bundle():
     from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-    if AIRFLOW_V_3_0_PLUS:
-        from sqlalchemy import func, select
+    if not AIRFLOW_V_3_0_PLUS:
+        yield
+        return
 
-        from airflow.models.dagbundle import DagBundleModel
-        from airflow.utils.session import create_session
+    from sqlalchemy import delete, func, select
 
-        with create_session() as session:
-            if (
-                session.scalar(
-                    select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name == "testing")
-                )
-                == 0
-            ):
-                testing = DagBundleModel(name="testing")
-                session.add(testing)
+    from airflow.models.dagbundle import DagBundleModel
+    from airflow.utils.session import create_session
+
+    created = False
+    with create_session() as session:
+        if (
+            session.scalar(
+                select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name == "testing")
+            )
+            == 0
+        ):
+            testing = DagBundleModel(name="testing")
+            session.add(testing)
+            created = True
+
+    try:
+        yield
+    finally:
+        if created:
+            with create_session() as session:
+                session.execute(delete(DagBundleModel).where(DagBundleModel.name == "testing"))

Review Comment:
   I am pretty sure this is one of the "quick fixes" I tried before creating this issue and this will blow up.  We'll see when the CI finishes.  I didn't dive too deep into it once I saw it wasn't working, but I suspect that this fails because (some of??) the DagModels in tests are being created after this fixture fires and _those_ aren't being cleaned, which causes an F-Key Reference error.  
   
   If the CI is failing, I'd suggest you start there.  If the CI passes then I'm mistaken and you can resolve this with my apologies. 



##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -1420,11 +1425,17 @@ def cleanup(self):
                         )
                         self.session.execute(delete(DagRun).where(DagRun.dag_id.in_(dag_ids)))
                         self.session.execute(delete(TaskInstance).where(TaskInstance.dag_id.in_(dag_ids)))
-                    self.session.execute(delete(XCom).where(XCom.dag_id.in_(dag_ids)))
-                    self.session.execute(delete(DagModel).where(DagModel.dag_id.in_(dag_ids)))
-                    self.session.execute(delete(TaskMap).where(TaskMap.dag_id.in_(dag_ids)))
-                    self.session.execute(delete(AssetEvent).where(AssetEvent.source_dag_id.in_(dag_ids)))
+                    if dag_ids:
+                        self.session.execute(delete(XCom).where(XCom.dag_id.in_(dag_ids)))
+                        self.session.execute(delete(DagModel).where(DagModel.dag_id.in_(dag_ids)))
+                        self.session.execute(delete(TaskMap).where(TaskMap.dag_id.in_(dag_ids)))
+                        self.session.execute(delete(AssetEvent).where(AssetEvent.source_dag_id.in_(dag_ids)))
+                    if AIRFLOW_V_3_0_PLUS and bundle_names:
+                        self.session.execute(
+                            delete(DagBundleModel).where(DagBundleModel.name.in_(bundle_names))
+                        )
                     self.session.commit()
+                    self.created_bundle_names.difference_update(bundle_names)

Review Comment:
   Similar to above, you can likely drop this.  We should be deleting all bundle_names, right?  The SQL above is batch-deleting ll of them, then here you remove all of them from the set.  MAYBE if you wanted to do this then validate that it is now an empty set, that could be a reason for it, but as of now it's just emptying the set before teardown.



##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -2924,40 +2935,60 @@ def mock_xcom_backend():
 def testing_dag_bundle():
     from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-    if AIRFLOW_V_3_0_PLUS:
-        from sqlalchemy import func, select
+    if not AIRFLOW_V_3_0_PLUS:
+        yield
+        return
 
-        from airflow.models.dagbundle import DagBundleModel
-        from airflow.utils.session import create_session
+    from sqlalchemy import delete, func, select
 
-        with create_session() as session:
-            if (
-                session.scalar(
-                    select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name == "testing")
-                )
-                == 0
-            ):
-                testing = DagBundleModel(name="testing")
-                session.add(testing)
+    from airflow.models.dagbundle import DagBundleModel
+    from airflow.utils.session import create_session
+
+    created = False
+    with create_session() as session:
+        if (
+            session.scalar(
+                select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name == "testing")
+            )
+            == 0
+        ):
+            testing = DagBundleModel(name="testing")
+            session.add(testing)
+            created = True
+
+    try:
+        yield
+    finally:
+        if created:
+            with create_session() as session:
+                session.execute(delete(DagBundleModel).where(DagBundleModel.name == "testing"))
 
 
 @pytest.fixture
 def testing_team():
     from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-    if AIRFLOW_V_3_0_PLUS:
-        from sqlalchemy import select
+    if not AIRFLOW_V_3_0_PLUS:
+        yield None
+        return
 
-        from airflow.models.team import Team
-        from airflow.utils.session import create_session
+    from sqlalchemy import delete, select
 
-        with create_session() as session:
-            team = session.scalar(select(Team).where(Team.name == "testing"))
-            if not team:
-                team = Team(name="testing")
-                session.add(team)
-                session.flush()
-            yield team
+    from airflow.models.team import Team
+    from airflow.utils.session import create_session
+
+    created = False
+    with create_session() as session:
+        team = session.scalar(select(Team).where(Team.name == "testing"))
+        if not team:
+            team = Team(name="testing")
+            session.add(team)
+            session.flush()
+            created = True
+        yield team
+        if created:
+            session.rollback()
+            session.execute(delete(Team).where(Team.name == "testing"))
 

Review Comment:
   SQL is one of my weak points, but this smells funny.  Rolling back before deleting feels wrong here?  100% willing to be wrong here and learn something from you if this is an area you are comfortable with. 
   
   If I understand this right, we're creating the Team and yielding it, then after the test finishes with it we're rolling it back, then deleting it?  I'm guessing the rollback is in case the test modified it during the yield?  Which maybe makes sense, but if we're deleting it anyway, does it matter if we roll it back first?  If the session is in a good/unmodified state then the rollback is unnecessary.  If the session is in a bad/modified state, then the test has already failed and we should nuke it anyway, no?
   
   Again, this is a blind spot for me, I cold absolutely be missing something.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix test fixture cleanup for bundles and teams [airflow]

Posted by "anishgirianish (via GitHub)" <gi...@apache.org>.
anishgirianish commented on code in PR #68376:
URL: https://github.com/apache/airflow/pull/68376#discussion_r3396936025


##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1181,10 +1181,10 @@ def _fetch_remote_logging_conn(conn_id: str, client: Client) -> Connection | Non
         from airflow.sdk.definitions.connection import Connection
 

Review Comment:
   I think this is not related to the issue attached to this pr. And I think changes to this file can be dropped



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org