sparkforensics-operator 0.1.0__py3-none-any.whl

sparkforensics_operator/__init__.py

@@ -0,0 +1,29 @@
+from sparkforensics_operator.callback import spark_forensics_callback
+from sparkforensics_operator.exceptions import ThresholdBreached
+from sparkforensics_operator.hooks.analyze.base import AnalyzeHook
+from sparkforensics_operator.hooks.analyze.subprocess import SubprocessAnalyzeHook
+from sparkforensics_operator.hooks.log_source.base import LogSourceHook
+from sparkforensics_operator.hooks.log_source.filesystem import FilesystemLogSourceHook
+from sparkforensics_operator.hooks.log_source.history_server import HistoryServerLogSourceHook
+from sparkforensics_operator.hooks.log_source.sftp import SFTPLogSourceHook
+from sparkforensics_operator.hooks.log_source.tunnel import SSHTunneledLogSourceHook
+from sparkforensics_operator.hooks.log_source.xcom import XComLogSourceHook
+from sparkforensics_operator.notify import Notifier
+from sparkforensics_operator.operator import SparkForensicsOperator
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "SparkForensicsOperator",
+    "spark_forensics_callback",
+    "ThresholdBreached",
+    "LogSourceHook",
+    "HistoryServerLogSourceHook",
+    "FilesystemLogSourceHook",
+    "XComLogSourceHook",
+    "SFTPLogSourceHook",
+    "SSHTunneledLogSourceHook",
+    "AnalyzeHook",
+    "SubprocessAnalyzeHook",
+    "Notifier",
+]

sparkforensics_operator/_compat.py

@@ -0,0 +1,23 @@
+"""
+Isolates the one real Airflow-2-vs-3 import difference this package depends
+on. Airflow 3's Task SDK moved BaseOperator/BaseOperatorLink to the
+airflow.sdk namespace; airflow.sdk does not exist on Airflow 2.x. Everywhere
+else in this package, import BaseOperator/BaseOperatorLink from here, never
+straight from airflow.
+
+Within the Airflow 2.x branch there is a second split: BaseOperatorLink
+moved out of airflow.models.baseoperator and into its own
+airflow.models.baseoperatorlink module in Airflow 2.8.0. On Airflow
+2.6.0/2.7.0, airflow.models.baseoperatorlink does not exist yet and
+BaseOperatorLink must still be imported from airflow.models.baseoperator.
+"""
+try:
+    from airflow.sdk import BaseOperator, BaseOperatorLink
+except ImportError:  # Airflow 2.x has no airflow.sdk package
+    from airflow.models.baseoperator import BaseOperator
+    try:
+        from airflow.models.baseoperatorlink import BaseOperatorLink  # Airflow >= 2.8
+    except ImportError:  # Airflow 2.6/2.7
+        from airflow.models.baseoperator import BaseOperatorLink
+
+__all__ = ["BaseOperator", "BaseOperatorLink"]

