You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "stevenschaerer (via GitHub)" <gi...@apache.org> on 2024/04/02 20:56:10 UTC

[PR] Use async db calls in WorkflowTrigger [airflow]

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

   This closes: #38672
   
   I opted to slightly extend some existing unit tests rather than write mostly identical new ones. If the `@sync_to_async` decorator and the `await`s before `self._get_count` are removed the tests will fail because the trigger and dummy tasks are executed sequentially.
   
   Should I also remove `TaskStateTrigger` in this PR? It is not used anymore outside of tests.


-- 
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] Use async db calls in WorkflowTrigger [airflow]

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


##########
tests/triggers/test_external_task.py:
##########
@@ -438,3 +464,12 @@ def test_serialization(self):
             "execution_dates": [timezone.datetime(2022, 1, 1)],
             "poll_interval": 5,
         }
+
+
+def mocked_get_count(*args, **kwargs):
+    time.sleep(0.3)

Review Comment:
   done



-- 
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] Use async db calls in WorkflowTrigger [airflow]

Posted by "Taragolis (via GitHub)" <gi...@apache.org>.
Taragolis merged PR #38689:
URL: https://github.com/apache/airflow/pull/38689


-- 
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] Use async db calls in WorkflowTrigger [airflow]

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

   > I don’t think we can just remove TaskStateTrigger, but let’s deprecate it for future removal.
   > 
   > It would also be a good idea to implement an version of `_get_count` to use the async SQLAlchemy API if available, and fallback to the `sync_to_async` version otherwise. Most people likely won’t have the async component (`pip install sqlalchemy[asyncio]`) but we should start to slowly add support to it in triggerers.
   
   I have deprecated `TaskStateTrigger`. Following the other comments I'll leave using async SQLAlchemy for a separate PR.


-- 
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] Use async db calls in WorkflowTrigger [airflow]

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

   I don’t think we can just remove TaskStateTrigger, but let’s deprecate it for future removal.
   
   It would also be a good idea to implement an version of `_get_count` to use the async SQLAlchemy API if available, and fallback to the `sync_to_async` version otherwise. Most people likely won’t have the async component (`pip install sqlalchemy[asyncio]`) but we should start to slowly add support to it in triggerers.


-- 
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] Use async db calls in WorkflowTrigger [airflow]

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


##########
airflow/triggers/external_task.py:
##########
@@ -98,44 +99,43 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
         while True:
             if self.failed_states:
-                failed_count = _get_count(
-                    self.execution_dates,
-                    self.external_task_ids,
-                    self.external_task_group_id,
-                    self.external_dag_id,
-                    self.failed_states,
-                )
+                failed_count = await self._get_count(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = _get_count(
-                    self.execution_dates,
-                    self.external_task_ids,
-                    self.external_task_group_id,
-                    self.external_dag_id,
-                    self.skipped_states,
-                )
+                skipped_count = await self._get_count(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = _get_count(
-                self.execution_dates,
-                self.external_task_ids,
-                self.external_task_group_id,
-                self.external_dag_id,
-                self.allowed_states,
-            )
+            allowed_count = await self._get_count(self.allowed_states)
             if allowed_count == len(self.execution_dates):
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    @sync_to_async
+    def _get_count(self, states: typing.Iterable[str] | None) -> int:
+        """
+        Get the count of records against dttm filter and states. Async wrapper for _get_count.
+
+        :param states: task or dag states
+        :return The count of records.
+        """
+        return _get_count(
+            dttm_filter=self.execution_dates,
+            external_task_ids=self.external_task_ids,
+            external_task_group_id=self.external_task_group_id,
+            external_dag_id=self.external_dag_id,
+            states=states,
+        )
+
 
+@deprecated(reason=("TaskStateTrigger has been deprecated and will be removed in future."))

Review Comment:
   Just found why tests around this class not failed
   
   ```suggestion
   @deprecated(
       reason="TaskStateTrigger has been deprecated and will be removed in future.",
       category=RemovedInAirflow3Warning,
   )
   ```
   
   And this is required to [deal with warnings](https://github.com/apache/airflow/blob/main/contributing-docs/testing/unit_tests.rst#handling-warnings) into the `tests/triggers/test_external_task.py::test_external_task.py::TestTaskStateTrigger`



-- 
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] Use async db calls in WorkflowTrigger [airflow]

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


##########
tests/triggers/test_external_task.py:
##########
@@ -438,3 +464,12 @@ def test_serialization(self):
             "execution_dates": [timezone.datetime(2022, 1, 1)],
             "poll_interval": 5,
         }
+
+
+def mocked_get_count(*args, **kwargs):
+    time.sleep(0.3)
+    return 1
+
+
+async def dummy_async_fun():

Review Comment:
   renamed to `fake`



-- 
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] Use async db calls in WorkflowTrigger [airflow]

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

   Also some async features are only available in sqlalchemy 2.0, so maybe we need to focus on this migration first.


