diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/variable.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/variable.py index fd49a5eae46d6..b9403a878d277 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/variable.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/variable.py @@ -34,3 +34,9 @@ class VariablePostBody(StrictBaseModel): value: str | None = Field(alias="val") description: str | None = Field(default=None) + + +class VariableKeysResponse(StrictBaseModel): + """Variable keys schema for list responses.""" + + keys: list[str] diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py index 06f07aee82389..864bb02b1ca8e 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/__init__.py @@ -51,6 +51,7 @@ authenticated_router.include_router( task_reschedules.router, prefix="/task-reschedules", tags=["Task Reschedules"] ) +authenticated_router.include_router(variables.keys_router, prefix="/variables", tags=["Variables"]) authenticated_router.include_router(variables.router, prefix="/variables", tags=["Variables"]) authenticated_router.include_router(xcoms.router, prefix="/xcoms", tags=["XComs"]) authenticated_router.include_router(hitl.router, prefix="/hitlDetails", tags=["Human in the Loop"]) 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 a0a7cb56045c1..b0cb7d5eaf02f 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 @@ -20,9 +20,12 @@ import logging from typing import Annotated -from fastapi import APIRouter, Depends, HTTPException, Path, Request, status +from fastapi import APIRouter, Depends, HTTPException, Path, Query, Request, status +from sqlalchemy import select +from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.execution_api.datamodels.variable import ( + VariableKeysResponse, VariablePostBody, VariableResponse, ) @@ -59,6 +62,8 @@ async def has_variable_access( dependencies=[Depends(has_variable_access)], ) +keys_router = APIRouter() + log = logging.getLogger(__name__) @@ -120,3 +125,25 @@ def delete_variable( ): """Delete an Airflow Variable.""" Variable.delete(key=variable_key, team_name=team_name) + + +@keys_router.get( + "/keys", + responses={ + status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, + }, +) +def get_variable_keys( + session: SessionDep, + team_name: Annotated[str | None, Depends(get_team_name_dep)] = None, + prefix: Annotated[str | None, Query()] = None, +) -> VariableKeysResponse: + """Get Airflow Variable keys, optionally filtered by prefix.""" + stmt = select(Variable.key) + if prefix is not None: + stmt = stmt.where(Variable.key.startswith(prefix)) + if team_name is not None: + stmt = stmt.where(Variable.team_name == team_name) + + keys = session.scalars(stmt).all() + return VariableKeysResponse(keys=list(keys)) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index dfa27f53ebd91..31ce8b8ff8dcf 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -41,11 +41,16 @@ RemoveUpstreamMapIndexesField, ) from airflow.api_fastapi.execution_api.versions.v2026_04_17 import AddStateEndpoints, AddTeamNameField +from airflow.api_fastapi.execution_api.versions.v2026_04_28 import AddVariableKeysEndpoint from airflow.api_fastapi.execution_api.versions.v2026_06_16 import AddRetryPolicyFields bundle = VersionBundle( HeadVersion(), Version("2026-06-16", AddRetryPolicyFields), + Version( + "2026-04-28", + AddVariableKeysEndpoint, + ), Version( "2026-04-17", AddTeamNameField, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_28.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_28.py new file mode 100644 index 0000000000000..0bc300a499837 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_28.py @@ -0,0 +1,28 @@ +# 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. + +from __future__ import annotations + +from cadwyn import VersionChange, endpoint + + +class AddVariableKeysEndpoint(VersionChange): + """Add GET /variables/keys endpoint for listing variable keys with optional prefix filter.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (endpoint("/variables/keys", ["GET"]).didnt_exist,) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index aa9f07411f87d..d7c0a9d2b59fc 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -50,6 +50,7 @@ GetTaskStates, GetTICount, GetVariable, + GetVariableKeys, GetXCom, GetXComCount, GetXComSequenceItem, @@ -61,6 +62,7 @@ PrevSuccessfulDagRunResult, PutVariable, TaskStatesResult, + VariableKeysResult, VariableResult, XComCountResponse, XComResult, @@ -128,6 +130,7 @@ class DagFileParsingResult(BaseModel): DagFileParsingResult | GetConnection | GetVariable + | GetVariableKeys | PutVariable | GetTaskStates | GetTICount @@ -147,6 +150,7 @@ class DagFileParsingResult(BaseModel): DagFileParseRequest | ConnectionResult | VariableResult + | VariableKeysResult | TaskStatesResult | PreviousDagRunResult | PreviousTIResult @@ -628,6 +632,10 @@ def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int dump_opts = {"exclude_unset": True} else: resp = var + elif isinstance(msg, GetVariableKeys): + from airflow.sdk.execution_time.request_handlers import handle_get_variable_keys + + resp, dump_opts = handle_get_variable_keys(self.client, msg) elif isinstance(msg, PutVariable): self.client.variables.set(msg.key, msg.value, msg.description) elif isinstance(msg, DeleteVariable): diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 9a22e61ad8a03..a7e72d17f7d1f 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -71,6 +71,7 @@ GetTaskStates, GetTICount, GetVariable, + GetVariableKeys, GetXCom, MaskSecret, OKResponse, @@ -79,6 +80,7 @@ TaskStatesResult, TICount, UpdateHITLDetail, + VariableKeysResult, VariableResult, XComResult, _new_encoder, @@ -87,6 +89,7 @@ from airflow.sdk.execution_time.request_handlers import ( handle_get_connection, handle_get_variable, + handle_get_variable_keys, handle_mask_secret, ) from airflow.sdk.execution_time.supervisor import WatchedSubprocess, make_buffered_socket_reader @@ -299,6 +302,7 @@ def from_api_response(cls, response: HITLDetailResponse) -> HITLDetailResponseRe | messages.TriggerStateSync | ConnectionResult | VariableResult + | VariableKeysResult | XComResult | DagRunStateResult | DRCount @@ -320,6 +324,7 @@ def from_api_response(cls, response: HITLDetailResponse) -> HITLDetailResponseRe | GetConnection | DeleteVariable | GetVariable + | GetVariableKeys | PutVariable | DeleteXCom | GetXCom @@ -516,6 +521,8 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, r resp = self.client.variables.delete(msg.key) elif isinstance(msg, GetVariable): resp, dump_opts = handle_get_variable(self.client, msg) + elif isinstance(msg, GetVariableKeys): + resp, dump_opts = handle_get_variable_keys(self.client, msg) elif isinstance(msg, PutVariable): self.client.variables.set(msg.key, msg.value, msg.description) elif isinstance(msg, DeleteXCom): 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 fe7611636358d..391530dc8fb5b 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 @@ -244,6 +244,34 @@ def test_post_variable_access_denied(self, client, caplog): assert any(msg.startswith("Checking write access for task instance") for msg in caplog.messages) +class TestGetVariableKeys: + @pytest.mark.parametrize( + ("prefix", "expected_keys"), + [ + pytest.param(None, {"prod_db_url", "prod_api_key", "dev_debug"}, id="no-prefix"), + pytest.param("prod_", {"prod_db_url", "prod_api_key"}, id="with-prefix"), + pytest.param("staging_", set(), id="no-match"), + ], + ) + def test_get_variable_keys(self, client, session, prefix, expected_keys): + Variable.set(key="prod_db_url", value="postgres://...", session=session) + Variable.set(key="prod_api_key", value="secret", session=session) + Variable.set(key="dev_debug", value="true", session=session) + session.commit() + + params = {"prefix": prefix} if prefix is not None else {} + response = client.get("/execution/variables/keys", params=params) + + assert response.status_code == 200 + assert set(response.json()["keys"]) == expected_keys + + def test_get_variable_keys_empty_db(self, client): + response = client.get("/execution/variables/keys") + + assert response.status_code == 200 + assert response.json() == {"keys": []} + + class TestDeleteVariable: @pytest.mark.parametrize( ("keys_to_create", "key_to_delete"), diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_28/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_28/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_28/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_28/test_variables.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_28/test_variables.py new file mode 100644 index 0000000000000..19654d255a32a --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_28/test_variables.py @@ -0,0 +1,35 @@ +# 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. + +from __future__ import annotations + +import pytest + +pytestmark = pytest.mark.db_test + + +@pytest.fixture +def old_ver_client(client): + """Last released execution API before `GET /variables/keys` was added.""" + client.headers["Airflow-API-Version"] = "2026-04-17" + return client + + +def test_variable_keys_endpoint_not_available_in_previous_version(old_ver_client): + response = old_ver_client.get("/execution/variables/keys") + + assert response.status_code == 404 diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 54927794bf17e..bfa90d72e9b8c 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -69,6 +69,7 @@ TITerminalStatePayload, TriggerDAGRunPayload, ValidationError as RemoteValidationError, + VariableKeysResponse, VariablePostBody, VariableResponse, XComResponse, @@ -477,6 +478,14 @@ def delete( # decouple from the server response string return OKResponse(ok=True) + def keys(self, prefix: str | None = None) -> VariableKeysResponse: + """List variable keys from the API server, optionally filtered by key prefix.""" + params: dict[str, str] = {} + if prefix is not None: + params["prefix"] = prefix + resp = self.client.get("variables/keys", params=params) + return VariableKeysResponse.model_validate_json(resp.read()) + class XComOperations: __slots__ = ("client",) diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index b5b100154c389..e9048b12b2d01 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -442,6 +442,17 @@ class ValidationError(BaseModel): ctx: Annotated[dict[str, Any] | None, Field(title="Context")] = None +class VariableKeysResponse(BaseModel): + """ + Variable keys schema for list responses. + """ + + model_config = ConfigDict( + extra="forbid", + ) + keys: Annotated[list[str], Field(title="Keys")] + + class VariablePostBody(BaseModel): """ Request body schema for creating variables. diff --git a/task-sdk/src/airflow/sdk/definitions/variable.py b/task-sdk/src/airflow/sdk/definitions/variable.py index 2e4c9aae3ca0f..eed381cd7b201 100644 --- a/task-sdk/src/airflow/sdk/definitions/variable.py +++ b/task-sdk/src/airflow/sdk/definitions/variable.py @@ -67,6 +67,23 @@ def set(cls, key: str, value: Any, description: str | None = None, serialize_jso except AirflowRuntimeError as e: log.exception(e) + @classmethod + def keys(cls, prefix: str | None = None) -> list[str]: + """ + Return Variable keys that start with the given prefix. + + The keys are fetched lazily on first access (iteration, indexing, len, etc.) + and cached for subsequent access. Only keys stored in the metadata database + are returned — secrets backends are not consulted. + + :param prefix: Optional key prefix to filter by. If None, all keys are returned. + """ + import lazy_object_proxy + + from airflow.sdk.execution_time.context import _get_variable_keys + + return lazy_object_proxy.Proxy(lambda: _get_variable_keys(prefix=prefix)) + @classmethod def delete(cls, key: str) -> None: from airflow.sdk.exceptions import AirflowRuntimeError diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 87c7881333ad4..0247ba0ce3e7b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -545,6 +545,11 @@ def from_variable_response(cls, variable_response: VariableResponse) -> Variable return cls(**variable_response.model_dump(exclude_defaults=True), type="VariableResult") +class VariableKeysResult(BaseModel): + keys: list[str] + type: Literal["VariableKeysResult"] = "VariableKeysResult" + + class DagRunResult(DagRun): type: Literal["DagRunResult"] = "DagRunResult" @@ -728,6 +733,7 @@ def from_api_response(cls, dag_response: DagResponse) -> DagResult: | TaskBreadcrumbsResult | TaskStatesResult | VariableResult + | VariableKeysResult | XComCountResponse | XComResult | XComSequenceIndexResult @@ -862,6 +868,11 @@ class GetVariable(BaseModel): type: Literal["GetVariable"] = "GetVariable" +class GetVariableKeys(BaseModel): + prefix: str | None = None + type: Literal["GetVariableKeys"] = "GetVariableKeys" + + class PutVariable(BaseModel): key: str value: str | None @@ -1061,6 +1072,7 @@ class GetDag(BaseModel): | GetTaskBreadcrumbs | GetTaskStates | GetVariable + | GetVariableKeys | GetXCom | GetXComCount | GetXComSequenceItem diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 66c1f3aa8b7eb..757e56f56db69 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -278,6 +278,17 @@ def _get_variable(key: str, deserialize_json: bool) -> Any: ) +def _get_variable_keys(prefix: str | None = None) -> list[str]: + from airflow.sdk.execution_time.comms import GetVariableKeys, VariableKeysResult + from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS + + msg = SUPERVISOR_COMMS.send(GetVariableKeys(prefix=prefix)) + if not isinstance(msg, VariableKeysResult): + return [] + + return msg.keys + + def _set_variable(key: str, value: Any, description: str | None = None, serialize_json: bool = False) -> None: # TODO: This should probably be moved to a separate module like `airflow.sdk.execution_time.comms` # or `airflow.sdk.execution_time.variable` diff --git a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py index eed3e840e395b..3b528c2e71a15 100644 --- a/task-sdk/src/airflow/sdk/execution_time/request_handlers.py +++ b/task-sdk/src/airflow/sdk/execution_time/request_handlers.py @@ -31,13 +31,16 @@ from airflow.sdk.api.datamodels._generated import ( ConnectionResponse, + VariableKeysResponse, VariableResponse, ) from airflow.sdk.execution_time.comms import ( ConnectionResult, GetConnection, GetVariable, + GetVariableKeys, MaskSecret, + VariableKeysResult, VariableResult, ) from airflow.sdk.log import mask_secret @@ -70,6 +73,16 @@ def handle_get_variable(client: Client, msg: GetVariable) -> tuple[BaseModel | N return var, {} +def handle_get_variable_keys( + client: Client, msg: GetVariableKeys +) -> tuple[BaseModel | None, dict[str, bool]]: + """Fetch variable keys filtered by prefix.""" + result = client.variables.keys(prefix=msg.prefix) + if not isinstance(result, VariableKeysResponse): + return result, {} + return VariableKeysResult(keys=result.keys), {"exclude_unset": True} + + def handle_mask_secret(msg: MaskSecret) -> None: """Register a value with the secrets masker.""" mask_secret(msg.value, msg.name) diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 375c5a9e30b8e..356fe9824e5c8 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -90,6 +90,7 @@ GetTaskStates, GetTICount, GetVariable, + GetVariableKeys, GetXCom, GetXComCount, GetXComSequenceItem, @@ -124,6 +125,7 @@ from airflow.sdk.execution_time.request_handlers import ( handle_get_connection, handle_get_variable, + handle_get_variable_keys, handle_mask_secret, ) @@ -1423,6 +1425,8 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: resp, dump_opts = handle_get_connection(self.client, msg) elif isinstance(msg, GetVariable): resp, dump_opts = handle_get_variable(self.client, msg) + elif isinstance(msg, GetVariableKeys): + resp, dump_opts = handle_get_variable_keys(self.client, msg) elif isinstance(msg, GetXCom): xcom = self.client.xcoms.get( msg.dag_id, msg.run_id, msg.task_id, msg.key, msg.map_index, msg.include_prior_dates diff --git a/task-sdk/tests/task_sdk/definitions/test_variables.py b/task-sdk/tests/task_sdk/definitions/test_variables.py index 3717f834735ff..beaefab83d917 100644 --- a/task-sdk/tests/task_sdk/definitions/test_variables.py +++ b/task-sdk/tests/task_sdk/definitions/test_variables.py @@ -25,7 +25,7 @@ from airflow.sdk import Variable from airflow.sdk.configuration import initialize_secrets_backends -from airflow.sdk.execution_time.comms import PutVariable, VariableResult +from airflow.sdk.execution_time.comms import GetVariableKeys, PutVariable, VariableKeysResult, VariableResult from airflow.sdk.execution_time.secrets import DEFAULT_SECRETS_SEARCH_PATH_WORKERS from tests_common.test_utils.config import conf_vars @@ -90,6 +90,53 @@ def test_var_set(self, key, value, description, serialize_json, mock_supervisor_ ) +class TestVariableKeys: + @pytest.mark.parametrize( + ("prefix", "keys"), + [ + pytest.param( + None, + ["prod_db", "prod_api", "dev_debug"], + id="all", + ), + pytest.param( + "prod_", + ["prod_db", "prod_api"], + id="with-prefix", + ), + pytest.param( + "nonexistent_", + [], + id="empty-result", + ), + ], + ) + def test_keys(self, prefix, keys, mock_supervisor_comms): + mock_supervisor_comms.send.return_value = VariableKeysResult(keys=keys) + + results = Variable.keys(prefix=prefix) + + # keys() is lazy — no API call until the proxy is accessed + mock_supervisor_comms.send.assert_not_called() + + materialized = list(results) + + mock_supervisor_comms.send.assert_called_once_with(msg=GetVariableKeys(prefix=prefix)) + assert materialized == keys + + def test_keys_cached_after_first_access(self, mock_supervisor_comms): + mock_supervisor_comms.send.return_value = VariableKeysResult(keys=["a", "b"]) + + results = Variable.keys(prefix="x_") + + # Multiple accesses should only trigger the API call once + list(results) + list(results) + len(results) + + mock_supervisor_comms.send.assert_called_once_with(msg=GetVariableKeys(prefix="x_")) + + class TestVariableFromSecrets: def test_var_get_from_secrets_found(self, mock_supervisor_comms, tmp_path): """Tests getting a variable from secrets backend.""" diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 3695af1fff592..90fc42fca5c9e 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -97,6 +97,7 @@ GetTaskStates, GetTICount, GetVariable, + GetVariableKeys, GetXCom, GetXComCount, GetXComSequenceItem, @@ -127,6 +128,7 @@ TriggerDagRun, UpdateHITLDetail, ValidateInletsAndOutlets, + VariableKeysResult, VariableResult, XComCountResponse, XComResult, @@ -1548,6 +1550,19 @@ class RequestTestCase: ), expected_body={"ok": True, "type": "OKResponse"}, ), + RequestTestCase( + message=GetVariableKeys(prefix="test_"), + test_id="get_variable_keys", + client_mock=ClientMock( + method_path="variables.keys", + kwargs={"prefix": "test_"}, + response=VariableKeysResult(keys=["test_key"]), + ), + expected_body={ + "keys": ["test_key"], + "type": "VariableKeysResult", + }, + ), RequestTestCase( message=DeferTask(next_method="execute_callback", classpath="my-classpath"), test_id="patch_task_instance_to_deferred",