sparkforensics_operator/callback.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import logging
+from typing import Callable
+
+from sparkforensics_operator.exceptions import ThresholdBreached
+from sparkforensics_operator.operator import run_spark_forensics
+
+log = logging.getLogger("airflow.task")
+
+
+def spark_forensics_callback(
+    *,
+    log_source,
+    backend,
+    report_dest: str,
+    max_runtime_ms: int | None = None,
+    max_spill_gb: float | None = None,
+    max_skew_ratio: float | None = None,
+    max_failed_task_rate_pct: float | None = None,
+    min_efficiency_pct: float | None = None,
+    on_threshold_breach: str = "fail",
+    notifier=None,
+    aws_conn_id: str | None = None,
+) -> Callable[[dict], None]:
+    """Builds an on_success_callback for the upstream Spark task, sharing
+    SparkForensicsOperator's exact execute() logic via run_spark_forensics.
+
+    Note (see the plan's "Global constraints"): Airflow itself catches and
+    only logs any exception this callback raises, it never retries and
+    never fails the upstream task. Set on_threshold_breach="fail" here only
+    to get that log-and-continue behavior on breach; it will not fail the
+    DAG the way the standalone-operator trigger shape does.
+    """
+    thresholds = {
+        "max_runtime_ms": max_runtime_ms,
+        "max_spill_gb": max_spill_gb,
+        "max_skew_ratio": max_skew_ratio,
+        "max_failed_task_rate_pct": max_failed_task_rate_pct,
+        "min_efficiency_pct": min_efficiency_pct,
+    }
+
+    def _callback(context: dict) -> None:
+        ti = context.get("ti")
+        try:
+            destination = run_spark_forensics(
+                context,
+                log_source=log_source,
+                backend=backend,
+                report_dest=report_dest,
+                thresholds=thresholds,
+                on_threshold_breach=on_threshold_breach,
+                notifier=notifier,
+                log=log,
+                aws_conn_id=aws_conn_id,
+            )
+        except ThresholdBreached as e:
+            # on_threshold_breach="fail" (the default) makes run_spark_forensics
+            # raise ThresholdBreached *after* the report has already been
+            # persisted successfully. Without this, the callback path would
+            # never record the report's location on the one path (a breach)
+            # where a user most wants the report link. run_spark_forensics
+            # attaches the already-known-good destination to the exception
+            # for exactly this reason (see operator.py); still push it to
+            # XCom before letting the breach propagate as usual.
+            destination = getattr(e, "destination", None)
+            if ti and destination is not None:
+                ti.xcom_push(key="return_value", value=destination)
+            raise
+
+        # No operator auto-pushes this callback's return value to XCom, so
+        # push it manually under the same key ("return_value") an operator's
+        # own return value would use (see links.py's _XCOM_RETURN_KEY),
+        # keeping the persisted report's location discoverable via
+        # ti.xcom_pull() on the upstream task. Guard against a malformed/test
+        # context missing "ti": this is a notification-adjacent side effect
+        # and must not turn an already-successful run into a new failure.
+        if ti:
+            ti.xcom_push(key="return_value", value=destination)
+
+    return _callback

sparkforensics_operator/exceptions.py

@@ -0,0 +1,10 @@
+from airflow.exceptions import AirflowFailException
+
+
+class ThresholdBreached(AirflowFailException):
+    """Raised when a configured SparkForensics threshold is violated and
+    on_threshold_breach="fail". Extends AirflowFailException (not
+    AirflowException): a breach is a deterministic verdict from the
+    already-completed analysis, so retrying would just re-run the same
+    fetch+analyze cycle to reach the same answer, and should not consume
+    the task's configured retries."""

sparkforensics_operator/hooks/analyze/base.py

@@ -0,0 +1,18 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from airflow.hooks.base import BaseHook
+
+from sparkforensics_operator.report import Report
+
+
+class AnalyzeHook(BaseHook, ABC):
+    """Runs sparkforensics analysis over a local event log and returns a
+    Report. v1 ships one implementation (subprocess against the
+    sparkforensics-analyze CLI), see the plan's "Future: http/MCP
+    AnalyzeHook backend" section for why an HTTP backend was deferred."""
+
+    @abstractmethod
+    def analyze(self, log_path: Path, thresholds: dict) -> Report:
+        """thresholds keys: max_runtime_ms, max_spill_gb, max_skew_ratio,
+        max_failed_task_rate_pct, min_efficiency_pct (all optional)."""

sparkforensics_operator/hooks/analyze/subprocess.py

