This is an automated email from the ASF dual-hosted git repository. kaxilnaik pushed a commit to branch v3-1-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit dd2a83ac56c2f0e3604fdb8ccfd565a5135fff7f Author: Duc Nguyen <[email protected]> AuthorDate: Sun Sep 14 00:49:13 2025 +0700 Allow slash in Variable keys from TaskSDK read/write(#55324) (cherry picked from commit 893f6279e6bb68ed875b2959300f26a58dc98a4a) --- .../api_fastapi/execution_api/routes/variables.py | 15 ++++- .../execution_api/versions/head/test_variables.py | 75 +++++++++++++++------- 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py index e1685987866..d2f5d21349c 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/variables.py @@ -55,7 +55,7 @@ log = logging.getLogger(__name__) @router.get( - "/{variable_key}", + "/{variable_key:path}", responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the variable"}, @@ -63,6 +63,9 @@ log = logging.getLogger(__name__) ) def get_variable(variable_key: str) -> VariableResponse: """Get an Airflow Variable.""" + if not variable_key: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Not Found") + try: variable_value = Variable.get(variable_key) except KeyError: @@ -78,7 +81,7 @@ def get_variable(variable_key: str) -> VariableResponse: @router.put( - "/{variable_key}", + "/{variable_key:path}", status_code=status.HTTP_201_CREATED, responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, @@ -87,12 +90,15 @@ def get_variable(variable_key: str) -> VariableResponse: ) def put_variable(variable_key: str, body: VariablePostBody): """Set an Airflow Variable.""" + if not variable_key: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Not Found") + Variable.set(key=variable_key, value=body.value, description=body.description) return {"message": "Variable successfully set"} @router.delete( - "/{variable_key}", + "/{variable_key:path}", status_code=status.HTTP_204_NO_CONTENT, responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, @@ -101,4 +107,7 @@ def put_variable(variable_key: str, body: VariablePostBody): ) def delete_variable(variable_key: str): """Delete an Airflow Variable.""" + if not variable_key: + raise HTTPException(status.HTTP_404_NOT_FOUND, detail="Not Found") + Variable.delete(key=variable_key) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py index 0ada165ed5b..0e112ad330a 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_variables.py @@ -65,17 +65,24 @@ def access_denied(client): class TestGetVariable: - def test_variable_get_from_db(self, client, session): - Variable.set(key="var1", value="value", session=session) + @pytest.mark.parametrize( + "key, value", + [ + ("var1", "value"), + ("var2/with_slash", "slash_value"), + ], + ) + def test_variable_get_from_db(self, client, session, key, value): + Variable.set(key=key, value=value, session=session) session.commit() - response = client.get("/execution/variables/var1") + response = client.get(f"/execution/variables/{key}") assert response.status_code == 200 - assert response.json() == {"key": "var1", "value": "value"} + assert response.json() == {"key": key, "value": value} # Remove connection - Variable.delete(key="var1", session=session) + Variable.delete(key=key, session=session) session.commit() @mock.patch.dict( @@ -88,13 +95,20 @@ class TestGetVariable: assert response.status_code == 200 assert response.json() == {"key": "key1", "value": "VALUE"} - def test_variable_get_not_found(self, client): - response = client.get("/execution/variables/non_existent_var") + @pytest.mark.parametrize( + "key", + [ + "non_existent_var", + "non/existent/slash/var", + ], + ) + def test_variable_get_not_found(self, client, key): + response = client.get(f"/execution/variables/{key}") assert response.status_code == 404 assert response.json() == { "detail": { - "message": "Variable with key 'non_existent_var' not found", + "message": f"Variable with key '{key}' not found", "reason": "not_found", } } @@ -117,14 +131,18 @@ class TestGetVariable: class TestPutVariable: @pytest.mark.parametrize( - "payload", + "key, payload", [ - pytest.param({"value": "{}", "description": "description"}, id="valid-payload"), - pytest.param({"value": "{}"}, id="missing-description"), + pytest.param("var_create", {"value": "{}", "description": "description"}, id="valid-payload"), + pytest.param("var_create", {"value": "{}"}, id="missing-description"), + pytest.param( + "var_create/with_slash", + {"value": "slash_value", "description": "Variable with slash"}, + id="slash-key", + ), ], ) - def test_should_create_variable(self, client, payload, session): - key = "var_create" + def test_should_create_variable(self, client, key, payload, session): response = client.put( f"/execution/variables/{key}", json=payload, @@ -132,7 +150,7 @@ class TestPutVariable: assert response.status_code == 201, response.json() assert response.json()["message"] == "Variable successfully set" - var_from_db = session.query(Variable).where(Variable.key == "var_create").first() + var_from_db = session.query(Variable).where(Variable.key == key).first() assert var_from_db is not None assert var_from_db.key == key assert var_from_db.val == payload["value"] @@ -179,8 +197,14 @@ class TestPutVariable: assert response.json()["detail"][0]["type"] == "extra_forbidden" assert response.json()["detail"][0]["msg"] == "Extra inputs are not permitted" - def test_overwriting_existing_variable(self, client, session): - key = "var_create" + @pytest.mark.parametrize( + "key", + [ + "var_create", + "var_create/with_slash", + ], + ) + def test_overwriting_existing_variable(self, client, session, key): Variable.set(key=key, value="value", session=session) session.commit() @@ -218,19 +242,26 @@ class TestPutVariable: class TestDeleteVariable: - def test_should_delete_variable(self, client, session): - for i in range(1, 3): - Variable.set(key=f"key{i}", value=i) + @pytest.mark.parametrize( + "keys_to_create, key_to_delete", + [ + (["key1", "key2"], "key1"), + (["key3/with_slash", "key4"], "key3/with_slash"), + ], + ) + def test_should_delete_variable(self, client, session, keys_to_create, key_to_delete): + for i, key in enumerate(keys_to_create, 1): + Variable.set(key=key, value=str(i)) vars = session.query(Variable).all() - assert len(vars) == 2 + assert len(vars) == len(keys_to_create) - response = client.delete("/execution/variables/key1") + response = client.delete(f"/execution/variables/{key_to_delete}") assert response.status_code == 204 vars = session.query(Variable).all() - assert len(vars) == 1 + assert len(vars) == len(keys_to_create) - 1 def test_should_not_delete_variable(self, client, session): Variable.set(key="key", value="value")