-- 
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] Use async db calls in WorkflowTrigger [airflow]

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

   > I don’t think we can just remove TaskStateTrigger, but let’s deprecate it for future removal.
   > 
   > 
   > 
   > It would also be a good idea to implement an version of `_get_count` to use the async SQLAlchemy API if available, and fallback to the `sync_to_async` version otherwise. Most people likely won’t have the async component (`pip install sqlalchemy[asyncio]`) but we should start to slowly add support to it in triggerers.
   
   I'm working on a draft AIP to add this support, I tested it a few months ago and the results were promising: https://github.com/apache/airflow/pull/36504


-- 
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] Use async db calls in WorkflowTrigger [airflow]

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


##########
tests/triggers/test_external_task.py:
##########
@@ -438,3 +464,12 @@ def test_serialization(self):
             "execution_dates": [timezone.datetime(2022, 1, 1)],
             "poll_interval": 5,
         }
+
+
+def mocked_get_count(*args, **kwargs):
+    time.sleep(0.3)

Review Comment:
   ```suggestion
       time.sleep(0.001)
   ```
   
   Try to use less possible delays in tests as possible.



##########
tests/triggers/test_external_task.py:
##########
@@ -438,3 +464,12 @@ def test_serialization(self):
             "execution_dates": [timezone.datetime(2022, 1, 1)],
             "poll_interval": 5,
         }
+
+
+def mocked_get_count(*args, **kwargs):
+    time.sleep(0.3)
+    return 1
+
+
+async def dummy_async_fun():

Review Comment:
   We avoid to use `dummy` into the Apache Airflow repo (still exists in some places) so better use here `fake`, `mocked`, and etc.



-- 
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] Use async db calls in WorkflowTrigger [airflow]

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


##########
airflow/triggers/external_task.py:
##########
@@ -98,44 +99,43 @@ async def run(self) -> typing.AsyncIterator[TriggerEvent]:
         """Check periodically tasks, task group or dag status."""
         while True:
             if self.failed_states:
-                failed_count = _get_count(
-                    self.execution_dates,
-                    self.external_task_ids,
-                    self.external_task_group_id,
-                    self.external_dag_id,
-                    self.failed_states,
-                )
+                failed_count = await self._get_count(self.failed_states)
                 if failed_count > 0:
                     yield TriggerEvent({"status": "failed"})
                     return
                 else:
                     yield TriggerEvent({"status": "success"})
                     return
             if self.skipped_states:
-                skipped_count = _get_count(
-                    self.execution_dates,
-                    self.external_task_ids,
-                    self.external_task_group_id,
-                    self.external_dag_id,
-                    self.skipped_states,
-                )
+                skipped_count = await self._get_count(self.skipped_states)
                 if skipped_count > 0:
                     yield TriggerEvent({"status": "skipped"})
                     return
-            allowed_count = _get_count(
-                self.execution_dates,
-                self.external_task_ids,
-                self.external_task_group_id,
-                self.external_dag_id,
-                self.allowed_states,
-            )
+            allowed_count = await self._get_count(self.allowed_states)
             if allowed_count == len(self.execution_dates):
                 yield TriggerEvent({"status": "success"})
                 return
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
+    @sync_to_async
+    def _get_count(self, states: typing.Iterable[str] | None) -> int:
+        """
+        Get the count of records against dttm filter and states. Async wrapper for _get_count.
+
+        :param states: task or dag states
+        :return The count of records.
+        """
+        return _get_count(
+            dttm_filter=self.execution_dates,
+            external_task_ids=self.external_task_ids,
+            external_task_group_id=self.external_task_group_id,
+            external_dag_id=self.external_dag_id,
+            states=states,
+        )
+
 
+@deprecated(reason=("TaskStateTrigger has been deprecated and will be removed in future."))

Review Comment:
   Thanks! This is done.



-- 
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