mainframe-migration-toolkit 0.2.0__py3-none-any.whl

mainframe_toolkit/io.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+from typing import Any, Iterator, Mapping
+
+from .errors import OptionalDependencyError
+
+
+def _spark_options(options: Mapping[str, Any] | None) -> dict[str, str]:
+    return {str(key): str(value).lower() if isinstance(value, bool) else str(value) for key, value in (options or {}).items()}
+
+
+def _require_spark() -> None:
+    try:
+        import pyspark  # noqa: F401
+    except ImportError as exc:
+        raise OptionalDependencyError("install mainframe-migration-toolkit[spark]") from exc
+
+
+def s3_uri(bucket_or_uri: str, key: str | None = None) -> str:
+    if bucket_or_uri.startswith(("s3://", "s3a://")):
+        base = "s3a://" + bucket_or_uri.split("://", 1)[1].rstrip("/")
+    else:
+        base = f"s3a://{bucket_or_uri.strip('/')}"
+    return f"{base}/{key.lstrip('/')}" if key else base
+
+
+def configure_s3(
+    spark: Any,
+    *,
+    endpoint: str | None = None,
+    region: str | None = None,
+    path_style_access: bool | None = None,
+    access_key: str | None = None,
+    secret_key: str | None = None,
+    session_token: str | None = None,
+) -> Any:
+    """Configure the Hadoop S3A connector without persisting credentials."""
+    _require_spark()
+    try:
+        (
+            spark._jvm.java.lang.Thread.currentThread()
+            .getContextClassLoader()
+            .loadClass("org.apache.hadoop.fs.s3a.S3AFileSystem")
+        )
+    except Exception as exc:
+        raise OptionalDependencyError(
+            "Spark is missing the hadoop-aws connector required for s3a:// paths"
+        ) from exc
+    configuration = spark.sparkContext._jsc.hadoopConfiguration()
+    configuration.set("fs.s3a.impl", "org.apache.hadoop.fs.s3a.S3AFileSystem")
+    if endpoint:
+        configuration.set("fs.s3a.endpoint", endpoint)
+    if region:
+        configuration.set("fs.s3a.endpoint.region", region)
+    if path_style_access is not None:
+        configuration.set("fs.s3a.path.style.access", str(path_style_access).lower())
+    credentials = (access_key, secret_key, session_token)
+    if any(credentials) and not (access_key and secret_key):
+        raise ValueError("access_key and secret_key must be supplied together")
+    if access_key and secret_key:
+        configuration.set("fs.s3a.access.key", access_key)
+        configuration.set("fs.s3a.secret.key", secret_key)
+        if session_token:
+            configuration.set("fs.s3a.session.token", session_token)
+    return spark
+
+
+def read_csv(
+    spark: Any,
+    path: str | Path,
+    *,
+    delimiter: str = ",",
+    header: bool = True,
+    schema: Any = None,
+    infer_schema: bool = False,
+    encoding: str = "UTF-8",
+    options: Mapping[str, Any] | None = None,
+) -> Any:
+    _require_spark()
+    reader_options = {
+        "sep": delimiter,
+        "header": str(header).lower(),
+        "inferSchema": str(infer_schema).lower(),
+        "encoding": encoding,
+    }
+    reader_options.update(_spark_options(options))
+    reader = spark.read.options(**reader_options)
+    if schema is not None:
+        reader = reader.schema(schema)
+    return reader.csv(str(path))
+
+
+def write_csv(
+    dataframe: Any,
+    path: str | Path,
+    *,
+    delimiter: str = ",",
+    header: bool = True,
+    mode: str = "overwrite",
+    coalesce: int | None = None,
+    options: Mapping[str, Any] | None = None,
+) -> None:
+    _require_spark()
+    target = dataframe.coalesce(coalesce) if coalesce else dataframe
+    writer_options = {"sep": delimiter, "header": str(header).lower()}
+    writer_options.update(_spark_options(options))
+    target.write.mode(mode).options(**writer_options).csv(str(path))
+
+
+def read_parquet(
+    spark: Any,
+    path: str | Path,
+    *,
+    options: Mapping[str, Any] | None = None,
+) -> Any:
+    _require_spark()
+    return spark.read.options(**_spark_options(options)).parquet(str(path))
+
+
+def write_parquet(
+    dataframe: Any,
+    path: str | Path,
+    *,
+    mode: str = "overwrite",
+    partition_by: list[str] | tuple[str, ...] | None = None,
+    options: Mapping[str, Any] | None = None,
+) -> None:
+    _require_spark()
+    writer = dataframe.write.mode(mode).options(**_spark_options(options))
+    if partition_by:
+        writer = writer.partitionBy(*partition_by)
+    writer.parquet(str(path))
+
+
+def read_dataset(
+    spark: Any,
+    path: str | Path,
+    *,
+    format: str | None = None,
+    options: Mapping[str, Any] | None = None,
+    **kwargs: Any,
+) -> Any:
+    selected = (format or Path(str(path)).suffix.lstrip(".") or "parquet").lower()
+    if selected == "csv":
+        return read_csv(spark, path, options=options, **kwargs)
+    if selected == "parquet":
+        return read_parquet(spark, path, options=options)
+    raise ValueError("format must be csv or parquet")
+
+
+def write_dataset(
+    dataframe: Any,
+    path: str | Path,
+    *,
+    format: str | None = None,
+    options: Mapping[str, Any] | None = None,
+    **kwargs: Any,
+) -> None:
+    selected = (format or Path(str(path)).suffix.lstrip(".") or "parquet").lower()
+    if selected == "csv":
+        write_csv(dataframe, path, options=options, **kwargs)
+        return
+    if selected == "parquet":
+        write_parquet(dataframe, path, options=options, **kwargs)
+        return
+    raise ValueError("format must be csv or parquet")
+
+
+def read_s3_csv(spark: Any, bucket_or_uri: str, key: str | None = None, **kwargs: Any) -> Any:
+    return read_csv(spark, s3_uri(bucket_or_uri, key), **kwargs)
+
+
+def write_s3_csv(dataframe: Any, bucket_or_uri: str, key: str | None = None, **kwargs: Any) -> None:
+    write_csv(dataframe, s3_uri(bucket_or_uri, key), **kwargs)
+
+
+def read_s3_parquet(spark: Any, bucket_or_uri: str, key: str | None = None, **kwargs: Any) -> Any:
+    return read_parquet(spark, s3_uri(bucket_or_uri, key), **kwargs)
+
+
+def write_s3_parquet(dataframe: Any, bucket_or_uri: str, key: str | None = None, **kwargs: Any) -> None:
+    write_parquet(dataframe, s3_uri(bucket_or_uri, key), **kwargs)
+
+
+def read_delimited_rows(
+    path: str | Path,
+    *,
+    delimiter: str = ",",
+    encoding: str = "utf-8",
+    header: bool = True,
+    fieldnames: list[str] | None = None,
+) -> Iterator[dict[str, str]]:
+    """Stream a local delimited file for control data and unit tests."""
+    with Path(path).open(encoding=encoding, newline="") as stream:
+        if header:
+            yield from csv.DictReader(stream, delimiter=delimiter)
+        else:
+            if not fieldnames:
+                raise ValueError("fieldnames are required when header is false")
+            yield from csv.DictReader(stream, delimiter=delimiter, fieldnames=fieldnames)

