streamwright-spark 0.1.0__py3-none-any.whl

streamwright/spark/__init__.py

@@ -0,0 +1 @@
+"""SparkJob engine for StreamWright."""

streamwright/spark/cli/__init__.py

@@ -0,0 +1,14 @@
+"""The `sw spark` command line."""
+
+import argparse
+
+from streamwright.spark.cli.main import main
+
+__all__ = ["main", "register"]
+
+
+def register(subparsers):
+    parser = subparsers.add_parser("spark", help="validate, run and inspect SparkJob pipelines", add_help=False)
+    parser.add_argument("-h", "--help", action="store_true")
+    parser.add_argument("arguments", nargs=argparse.REMAINDER)
+    parser.set_defaults(command_handler=lambda args: main(["--help"] if args.help else args.arguments))

streamwright/spark/cli/__main__.py

@@ -0,0 +1,7 @@
+"""python -m streamwright.spark.cli"""
+
+import sys
+
+from streamwright.spark.cli.main import main
+
+sys.exit(main())

streamwright/spark/cli/aws_secrets.py

@@ -0,0 +1,58 @@
+"""AWS Secrets Manager access, imported only by the driver when configured."""
+
+import json
+import re
+
+import boto3
+
+
+def read_secrets(source, job):
+    region = ((job.get("runtime") or {}).get("platform") or {}).get("region")
+    client = boto3.client("secretsmanager", region_name=region)
+    data = {}
+
+    def fetch_value(secret_id):
+        reply = client.get_secret_value(SecretId=secret_id)
+        raw = reply.get("SecretString")
+        if raw is None and "SecretBinary" in reply:
+            raw = reply["SecretBinary"].decode("utf-8")
+        return raw or ""
+
+    mapping = source.get("mapping")
+    if isinstance(mapping, dict):
+        for variable, secret_id in mapping.items():
+            raw = fetch_value(secret_id)
+            try:
+                parsed = json.loads(raw)
+                if isinstance(parsed, dict):
+                    if variable in parsed:
+                        data[variable] = parsed[variable]
+                    elif len(parsed) == 1:
+                        data[variable] = next(iter(parsed.values()))
+                    else:
+                        data.update(parsed)
+                else:
+                    data[variable] = raw
+            except (ValueError, TypeError):
+                data[variable] = raw
+
+    secret_ids = source.get("ids") or source.get("id")
+    if isinstance(secret_ids, str):
+        secret_ids = [secret_ids]
+    elif not isinstance(secret_ids, (list, tuple)):
+        secret_ids = []
+    for secret_id in secret_ids:
+        raw = fetch_value(secret_id)
+        try:
+            parsed = json.loads(raw)
+            if isinstance(parsed, dict):
+                data.update(parsed)
+            else:
+                leaf = secret_id.split("/")[-1].split(":")[-1]
+                leaf = re.sub(r"-[A-Za-z0-9]{6}$", "", leaf)
+                data[leaf] = raw
+        except (ValueError, TypeError):
+            leaf = secret_id.split("/")[-1].split(":")[-1]
+            leaf = re.sub(r"-[A-Za-z0-9]{6}$", "", leaf)
+            data[leaf] = raw
+    return data

streamwright/spark/cli/common.py

@@ -0,0 +1,73 @@
+"""Pieces shared by the sw spark commands."""
+
+import sys
+
+from streamwright.core.jobs.inputs import InputError as InputsError
+from streamwright.core.jobs.inputs import parse_set
+from streamwright.core.sdk.contract import resolve_connection
+from streamwright.spark.contracts import Connection
+from streamwright.spark.runtime import registry
+
+__all__ = ["InputsError", "add_job_arguments", "connections_of", "err", "load_plugins", "parse_set"]
+
+
+def err(text):
+    sys.stderr.write(text + "\n")
+
+
+def add_job_arguments(parser, many=False):
+    if many:
+        parser.add_argument("job", nargs="*", metavar="PATH", help="path to the job definition YAML file(s)")
+        parser.add_argument(
+            "--job",
+            "-j",
+            action="append",
+            default=[],
+            dest="flag_jobs",
+            help="path to the job definition YAML file (can be specified multiple times)",
+        )
+        parser.add_argument("--file", "-f", action="append", default=[], dest="flag_files", help="alias for --job")
+        parser.add_argument("--pipeline", action="append", default=[], dest="flag_pipelines", help="alias for --job")
+    else:
+        parser.add_argument("job", nargs="?", metavar="JOB", help="path to the job definition YAML file or directory")
+        parser.add_argument("--job", "-j", dest="flag_job", help="path to the job definition YAML file or directory")
+        parser.add_argument("--file", "-f", dest="flag_file", help="alias for --job")
+        parser.add_argument("--pipeline", dest="flag_pipeline", help="alias for --job")
+    parser.add_argument("--config", help="a YAML file with a `config:` mapping (values for spec.config)")
+    parser.add_argument(
+        "--set",
+        action="append",
+        default=[],
+        metavar="NAME=VALUE",
+        help="override a spec.config value (e.g. --set key=value)",
+    )
+    parser.add_argument("--overlay", action="append", default=[], help="YAML merged over the job; a null deletes a key")
+    parser.add_argument(
+        "--allow-connector", action="append", default=None, metavar="PROVIDER", help="only load these providers"
+    )
+
+
+def load_plugins(allow=None):
+    """({provider: adapter}, {platform: launcher}, problems)."""
+    adapters, problems = registry.load_all("adapter", allow)
+    launchers, more = registry.load_all("launcher")
+    return adapters, launchers, problems + more
+
+
+def connections_of(job, *, adapters=None, declared=None, config=None, secrets=None):
+    """
+    {name: Connection} of the connections the Spark engine runs natively (a Spark adapter): `provider` apart, the rest
+    is the auth block. A provider without a Spark adapter is reached only through call steps, which resolve its
+    connection in the source engine.
+    """
+    connections = {}
+    for name, connection in (job.get("connections") or {}).items():
+        provider = connection["provider"]
+        adapter = (adapters or {}).get(provider)
+        if adapter is None:
+            continue
+        original = (declared or job)["connections"][name]
+        auth = {key: value for key, value in original.items() if key != "provider"}
+        values = resolve_connection(adapter.definition, auth, config=config or {}, secrets=secrets or {})
+        connections[name] = Connection(name, provider, values)
+    return connections

