You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "amoghrajesh (via GitHub)" <gi...@apache.org> on 2025/01/02 09:49:55 UTC

[PR] AIP-72: Gracefully handle not-found XCOMs in task sdk API client [airflow]

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

   <!--
    Licensed to the Apache Software Foundation (ASF) under one
    or more contributor license agreements.  See the NOTICE file
    distributed with this work for additional information
    regarding copyright ownership.  The ASF licenses this file
    to you under the Apache License, Version 2.0 (the
    "License"); you may not use this file except in compliance
    with the License.  You may obtain a copy of the License at
   
      http://www.apache.org/licenses/LICENSE-2.0
   
    Unless required by applicable law or agreed to in writing,
    software distributed under the License is distributed on an
    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    KIND, either express or implied.  See the License for the
    specific language governing permissions and limitations
    under the License.
    -->
   
   <!--
   Thank you for contributing! Please make sure that your code changes
   are covered with tests. And in case of new features or big changes
   remember to adjust the documentation.
   
   Feel free to ping committers for the review!
   
   In case of an existing issue, reference it using one of the following:
   
   closes: #ISSUE
   related: #ISSUE
   
   How to write a good git commit message:
   http://chris.beams.io/posts/git-commit/
   -->
   
   closes: #45341 
   
   ### Testing results
   DAG used:
   ```
   from airflow import DAG
   from airflow.providers.standard.operators.python import PythonOperator
   
   def push_to_xcom(**kwargs):
       value = "Hello, XCom!"
       return value
   
   
   def pull_from_xcom(**kwargs):
       ti = kwargs['ti']
       xcom_value = ti.xcom_pull(task_ids='invalid_id')
       print(f"Retrieved XCom Value: {xcom_value}")
   
   
   with DAG(
       'xcom_example',
       schedule=None,
       catchup=False,
   ) as dag:
   
       push_xcom_task = PythonOperator(
           task_id='push_xcom_task',
           python_callable=push_to_xcom,
       )
   
       pull_xcom_task = PythonOperator(
           task_id='pull_xcom_task',
           python_callable=pull_from_xcom,
       )
   
       push_xcom_task >> pull_xcom_task
   ```
   
   #### Legacy
   Task 1 XCom pushed:
   <img width="1727" alt="image" src="https://github.com/user-attachments/assets/b4ab2d7a-72e4-4e13-ba0d-2b3faa7aee7f" />
   
   Task 2 XCom consumed as None when absent:
   <img width="1727" alt="image" src="https://github.com/user-attachments/assets/461744d5-6c77-4630-b98f-87d7d0942b67" />
   
   
   
   #### Task SDK
   Task 1 XCom pushed:
   <img width="1727" alt="image" src="https://github.com/user-attachments/assets/c3967e43-8c56-4779-ab00-c724c2d18ab3" />
   
   Task 2 XCom consumed as None when absent:
   <img width="1727" alt="image" src="https://github.com/user-attachments/assets/1a3c30ca-e0e8-47d4-a555-c4311b4ca094" />
   
   
   <!-- Please keep an empty line above the dashes. -->
   ---
   **^ Add meaningful description above**
   Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
   In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
   In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
   In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
   


-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   Okay, I misunderstood on the type of exception. Let me add that and add a test case, for another case



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   I added a case for 500 (which is a valid response from the api) and tested it out by hacking the API to send that to see if the scheduler crashes, it doesn't all good.



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   (And it mayb be a good idea to have a test case for non-404 exceptions)



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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

   Hi @uranusjr are you ok with the reply on this comment? https://github.com/apache/airflow/pull/45344#discussion_r1900781515


-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   Uuups, review problem. MUST re-raise other!



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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

   > @amoghrajesh - can you please rebase that one -> we found and issue with @jscheffl with the new caching scheme - fixed in #45347 that would run "main" version of the tests. I am asking in all affected PRs to rebase.
   
   Oh ok, let me rebase this one


-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   No I mean with the current code, if a ServerResponseError that’s not 404 is raised by the XCom API, this would silently suppress that exception and fall over the the line after this. Since `resp` is not defined if the `try` block is unsuccessful, this would result in this function failing with a very unhelpful `name 'resp' is not defined`. A `raise` statement should be inserted after line 239 to retain the original exception.



-- 
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] AIP-72: Gracefully handle not-found XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -726,7 +727,12 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = var_result.model_dump_json().encode()
         elif isinstance(msg, GetXCom):
             xcom = self.client.xcoms.get(msg.dag_id, msg.run_id, msg.task_id, msg.key, msg.map_index)
-            xcom_result = XComResult.from_xcom_response(xcom)
+            if isinstance(xcom, XComResponse):
+                xcom_result = XComResult.from_xcom_response(xcom)
+            else:
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence making this resp with key as `key` and value as None, so that the message is sent back to task runner.
+                xcom_result = XComResult.from_xcom_response(XComResponse(key=msg.key, value=None))

Review Comment:
   This branch indicates that the "xcom" is not found (The server responded 404).
   Instead of complicating the client which is the sdk, I chose to send a message to the task runner (it will receive it through get_message) with an empty "value". If the value is empty, the `xcom_pull` in task runner will return the `default` which is `None` anyways: https://github.com/apache/airflow/blob/main/task_sdk/src/airflow/sdk/execution_time/task_runner.py#L198-L204 which is the behaviour I am intending.



