amoghrajesh commented on code in PR #72546:
URL: https://github.com/apache/airflow/pull/72546#discussion_r4034053225
##########
providers/standard/src/airflow/providers/standard/operators/bash.py:
##########
@@ -261,5 +328,95 @@ def _is_inline_command(cls, bash_command: str) -> bool:
"""Return True if the bash command is an inline string. False if it's
a bash script file."""
return not bash_command.endswith(tuple(cls.template_ext))
+ @contextlib.contextmanager
+ def _xcom_workspace(self) -> Iterator[tuple[Path, Path | None]]:
+ """
+ Create the ``$AIRFLOW_XCOM_DIR`` workspace for a single run.
+
+ Kept separate from the ``working_directory()`` script cwd tempdir so a
rendered ``.sh``
+ file can never collide with an XCom key.
+ """
+ with tempfile.TemporaryDirectory(prefix="airflow-bash-xcom-") as
workspace:
+ xcom_dir = Path(workspace) / "xcom"
+ xcom_dir.mkdir()
+ bin_dir: Path | None = None
+ if self.xcom_helper_name:
+ bin_dir = Path(workspace) / "bin"
+ bin_dir.mkdir()
+ helper_path = bin_dir / self.xcom_helper_name
+ helper_path.write_text(_XCOM_HELPER_SHIM)
+ helper_path.chmod(0o755)
+ yield xcom_dir, bin_dir
+
+ def _read_xcom_file(self, path: Path) -> tuple[str | None, Any, str |
None]:
+ """Read one file under ``$AIRFLOW_XCOM_DIR`` and return ``(key, value,
error)``."""
+ size = path.stat().st_size
+ if size > self.max_xcom_file_size:
+ return (
+ None,
+ None,
+ f"XCom file {path} is {size} bytes, exceeding
max_xcom_file_size "
+ f"({self.max_xcom_file_size} bytes); it was not pushed.",
+ )
+ text = path.read_bytes().decode(self.output_encoding)
+ value = text[:-1] if text.endswith("\n") else text
+ if path.suffix != ".json":
+ return path.name, value, None
+ try:
+ return path.stem, json.loads(text), None
+ except json.JSONDecodeError as exc:
+ # Keep the ".json" suffix on the key so downstream consumers can
tell this value
+ # was not parsed, and never silently drop the diagnostic content.
+ return (
+ path.name,
+ value,
+ f"XCom file {path} contains invalid JSON ({exc}); the raw
content was pushed "
+ f"under {path.name!r} instead.",
+ )
+
+ def _read_xcom_dir(self, xcom_dir: Path, errors: list[str]) -> dict[str,
Any]:
+ result: dict[str, Any] = {}
+ for entry in sorted(os.scandir(xcom_dir), key=lambda e: e.name):
+ # is_symlink() is checked first because is_dir()/is_file() follow
symlinks.
+ if entry.is_symlink() or not (entry.is_dir() or entry.is_file()):
+ errors.append(
+ f"XCom directory entry {entry.path!r} is not a regular
file; it was not pushed."
+ )
+ continue
+ if entry.is_dir():
+ errors.append(
+ f"XCom directory entry {entry.path!r} is a subdirectory;
subdirectories are not "
+ "supported. Use a '.json' file (or `xcom push --json`) for
structured values."
+ )
+ continue
+ key, value, error = self._read_xcom_file(Path(entry.path))
+ if error is not None:
+ errors.append(error)
+ if key is not None:
+ result[key] = value
+ return result
+
+ def _collect_and_push_xcom(
+ self, context: Context, xcom_dir: Path, exception_in_flight: bool
+ ) -> dict[str, Any]:
+ """
+ Read ``$AIRFLOW_XCOM_DIR`` and push every entry found in it.
+
+ Called from a ``finally`` block, so collection errors are gathered
rather than raised
+ inline, and are downgraded to a warning when the command already
failed or skipped --
+ an original Bash failure must never be masked by a bad XCom file.
+ """
+ errors: list[str] = []
+ entries = self._read_xcom_dir(xcom_dir, errors=errors)
+ for key, value in entries.items():
+ context["ti"].xcom_push(key=key, value=value)
+ if errors:
+ message = "; ".join(errors)
+ if exception_in_flight:
+ self.log.warning("Errors while collecting the XCom directory:
%s", message)
+ else:
+ raise ValueError(f"Errors while collecting the XCom directory:
{message}")
Review Comment:
This fails a task whose command passed. If someone runs `mkdir
"$AIRFLOW_XCOM_DIR/tmp"` for scratch space, or if one file is a byte over
`max_xcom_file_size`, their green task goes red. Please log a warning and push
the good entries instead of raising.
##########
providers/standard/src/airflow/providers/standard/operators/bash.py:
##########
@@ -261,5 +328,95 @@ def _is_inline_command(cls, bash_command: str) -> bool:
"""Return True if the bash command is an inline string. False if it's
a bash script file."""
return not bash_command.endswith(tuple(cls.template_ext))
+ @contextlib.contextmanager
+ def _xcom_workspace(self) -> Iterator[tuple[Path, Path | None]]:
+ """
+ Create the ``$AIRFLOW_XCOM_DIR`` workspace for a single run.
+
+ Kept separate from the ``working_directory()`` script cwd tempdir so a
rendered ``.sh``
+ file can never collide with an XCom key.
+ """
+ with tempfile.TemporaryDirectory(prefix="airflow-bash-xcom-") as
workspace:
+ xcom_dir = Path(workspace) / "xcom"
+ xcom_dir.mkdir()
+ bin_dir: Path | None = None
+ if self.xcom_helper_name:
+ bin_dir = Path(workspace) / "bin"
+ bin_dir.mkdir()
+ helper_path = bin_dir / self.xcom_helper_name
+ helper_path.write_text(_XCOM_HELPER_SHIM)
+ helper_path.chmod(0o755)
+ yield xcom_dir, bin_dir
+
+ def _read_xcom_file(self, path: Path) -> tuple[str | None, Any, str |
None]:
+ """Read one file under ``$AIRFLOW_XCOM_DIR`` and return ``(key, value,
error)``."""
+ size = path.stat().st_size
+ if size > self.max_xcom_file_size:
+ return (
+ None,
+ None,
+ f"XCom file {path} is {size} bytes, exceeding
max_xcom_file_size "
+ f"({self.max_xcom_file_size} bytes); it was not pushed.",
+ )
+ text = path.read_bytes().decode(self.output_encoding)
Review Comment:
This runs inside the `finally` block. If a script writes a binary file here,
decode raises `UnicodeDecodeError`, and that replaces the real
`AirflowException` from the failed command. Your own docstring says a bad xcom
file must never mask a Bash failure. This is that case. Wrap the read in a try
and add the failure to errors, like you already do for bad JSON.
Same problem with `path.stat()` on line 353 and `os.scandir()` on line 379.
A file the script deletes or chmod 000s after the scan raises `OSError` out of
the finally.
##########
providers/standard/src/airflow/providers/standard/operators/bash.py:
##########
@@ -214,18 +251,48 @@ def execute(self, context: Context):
env = self.get_env(context)
self._is_inline_cmd = self._is_inline_command(bash_command=cast("str",
self.bash_command))
- if self._is_inline_cmd:
- result = self._run_inline_command(bash_path=bash_path, env=env)
- else:
- result = self._run_rendered_script_file(bash_path=bash_path,
env=env)
-
- if result.exit_code in self.skip_on_exit_code:
- raise AirflowSkipException(f"Bash command returned exit code
{result.exit_code}. Skipping.")
- if result.exit_code != 0:
- raise AirflowException(
- f"Bash command failed. The command returned a non-zero exit
code {result.exit_code}."
- )
+ pushed_xcoms: dict[str, Any] = {}
+ with contextlib.ExitStack() as stack:
+ xcom_dir: Path | None = None
+ if self.do_xcom_push:
+ xcom_dir, bin_dir = stack.enter_context(self._xcom_workspace())
+ # get_env() returns self.env itself when append_env is False,
so copy before
+ # adding per-run vars -- otherwise these temp paths leak into
the operator's
+ # own env attribute and survive into later runs.
+ env = dict(env)
+ env["AIRFLOW_XCOM_DIR"] = os.fspath(xcom_dir)
+ if bin_dir is not None:
+ env["PATH"] = f"{bin_dir}{os.pathsep}{env.get('PATH') or
os.defpath}"
+
+ exception_in_flight = False
+ try:
+ if self._is_inline_cmd:
+ result = self._run_inline_command(bash_path=bash_path,
env=env)
+ else:
+ result =
self._run_rendered_script_file(bash_path=bash_path, env=env)
+
+ if result.exit_code in self.skip_on_exit_code:
+ raise AirflowSkipException(
+ f"Bash command returned exit code {result.exit_code}.
Skipping."
+ )
+ if result.exit_code != 0:
+ raise AirflowException(
+ f"Bash command failed. The command returned a non-zero
exit code {result.exit_code}."
+ )
+ except BaseException:
+ exception_in_flight = True
+ raise
+ finally:
+ # Collecting here (rather than after the exit-code checks
above) means XComs are
+ # pushed whether the command succeeded, failed, or was skipped.
+ if xcom_dir is not None:
+ pushed_xcoms = self._collect_and_push_xcom(
+ context=context, xcom_dir=xcom_dir,
exception_in_flight=exception_in_flight
+ )
+
+ if "return_value" in pushed_xcoms:
Review Comment:
`return_value` as a magic filename causes three problems:
1. It returns before `self.output_processor()` call on line 296, so a user
with both set gets their processor silently ignored.
2. On success the key is pushed twice: once on line 412, once by the
runner's `_push_xcom_if_needed`.
3. On failure line 412 pushes it anyway, so a failed task ends up with a
return_value xcom. A downstream `xcom_pull(task_ids="failed_task")` then reads
it. The docs say this only applies when the command succeeds.
Simplest fix: drop the special case and treat `return_value` as an ordinary
key. Costs users one line of `output_processor` and removes all three problems.
##########
providers/standard/src/airflow/providers/standard/operators/bash.py:
##########
@@ -162,6 +195,8 @@ def __init__(
skip_on_exit_code: int | Container[int] | None = 99,
cwd: str | None = None,
output_processor: Callable[[str], Any] = lambda result: result,
+ xcom_helper_name: str | None = "xcom",
Review Comment:
Default this to `None`. As of now, every existing `BashOperator` gets a new
xcom executable prepended to PATH on upgrade because `do_xcom_push` is already
True by default. Anyone who calls their own xcom script from a Bash task
silently gets this shim instead. The `$AIRFLOW_XCOM_DIR` folder is harmless
always to provide. The PATH write is not.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