streamwright/spark/cli/inputs.py

@@ -0,0 +1,40 @@
+"""Config and secret values for a run, through core input handling."""
+
+from importlib import import_module
+
+from streamwright.core.config import inputs as core_inputs
+from streamwright.core.jobs import inputs as job_inputs
+from streamwright.spark.cli.common import err, parse_set
+
+__all__ = ["InputError", "config_values", "secret_values"]
+
+InputError = core_inputs.InputError
+
+
+def config_values(job, args):
+    provided = {}
+    if args.config:
+        provided.update(job_inputs.read_config_values(args.config))
+    provided.update(parse_set(args.set))
+    return job_inputs.effective_config(job, provided)
+
+
+def secret_values(job, args):
+    """STREAMWRIGHT_SECRET_<NAME> variables, --secrets FILE, or runtime.secrets (AWS Secrets Manager); never printed."""
+    declared = ((job.get("spec") or {}).get("secrets")) or {}
+    provided = {}
+    if getattr(args, "secrets", None):
+        provided.update(core_inputs.read_values_file(args.secrets, err))
+    for name, value in core_inputs.secrets_from_env(list(declared)).items():
+        provided.setdefault(name, value)
+    source = (job.get("runtime") or {}).get("secrets") or {}
+    if source.get("source") == "aws_secretsmanager" and len(set(provided) & set(declared)) < len(declared):
+        try:
+            aws_secrets = import_module("streamwright.spark.cli.aws_secrets")
+        except ImportError as exc:
+            if exc.name != "boto3":
+                raise
+            raise InputError('boto3 is not installed; pip install "streamwright-spark[emr]"') from exc
+        for name, value in aws_secrets.read_secrets(source, job).items():
+            provided.setdefault(name, value)
+    return job_inputs.secret_values(job, provided)

streamwright/spark/cli/lock.py

@@ -0,0 +1,65 @@
+"""sw spark lock: write or check the version lock."""
+
+import os
+
+from streamwright.spark.cli.common import err, load_plugins
+from streamwright.spark.job import check_job, load_job
+from streamwright.spark.launchers.common import sha256_file
+from streamwright.spark.runtime import lock as lockmod
+
+__all__ = ["add_arguments", "run"]
+
+
+def add_arguments(parser):
+    parser.add_argument("--spark", required=True, help="the Spark version the job runs on, e.g. 3.5.3")
+    parser.add_argument(
+        "--out", help="output path for the generated lock file (defaults to runtime.artifacts.lock or <name>.lock.json)"
+    )
+    parser.add_argument("--emr-release", help="EMR release label (e.g. emr-7.2.0) to record in the lock file")
+    parser.add_argument(
+        "--env-archive", help="path to packaged virtual environment archive (.tar.gz) for remote executors"
+    )
+    parser.add_argument(
+        "--jar",
+        action="append",
+        default=[],
+        metavar="GROUP:ARTIFACT=VERSION",
+        help="JAR dependency override in GROUP:ARTIFACT=VERSION format",
+    )
+    parser.add_argument("--check", action="store_true", help="verify an existing lock instead of writing one")
+
+
+def run(args):
+    adapters, launchers, _ = load_plugins(args.allow_connector)
+    job = load_job(args.job, args.overlay)
+    problems, _ = check_job(job, adapters, launchers)
+    if problems:
+        for text in problems:
+            err("error: " + text)
+        return 2
+    used = sorted({c["provider"] for c in job["connections"].values()})
+    jars = dict(item.split("=", 1) for item in args.jar if "=" in item)
+    try:
+        python_env = None
+        if args.env_archive:
+            python_env = {"archive": os.path.basename(args.env_archive), "sha256": sha256_file(args.env_archive)}
+        lock = lockmod.build_lock(
+            job["name"],
+            [adapters[p] for p in used],
+            args.spark,
+            extra_jars=jars,
+            emr_release=args.emr_release,
+            python_env=python_env,
+        )
+    except lockmod.LockError as exc:
+        err("error: %s" % exc)
+        return 2
+    out = args.out or ((job.get("runtime") or {}).get("artifacts") or {}).get("lock") or "%s.lock.json" % job["name"]
+    if args.check:
+        problems = lockmod.check_lock(lockmod.load(out), args.spark)
+        for text in problems:
+            err("lock: " + text)
+        return 2 if problems else 0
+    lockmod.dump(lock, out)
+    print("wrote %s (fingerprint %s)" % (out, lockmod.fingerprint(lock)[:12]))
+    return 0