##########
task_sdk/tests/execution_time/test_supervisor.py:
##########
@@ -762,14 +764,15 @@ def watched_subprocess(self, mocker):
         )
 
     @pytest.mark.parametrize(
-        ["message", "expected_buffer", "client_attr_path", "method_arg", "mock_response"],
+        ["message", "expected_buffer", "client_attr_path", "method_arg", "mock_response", "decoded_buffer"],

Review Comment:
   The `decoded_buffer` parameter overrides the comparison in https://github.com/apache/airflow/pull/45344/files#diff-f55b9ba6be46de22b2838085bf1e0aa4193797c3ff7d0cdc4e06a626c96df1baR970-R971. The idea is that when the supervisor a response to the buffer, if `decoded_buffer` is present, it is used for comparison instead of `mock_response`. This is logical because for cases like these when the client responds 404 and we do not pass that all the way to task runner but transform it to something else like in this case (client responded 404, but we sent XcomResponse with value=None).



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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

   @amoghrajesh - can you please rebase that one -> we found and issue with @jscheffl with the new caching scheme - fixed in https://github.com/apache/airflow/pull/45347 that would run "main" version of the tests. I am asking in all affected PRs to rebase.


-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   We have an event hook defined here: https://github.com/apache/airflow/blob/main/task_sdk/src/airflow/sdk/api/client.py#L300
   With handlers here: https://github.com/apache/airflow/blob/main/task_sdk/src/airflow/sdk/api/client.py#L92
   
   The response hook makes sure that only a ServerResponseError is raised, if there is an error:
   https://github.com/apache/airflow/blob/main/task_sdk/src/airflow/sdk/api/client.py#L356-L392



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -726,7 +727,12 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = var_result.model_dump_json().encode()
         elif isinstance(msg, GetXCom):
             xcom = self.client.xcoms.get(msg.dag_id, msg.run_id, msg.task_id, msg.key, msg.map_index)
-            xcom_result = XComResult.from_xcom_response(xcom)
+            if isinstance(xcom, XComResponse):
+                xcom_result = XComResult.from_xcom_response(xcom)
+            else:
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence making this resp with key as `key` and value as None, so that the message is sent back to task runner.
+                xcom_result = XComResult.from_xcom_response(XComResponse(key=msg.key, value=None))

Review Comment:
   Let me actually make that change



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -726,7 +727,12 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = var_result.model_dump_json().encode()
         elif isinstance(msg, GetXCom):
             xcom = self.client.xcoms.get(msg.dag_id, msg.run_id, msg.task_id, msg.key, msg.map_index)
-            xcom_result = XComResult.from_xcom_response(xcom)
+            if isinstance(xcom, XComResponse):
+                xcom_result = XComResult.from_xcom_response(xcom)
+            else:
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence making this resp with key as `key` and value as None, so that the message is sent back to task runner.
+                xcom_result = XComResult.from_xcom_response(XComResponse(key=msg.key, value=None))

Review Comment:
   Why are you not "just" moving the fallback handling into task_sdk/src/airflow/sdk/api/client.py but return an ErrorReponse which then later needs to be converted? Is this not over-complicating?



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/api/client.py:
##########
@@ -219,7 +219,24 @@ def get(
         params = {}
         if map_index is not None:
             params.update({"map_index": map_index})
-        resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        try:
+            resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}", params=params)
+        except ServerResponseError as e:
+            if e.response.status_code == HTTPStatus.NOT_FOUND:
+                log.error(
+                    "XCom not found",
+                    dag_id=dag_id,
+                    run_id=run_id,
+                    task_id=task_id,
+                    key=key,
+                    map_index=map_index,
+                    detail=e.detail,
+                    status_code=e.response.status_code,
+                )
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence returning with key as `key` and value as `None`, so that the message is sent back to task runner
+                # and the default value of None in xcom_pull is used.
+                return XComResponse(key=key, value=None)

Review Comment:
   Should this re-raise the exception for other errors?



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -726,7 +727,12 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = var_result.model_dump_json().encode()
         elif isinstance(msg, GetXCom):
             xcom = self.client.xcoms.get(msg.dag_id, msg.run_id, msg.task_id, msg.key, msg.map_index)
-            xcom_result = XComResult.from_xcom_response(xcom)
+            if isinstance(xcom, XComResponse):
+                xcom_result = XComResult.from_xcom_response(xcom)
+            else:
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence making this resp with key as `key` and value as None, so that the message is sent back to task runner.
+                xcom_result = XComResult.from_xcom_response(XComResponse(key=msg.key, value=None))

Review Comment:
   I actually handled it that way and its much simpler and cleaner.



##########
task_sdk/tests/execution_time/test_supervisor.py:
##########
@@ -762,14 +764,15 @@ def watched_subprocess(self, mocker):
         )
 
     @pytest.mark.parametrize(
-        ["message", "expected_buffer", "client_attr_path", "method_arg", "mock_response"],
+        ["message", "expected_buffer", "client_attr_path", "method_arg", "mock_response", "decoded_buffer"],

Review Comment:
   Not applicable, handled differently



-- 
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] AIP-72: Gracefully handle 'not-found' XCOMs in task sdk API client [airflow]

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


##########
task_sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -726,7 +727,12 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger):
             resp = var_result.model_dump_json().encode()
         elif isinstance(msg, GetXCom):
             xcom = self.client.xcoms.get(msg.dag_id, msg.run_id, msg.task_id, msg.key, msg.map_index)
-            xcom_result = XComResult.from_xcom_response(xcom)
+            if isinstance(xcom, XComResponse):
+                xcom_result = XComResult.from_xcom_response(xcom)
+            else:
+                # Airflow 2.x just ignores the absence of an XCom and moves on with a return value of None
+                # Hence making this resp with key as `key` and value as None, so that the message is sent back to task runner.
+                xcom_result = XComResult.from_xcom_response(XComResponse(key=msg.key, value=None))

Review Comment:
   I briefly gave it some thought and thought that the client shouldn't have to handle such a scenario. But now that I think about it, that will significantly simplify it



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