mainframe_toolkit/jcl.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import re
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class DDStatement:
+    name: str
+    raw: str
+    datasets: list[str] = field(default_factory=list)
+    disposition: str | None = None
+    dummy: bool = False
+    inline: bool = False
+
+
+@dataclass
+class JCLStep:
+    order: int
+    name: str
+    program: str | None
+    procedure: str | None
+    condition: str | None
+    parameters: str | None
+    raw: str
+    dd: list[DDStatement] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class JCLJob:
+    name: str | None
+    steps: list[JCLStep]
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"name": self.name, "steps": [step.to_dict() for step in self.steps]}
+
+
+def parse_jcl(path_or_text: str | Path, *, from_string: bool = False) -> JCLJob:
+    if from_string:
+        text = str(path_or_text)
+    else:
+        text = Path(path_or_text).read_text(encoding="utf-8", errors="replace")
+    statements = _logical_statements(text)
+    job_name: str | None = None
+    steps: list[JCLStep] = []
+    current: JCLStep | None = None
+    for name, operation, body in statements:
+        upper_operation = operation.upper()
+        if upper_operation == "JOB":
+            job_name = name.upper() if name else None
+            continue
+        if upper_operation == "EXEC":
+            program_match = re.search(r"(?:^|,)\s*PGM\s*=\s*([A-Z0-9@$#_-]+)", body, flags=re.IGNORECASE)
+            proc_match = re.search(r"(?:^|,)\s*PROC\s*=\s*([A-Z0-9@$#_-]+)", body, flags=re.IGNORECASE)
+            if not program_match and not proc_match:
+                proc_match = re.match(r"\s*([A-Z0-9@$#_-]+)", body, flags=re.IGNORECASE)
+            current = JCLStep(
+                order=len(steps) + 1,
+                name=(name or f"STEP{len(steps) + 1:03d}").upper(),
+                program=program_match.group(1).upper() if program_match else None,
+                procedure=proc_match.group(1).upper() if proc_match else None,
+                condition=_parameter(body, "COND"),
+                parameters=_parameter(body, "PARM"),
+                raw=body.strip(),
+            )
+            steps.append(current)
+            continue
+        if upper_operation == "DD" and current is not None:
+            dd_name = name.upper() if name else (current.dd[-1].name if current.dd else "UNNAMED")
+            statement = DDStatement(
+                name=dd_name,
+                raw=body.strip(),
+                datasets=[match.rstrip(",") for match in re.findall(r"\bDSN(?:AME)?\s*=\s*([^,\s]+)", body, flags=re.IGNORECASE)],
+                disposition=_parameter(body, "DISP"),
+                dummy=bool(re.search(r"\bDUMMY\b", body, flags=re.IGNORECASE)),
+                inline=bool(re.match(r"\s*(?:\*|DATA\b)", body, flags=re.IGNORECASE)),
+            )
+            current.dd.append(statement)
+    return JCLJob(job_name, steps)
+
+
+def _logical_statements(text: str) -> list[tuple[str, str, str]]:
+    result: list[tuple[str, str, str]] = []
+    active: tuple[str, str, str] | None = None
+    in_inline = False
+    for raw_line in text.splitlines():
+        if in_inline:
+            if raw_line.strip() == "/*":
+                in_inline = False
+                continue
+            if not raw_line.startswith("//"):
+                continue
+            in_inline = False
+        if raw_line.startswith("//*") or raw_line.strip() in {"", "//"}:
+            continue
+        if not raw_line.startswith("//"):
+            continue
+        content = raw_line[2:]
+        match = re.match(r"([^\s]*)\s+([A-Za-z0-9]+)\s*(.*)$", content)
+        if match and match.group(2).upper() in {"JOB", "EXEC", "DD", "PROC", "PEND"}:
+            if active is not None:
+                result.append(active)
+            active = (match.group(1), match.group(2), match.group(3))
+        elif active is not None:
+            active = (active[0], active[1], f"{active[2]} {content.strip()}")
+        if active is not None and active[1].upper() == "DD" and re.match(r"\s*(?:\*|DATA\b)", active[2], flags=re.IGNORECASE):
+            result.append(active)
+            active = None
+            in_inline = True
+    if active is not None:
+        result.append(active)
+    return result
+
+
+def _parameter(body: str, name: str) -> str | None:
+    match = re.search(
+        rf"(?:^|,)\s*{re.escape(name)}\s*=\s*(\([^)]*\)|'[^']*'|\"[^\"]*\"|[^,\s]+)",
+        body,
+        flags=re.IGNORECASE,
+    )
+    return match.group(1).strip() if match else None

