Merge pull request #160 from forcedotcom/jo_lineage_gen

generate lineage

Author: joroscoSF

README.md

@@ -245,6 +245,29 @@ If none of these exist, `find_file_path` raises `FileNotFoundError` with the lis
 
 `$LIBRARY_PATH` is set automatically to the root of the package at runtime inside Data Cloud.
 
+### Generating lineage
+
+`generate_lineage` runs a script entry point locally and derives the field-level `dataLineage` section of `config.json` from the Spark plans of the DataFrames it writes:
+
+```python
+from datacustomcode.lineage import generate_lineage
+
+result = generate_lineage("payload/entrypoint.py", sf_cli_org="my-org")
+print(result.to_dict()["dataLineage"])
+# [{"targetField": "Output__dll.name__c",
+#   "sourceFields": ["Input__dll.first__c", "Input__dll.last__c"]}, ...]
+```
+
+The entry point runs exactly as it would locally (same config, reader, and credentials), except that writes are traced and validated against the target DLO instead of being performed. It runs in a separate Python process, which inherits your working directory, environment, and `sys.path`; errors raised by the script are re-raised from `generate_lineage`. Reads use the reader's `default_row_limit` (1000 by default; set it in your `config.yaml`).
+
+> **Review the output before you use it.** The generated lineage is a starting point for review, not a guarantee. It reflects only the code paths and data your local run exercised, and it can't see inside UDFs. Accurate and complete `dataLineage` in your package is your responsibility as the package author.
+
+- A target field's sources are the fields its *value* is computed from. Fields that only decide which rows, groups or branch are used aren't sources, however the logic is written: filter conditions, join keys, `groupBy` keys, sort keys, window `partitionBy`/`orderBy` keys, `when()`/`if()` conditions, and aggregate `FILTER (WHERE …)` clauses. For example, `groupBy("region").agg(sum("amount"))` and `sum("amount").over(Window.partitionBy("region"))` both trace to `amount` only.
+- Fields computed only from literals are listed in `noSources`. So are fields whose value doesn't come from any input field's value, such as `row_number()`, `rank()`, `count("*")`, or `sum(when(cond, 1))`.
+- Fields that can't be traced are listed in `unresolved` with a reason and get no lineage. This covers data that didn't come from `read_dlo`/`read_dmo` (for example a `toPandas()` → `createDataFrame` round trip), `mapInPandas`, and subqueries.
+- `pairCount` is the total number of (target, source) field pairs.
+- Only batch script packages are supported.
+
 
 ## CLI
 

src/datacustomcode/client.py