streamwright/spark/cli/main.py

@@ -0,0 +1,94 @@
+"""
+`sw spark ...`. Entry points `streamwright-spark` / `sw-spark`; attaching it to `streamwright` core is a lazy
+dispatch like `streamwright mcp` (see README). Exit codes: 0 ok, 1 failed, 2 usage/validation, 130 interrupted.
+
+    sw spark validate PATH... [--config F] [--set K=V] [--overlay F] [--analyze] [--graph] [--strict] [--lock F]
+    sw spark run JOB [--config F] [--set K=V] [--secrets F] [--overlay F] [--once|--drain|--follow [--interval 10s]]
+                     [--max-messages N] [--run-key K] [--conf K=V] [--lock F] [--lock-required] [--dry-run] [--detach]
+    sw spark lock JOB --spark V [--out F] [--emr-release L] [--env-archive F] [--jar G:A=V] [--check]
+    sw spark state init|status|batches|show|retry|abandon|purge|lease JOB [ARG]
+"""
+
+import argparse
+import sys
+
+from streamwright.spark.cli import lock, run, state, validate
+from streamwright.spark.cli.common import InputsError, add_job_arguments, err
+from streamwright.spark.cli.inputs import InputError
+from streamwright.spark.job import JobError
+from streamwright.spark.runtime.lock import LockError
+
+__all__ = ["main"]
+
+_COMMANDS = (
+    ("validate", validate, "check jobs without Spark, brokers or credentials", True),
+    ("plan", validate, "show Spark actions and step dependencies without running the job", True),
+    ("run", run, "run a job", False),
+    ("lock", lock, "write or check the version lock", False),
+    ("state", state, "inspect and repair bookkeeping", None),
+)
+
+
+def _parser():
+    parser = argparse.ArgumentParser(prog="sw spark", description="Spark jobs for StreamWright.")
+    sub = parser.add_subparsers(dest="command")
+    for name, module, help_text, many in _COMMANDS:
+        command = sub.add_parser(name, help=help_text)
+        if many is not None:
+            add_job_arguments(command, many)
+        module.add_arguments(command)
+        command.set_defaults(handler=_run_plan if name == "plan" else module.run)
+    return parser
+
+
+def _run_plan(args):
+    args.analyze = True
+    args.graph = True
+    return validate.run(args)
+
+
+def _normalize_job_args(args):
+    if hasattr(args, "flag_job"):
+        flag_val = args.flag_job or getattr(args, "flag_file", None) or getattr(args, "flag_pipeline", None)
+        if flag_val:
+            if hasattr(args, "arg") and args.job and args.arg is None:
+                args.arg = args.job
+            args.job = flag_val
+        elif not args.job:
+            err("error: job is required (pass JOB positional argument or --job PATH)")
+            return False
+
+    if hasattr(args, "flag_jobs"):
+        jobs = (
+            list(args.job or [])
+            + list(args.flag_jobs or [])
+            + list(getattr(args, "flag_files", None) or [])
+            + list(getattr(args, "flag_pipelines", None) or [])
+        )
+        if not jobs:
+            err("error: at least one job path is required (pass PATH positional argument or --job PATH)")
+            return False
+        args.job = jobs
+
+    return True
+
+
+def main(argv=None):
+    parser = _parser()
+    args = parser.parse_args(argv)
+    if not getattr(args, "handler", None):
+        parser.print_help()
+        return 2
+    if not _normalize_job_args(args):
+        return 2
+    try:
+        return args.handler(args)
+    except (JobError, InputError, InputsError, LockError) as exc:
+        err("error: %s" % exc)
+        return 2
+    except KeyboardInterrupt:
+        return 130
+
+
+if __name__ == "__main__":
+    sys.exit(main())