-
Notifications
You must be signed in to change notification settings - Fork 17.3k
Improve error handling for SUPERVISOR_COMMS access outside task context #61630
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
afd24cf
3ee1138
f2f1f3c
c5b7eb2
b5ac336
5631eed
f5aedfd
f3ac663
e8ba49a
41abed9
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -332,9 +332,26 @@ def _get_variable(key: str, deserialize_json: bool) -> Any: | |
| ) | ||
|
|
||
| # If no backend found the variable, raise a not found error (mirrors _get_connection) | ||
| from airflow.sdk.exceptions import AirflowRuntimeError, ErrorType | ||
| from airflow.sdk.execution_time import task_runner | ||
| from airflow.sdk.execution_time.comms import ErrorResponse | ||
|
|
||
| if not hasattr(task_runner, "SUPERVISOR_COMMS"): | ||
| raise AirflowRuntimeError( | ||
| ErrorResponse( | ||
| error=ErrorType.VARIABLE_NOT_FOUND, | ||
| detail={ | ||
| "message": ( | ||
| f"Variable '{key}' not found. Note: SUPERVISOR_COMMS is not available, " | ||
| "which means this code is running outside a task execution context " | ||
| "(e.g., at the top level of a DAG file). " | ||
| "Consider using environment variables (AIRFLOW_VAR_<key>), " | ||
| "Jinja templates ({{ var.value.<key> }}), " | ||
| "or move the Variable.get() call inside a task function." | ||
| ) | ||
| }, | ||
| ) | ||
| ) | ||
|
|
||
| raise AirflowRuntimeError( | ||
| ErrorResponse(error=ErrorType.VARIABLE_NOT_FOUND, detail={"message": f"Variable {key} not found"}) | ||
| ) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Related:
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good catch, that one slipped through. Fixed. |
||
|
|
@@ -345,17 +362,31 @@ def _get_variable(key: str, deserialize_json: bool) -> Any: | |
|
|
||
| def _get_variable_keys(prefix: str | None = None) -> list[str]: | ||
| from airflow.sdk.exceptions import AirflowRuntimeError | ||
| from airflow.sdk.execution_time import task_runner | ||
| from airflow.sdk.execution_time.comms import ( | ||
| ErrorResponse, | ||
| GetVariableKeys, | ||
| VariableKeysResult, | ||
| ) | ||
| from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS | ||
|
|
||
| if not hasattr(task_runner, "SUPERVISOR_COMMS"): | ||
| raise AirflowRuntimeError( | ||
| ErrorResponse( | ||
| error=ErrorType.GENERIC_ERROR, | ||
| detail={ | ||
| "message": ( | ||
| "Variable.keys() requires a task execution context (SUPERVISOR_COMMS is not available). " | ||
| "This typically happens when calling Variable.keys() at the top level of a DAG file " | ||
| "or outside of a running task. Variable.keys() can only be used inside a task." | ||
| ) | ||
| }, | ||
| ) | ||
| ) | ||
|
|
||
| all_keys: list[str] = [] | ||
| offset = 0 | ||
| while True: | ||
| msg = SUPERVISOR_COMMS.send( | ||
| msg = task_runner.SUPERVISOR_COMMS.send( | ||
| GetVariableKeys(prefix=prefix, limit=_VARIABLE_KEYS_PAGE_SIZE, offset=offset) | ||
| ) | ||
| if isinstance(msg, ErrorResponse): | ||
|
|
@@ -377,11 +408,25 @@ def _set_variable(key: str, value: Any, description: str | None = None, serializ | |
| # keep Task SDK as a separate package than execution time mods. | ||
| import json | ||
|
|
||
| from airflow.sdk.execution_time import task_runner | ||
| from airflow.sdk.execution_time.cache import SecretCache | ||
| from airflow.sdk.execution_time.comms import PutVariable | ||
| from airflow.sdk.execution_time.comms import ErrorResponse, PutVariable | ||
| from airflow.sdk.execution_time.secrets.execution_api import ExecutionAPISecretsBackend | ||
| from airflow.sdk.execution_time.supervisor import ensure_secrets_backend_loaded | ||
| from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS | ||
|
|
||
| if not hasattr(task_runner, "SUPERVISOR_COMMS"): | ||
| raise AirflowRuntimeError( | ||
| ErrorResponse( | ||
| error=ErrorType.GENERIC_ERROR, | ||
| detail={ | ||
| "message": ( | ||
| "Variable.set() requires a task execution context (SUPERVISOR_COMMS is not available). " | ||
| "This typically happens when calling Variable.set() at the top level of a DAG file " | ||
| "or outside of a running task. Variable.set() can only be used inside a task." | ||
| ) | ||
| }, | ||
| ) | ||
| ) | ||
|
|
||
| # check for write conflicts on the worker | ||
| for secrets_backend in ensure_secrets_backend_loaded(): | ||
|
|
@@ -412,7 +457,7 @@ def _set_variable(key: str, value: Any, description: str | None = None, serializ | |
| except Exception as e: | ||
| log.exception(e) | ||
|
|
||
| SUPERVISOR_COMMS.send(PutVariable(key=key, value=value, description=description)) | ||
| task_runner.SUPERVISOR_COMMS.send(PutVariable(key=key, value=value, description=description)) | ||
|
|
||
| # Invalidate cache after setting the variable | ||
| SecretCache.invalidate_variable(key) | ||
|
|
@@ -424,11 +469,25 @@ def _delete_variable(key: str) -> None: | |
| # A reason to not move it to `airflow.sdk.execution_time.comms` is that it | ||
| # will make that module depend on Task SDK, which is not ideal because we intend to | ||
| # keep Task SDK as a separate package than execution time mods. | ||
| from airflow.sdk.execution_time import task_runner | ||
| from airflow.sdk.execution_time.cache import SecretCache | ||
| from airflow.sdk.execution_time.comms import DeleteVariable | ||
| from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS | ||
| from airflow.sdk.execution_time.comms import DeleteVariable, ErrorResponse | ||
|
|
||
| if not hasattr(task_runner, "SUPERVISOR_COMMS"): | ||
| raise AirflowRuntimeError( | ||
| ErrorResponse( | ||
| error=ErrorType.GENERIC_ERROR, | ||
| detail={ | ||
| "message": ( | ||
| "Variable.delete() requires a task execution context (SUPERVISOR_COMMS is not available). " | ||
| "This typically happens when calling Variable.delete() at the top level of a DAG file " | ||
| "or outside of a running task. Variable.delete() can only be used inside a task." | ||
| ) | ||
| }, | ||
| ) | ||
| ) | ||
|
|
||
| msg = SUPERVISOR_COMMS.send(DeleteVariable(key=key)) | ||
| msg = task_runner.SUPERVISOR_COMMS.send(DeleteVariable(key=key)) | ||
| if TYPE_CHECKING: | ||
| assert isinstance(msg, OKResponse) | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This OTLP-noise filter is unrelated to the
SUPERVISOR_COMMSerror-handling scope of the PR. It's a legitimate CI-flake fix (other loggers pollutingcaplog.messages), but bundling it here makes the change harder to bisect later if it ever causes regressions.Either split it into its own PR or call it out explicitly in the description so reviewers know it's intentional scope creep.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fair, the scope creep is real. Kept it in this PR since it surfaced while debugging the CI flake that was blocking the merge, but called it out explicitly under "Out-of-scope change included" in the description, referencing commit f2f1f3c so it stays discoverable if it ever needs to be bisected.