@@ -0,0 +1,72 @@
+import subprocess
+import tempfile
+from pathlib import Path
+
+from airflow.exceptions import AirflowException
+
+from sparkforensics_operator.report import THRESHOLD_CLI_FLAGS, Report, parse_report_json, parse_threshold_results
+
+from .base import AnalyzeHook
+
+# Exit codes 0 (success), 1 (threshold violated), and 3 (thresholds inconclusive) per the CLI's documented exit-code contract; anything else (OOM-kill, wrapper failure, etc.) means --out file is not trustworthy.
+_USABLE_OUT_FILE_EXIT_CODES = {0, 1, 3}
+
+
+class SubprocessAnalyzeHook(AnalyzeHook):
+    """Shells out to `sparkforensics-analyze <log_path> --format json --out
+    <tmpfile> [threshold flags]`. Requires Node.js (engines
+    '>=22.18.0 <23.0.0 || >=23.6.0') and the sparkforensics npm package
+    installed on the worker, with sparkforensics-analyze resolvable on
+    PATH (or pass analyze_bin=<full path>)."""
+
+    def __init__(self, analyze_bin: str = "sparkforensics-analyze", timeout: int = 900):
+        super().__init__()
+        self.analyze_bin = analyze_bin
+        self.timeout = timeout
+
+    def _threshold_args(self, thresholds: dict) -> list:
+        args = []
+        for key, flag in THRESHOLD_CLI_FLAGS.items():
+            value = thresholds.get(key)
+            if value is not None:
+                args.extend([flag, str(value)])
+        return args
+
+    def analyze(self, log_path: Path, thresholds: dict) -> Report:
+        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as out_file:
+            out_path = Path(out_file.name)
+
+        args = [self.analyze_bin, str(log_path), "--format", "json", "--out", str(out_path)]
+        args.extend(self._threshold_args(thresholds))
+
+        try:
+            try:
+                proc = subprocess.run(args, capture_output=True, text=True, timeout=self.timeout)
+            except FileNotFoundError as e:
+                raise AirflowException(
+                    f"sparkforensics-analyze binary not found: {self.analyze_bin!r}. "
+                    "Install the sparkforensics npm package on this worker, or pass "
+                    "analyze_bin=<full path to sparkforensics-analyze>."
+                ) from e
+            except subprocess.TimeoutExpired as e:
+                raise AirflowException(
+                    f"sparkforensics-analyze timed out after {self.timeout}s analyzing {log_path}."
+                ) from e
+
+            if proc.returncode == 2:
+                raise AirflowException(
+                    f"sparkforensics-analyze failed to parse the event log or was given "
+                    f"bad arguments (exit 2): {proc.stderr.strip()}"
+                )
+            if proc.returncode not in _USABLE_OUT_FILE_EXIT_CODES:
+                raise AirflowException(
+                    f"sparkforensics-analyze exited with unexpected code {proc.returncode}: "
+                    f"{proc.stderr.strip()}"
+                )
+
+            report = parse_report_json(out_path)
+            report.threshold_results = parse_threshold_results(thresholds, proc.stderr)
+            report.exit_code = proc.returncode
+            return report
+        finally:
+            out_path.unlink(missing_ok=True)

sparkforensics_operator/hooks/log_source/_dest_root.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import shutil
+import tempfile
+from pathlib import Path
+
+
+def dest_for(dest_root: Path, source_path: str | Path) -> Path:
+    """The dest_root/<basename> convention shared by every LogSourceHook that
+    stages a single file: FilesystemLogSourceHook (local copy) and
+    SFTPLogSourceHook (remote-to-local fetch)."""
+    return dest_root / Path(source_path).name
+
+
+def make_dest_root(dest_dir: str | None) -> tuple[Path, Path | None]:
+    """Resolves the directory a hook stages files into: dest_dir if the
+    caller gave one (a shared, caller-managed directory — the second tuple
+    element is None so cleanup() never touches it), else a fresh private
+    temp dir this hook owns (returned as both elements, so the caller can
+    rmtree it later, on error or in cleanup())."""
+    if dest_dir is not None:
+        root = Path(dest_dir)
+        return root, None
+    owned = Path(tempfile.mkdtemp(prefix="sparkforensics-"))
+    return owned, owned
+
+
+def remove_if_owned(owned_temp_root: Path | None) -> None:
+    if owned_temp_root is not None:
+        shutil.rmtree(owned_temp_root, ignore_errors=True)