mainframe_toolkit/pipeline.py

@@ -0,0 +1,322 @@
+from __future__ import annotations
+
+import json
+import re
+import time
+import traceback
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable, Iterable, Mapping
+
+from .io import write_dataset
+
+
+class PipelineRunError(RuntimeError):
+    def __init__(self, program: str, cause: BaseException):
+        super().__init__(f"program {program} failed: {cause}")
+        self.program = program
+        self.cause = cause
+
+
+class ProgramReturnCodeError(RuntimeError):
+    pass
+
+
+class DatasetRegistry:
+    """In-memory handoff between converted programs."""
+
+    def __init__(self, initial: Mapping[str, Any] | None = None):
+        self._datasets: dict[str, Any] = dict(initial or {})
+        self._versions: dict[str, int] = {name: 1 for name in self._datasets}
+
+    def get(self, name: str) -> Any:
+        try:
+            return self._datasets[name]
+        except KeyError as exc:
+            available = ", ".join(sorted(self._datasets)) or "<none>"
+            raise KeyError(f"dataset {name!r} is unavailable; available: {available}") from exc
+
+    def put(self, name: str, value: Any) -> int:
+        self._datasets[name] = value
+        self._versions[name] = self._versions.get(name, 0) + 1
+        return self._versions[name]
+
+    def contains(self, name: str) -> bool:
+        return name in self._datasets
+
+    def snapshot(self) -> dict[str, Any]:
+        return dict(self._datasets)
+
+    def version(self, name: str) -> int:
+        return self._versions.get(name, 0)
+
+
+@dataclass
+class PipelineContext:
+    spark: Any
+    job_name: str
+    audit_dir: str | Path
+    datasets: DatasetRegistry = field(default_factory=DatasetRegistry)
+    run_id: str = field(default_factory=lambda: _default_run_id())
+    audit_format: str = "parquet"
+    audit_options: Mapping[str, Any] = field(default_factory=dict)
+    persist_intermediates: bool = True
+    metadata: dict[str, Any] = field(default_factory=dict)
+    _program: str | None = field(default=None, init=False, repr=False)
+    _step: int = field(default=0, init=False, repr=False)
+    _published: list[dict[str, Any]] = field(default_factory=list, init=False, repr=False)
+
+    @property
+    def run_dir(self) -> Path:
+        return Path(self.audit_dir) / _safe_name(self.job_name) / _safe_name(self.run_id)
+
+    def input(self, name: str) -> Any:
+        return self.datasets.get(name)
+
+    def publish(
+        self,
+        name: str,
+        value: Any,
+        *,
+        audit: bool = True,
+        format: str | None = None,
+        options: Mapping[str, Any] | None = None,
+    ) -> Any:
+        if self._program is None:
+            raise RuntimeError("publish must be called while a program is executing")
+        if self.persist_intermediates and _is_spark_dataframe(value):
+            value = value.persist()
+        version = self.datasets.put(name, value)
+        item: dict[str, Any] = {
+            "dataset": name,
+            "version": version,
+            "program": self._program,
+            "audited": audit,
+        }
+        if audit:
+            item["audit_path"] = self._write_audit(
+                name,
+                value,
+                format=format or self.audit_format,
+                options=options,
+            )
+        self._published.append(item)
+        return value
+
+    def _write_audit(
+        self,
+        name: str,
+        value: Any,
+        *,
+        format: str,
+        options: Mapping[str, Any] | None,
+    ) -> str:
+        selected = format.lower()
+        base = self.run_dir / f"{self._step:03d}_{_safe_name(self._program or 'program')}" / _safe_name(name)
+        if _is_spark_dataframe(value):
+            write_dataset(
+                value,
+                base,
+                format=selected,
+                mode="overwrite",
+                options={**self.audit_options, **dict(options or {})},
+            )
+            return str(base)
+
+        base.mkdir(parents=True, exist_ok=True)
+        output = base / "part-00000.jsonl"
+        rows: Iterable[Any]
+        if isinstance(value, Mapping):
+            rows = [value]
+        elif isinstance(value, (str, bytes)):
+            rows = [{"value": value.decode() if isinstance(value, bytes) else value}]
+        else:
+            try:
+                rows = iter(value)
+            except TypeError:
+                rows = [{"value": value}]
+        with output.open("w", encoding="utf-8") as stream:
+            for row in rows:
+                stream.write(json.dumps(row, ensure_ascii=False, default=str) + "\n")
+        return str(output)
+
+    def _begin_program(self, name: str, step: int) -> int:
+        self._program = name
+        self._step = step
+        before = len(self._published)
+        return before
+
+    def _finish_program(self) -> None:
+        self._program = None
+
+
+@dataclass(frozen=True)
+class ProgramResult:
+    outputs: Mapping[str, Any] = field(default_factory=dict)
+    return_code: int = 0
+    messages: tuple[str, ...] = field(default_factory=tuple)
+
+
+ProgramFunction = Callable[[PipelineContext], Mapping[str, Any] | ProgramResult | None]
+ProgramCondition = Callable[[PipelineContext, Mapping[str, int]], bool]
+
+
+@dataclass(frozen=True)
+class PipelineProgram:
+    name: str
+    function: ProgramFunction
+    audit_outputs: bool = True
+    audit_format: str | None = None
+    allowed_return_codes: tuple[int, ...] = (0, 4)
+    condition: ProgramCondition | None = None
+
+    def execute(self, context: PipelineContext) -> Mapping[str, Any] | ProgramResult | None:
+        return self.function(context)
+
+
+class PipelineRunner:
+    def __init__(self, context: PipelineContext, programs: Iterable[PipelineProgram]):
+        self.context = context
+        self.programs = list(programs)
+        names = [program.name for program in self.programs]
+        if len(names) != len(set(names)):
+            raise ValueError("program names must be unique")
+
+    def run(self, *, through: str | None = None) -> dict[str, Any]:
+        selected = self._selected_programs(through)
+        manifest: dict[str, Any] = {
+            "job": self.context.job_name,
+            "run_id": self.context.run_id,
+            "started_at": _now(),
+            "through": through,
+            "status": "running",
+            "programs": [],
+            "metadata": self.context.metadata,
+        }
+        self._write_manifest(manifest)
+        return_codes: dict[str, int] = {}
+        for index, program in enumerate(selected, 1):
+            try:
+                should_run = program.condition is None or program.condition(self.context, return_codes)
+            except BaseException as exc:
+                record = {
+                    "step": index,
+                    "program": program.name,
+                    "status": "failed",
+                    "phase": "condition",
+                    "error_type": type(exc).__name__,
+                    "error": str(exc),
+                    "traceback": "".join(traceback.format_exception(exc)),
+                }
+                manifest["programs"].append(record)
+                manifest.update(status="failed", failed_program=program.name, finished_at=_now())
+                self._write_manifest(manifest)
+                raise PipelineRunError(program.name, exc) from exc
+            if not should_run:
+                manifest["programs"].append(
+                    {
+                        "step": index,
+                        "program": program.name,
+                        "status": "skipped",
+                        "reason": "JCL condition evaluated to false",
+                    }
+                )
+                self._write_manifest(manifest)
+                continue
+            started = time.perf_counter()
+            published_before = self.context._begin_program(program.name, index)
+            record: dict[str, Any] = {
+                "step": index,
+                "program": program.name,
+                "status": "running",
+                "started_at": _now(),
+            }
+            manifest["programs"].append(record)
+            self._write_manifest(manifest)
+            try:
+                execution_result = program.execute(self.context)
+                if isinstance(execution_result, ProgramResult):
+                    return_code = execution_result.return_code
+                    outputs = execution_result.outputs
+                    record["messages"] = list(execution_result.messages)
+                else:
+                    return_code = 0
+                    outputs = execution_result
+                record["return_code"] = return_code
+                return_codes[program.name] = return_code
+                if return_code not in program.allowed_return_codes:
+                    raise ProgramReturnCodeError(
+                        f"return code {return_code} is not allowed; expected one of {program.allowed_return_codes}"
+                    )
+                if outputs is not None:
+                    if not isinstance(outputs, Mapping):
+                        raise TypeError("program must return a dataset mapping or None")
+                    already = {
+                        item["dataset"]
+                        for item in self.context._published[published_before:]
+                    }
+                    for name, value in outputs.items():
+                        if name not in already:
+                            self.context.publish(
+                                name,
+                                value,
+                                audit=program.audit_outputs,
+                                format=program.audit_format,
+                            )
+                record["status"] = "success"
+                record["outputs"] = self.context._published[published_before:]
+            except BaseException as exc:
+                record.update(
+                    status="failed",
+                    error_type=type(exc).__name__,
+                    error=str(exc),
+                    traceback="".join(traceback.format_exception(exc)),
+                )
+                manifest.update(status="failed", failed_program=program.name, finished_at=_now())
+                record["duration_ms"] = round((time.perf_counter() - started) * 1000, 3)
+                self.context._finish_program()
+                self._write_manifest(manifest)
+                raise PipelineRunError(program.name, exc) from exc
+            record["duration_ms"] = round((time.perf_counter() - started) * 1000, 3)
+            record["finished_at"] = _now()
+            self.context._finish_program()
+            self._write_manifest(manifest)
+        manifest.update(status="success", finished_at=_now())
+        self._write_manifest(manifest)
+        return manifest
+
+    def _selected_programs(self, through: str | None) -> list[PipelineProgram]:
+        if through is None:
+            return self.programs
+        for index, program in enumerate(self.programs):
+            if program.name.casefold() == through.casefold():
+                return self.programs[: index + 1]
+        raise ValueError(f"unknown checkpoint program {through!r}")
+
+    def _write_manifest(self, manifest: Mapping[str, Any]) -> None:
+        path = self.context.run_dir / "run-manifest.json"
+        path.parent.mkdir(parents=True, exist_ok=True)
+        temporary = path.with_suffix(".tmp")
+        with temporary.open("w", encoding="utf-8") as stream:
+            json.dump(manifest, stream, ensure_ascii=False, indent=2, default=str)
+            stream.write("\n")
+        temporary.replace(path)
+
+
+def _is_spark_dataframe(value: Any) -> bool:
+    return hasattr(value, "write") and hasattr(value, "schema") and hasattr(value, "sparkSession")
+
+
+def _safe_name(value: str) -> str:
+    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._") or "unnamed"
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _default_run_id() -> str:
+    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
+    return f"{stamp}-{uuid.uuid4().hex[:8]}"