From 9316ed6df85234ba7f36437152329d9b1a27424d Mon Sep 17 00:00:00 2001
From: Shahar Epstein <60007259+shahar1@users.noreply.github.com>
Date: Sat, 21 Dec 2024 19:41:16 +0200
Subject: [PATCH] Allow fetching XCom with forward slash from the API and
 escape it in the UI (#45134)

---
 airflow/api_connexion/openapi/v1.yaml         |  1 +
 airflow/www/static/js/api/useTaskXcom.ts      | 18 ++++++++-------
 newsfragments/45134.bugfix.rst                |  1 +
 .../endpoints/test_xcom_endpoint.py           | 23 +++++++++++++++----
 4 files changed, 31 insertions(+), 12 deletions(-)
 create mode 100644 newsfragments/45134.bugfix.rst

diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index fcf41462d233c..0ead649f4215a 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -5528,6 +5528,7 @@ components:
       name: xcom_key
       schema:
         type: string
+        format: path
       required: true
       description: The XCom key.
 
diff --git a/airflow/www/static/js/api/useTaskXcom.ts b/airflow/www/static/js/api/useTaskXcom.ts
index 403233285eb11..d8758e60b9056 100644
--- a/airflow/www/static/js/api/useTaskXcom.ts
+++ b/airflow/www/static/js/api/useTaskXcom.ts
@@ -57,14 +57,16 @@ export const useTaskXcomEntry = ({
 }: TaskXcomProps) =>
   useQuery(
     ["taskXcom", dagId, dagRunId, taskId, mapIndex, xcomKey, tryNumber],
-    () =>
-      axios.get<AxiosResponse, API.XCom>(
-        getMetaValue("task_xcom_entry_api")
-          .replace("_DAG_RUN_ID_", dagRunId)
-          .replace("_TASK_ID_", taskId)
-          .replace("_XCOM_KEY_", xcomKey),
-        { params: { map_index: mapIndex, stringify: false } }
-      ),
+    () => {
+      const taskXcomEntryApiUrl = getMetaValue("task_xcom_entry_api")
+        .replace("_DAG_RUN_ID_", dagRunId)
+        .replace("_TASK_ID_", taskId)
+        .replace("_XCOM_KEY_", encodeURIComponent(xcomKey));
+
+      return axios.get<AxiosResponse, API.XCom>(taskXcomEntryApiUrl, {
+        params: { map_index: mapIndex, stringify: false },
+      });
+    },
     {
       enabled: !!xcomKey,
     }
diff --git a/newsfragments/45134.bugfix.rst b/newsfragments/45134.bugfix.rst
new file mode 100644
index 0000000000000..09aaae23a3487
--- /dev/null
+++ b/newsfragments/45134.bugfix.rst
@@ -0,0 +1 @@
+(v2 API & UI) Allow fetching XCom with forward slash from the API and escape it in the UI
diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py
index 4f0072860fd3e..ba90e28f3ac2d 100644
--- a/tests/api_connexion/endpoints/test_xcom_endpoint.py
+++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py
@@ -226,52 +226,67 @@ def _create_xcom_entry(
         )
 
     @pytest.mark.parametrize(
-        "allowed, query, expected_status_or_value",
+        "allowed, query, expected_status_or_value, key",
         [
             pytest.param(
                 True,
                 "?deserialize=true",
                 "real deserialized TEST_VALUE",
+                "key",
                 id="true",
             ),
             pytest.param(
                 False,
                 "?deserialize=true",
                 400,
+                "key",
                 id="disallowed",
             ),
             pytest.param(
                 True,
                 "?deserialize=false",
                 "orm deserialized TEST_VALUE",
+                "key",
                 id="false-irrelevant",
             ),
             pytest.param(
                 False,
                 "?deserialize=false",
                 "orm deserialized TEST_VALUE",
+                "key",
                 id="false",
             ),
             pytest.param(
                 True,
                 "",
                 "orm deserialized TEST_VALUE",
+                "key",
                 id="default-irrelevant",
             ),
             pytest.param(
                 False,
                 "",
                 "orm deserialized TEST_VALUE",
+                "key",
                 id="default",
             ),
+            pytest.param(
+                False,
+                "",
+                "orm deserialized TEST_VALUE",
+                "key/with/slashes",
+                id="key-with-slashes",
+            ),
         ],
     )
     @conf_vars({("core", "xcom_backend"): "tests.api_connexion.endpoints.test_xcom_endpoint.CustomXCom"})
-    def test_custom_xcom_deserialize(self, allowed: bool, query: str, expected_status_or_value: int | str):
+    def test_custom_xcom_deserialize(
+        self, allowed: bool, query: str, expected_status_or_value: int | str, key: str
+    ):
         XCom = resolve_xcom_backend()
-        self._create_xcom_entry("dag", "run", utcnow(), "task", "key", backend=XCom)
+        self._create_xcom_entry("dag", "run", utcnow(), "task", key, backend=XCom)
 
-        url = f"/api/v1/dags/dag/dagRuns/run/taskInstances/task/xcomEntries/key{query}"
+        url = f"/api/v1/dags/dag/dagRuns/run/taskInstances/task/xcomEntries/{key}{query}"
         with mock.patch("airflow.api_connexion.endpoints.xcom_endpoint.XCom", XCom):
             with conf_vars({("api", "enable_xcom_deserialize_support"): str(allowed)}):
                 response = self.client.get(url, environ_overrides={"REMOTE_USER": "test"})