sparkforensics_operator/hooks/log_source/_path_template.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def _safe_path_component(value: object) -> str | None:
+    """Reduce an attacker-influenceable context string (e.g. run_id, which is
+    settable via Airflow's trigger API) to its filesystem basename before it
+    is interpolated into path_template, the same way history_server.py
+    sanitizes app_id. None is preserved as None rather than becoming the
+    literal string "None"."""
+    if value is None:
+        return None
+    return Path(str(value)).name
+
+
+def _template_vars(context: dict) -> dict:
+    dag = context.get("dag")
+    task = context.get("task")
+    return {
+        "ds": _safe_path_component(context.get("ds")),
+        "run_id": _safe_path_component(context.get("run_id")),
+        # Not sanitized: logical_date is datetime-like and path_template may
+        # apply a strftime format spec to it (e.g. "{logical_date:%Y-%m-%d}");
+        # Path(...).name would break that formatting and, unlike a plain
+        # string substitution, .format()'s format-spec mini-language for a
+        # datetime doesn't accept arbitrary attacker strings.
+        "logical_date": context.get("logical_date"),
+        "dag_id": _safe_path_component(dag.dag_id) if dag is not None else None,
+        "task_id": _safe_path_component(task.task_id) if task is not None else None,
+    }

sparkforensics_operator/hooks/log_source/_rolling_log.py

@@ -0,0 +1,4 @@
+import re
+
+# Mirrors isRollingLogDirectory in sparkforensics's src/cli/collect-run.ts.
+_ROLLING_ENTRY_RE = re.compile(r"^events_\d+_")

sparkforensics_operator/hooks/log_source/base.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from airflow.hooks.base import BaseHook
+
+
+class LogSourceHook(BaseHook, ABC):
+    """Fetches a Spark job's event log to a local path. Implementations
+    differ only in where the log comes from: a Spark History Server, a
+    filesystem/HDFS path pattern, or an XCom value pushed by the upstream
+    Spark task.
+    """
+
+    @abstractmethod
+    def fetch(self, context: dict) -> Path:
+        """Return a local path to the event log: either a single file, or a
+        directory of rolling-log segments (files named events_<n>_...)."""
+
+    def cleanup(self, path: Path) -> None:
+        """Called after analysis, on success or failure. No-op by default:
+        override only to remove a path this hook created for itself (e.g. a
+        private temp dir); never one it merely handed back to the caller,
+        such as a user-owned filesystem path or another task's
+        XCom-referenced file."""

sparkforensics_operator/hooks/log_source/filesystem.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+from airflow.exceptions import AirflowException
+
+from ._dest_root import dest_for
+from ._path_template import _template_vars
+from .base import LogSourceHook
+
+
+class FilesystemLogSourceHook(LogSourceHook):
+    """Reads an event log from a configured filesystem/HDFS path pattern.
+    "HDFS" here means an already-mounted path (NFS gateway, WebHDFS FUSE
+    mount, or similar), this hook does no Hadoop-client I/O of its own, it
+    just reads/copies a local path, the same way FilesystemLogSourceHook's
+    output is handed straight to sparkforensics-analyze's own local-file
+    reader."""
+
+    def __init__(self, path_template: str, dest_dir: str | None = None):
+        super().__init__()
+        self.path_template = path_template
+        self.dest_dir = dest_dir
+
+    def fetch(self, context: dict) -> Path:
+        rendered = self.path_template.format(**_template_vars(context))
+        source = Path(rendered)
+        if not source.exists():
+            raise AirflowException(f"Configured log path does not exist: {source}")
+        if self.dest_dir is None:
+            return source
+
+        dest_root = Path(self.dest_dir)
+        dest_root.mkdir(parents=True, exist_ok=True)
+        dest = dest_for(dest_root, source)
+        if source.is_dir():
+            shutil.copytree(source, dest, dirs_exist_ok=True)
+        else:
+            shutil.copy2(source, dest)
+        return dest
+
+    def cleanup(self, path: Path) -> None:
+        if self.dest_dir is None:
+            return  # fetch() returned the caller's own path; never delete it.
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)