@@ -376,7 +376,13 @@ def __new__(
             instance._spark_named_credential = spark_named_credential
             # Initialize Readers and Writers from config
             # and/or provided reader and writer
-            if reader is None or writer is None:
+            reader_forced = (
+                config.reader_config is not None and config.reader_config.force
+            )
+            writer_forced = (
+                config.writer_config is not None and config.writer_config.force
+            )
+            if reader is None or writer is None or reader_forced or writer_forced:
                 # We need a spark because we will initialize readers and writers.
                 # Reuse the process-wide session if one client already built it,
                 # so a Client and a StreamingClient share a single connection.
@@ -404,19 +410,15 @@ def __new__(
                 raise ValueError(
                     "Reader config is required when reader is not provided"
                 )
-            elif reader is None or (
-                config.reader_config is not None and config.reader_config.force
-            ):
+            elif reader is None or reader_forced:
                 reader_init = config.reader_config.to_object(spark)  # type: ignore
             else:
                 reader_init = reader
             if config.writer_config is None and writer is None:
                 raise ValueError(
                     "Writer config is required when writer is not provided"
                 )
-            elif writer is None or (
-                config.writer_config is not None and config.writer_config.force
-            ):
+            elif writer is None or writer_forced:
                 writer_init = config.writer_config.to_object(spark)  # type: ignore
             else:
                 writer_init = writer

src/datacustomcode/lineage/__init__.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Generate config.json ``dataLineage`` by tracing a script's Spark plans."""
+
+from __future__ import annotations
+
+import builtins
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from typing import (
+    Any,
+    Dict,
+    NoReturn,
+    Optional,
+    Sequence,
+)
+
+from datacustomcode.lineage.capture import (
+    LineageCaptureReader,
+    LineageCaptureWriter,
+    LineageCollector,
+    LineageResult,
+    install_capture,
+)
+from datacustomcode.lineage.tracer import LineageTraceError
+
+__all__ = [
+    "LineageCaptureReader",
+    "LineageCaptureWriter",
+    "LineageCollector",
+    "LineageResult",
+    "LineageTraceError",
+    "generate_lineage",
+    "install_capture",
+]
+
+
+def generate_lineage(
+    entrypoint: str,
+    config_file: Optional[str] = None,
+    dependencies: Sequence[str] = (),
+    profile: str = "default",
+    sf_cli_org: Optional[str] = None,
+) -> LineageResult:
+    """Run a script entrypoint and return the field-level lineage it produces.
+
+    The entrypoint runs exactly as with ``run`` (same config resolution, reader
+    and credentials), except that writes are traced instead of performed. It runs
+    in a fresh Python process, so no client or config from the caller's process
+    leaks into the capture. The process inherits this one's working directory,
+    environment, ``sys.path`` and stdout/stderr.
+
+    Args:
+        entrypoint: The script entrypoint to run.
+        config_file: Optional SDK config file, as for ``run``.
+        dependencies: Modules to import before running, as for ``run``.
+        profile: The credentials profile to use.
+        sf_cli_org: Optional SF CLI org alias or username.
+
+    Returns:
+        The captured :class:`LineageResult`.
+
+    Raises:
+        ValueError: If the package is not a batch script.
+        Exception: Whatever the run raised, re-raised as the same built-in
+            exception type or :class:`LineageTraceError`, else ``RuntimeError``.
+    """
+    with tempfile.TemporaryDirectory(prefix="datacustomcode-lineage-") as tmp:
+        request_file = os.path.join(tmp, "request.json")
+        response_file = os.path.join(tmp, "response.json")
+        with open(request_file, "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "args": {
+                        "entrypoint": entrypoint,
+                        "config_file": config_file,
+                        "dependencies": list(dependencies),
+                        "profile": profile,
+                        "sf_cli_org": sf_cli_org,
+                    },
+                    "response_file": response_file,
+                },
+                f,
+            )
+        # The worker sees the same modules this process does
+        env = dict(os.environ)
+        env["PYTHONPATH"] = os.pathsep.join(path for path in sys.path if path)
+        completed = subprocess.run(
+            [sys.executable, "-m", "datacustomcode.lineage._worker", request_file],
+            env=env,
+            check=False,
+        )
+        if not os.path.exists(response_file):
+            raise RuntimeError(
+                f"Lineage generation exited with code {completed.returncode} "
+                "without a result"
+            )
+        with open(response_file, encoding="utf-8") as f:
+            response = json.load(f)
+    if "error" in response:
+        _raise(response["error"])
+    return LineageResult.from_dict(response["result"])
+
+
+def _raise(error: Dict[str, Any]) -> NoReturn:
+    """Re-raise an error reported by the worker as the closest local type."""
+    message = error["message"]
+    if error["type"] == LineageTraceError.__name__:
+        raise LineageTraceError(message)
+    exc_type = getattr(builtins, error["type"], None)
+    if (
+        error["module"] == "builtins"
+        and isinstance(exc_type, type)
+        and issubclass(exc_type, Exception)
+    ):
+        try:
+            exc = exc_type(message)
+        except TypeError:  # e.g. UnicodeDecodeError takes more arguments
+            pass
+        else:
+            raise exc
+    raise RuntimeError(f"{error['module']}.{error['type']}: {message}")

src/datacustomcode/lineage/_worker.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Runs one lineage capture in a fresh process, for ``generate_lineage``.
+
+Usage: ``python -m datacustomcode.lineage._worker <request.json>``. The request
+holds the ``run_capture`` arguments and the path of a response file, which gets
+either ``{"result": LineageResult.to_dict()}`` or
+``{"error": {"module", "type", "message"}}``. The response goes to a file rather
+than stdout because the entrypoint prints freely to stdout.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import traceback
+from typing import (
+    List,
+    Optional,
+    Sequence,
+)
+
+from datacustomcode.lineage.capture import (
+    LineageCollector,
+    LineageResult,
+    install_capture,
+)
+
+
+def run_capture(
+    entrypoint: str,
+    config_file: Optional[str],
+    dependencies: Sequence[str],
+    profile: str,
+    sf_cli_org: Optional[str],
+) -> LineageResult:
+    """Run the entrypoint with capture installed. Call only in a fresh process."""
+    from datacustomcode.run import (
+        execute_entrypoint,
+        import_dependencies,
+        prepare_run,
+    )
+
+    # Dependencies are imported after capture is installed, so a Client they
+    # create at import time is captured too
+    context = prepare_run(entrypoint, config_file, [], profile, sf_cli_org)
+    if context.package_type != "script":
+        raise ValueError(
+            f"Lineage can only be generated for script packages, not "
+            f"{context.package_type} packages"
+        )
+    if context.config_json.get("streamingSource"):
+        raise ValueError("Lineage generation does not support streaming scripts")
+
+    collector = LineageCollector()
+    install_capture(collector)
+    import_dependencies(list(dependencies))
+    execute_entrypoint(entrypoint, context)
+    return collector.result()
+
+
+def main(argv: List[str]) -> int:
+    with open(argv[0], encoding="utf-8") as f:
+        request = json.load(f)
+    try:
+        result = run_capture(**request["args"])
+        response: dict = {"result": result.to_dict()}
+    except Exception as exc:
+        # The caller re-raises from the response; the traceback is for the
+        # customer, pointing into their entrypoint
+        traceback.print_exc()
+        response = {
+            "error": {
+                "module": type(exc).__module__,
+                "type": type(exc).__name__,
+                "message": str(exc),
+            }
+        }
+    with open(request["response_file"], "w", encoding="utf-8") as f:
+        json.dump(response, f)
+    return 1 if "error" in response else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))