sparkforensics_operator/hooks/log_source/history_server.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import time
+import zipfile
+from pathlib import Path
+from urllib.parse import quote
+
+import requests
+from airflow.exceptions import AirflowException
+
+from ._dest_root import make_dest_root, remove_if_owned
+from ._rolling_log import _ROLLING_ENTRY_RE
+from .base import LogSourceHook
+
+
+class HistoryServerLogSourceHook(LogSourceHook):
+    """Downloads a Spark job's event log from a Spark History Server's REST
+    API: GET {base_url}/api/v1/applications/{app_id}[/{attempt_id}]/logs,
+    which always returns a zip (one entry for a single event-log file,
+    multiple events_<n>_... entries for a rolling log)."""
+
+    def __init__(
+        self,
+        base_url: str,
+        app_id: str,
+        attempt_id: str | None = None,
+        dest_dir: str | None = None,
+        timeout: int = 300,
+    ):
+        super().__init__()
+        self.base_url = base_url.rstrip("/")
+        self.app_id = app_id
+        self.attempt_id = attempt_id
+        self.dest_dir = dest_dir
+        self.timeout = timeout
+        self._owned_temp_root: Path | None = None
+
+    def _build_url(self) -> str:
+        segments = ["api", "v1", "applications", quote(self.app_id, safe="")]
+        if self.attempt_id:
+            segments.append(quote(self.attempt_id, safe=""))
+        segments.append("logs")
+        return f"{self.base_url}/{'/'.join(segments)}"
+
+    def fetch(self, context: dict) -> Path:
+        url = self._build_url()
+        response = requests.get(url, timeout=self.timeout, stream=True)
+        if response.status_code != 200:
+            raise AirflowException(
+                f"Spark History Server log download failed ({response.status_code}) for {url}"
+            )
+
+        dest_root, self._owned_temp_root = make_dest_root(self.dest_dir)
+
+        try:
+            dest_root.mkdir(parents=True, exist_ok=True)
+
+            # self.app_id is only safe as a URL segment (quoted above); sanitize
+            # it before reusing it as a filesystem path component so a
+            # crafted/unexpected app_id (e.g. containing "../") can't escape
+            # dest_root.
+            safe_app_id = Path(self.app_id).name
+            zip_path = dest_root / f"{safe_app_id}.zip"
+            # self.timeout only bounds each individual socket read/connect (see
+            # the requests.get call above); a slow-trickling connection could
+            # otherwise keep the download running far longer than self.timeout
+            # implies. Track an explicit wall-clock deadline so self.timeout
+            # bounds the entire transfer, not just each chunk read.
+            deadline = time.monotonic() + self.timeout
+            with open(zip_path, "wb") as fh:
+                for chunk in response.iter_content(chunk_size=1024 * 1024):
+                    fh.write(chunk)
+                    if time.monotonic() > deadline:
+                        response.close()
+                        raise AirflowException(
+                            f"Spark History Server log download for app {self.app_id} "
+                            f"exceeded {self.timeout}s (total transfer time, not just a "
+                            "single read/connect)."
+                        )
+
+            extract_dir = dest_root / safe_app_id
+            with zipfile.ZipFile(zip_path) as zf:
+                names = zf.namelist()
+                zf.extractall(extract_dir)
+            zip_path.unlink()
+
+            if len(names) == 1:
+                return extract_dir / names[0]
+            if any(_ROLLING_ENTRY_RE.match(Path(name).name) for name in names):
+                return extract_dir
+            raise AirflowException(
+                f"Unexpected Spark History Server log archive contents for app {self.app_id}: {names}"
+            )
+        except Exception:
+            remove_if_owned(self._owned_temp_root)
+            raise
+
+    def cleanup(self, path: Path) -> None:
+        remove_if_owned(self._owned_temp_root)

sparkforensics_operator/hooks/log_source/sftp.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import concurrent.futures
+from pathlib import Path
+
+from airflow.exceptions import AirflowException
+
+from ._dest_root import dest_for, make_dest_root, remove_if_owned
+from ._path_template import _template_vars
+from ._rolling_log import _ROLLING_ENTRY_RE
+from .base import LogSourceHook
+
+
+class SFTPLogSourceHook(LogSourceHook):
+    """Reads an event log from an on-prem filesystem/HDFS path pattern the
+    cloud Airflow worker can't mount directly, over the same SSH connection
+    already configured for SSHOperator (via SFTPHook). The SSH-reachable
+    equivalent of FilesystemLogSourceHook: unlike that hook, every fetch is
+    a remote-to-local copy. Same cleanup() convention as
+    HistoryServerLogSourceHook (not FilesystemLogSourceHook): the path is
+    auto-cleaned when dest_dir is None (a private temp dir this hook
+    owns), but is a no-op when dest_dir is set (a caller-managed shared
+    directory)."""
+
+    def __init__(self, ssh_conn_id: str, path_template: str, dest_dir: str | None = None):
+        super().__init__()
+        self.ssh_conn_id = ssh_conn_id
+        self.path_template = path_template
+        self.dest_dir = dest_dir
+        self._owned_temp_root: Path | None = None
+
+    def fetch(self, context: dict) -> Path:
+        from airflow.providers.sftp.hooks.sftp import SFTPHook
+
+        remote_path = self.path_template.format(**_template_vars(context))
+        dest_root, self._owned_temp_root = make_dest_root(self.dest_dir)
+
+        try:
+            sftp_hook = SFTPHook(ssh_conn_id=self.ssh_conn_id)
+            if not sftp_hook.path_exists(remote_path):
+                raise AirflowException(
+                    f"Configured log path does not exist (checked over SFTP via "
+                    f"ssh_conn_id={self.ssh_conn_id!r}): {remote_path}"
+                )
+
+            dest_root.mkdir(parents=True, exist_ok=True)
+
+            if sftp_hook.isdir(remote_path):
+                entries = sftp_hook.list_directory(remote_path) or []
+                rolling_entries = [e for e in entries if _ROLLING_ENTRY_RE.match(Path(e).name)]
+                if not rolling_entries:
+                    raise AirflowException(
+                        f"Unexpected SFTP log directory contents (not a rolling-log "
+                        f"layout) at {remote_path!r} (checked over SFTP via "
+                        f"ssh_conn_id={self.ssh_conn_id!r}): {entries}"
+                    )
+                source_name = Path(remote_path).name
+                if not source_name:
+                    raise AirflowException(
+                        f"path_template resolved to a remote directory with no name "
+                        f"component to stage locally under dest_dir: {remote_path!r}"
+                    )
+                # Isolate per-source, same convention as FilesystemLogSourceHook's
+                # dest_root / source.name: a caller-supplied dest_dir shared across
+                # runs must not mix one run's events_<n>_* segments with another's.
+                source_dir = dest_root / source_name
+                source_dir.mkdir(parents=True, exist_ok=True)
+
+                def _download(entry: str) -> None:
+                    # Re-derive the basename that already passed the
+                    # _ROLLING_ENTRY_RE check above and use it (not the raw
+                    # entry) for the local destination, so a malformed
+                    # directory entry can't escape source_dir.
+                    local_name = Path(entry).name
+                    remote_file = f"{remote_path.rstrip('/')}/{entry}"
+                    sftp_hook.retrieve_file(remote_file, str(source_dir / local_name))
+
+                # Segments are independent files over the same SSH connection;
+                # overlapping the transfers cuts wall-clock time for
+                # multi-segment rolling logs versus one-at-a-time downloads.
+                with concurrent.futures.ThreadPoolExecutor(
+                    max_workers=min(8, len(rolling_entries))
+                ) as pool:
+                    list(pool.map(_download, rolling_entries))
+                return source_dir
+
+            source_name = Path(remote_path).name
+            if not source_name:
+                raise AirflowException(
+                    f"path_template resolved to a remote path with no filename "
+                    f"component to stage locally under dest_dir: {remote_path!r}"
+                )
+            local_path = dest_for(dest_root, remote_path)
+            sftp_hook.retrieve_file(remote_path, str(local_path))
+            return local_path
+        except AirflowException:
+            remove_if_owned(self._owned_temp_root)
+            raise
+        except Exception as e:
+            remove_if_owned(self._owned_temp_root)
+            raise AirflowException(
+                f"SFTP log fetch failed (ssh_conn_id={self.ssh_conn_id!r}, "
+                f"remote_path={remote_path!r}): {e}"
+            ) from e
+
+    def cleanup(self, path: Path) -> None:
+        remove_if_owned(self._owned_temp_root)