From f9d157b6f395684dc530e5eda3dcf254abcc8a02 Mon Sep 17 00:00:00 2001 From: Jesus Orosco Date: Fri, 9 Oct 2026 09:54:24 -0700 Subject: [PATCH 1/6] adding some lineage generating features --- README.md | 21 ++ src/datacustomcode/lineage/__init__.py | 82 ++++++ src/datacustomcode/lineage/capture.py | 280 +++++++++++++++++++ src/datacustomcode/lineage/tracer.py | 372 +++++++++++++++++++++++++ src/datacustomcode/run.py | 53 +++- tests/lineage/__init__.py | 1 + tests/lineage/conftest.py | 20 ++ tests/lineage/test_generate.py | 236 ++++++++++++++++ tests/lineage/test_tracer.py | 225 +++++++++++++++ tests/test_run.py | 29 ++ 10 files changed, 1318 insertions(+), 1 deletion(-) create mode 100644 src/datacustomcode/lineage/__init__.py create mode 100644 src/datacustomcode/lineage/capture.py create mode 100644 src/datacustomcode/lineage/tracer.py create mode 100644 tests/lineage/__init__.py create mode 100644 tests/lineage/conftest.py create mode 100644 tests/lineage/test_generate.py create mode 100644 tests/lineage/test_tracer.py diff --git a/README.md b/README.md index a9f5ae8..56dcacc 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,27 @@ 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. Reads use the reader's `default_row_limit` (1000 by default; set it in your `config.yaml`). + +- A target field's sources are the fields its *value* is computed from. Filter conditions and join keys aren't sources. +- Fields computed only from literals are listed in `noSources`. +- 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 diff --git a/src/datacustomcode/lineage/__init__.py b/src/datacustomcode/lineage/__init__.py new file mode 100644 index 0000000..9ca831b --- /dev/null +++ b/src/datacustomcode/lineage/__init__.py @@ -0,0 +1,82 @@ +# 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 + +from typing import Optional, Sequence + +from datacustomcode.lineage.capture import ( + LineageCaptureReader, + LineageCaptureWriter, + LineageCollector, + LineageResult, + capture_lineage, +) +from datacustomcode.lineage.tracer import LineageTraceError + +__all__ = [ + "LineageCaptureReader", + "LineageCaptureWriter", + "LineageCollector", + "LineageResult", + "LineageTraceError", + "capture_lineage", + "generate_lineage", +] + + +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. + + 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. + """ + from datacustomcode.run import execute_entrypoint, prepare_run + + context = prepare_run( + entrypoint, config_file, list(dependencies), 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() + with capture_lineage(collector): + execute_entrypoint(entrypoint, context) + return collector.result() diff --git a/src/datacustomcode/lineage/capture.py b/src/datacustomcode/lineage/capture.py new file mode 100644 index 0000000..7a884c5 --- /dev/null +++ b/src/datacustomcode/lineage/capture.py @@ -0,0 +1,280 @@ +# 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. +"""Readers and writers that record lineage instead of writing data. + +While :func:`capture_lineage` is active, the global reader/writer config is +replaced (with ``force=True``) by :class:`LineageCaptureReader` and +:class:`LineageCaptureWriter`. They wrap the reader/writer the config had +already resolved, so the customer's reader choice and credentials still apply: +reads go through the delegate and are registered as lineage sources, writes are +validated through the delegate (when it supports it), traced, and never written. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterator, + List, + Optional, + Union, +) + +from datacustomcode.config import AccessLayerObjectConfig, config +from datacustomcode.io.reader.base import BaseDataCloudReader +from datacustomcode.io.writer.base import ( + MERGE_RECORD_TYPE_COLUMN, + BaseDataCloudWriter, + WriteMode, +) +from datacustomcode.lineage.tracer import SourceRegistry, TraceResult + +if TYPE_CHECKING: + from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession + from pyspark.sql.types import AtomicType, StructType + + +@dataclass +class LineageResult: + """Field-level lineage captured from one run of an entrypoint. + + Attributes: + data_lineage: Target field to its sorted source fields. + unresolved: Target field to the reason its lineage could not be traced. + no_sources: Target fields computed from no source field (e.g. literals). + targets: Objects written. + sources: Objects read. + """ + + data_lineage: Dict[str, List[str]] = field(default_factory=dict) + unresolved: Dict[str, str] = field(default_factory=dict) + no_sources: List[str] = field(default_factory=list) + targets: List[str] = field(default_factory=list) + sources: List[str] = field(default_factory=list) + + @property + def pair_count(self) -> int: + """Number of (target, source) field pairs in ``data_lineage``.""" + return sum(len(sources) for sources in self.data_lineage.values()) + + def to_dict(self) -> Dict[str, Any]: + """Deterministic JSON-ready form; ``dataLineage`` matches config.json.""" + return { + "dataLineage": [ + {"targetField": target, "sourceFields": self.data_lineage[target]} + for target in sorted(self.data_lineage) + ], + "unresolved": [ + {"targetField": target, "reason": self.unresolved[target]} + for target in sorted(self.unresolved) + ], + "noSources": sorted(self.no_sources), + "targets": sorted(self.targets), + "sources": sorted(self.sources), + "pairCount": self.pair_count, + } + + +class LineageCollector: + """Accumulates sources read and target fields written during a run.""" + + def __init__(self) -> None: + self.registry = SourceRegistry() + self._targets: Dict[str, Dict[str, TraceResult]] = {} + + def record_read(self, object_name: str, dataframe: PySparkDataFrame) -> None: + self.registry.register(dataframe, object_name) + + def record_write(self, object_name: str, dataframe: PySparkDataFrame) -> None: + fields = self._targets.setdefault(object_name, {}) + for column, traced in self.registry.trace(dataframe).items(): + if column == MERGE_RECORD_TYPE_COLUMN: + continue + # Repeat writes to the same target: a field reads from the union of + # every write, and stays unresolved if any write was unresolved + previous = fields.get(column) + if previous is None or (traced.unresolved and not previous.unresolved): + fields[column] = traced + elif not previous.unresolved: + fields[column] = TraceResult(sources=previous.sources | traced.sources) + + def result(self) -> LineageResult: + result = LineageResult( + targets=list(self._targets), sources=list(self.registry.objects) + ) + for object_name, fields in self._targets.items(): + for column, traced in fields.items(): + target = f"{object_name}.{column}" + if traced.reason is not None: + result.unresolved[target] = traced.reason + elif traced.sources: + result.data_lineage[target] = sorted(traced.sources) + else: + result.no_sources.append(target) + return result + + +@dataclass +class _CaptureSession: + collector: LineageCollector + reader_config: AccessLayerObjectConfig + writer_config: AccessLayerObjectConfig + + +# The capture in progress, read by the capture classes when Client builds them +_active: Dict[str, Optional[_CaptureSession]] = {"session": None} + + +def _active_session() -> _CaptureSession: + session = _active["session"] + if session is None: + raise RuntimeError( + "Lineage capture readers and writers can only be created inside " + "capture_lineage()" + ) + return session + + +class _Delegating: + """Falls back to the delegate for attributes the capture class lacks.""" + + _delegate: Any + + def __getattr__(self, item: str) -> Any: + if item == "_delegate": + raise AttributeError(item) + return getattr(self._delegate, item) + + +class LineageCaptureReader(_Delegating, BaseDataCloudReader): + """Reads through the configured reader and registers lineage sources.""" + + CONFIG_NAME = "LineageCaptureReader" + + def __init__(self, spark: SparkSession) -> None: + super().__init__(spark) + session = _active_session() + self._collector = session.collector + self._delegate = session.reader_config.to_object(spark) + + def read_dlo( + self, + name: str, + schema: Union[AtomicType, StructType, str, None] = None, + ) -> PySparkDataFrame: + dataframe = ( + self._delegate.read_dlo(name) + if schema is None + else self._delegate.read_dlo(name, schema) + ) + self._collector.record_read(name, dataframe) + return dataframe # type: ignore[no-any-return] + + def read_dmo( + self, + name: str, + schema: Union[AtomicType, StructType, str, None] = None, + ) -> PySparkDataFrame: + dataframe = ( + self._delegate.read_dmo(name) + if schema is None + else self._delegate.read_dmo(name, schema) + ) + self._collector.record_read(name, dataframe) + return dataframe # type: ignore[no-any-return] + + +class LineageCaptureWriter(_Delegating, BaseDataCloudWriter): + """Traces DataFrames handed to the writer; never writes or shows them.""" + + CONFIG_NAME = "LineageCaptureWriter" + + def __init__(self, spark: SparkSession) -> None: + super().__init__(spark) + session = _active_session() + self._collector = session.collector + self._delegate = session.writer_config.to_object(spark) + + def _validate_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + # Same column check `run` does, so lineage never names a missing field + validate = getattr( + self._delegate, "validate_dataframe_columns_against_dlo", None + ) + if validate is not None: + validate(dataframe, name) + + def write_to_dlo( + self, + name: str, + dataframe: PySparkDataFrame, + write_mode: WriteMode, + **kwargs: Any, + ) -> None: + self.auto_write_to_dlo(name, dataframe) + + def write_to_dmo( + self, + name: str, + dataframe: PySparkDataFrame, + write_mode: WriteMode, + **kwargs: Any, + ) -> None: + self.auto_write_to_dmo(name, dataframe) + + def auto_write_to_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: + self._validate_dlo(name, dataframe) + self._collector.record_write(name, dataframe) + + def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: + self._collector.record_write(name, dataframe) + + +@contextmanager +def capture_lineage(collector: LineageCollector) -> Iterator[LineageCollector]: + """Route every ``Client`` created inside the block through lineage capture. + + The global reader/writer config is swapped for forced capture configs that + wrap the currently configured reader/writer, and any existing ``Client`` + singleton is set aside so the entrypoint's ``Client()`` picks them up. Both + are restored on exit. + """ + from datacustomcode.client import Client, StreamingClient + + if config.reader_config is None or config.writer_config is None: + raise ValueError("A reader and writer must be configured to capture lineage") + if _active["session"] is not None: + raise RuntimeError("Lineage capture is already active") + + saved_reader, saved_writer = config.reader_config, config.writer_config + saved_clients = Client._instance, StreamingClient._instance + _active["session"] = _CaptureSession(collector, saved_reader, saved_writer) + config.reader_config = AccessLayerObjectConfig( + type_config_name=LineageCaptureReader.CONFIG_NAME, force=True + ) + config.writer_config = AccessLayerObjectConfig( + type_config_name=LineageCaptureWriter.CONFIG_NAME, force=True + ) + Client._instance = None + StreamingClient._instance = None + try: + yield collector + finally: + config.reader_config, config.writer_config = saved_reader, saved_writer + Client._instance, StreamingClient._instance = saved_clients + _active["session"] = None diff --git a/src/datacustomcode/lineage/tracer.py b/src/datacustomcode/lineage/tracer.py new file mode 100644 index 0000000..6a9e7ea --- /dev/null +++ b/src/datacustomcode/lineage/tracer.py @@ -0,0 +1,372 @@ +# 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. +"""Column-level lineage tracing over a DataFrame's analyzed logical plan. + +Every column in a Spark plan is an attribute with a unique ``exprId``. Walking the +analyzed plan bottom-up, each attribute is resolved to the set of *leaf columns* +it is computed from, where a leaf column is ``(leaf relation identity, position)``. + +Leaf identity (rather than ``exprId``) is what makes self-joins traceable: Spark +re-aliases one side of a self-join to fresh ``exprId`` values, but both sides +still sit on the same underlying RDD / local data. + +Only the columns that feed a column's *value* are sources. Filter predicates, +join conditions and sort keys gate or order rows and are deliberately ignored. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from typing import ( + TYPE_CHECKING, + Any, + Dict, + FrozenSet, + Hashable, + Iterator, + List, + Optional, + Tuple, +) + +if TYPE_CHECKING: + from pyspark.sql import DataFrame as PySparkDataFrame + +LeafColumn = Tuple[Hashable, int] +"""``(leaf relation identity, column position)``.""" + +UNTRACEABLE_SOURCE = "untraceable source" + +_OUTPUT_SEPARATOR = "\u0001" + +# Expressions that hide the columns they read in a nested plan the JSON does not +# carry, so their inputs cannot be traced +_OPAQUE_EXPRESSIONS = frozenset( + {"ScalarSubquery", "Exists", "ListQuery", "InSubquery", "LateralSubquery"} +) + + +class LineageTraceError(Exception): + """Raised when a DataFrame's plan cannot be traced at all.""" + + +@dataclass(frozen=True) +class ColumnTrace: + """Lineage of one output column. + + Attributes: + name: Output column name. + leaves: Leaf columns the value is computed from. Empty for literals. + opaque_reason: Set when the plan hides this column's inputs (e.g. a + ``mapInPandas`` or subquery); ``leaves`` is then incomplete. + """ + + name: str + leaves: FrozenSet[LeafColumn] = field(default_factory=frozenset) + opaque_reason: Optional[str] = None + + +@dataclass(frozen=True) +class _AttrTrace: + leaves: FrozenSet[LeafColumn] + opaque_reason: Optional[str] = None + + def merge(self, other: _AttrTrace) -> _AttrTrace: + return _AttrTrace( + self.leaves | other.leaves, self.opaque_reason or other.opaque_reason + ) + + +_EMPTY = _AttrTrace(frozenset()) + + +@dataclass(frozen=True) +class TraceResult: + """Lineage of one output column in terms of registered source fields. + + Attributes: + sources: ``"Object.field"`` names the value is computed from. + reason: Set when the column cannot be fully traced; ``sources`` must + then not be used. + """ + + sources: FrozenSet[str] = field(default_factory=frozenset) + reason: Optional[str] = None + + @property + def unresolved(self) -> bool: + return self.reason is not None + + +class SourceRegistry: + """Maps leaf columns of DataFrames read from Data Cloud to their fields.""" + + def __init__(self) -> None: + self._fields: Dict[LeafColumn, str] = {} + self.objects: List[str] = [] + + def register(self, df: PySparkDataFrame, object_name: str) -> None: + """Record that the columns of ``df`` are the fields of ``object_name``.""" + for column in trace_columns(df): + for leaf in column.leaves: + self._fields[leaf] = f"{object_name}.{column.name}" + if object_name not in self.objects: + self.objects.append(object_name) + + def trace(self, df: PySparkDataFrame) -> Dict[str, TraceResult]: + """Resolve each column of ``df`` to the registered fields it reads.""" + results = {} + for column in trace_columns(df): + if column.opaque_reason is not None: + results[column.name] = TraceResult(reason=column.opaque_reason) + elif any(leaf not in self._fields for leaf in column.leaves): + results[column.name] = TraceResult(reason=UNTRACEABLE_SOURCE) + else: + results[column.name] = TraceResult( + sources=frozenset(self._fields[leaf] for leaf in column.leaves) + ) + return results + + +def trace_columns(df: PySparkDataFrame) -> List[ColumnTrace]: + """Trace every output column of ``df`` back to its leaf columns. + + Args: + df: A classic (non Spark Connect) PySpark DataFrame. + + Returns: + One :class:`ColumnTrace` per column of ``df``, in column order. + + Raises: + LineageTraceError: If the plan is not accessible or is malformed. + """ + jdf = getattr(df, "_jdf", None) + if jdf is None: + raise LineageTraceError( + "Lineage tracing requires a classic PySpark DataFrame; Spark Connect " + "DataFrames are not supported." + ) + jvm = df.sparkSession._jvm + jplan = jdf.queryExecution().analyzed() + nodes = json.loads(jplan.toJSON()) + jnodes = _jvm_preorder(jplan) + if len(nodes) != len(jnodes): + raise LineageTraceError( + f"Plan JSON has {len(nodes)} nodes but the plan tree has {len(jnodes)}" + ) + + children = _children_from_preorder(nodes) + outputs = [_parse_output(jnode) for jnode in jnodes] + traces: Dict[int, _AttrTrace] = {} + + # In pre-order every child comes after its parent, so walking backwards + # resolves all children before the node that consumes them + for index in range(len(nodes) - 1, -1, -1): + _trace_node( + nodes[index], + jnodes[index], + outputs[index], + [outputs[c] for c in children[index]], + traces, + jvm, + ) + + root_ids = outputs[0] + if len(root_ids) != len(df.columns): + raise LineageTraceError( + f"Plan output has {len(root_ids)} columns but the DataFrame has " + f"{len(df.columns)}" + ) + return [ + ColumnTrace( + name=name, + leaves=traces.get(expr_id, _EMPTY).leaves, + opaque_reason=traces.get(expr_id, _EMPTY).opaque_reason, + ) + for name, expr_id in zip(df.columns, root_ids, strict=False) + ] + + +def _trace_node( + node: dict, + jnode: Any, + output: List[int], + child_outputs: List[List[int]], + traces: Dict[int, _AttrTrace], + jvm: Any, +) -> None: + node_class = node["class"].rsplit(".", 1)[-1] + + if not child_outputs: + key = _leaf_key(jnode, node_class, jvm) + for position, expr_id in enumerate(output): + traces[expr_id] = _AttrTrace(frozenset({(key, position)})) + return + + if node_class == "Union": + # Union reuses child 0's exprIds; merge every child positionally into them + for position, expr_id in enumerate(output): + merged = _EMPTY + for child in child_outputs: + if position < len(child): + merged = merged.merge(traces.get(child[position], _EMPTY)) + traces[expr_id] = merged + elif node_class == "Expand": + for position, expr_id in enumerate(output): + merged = _EMPTY + for projection in node.get("projections", []): + if position < len(projection): + merged = merged.merge( + _trace_expression(projection[position], traces) + ) + traces[expr_id] = merged + elif node_class == "Generate": + generator = _trace_expression(node.get("generator", []), traces) + for attr_list in node.get("generatorOutput", []): + for attr in attr_list: + traces[_expr_id(attr)] = generator + else: + for expression in _expression_trees(node): + _trace_aliases(expression, traces) + + for expr_id in output: + if expr_id not in traces: + traces[expr_id] = _AttrTrace( + frozenset(), f"columns produced by {node_class} cannot be traced" + ) + + +def _trace_aliases(expression: List[dict], traces: Dict[int, _AttrTrace]) -> None: + """Record every Alias in a flattened expression tree as a new attribute.""" + for index, expr_node in enumerate(expression): + if expr_node["class"].endswith(".Alias"): + end = _subtree_end(expression, index) + traces[_expr_id(expr_node)] = _trace_expression( + expression[index + 1 : end], traces + ) + + +def _trace_expression( + expression: List[dict], traces: Dict[int, _AttrTrace] +) -> _AttrTrace: + """Union the traces of every attribute an expression (sub)tree references.""" + result = _EMPTY + for expr_node in expression: + expr_class = expr_node["class"].rsplit(".", 1)[-1] + if expr_class in _OPAQUE_EXPRESSIONS: + result = result.merge( + _AttrTrace(frozenset(), f"{expr_class} inputs cannot be traced") + ) + elif expr_class == "AttributeReference": + expr_id = _expr_id(expr_node) + if expr_id in traces: + result = result.merge(traces[expr_id]) + else: + result = result.merge( + _AttrTrace(frozenset(), f"unresolved attribute {expr_id}") + ) + return result + + +def _expression_trees(value: Any) -> Iterator[List[dict]]: + """Yield every flattened expression tree held in a plan node's fields.""" + if isinstance(value, dict): + for key, field_value in value.items(): + if key not in ("class", "num-children"): + yield from _expression_trees(field_value) + elif isinstance(value, list) and value: + if all(isinstance(item, dict) and "num-children" in item for item in value): + yield value + else: + for item in value: + if isinstance(item, list): + yield from _expression_trees(item) + + +def _subtree_end(nodes: List[dict], start: int) -> int: + """Index one past the pre-order subtree rooted at ``start``.""" + pending = 1 + index = start + while pending: + pending += nodes[index]["num-children"] - 1 + index += 1 + return index + + +def _children_from_preorder(nodes: List[dict]) -> List[List[int]]: + children: List[List[int]] = [[] for _ in nodes] + open_parents: List[List[int]] = [] # [parent index, children still expected] + for index, node in enumerate(nodes): + if open_parents: + parent = open_parents[-1] + children[parent[0]].append(index) + parent[1] -= 1 + if parent[1] == 0: + open_parents.pop() + if node["num-children"] > 0: + open_parents.append([index, node["num-children"]]) + return children + + +def _jvm_preorder(jplan: Any) -> List[Any]: + result = [] + stack = [jplan] + while stack: + jnode = stack.pop() + result.append(jnode) + jchildren = jnode.children() + stack.extend(jchildren.apply(i) for i in reversed(range(jchildren.size()))) + return result + + +def _parse_output(jnode: Any) -> List[int]: + """ExprIds of a plan node's output attributes, in order. + + One JVM round trip per node: attributes render as ``name#``. + """ + count = jnode.output().size() + if count == 0: + return [] + rendered = jnode.output().mkString(_OUTPUT_SEPARATOR).split(_OUTPUT_SEPARATOR) + if len(rendered) != count: + raise LineageTraceError("Could not parse plan output attributes") + expr_ids = [] + for attr in rendered: + digits = attr.rsplit("#", 1)[-1] + end = 0 + while end < len(digits) and digits[end].isdigit(): + end += 1 + if end == 0: + raise LineageTraceError(f"Could not parse plan attribute {attr!r}") + expr_ids.append(int(digits[:end])) + return expr_ids + + +def _expr_id(expr_node: dict) -> int: + return int(expr_node["exprId"]["id"]) + + +def _leaf_key(jnode: Any, node_class: str, jvm: Any) -> Hashable: + identity = jvm.java.lang.System.identityHashCode + if node_class == "LogicalRDD": + return ("rdd", jnode.rdd().id()) + if node_class == "LocalRelation": + return ("local", identity(jnode.data())) + if node_class == "LogicalRelation": + return ("relation", identity(jnode.relation())) + if node_class == "DataSourceV2Relation": + return ("v2", identity(jnode.table())) + return ("node", identity(jnode)) diff --git a/src/datacustomcode/run.py b/src/datacustomcode/run.py index c1b76d8..16a8ecc 100644 --- a/src/datacustomcode/run.py +++ b/src/datacustomcode/run.py @@ -12,6 +12,7 @@ # 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. +from dataclasses import dataclass, field import importlib import json import os @@ -102,6 +103,15 @@ def _update_config_options(profile: Optional[str], sf_cli_org: Optional[str]): ) +@dataclass(frozen=True) +class RunContext: + """State resolved by :func:`prepare_run` and consumed by + :func:`execute_entrypoint`.""" + + package_type: str + config_json: dict = field(default_factory=dict) + + def run_entrypoint( entrypoint: str, config_file: Union[str, None], @@ -121,6 +131,34 @@ def run_entrypoint( sf_cli_org: Optional SF CLI org alias or username. If provided, credentials are fetched via `sf org display` instead of from credentials.ini. """ + context = prepare_run(entrypoint, config_file, dependencies, profile, sf_cli_org) + execute_entrypoint(entrypoint, context, test_file) + + +def prepare_run( + entrypoint: str, + config_file: Union[str, None], + dependencies: List[str], + profile: str, + sf_cli_org: Optional[str] = None, +) -> RunContext: + """Resolve config and import dependencies for an entrypoint, without running it. + + This is the setup half of :func:`run_entrypoint`. Callers that need to adjust + the resolved global ``config`` before the entrypoint runs (e.g. lineage + generation swapping in capturing readers/writers) call this, make their + changes, then call :func:`execute_entrypoint`. + + Args: + entrypoint: The entrypoint script to run. + config_file: The config file to use. + dependencies: The dependencies to import. + profile: The credentials profile to use. + sf_cli_org: Optional SF CLI org alias or username. + + Returns: + The resolved :class:`RunContext`. + """ add_py_folder(entrypoint) # Read dataspace from config.json (required) @@ -184,8 +222,21 @@ def run_entrypoint( except (ModuleNotFoundError, AttributeError) as inner_exc: raise inner_exc from exc + return RunContext(package_type=package_type, config_json=config_json) + + +def execute_entrypoint( + entrypoint: str, context: RunContext, test_file: Optional[str] = None +) -> None: + """Run an entrypoint previously set up with :func:`prepare_run`. + + Args: + entrypoint: The entrypoint script to run. + context: The context returned by :func:`prepare_run`. + test_file: Optional test JSON file for function testing. + """ # Handle test file for functions - if test_file and package_type == "function": + if test_file and context.package_type == "function": run_function_with_test(entrypoint, test_file) else: runpy.run_path(entrypoint, init_globals=globals(), run_name="__main__") diff --git a/tests/lineage/__init__.py b/tests/lineage/__init__.py new file mode 100644 index 0000000..6ae5294 --- /dev/null +++ b/tests/lineage/__init__.py @@ -0,0 +1 @@ +# Package initialization file diff --git a/tests/lineage/conftest.py b/tests/lineage/conftest.py new file mode 100644 index 0000000..f262243 --- /dev/null +++ b/tests/lineage/conftest.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from pyspark.sql import SparkSession +import pytest + + +@pytest.fixture(scope="module") +def spark(): + # Arrow and case sensitivity match the SDK's default config.yaml + session = ( + SparkSession.builder.master("local[1]") + .appName("lineage-test") + .config("spark.ui.enabled", "false") + .config("spark.sql.caseSensitive", "true") + .config("spark.sql.execution.arrow.pyspark.enabled", "true") + .getOrCreate() + ) + session.sparkContext.setLogLevel("OFF") + yield session + session.stop() diff --git a/tests/lineage/test_generate.py b/tests/lineage/test_generate.py new file mode 100644 index 0000000..46665ad --- /dev/null +++ b/tests/lineage/test_generate.py @@ -0,0 +1,236 @@ +"""End-to-end tests for generate_lineage with an in-memory reader (no org).""" + +from __future__ import annotations + +import json +import textwrap + +import pandas as pd +import pytest +import yaml + +from datacustomcode.client import ( + Client, + StreamingClient, + _BaseClient, +) +from datacustomcode.config import config +from datacustomcode.io.reader.base import BaseDataCloudReader +from datacustomcode.io.writer.base import BaseDataCloudWriter +from datacustomcode.lineage import capture_lineage, generate_lineage +from datacustomcode.lineage.capture import LineageCollector +from datacustomcode.scan import SDK_CONFIG_DIR, SDK_CONFIG_FILE + +pytestmark = pytest.mark.spark + +TABLES = { + "orders__dll": pd.DataFrame( + {"id": [1, 2], "first": ["a", "b"], "last": ["x", "y"], "amount": [1.0, 2.0]} + ), + "customers__dlm": pd.DataFrame({"id": [1, 2], "region": ["us", "eu"]}), +} +VALIDATED: list[tuple[str, list[str]]] = [] + + +class InMemoryReader(BaseDataCloudReader): + CONFIG_NAME = "LineageTestInMemoryReader" + + def __init__(self, spark, dataspace=None): + super().__init__(spark) + + def read_dlo(self, name, schema=None): + return self.spark.createDataFrame(TABLES[name]) + + read_dmo = read_dlo + + +class RefusingWriter(BaseDataCloudWriter): + CONFIG_NAME = "LineageTestRefusingWriter" + + def __init__(self, spark, dataspace=None): + super().__init__(spark) + + def validate_dataframe_columns_against_dlo(self, dataframe, dlo_name): + VALIDATED.append((dlo_name, list(dataframe.columns))) + + def write_to_dlo(self, name, dataframe, write_mode): + raise AssertionError("lineage generation must never write") + + write_to_dmo = write_to_dlo + + +@pytest.fixture(autouse=True) +def isolated_client(spark): + saved = (config.reader_config, config.writer_config, config.spark_config) + saved_spark = _BaseClient._shared_spark + Client._instance = StreamingClient._instance = None + _BaseClient._shared_spark = spark + VALIDATED.clear() + yield + config.reader_config, config.writer_config, config.spark_config = saved + Client._instance = StreamingClient._instance = None + _BaseClient._shared_spark = saved_spark + + +@pytest.fixture +def config_file(tmp_path): + path = tmp_path / "config.yaml" + path.write_text( + yaml.dump( + { + "reader_config": {"type_config_name": InMemoryReader.CONFIG_NAME}, + "writer_config": {"type_config_name": RefusingWriter.CONFIG_NAME}, + } + ) + ) + return str(path) + + +def make_package(tmp_path, body: str, config_json=None, package_type=None) -> str: + payload = tmp_path / "payload" + payload.mkdir() + entrypoint = payload / "entrypoint.py" + entrypoint.write_text( + "from pyspark.sql import functions as F\n" + "from datacustomcode.client import Client\n" + "from datacustomcode.io.writer.base import WriteMode\n" + "client = Client()\n" + textwrap.dedent(body) + ) + (payload / "config.json").write_text( + json.dumps(config_json or {"dataspace": "default"}) + ) + if package_type: + sdk_dir = tmp_path / SDK_CONFIG_DIR + sdk_dir.mkdir() + (sdk_dir / SDK_CONFIG_FILE).write_text(json.dumps({"type": package_type})) + return str(entrypoint) + + +def test_generate_lineage_end_to_end(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + """ + orders = client.read_dlo("orders__dll") + lookup = orders.sparkSession.createDataFrame(orders.select("id").toPandas()) + df = orders.join(lookup, "id").select( + F.col("id").alias("id__c"), + F.concat_ws(" ", "first", "last").alias("name__c"), + F.lit("x").alias("const__c"), + lookup["id"].alias("lookup_id__c"), + ) + client.write_to_dlo("out__dll", df, WriteMode.APPEND) + """, + ) + + result = generate_lineage(entrypoint, config_file) + + assert result.to_dict() == { + "dataLineage": [ + {"targetField": "out__dll.id__c", "sourceFields": ["orders__dll.id"]}, + { + "targetField": "out__dll.name__c", + "sourceFields": ["orders__dll.first", "orders__dll.last"], + }, + ], + "unresolved": [ + {"targetField": "out__dll.lookup_id__c", "reason": "untraceable source"} + ], + "noSources": ["out__dll.const__c"], + "targets": ["out__dll"], + "sources": ["orders__dll"], + "pairCount": 3, + } + assert VALIDATED == [("out__dll", ["id__c", "name__c", "const__c", "lookup_id__c"])] + + +def test_repeat_writes_union_sources(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + """ + orders = client.read_dlo("orders__dll") + client.write_to_dlo( + "out__dll", orders.select(F.col("first").alias("v__c")), WriteMode.APPEND + ) + client.write_to_dlo( + "out__dll", orders.select(F.col("last").alias("v__c")), WriteMode.APPEND + ) + """, + ) + + lineage = generate_lineage(entrypoint, config_file).to_dict()["dataLineage"] + + assert lineage == [ + { + "targetField": "out__dll.v__c", + "sourceFields": ["orders__dll.first", "orders__dll.last"], + } + ] + + +def test_dmo_read_and_write(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + """ + df = client.read_dmo("customers__dlm").select( + F.upper("region").alias("region__c") + ) + client.write_to_dmo("out__dlm", df, WriteMode.APPEND) + """, + ) + + result = generate_lineage(entrypoint, config_file) + + assert result.data_lineage == {"out__dlm.region__c": ["customers__dlm.region"]} + assert VALIDATED == [] + + +def test_rejects_function_packages(tmp_path, config_file): + entrypoint = make_package(tmp_path, "", package_type="function") + + with pytest.raises(ValueError, match="only be generated for script packages"): + generate_lineage(entrypoint, config_file) + + +def test_rejects_streaming_scripts(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + "", + config_json={"dataspace": "default", "streamingSource": {"name": "a__dll"}}, + ) + + with pytest.raises(ValueError, match="does not support streaming"): + generate_lineage(entrypoint, config_file) + + +def test_script_error_propagates_and_restores_state(tmp_path, config_file): + entrypoint = make_package(tmp_path, "raise RuntimeError('boom')\n") + sentinel = object.__new__(Client) + Client._instance = sentinel + + with pytest.raises(RuntimeError, match="boom"): + generate_lineage(entrypoint, config_file) + + assert config.reader_config.type_config_name == InMemoryReader.CONFIG_NAME + assert config.writer_config.type_config_name == RefusingWriter.CONFIG_NAME + assert Client._instance is sentinel + + +def test_existing_client_is_not_reused_inside_capture(spark, config_file): + config.load(config_file) + outside = Client() + + with capture_lineage(LineageCollector()): + inside = Client() + assert inside is not outside + assert type(inside._writer).__name__ == "LineageCaptureWriter" + + assert Client() is outside + + +def test_capture_falls_through_to_delegate_attributes(config_file): + config.load(config_file) + + with capture_lineage(LineageCollector()): + writer = Client()._writer + assert writer.validate_dataframe_columns_against_dlo is not None + assert isinstance(writer._delegate, RefusingWriter) diff --git a/tests/lineage/test_tracer.py b/tests/lineage/test_tracer.py new file mode 100644 index 0000000..dc3dc47 --- /dev/null +++ b/tests/lineage/test_tracer.py @@ -0,0 +1,225 @@ +"""Tests for analyzed-plan column lineage tracing (real local Spark).""" + +from __future__ import annotations + +import time + +import pandas as pd +from pyspark.sql import Window, functions as F +from pyspark.sql.types import StringType +import pytest + +from datacustomcode.lineage.tracer import ( + UNTRACEABLE_SOURCE, + LineageTraceError, + SourceRegistry, + TraceResult, +) + +pytestmark = pytest.mark.spark + + +@pytest.fixture +def registry(): + return SourceRegistry() + + +@pytest.fixture +def orders(spark, registry): + df = spark.createDataFrame( + pd.DataFrame( + { + "id": [1, 2], + "first": ["a", "b"], + "last": ["x", "y"], + "amount": [1.0, 2.0], + "tags": [["t1", "t2"], ["t3"]], + } + ) + ) + registry.register(df, "orders__dll") + return df + + +@pytest.fixture +def customers(spark, registry): + df = spark.createDataFrame(pd.DataFrame({"id": [1, 2], "region": ["us", "eu"]})) + registry.register(df, "customers__dll") + return df + + +def sources(results: dict[str, TraceResult]) -> dict[str, set[str]]: + assert not [r for r in results.values() if r.unresolved], results + return {name: set(r.sources) for name, r in results.items()} + + +def test_passthrough(registry, orders): + assert sources(registry.trace(orders.select("id", "amount"))) == { + "id": {"orders__dll.id"}, + "amount": {"orders__dll.amount"}, + } + + +def test_rename_and_concat(registry, orders): + df = orders.select( + F.col("id").alias("order_id"), + F.concat_ws(" ", "first", "last").alias("full_name"), + ) + assert sources(registry.trace(df)) == { + "order_id": {"orders__dll.id"}, + "full_name": {"orders__dll.first", "orders__dll.last"}, + } + + +def test_filter_and_join_keys_are_not_sources(registry, orders, customers): + df = ( + orders.filter(F.col("amount") > 1) + .join(customers, orders.id == customers.id) + .select(orders["first"], customers["region"]) + ) + assert sources(registry.trace(df)) == { + "first": {"orders__dll.first"}, + "region": {"customers__dll.region"}, + } + + +def test_using_join_key_is_coalesced(registry, orders, customers): + df = orders.join(customers, "id", "full_outer").select("id") + assert sources(registry.trace(df)) == { + "id": {"orders__dll.id", "customers__dll.id"} + } + + +def test_python_udf_inputs(registry, orders): + shout = F.udf(lambda a, b: f"{a}{b}".upper(), StringType()) + df = orders.select(shout("first", "last").alias("loud")) + assert sources(registry.trace(df)) == { + "loud": {"orders__dll.first", "orders__dll.last"} + } + + +def test_aggregate(registry, orders): + df = orders.groupBy("first").agg(F.sum("amount").alias("total")) + assert sources(registry.trace(df)) == { + "first": {"orders__dll.first"}, + "total": {"orders__dll.amount"}, + } + + +def test_literal_has_no_sources(registry, orders): + results = registry.trace(orders.select(F.lit("const").alias("c"))) + assert results["c"] == TraceResult() + + +def test_union_is_positional(registry, orders, customers): + df = orders.select("id", "first").union(customers.select("id", "region")) + assert sources(registry.trace(df)) == { + "id": {"orders__dll.id", "customers__dll.id"}, + "first": {"orders__dll.first", "customers__dll.region"}, + } + + +def test_union_by_name(registry, orders, customers): + df = orders.select("id", F.col("first").alias("region")).unionByName( + customers.select("region", "id") + ) + assert sources(registry.trace(df)) == { + "id": {"orders__dll.id", "customers__dll.id"}, + "region": {"orders__dll.first", "customers__dll.region"}, + } + + +def test_window(registry, orders): + window = Window.partitionBy("first").orderBy("amount") + df = orders.select("id", F.sum("amount").over(window).alias("running")) + assert sources(registry.trace(df)) == { + "id": {"orders__dll.id"}, + "running": {"orders__dll.amount", "orders__dll.first"}, + } + + +def test_explode(registry, orders): + df = orders.select("id", F.explode("tags").alias("tag")) + assert sources(registry.trace(df)) == { + "id": {"orders__dll.id"}, + "tag": {"orders__dll.tags"}, + } + + +def test_rollup(registry, orders): + df = orders.rollup("first").agg(F.sum("amount").alias("total")) + assert sources(registry.trace(df)) == { + "first": {"orders__dll.first"}, + "total": {"orders__dll.amount"}, + } + + +def test_self_join(registry, orders): + left, right = orders.alias("l"), orders.alias("r") + df = left.join(right, F.col("l.id") == F.col("r.id")).select( + F.col("l.first").alias("left_first"), F.col("r.last").alias("right_last") + ) + assert sources(registry.trace(df)) == { + "left_first": {"orders__dll.first"}, + "right_last": {"orders__dll.last"}, + } + + +def test_dropped_then_readded(registry, orders): + df = orders.drop("first").withColumn("first", F.upper("last")) + assert sources(registry.trace(df.select("first"))) == { + "first": {"orders__dll.last"} + } + + +def test_to_pandas_round_trip_is_untraceable(spark, registry, orders): + df = spark.createDataFrame(orders.select("id").toPandas()) + assert registry.trace(df)["id"] == TraceResult(reason=UNTRACEABLE_SOURCE) + + +def test_unregistered_dataframe_is_untraceable(spark, registry, orders): + lookup = spark.createDataFrame(pd.DataFrame({"id": [1], "label": ["one"]})) + df = orders.join(lookup, "id").select("first", "label") + results = registry.trace(df) + assert results["first"].sources == {"orders__dll.first"} + assert results["label"] == TraceResult(reason=UNTRACEABLE_SOURCE) + + +def test_opaque_operator_is_unresolved(registry, orders): + def identity(batches): + yield from batches + + df = orders.select("id").mapInPandas(identity, "id long") + assert registry.trace(df)["id"].unresolved + + +def test_scalar_subquery_is_unresolved(spark, registry, orders): + orders.createOrReplaceTempView("orders_view") + df = spark.sql( + "SELECT id, (SELECT max(amount) FROM orders_view) AS top FROM orders_view" + ) + results = registry.trace(df) + assert results["id"].sources == {"orders__dll.id"} + assert results["top"].unresolved + + +def test_wide_passthrough_is_fast(spark): + columns = [f"field_{i:04d}__c" for i in range(512)] + wide = spark.createDataFrame(pd.DataFrame([list(range(512))], columns=columns)) + registry = SourceRegistry() + registry.register(wide, "wide__dll") + + start = time.monotonic() + results = sources(registry.trace(wide.select(*columns))) + elapsed = time.monotonic() - start + + assert results == {c: {f"wide__dll.{c}"} for c in columns} + assert elapsed < 1.0 + + +def test_spark_connect_dataframe_rejected(registry): + class ConnectLike: + pass + + with pytest.raises(LineageTraceError, match="Spark Connect"): + registry.trace(ConnectLike()) # type: ignore[arg-type] diff --git a/tests/test_run.py b/tests/test_run.py index 154b0bf..c09db5f 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -575,3 +575,32 @@ def test_streaming_source_set_from_streaming_source_field(self): def test_streaming_source_none_for_batch_without_field(self): content = self._run_capturing_streaming_source({"dataspace": "default"}) assert "streaming_source: None" in content + + +def test_prepare_run_resolves_without_executing(tmp_path): + from datacustomcode.run import prepare_run + + entrypoint = tmp_path / "entrypoint.py" + entrypoint.write_text("raise AssertionError('must not run')\n") + (tmp_path / "config.json").write_text(json.dumps({"dataspace": "ds1"})) + + context = prepare_run(str(entrypoint), None, [], "default") + + assert context.package_type == "script" + assert context.config_json == {"dataspace": "ds1"} + assert config.reader_config.options["dataspace"] == "ds1" + + +def test_execute_entrypoint_runs_prepared_script(tmp_path): + from datacustomcode.run import execute_entrypoint, prepare_run + + marker = tmp_path / "ran.txt" + entrypoint = tmp_path / "entrypoint.py" + entrypoint.write_text(f"open({str(marker)!r}, 'w').write(__name__)\n") + (tmp_path / "config.json").write_text(json.dumps({"dataspace": "default"})) + + context = prepare_run(str(entrypoint), None, [], "default") + assert not marker.exists() + execute_entrypoint(str(entrypoint), context) + + assert marker.read_text() == "__main__" From ded2ea21172b5c8faa11b5fa77541b3025aa9458 Mon Sep 17 00:00:00 2001 From: Jesus Orosco Date: Fri, 9 Oct 2026 15:02:11 -0700 Subject: [PATCH 2/6] remediate LLM findings --- src/datacustomcode/lineage/__init__.py | 13 ++- src/datacustomcode/lineage/capture.py | 154 ++++++++++++++----------- src/datacustomcode/lineage/tracer.py | 104 +++++++++++++---- src/datacustomcode/run.py | 20 +++- tests/lineage/test_generate.py | 83 +++++++++++-- tests/lineage/test_tracer.py | 80 +++++++++++++ 6 files changed, 349 insertions(+), 105 deletions(-) diff --git a/src/datacustomcode/lineage/__init__.py b/src/datacustomcode/lineage/__init__.py index 9ca831b..66f80bc 100644 --- a/src/datacustomcode/lineage/__init__.py +++ b/src/datacustomcode/lineage/__init__.py @@ -63,11 +63,15 @@ def generate_lineage( Raises: ValueError: If the package is not a batch script. """ - from datacustomcode.run import execute_entrypoint, prepare_run - - context = prepare_run( - entrypoint, config_file, list(dependencies), profile, sf_cli_org + from datacustomcode.run import ( + execute_entrypoint, + import_dependencies, + prepare_run, ) + + # Dependencies are imported inside the capture, 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 " @@ -78,5 +82,6 @@ def generate_lineage( collector = LineageCollector() with capture_lineage(collector): + import_dependencies(list(dependencies)) execute_entrypoint(entrypoint, context) return collector.result() diff --git a/src/datacustomcode/lineage/capture.py b/src/datacustomcode/lineage/capture.py index 7a884c5..2ee015c 100644 --- a/src/datacustomcode/lineage/capture.py +++ b/src/datacustomcode/lineage/capture.py @@ -14,12 +14,14 @@ # limitations under the License. """Readers and writers that record lineage instead of writing data. -While :func:`capture_lineage` is active, the global reader/writer config is -replaced (with ``force=True``) by :class:`LineageCaptureReader` and -:class:`LineageCaptureWriter`. They wrap the reader/writer the config had -already resolved, so the customer's reader choice and credentials still apply: -reads go through the delegate and are registered as lineage sources, writes are -validated through the delegate (when it supports it), traced, and never written. +While :func:`capture_lineage` is active, the reader and writer of every client +(already existing or created inside the block) are wrapped in place by +:class:`LineageCaptureReader` and :class:`LineageCaptureWriter`. The wrappers +delegate to the reader/writer the client actually resolved, whether from config +or passed to ``Client(reader=..., writer=...)``, so the customer's reader choice +and credentials still apply: reads go through the delegate and are registered as +lineage sources, writes are validated through the delegate (when it supports +it), traced, and never written. """ from __future__ import annotations @@ -36,7 +38,6 @@ Union, ) -from datacustomcode.config import AccessLayerObjectConfig, config from datacustomcode.io.reader.base import BaseDataCloudReader from datacustomcode.io.writer.base import ( MERGE_RECORD_TYPE_COLUMN, @@ -46,9 +47,11 @@ from datacustomcode.lineage.tracer import SourceRegistry, TraceResult if TYPE_CHECKING: - from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession + from pyspark.sql import DataFrame as PySparkDataFrame from pyspark.sql.types import AtomicType, StructType + from datacustomcode.client import _BaseClient + @dataclass class LineageResult: @@ -130,27 +133,6 @@ def result(self) -> LineageResult: return result -@dataclass -class _CaptureSession: - collector: LineageCollector - reader_config: AccessLayerObjectConfig - writer_config: AccessLayerObjectConfig - - -# The capture in progress, read by the capture classes when Client builds them -_active: Dict[str, Optional[_CaptureSession]] = {"session": None} - - -def _active_session() -> _CaptureSession: - session = _active["session"] - if session is None: - raise RuntimeError( - "Lineage capture readers and writers can only be created inside " - "capture_lineage()" - ) - return session - - class _Delegating: """Falls back to the delegate for attributes the capture class lacks.""" @@ -163,53 +145,51 @@ def __getattr__(self, item: str) -> Any: class LineageCaptureReader(_Delegating, BaseDataCloudReader): - """Reads through the configured reader and registers lineage sources.""" + """Reads through a client's reader and registers lineage sources.""" - CONFIG_NAME = "LineageCaptureReader" - - def __init__(self, spark: SparkSession) -> None: - super().__init__(spark) - session = _active_session() - self._collector = session.collector - self._delegate = session.reader_config.to_object(spark) + def __init__( + self, delegate: BaseDataCloudReader, collector: LineageCollector + ) -> None: + super().__init__(getattr(delegate, "spark", None)) # type: ignore[arg-type] + self._delegate = delegate + self._collector = collector def read_dlo( self, name: str, schema: Union[AtomicType, StructType, str, None] = None, ) -> PySparkDataFrame: - dataframe = ( + dataframe: PySparkDataFrame = ( self._delegate.read_dlo(name) if schema is None else self._delegate.read_dlo(name, schema) ) self._collector.record_read(name, dataframe) - return dataframe # type: ignore[no-any-return] + return dataframe def read_dmo( self, name: str, schema: Union[AtomicType, StructType, str, None] = None, ) -> PySparkDataFrame: - dataframe = ( + dataframe: PySparkDataFrame = ( self._delegate.read_dmo(name) if schema is None else self._delegate.read_dmo(name, schema) ) self._collector.record_read(name, dataframe) - return dataframe # type: ignore[no-any-return] + return dataframe class LineageCaptureWriter(_Delegating, BaseDataCloudWriter): """Traces DataFrames handed to the writer; never writes or shows them.""" - CONFIG_NAME = "LineageCaptureWriter" - - def __init__(self, spark: SparkSession) -> None: - super().__init__(spark) - session = _active_session() - self._collector = session.collector - self._delegate = session.writer_config.to_object(spark) + def __init__( + self, delegate: BaseDataCloudWriter, collector: LineageCollector + ) -> None: + super().__init__(getattr(delegate, "spark", None)) # type: ignore[arg-type] + self._delegate = delegate + self._collector = collector def _validate_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: # Same column check `run` does, so lineage never names a missing field @@ -245,36 +225,76 @@ def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: self._collector.record_write(name, dataframe) +@dataclass +class _WrappedClient: + client: _BaseClient + reader: BaseDataCloudReader + writer: BaseDataCloudWriter + + +class _CaptureSession: + def __init__(self, collector: LineageCollector) -> None: + self.collector = collector + self._wrapped: Dict[int, _WrappedClient] = {} + + def wrap(self, client: _BaseClient) -> None: + """Route a client's reads and writes through capture (idempotent).""" + if id(client) in self._wrapped: + return + reader, writer = client._reader, client._writer + self._wrapped[id(client)] = _WrappedClient(client, reader, writer) + client._reader = LineageCaptureReader(reader, self.collector) + client._writer = LineageCaptureWriter(writer, self.collector) + + def restore(self) -> None: + for wrapped in self._wrapped.values(): + wrapped.client._reader = wrapped.reader + wrapped.client._writer = wrapped.writer + self._wrapped.clear() + + +# The capture in progress; guards against nested captures +_active: Dict[str, Optional[_CaptureSession]] = {"session": None} + + @contextmanager def capture_lineage(collector: LineageCollector) -> Iterator[LineageCollector]: - """Route every ``Client`` created inside the block through lineage capture. + """Capture lineage from every ``Client`` used inside the block. - The global reader/writer config is swapped for forced capture configs that - wrap the currently configured reader/writer, and any existing ``Client`` - singleton is set aside so the entrypoint's ``Client()`` picks them up. Both - are restored on exit. + Clients that already exist (including ones a dependency created and kept a + reference to) are wrapped on entry, and ``Client.__new__`` is hooked so + clients created inside the block are wrapped as they are built, however + their reader and writer were supplied. Clients are unwrapped, the hook is + removed, and the ``Client`` singletons are restored on exit. """ - from datacustomcode.client import Client, StreamingClient + from datacustomcode.client import ( + Client, + StreamingClient, + _BaseClient, + ) - if config.reader_config is None or config.writer_config is None: - raise ValueError("A reader and writer must be configured to capture lineage") if _active["session"] is not None: raise RuntimeError("Lineage capture is already active") - saved_reader, saved_writer = config.reader_config, config.writer_config + session = _CaptureSession(collector) saved_clients = Client._instance, StreamingClient._instance - _active["session"] = _CaptureSession(collector, saved_reader, saved_writer) - config.reader_config = AccessLayerObjectConfig( - type_config_name=LineageCaptureReader.CONFIG_NAME, force=True - ) - config.writer_config = AccessLayerObjectConfig( - type_config_name=LineageCaptureWriter.CONFIG_NAME, force=True - ) - Client._instance = None - StreamingClient._instance = None + original_new = _BaseClient.__dict__["__new__"] + build = original_new.__func__ + + def capturing_new(cls: type, *args: Any, **kwargs: Any) -> _BaseClient: + client = build(cls, *args, **kwargs) + session.wrap(client) + return client # type: ignore[no-any-return] + + _active["session"] = session try: + for client in saved_clients: + if client is not None: + session.wrap(client) + _BaseClient.__new__ = staticmethod(capturing_new) # type: ignore[assignment] yield collector finally: - config.reader_config, config.writer_config = saved_reader, saved_writer + _BaseClient.__new__ = original_new # type: ignore[method-assign] + session.restore() Client._instance, StreamingClient._instance = saved_clients _active["session"] = None diff --git a/src/datacustomcode/lineage/tracer.py b/src/datacustomcode/lineage/tracer.py index 6a9e7ea..6eea9c4 100644 --- a/src/datacustomcode/lineage/tracer.py +++ b/src/datacustomcode/lineage/tracer.py @@ -39,6 +39,7 @@ Iterator, List, Optional, + Set, Tuple, ) @@ -49,6 +50,7 @@ """``(leaf relation identity, column position)``.""" UNTRACEABLE_SOURCE = "untraceable source" +AMBIGUOUS_SOURCE = "ambiguous source" _OUTPUT_SEPARATOR = "\u0001" @@ -72,19 +74,24 @@ class ColumnTrace: leaves: Leaf columns the value is computed from. Empty for literals. opaque_reason: Set when the plan hides this column's inputs (e.g. a ``mapInPandas`` or subquery); ``leaves`` is then incomplete. + direct: The column is its single leaf column passed through unchanged + (possibly renamed), not computed from it. """ name: str leaves: FrozenSet[LeafColumn] = field(default_factory=frozenset) opaque_reason: Optional[str] = None + direct: bool = False @dataclass(frozen=True) class _AttrTrace: leaves: FrozenSet[LeafColumn] opaque_reason: Optional[str] = None + direct: bool = False def merge(self, other: _AttrTrace) -> _AttrTrace: + # A merge combines inputs, so the result is never a direct pass-through return _AttrTrace( self.leaves | other.leaves, self.opaque_reason or other.opaque_reason ) @@ -112,17 +119,28 @@ def unresolved(self) -> bool: class SourceRegistry: - """Maps leaf columns of DataFrames read from Data Cloud to their fields.""" + """Maps leaf columns of DataFrames read from Data Cloud to their fields. + + Only direct columns name a field: a column a reader computes (e.g. + ``upper(name).alias("normalized")``) does not claim the leaf it reads. A leaf + claimed by two different fields (two direct columns over it, or two objects + read from the same relation) is ambiguous, and so is anything traced to it. + """ def __init__(self) -> None: self._fields: Dict[LeafColumn, str] = {} + self._ambiguous: Set[LeafColumn] = set() self.objects: List[str] = [] def register(self, df: PySparkDataFrame, object_name: str) -> None: - """Record that the columns of ``df`` are the fields of ``object_name``.""" + """Record that the direct columns of ``df`` are fields of ``object_name``.""" for column in trace_columns(df): - for leaf in column.leaves: - self._fields[leaf] = f"{object_name}.{column.name}" + if not column.direct: + continue + (leaf,) = column.leaves + field_name = f"{object_name}.{column.name}" + if self._fields.setdefault(leaf, field_name) != field_name: + self._ambiguous.add(leaf) if object_name not in self.objects: self.objects.append(object_name) @@ -132,6 +150,8 @@ def trace(self, df: PySparkDataFrame) -> Dict[str, TraceResult]: for column in trace_columns(df): if column.opaque_reason is not None: results[column.name] = TraceResult(reason=column.opaque_reason) + elif any(leaf in self._ambiguous for leaf in column.leaves): + results[column.name] = TraceResult(reason=AMBIGUOUS_SOURCE) elif any(leaf not in self._fields for leaf in column.leaves): results[column.name] = TraceResult(reason=UNTRACEABLE_SOURCE) else: @@ -190,14 +210,13 @@ def trace_columns(df: PySparkDataFrame) -> List[ColumnTrace]: f"Plan output has {len(root_ids)} columns but the DataFrame has " f"{len(df.columns)}" ) - return [ - ColumnTrace( - name=name, - leaves=traces.get(expr_id, _EMPTY).leaves, - opaque_reason=traces.get(expr_id, _EMPTY).opaque_reason, + columns = [] + for name, expr_id in zip(df.columns, root_ids, strict=False): + trace = traces.get(expr_id, _EMPTY) + columns.append( + ColumnTrace(name, trace.leaves, trace.opaque_reason, trace.direct) ) - for name, expr_id in zip(df.columns, root_ids, strict=False) - ] + return columns def _trace_node( @@ -213,7 +232,7 @@ def _trace_node( if not child_outputs: key = _leaf_key(jnode, node_class, jvm) for position, expr_id in enumerate(output): - traces[expr_id] = _AttrTrace(frozenset({(key, position)})) + traces[expr_id] = _AttrTrace(frozenset({(key, position)}), direct=True) return if node_class == "Union": @@ -234,10 +253,14 @@ def _trace_node( ) traces[expr_id] = merged elif node_class == "Generate": - generator = _trace_expression(node.get("generator", []), traces) - for attr_list in node.get("generatorOutput", []): - for attr in attr_list: - traces[_expr_id(attr)] = generator + generated = [ + _expr_id(attr) + for attr_list in node.get("generatorOutput", []) + for attr in attr_list + ] + per_output = _trace_generator(node.get("generator", []), len(generated), traces) + for expr_id, trace in zip(generated, per_output, strict=True): + traces[expr_id] = trace else: for expression in _expression_trees(node): _trace_aliases(expression, traces) @@ -253,10 +276,51 @@ def _trace_aliases(expression: List[dict], traces: Dict[int, _AttrTrace]) -> Non """Record every Alias in a flattened expression tree as a new attribute.""" for index, expr_node in enumerate(expression): if expr_node["class"].endswith(".Alias"): - end = _subtree_end(expression, index) - traces[_expr_id(expr_node)] = _trace_expression( - expression[index + 1 : end], traces - ) + child = expression[index + 1 : _subtree_end(expression, index)] + if ( + len(child) == 1 + and child[0]["class"].endswith(".AttributeReference") + and _expr_id(child[0]) in traces + ): + # A rename keeps the column a direct pass-through + traces[_expr_id(expr_node)] = traces[_expr_id(child[0])] + else: + traces[_expr_id(expr_node)] = _trace_expression(child, traces) + + +def _trace_generator( + generator: List[dict], output_count: int, traces: Dict[int, _AttrTrace] +) -> List[_AttrTrace]: + """Trace each output of a Generate node's generator expression.""" + whole = _trace_expression(generator, traces) + if output_count <= 1 or whole.opaque_reason is not None: + return [whole] * output_count + generator_class = generator[0]["class"].rsplit(".", 1)[-1] + if generator_class == "Stack": + # stack(n, e0, e1, ...) puts e_j in output j % output_count + per_output = [_EMPTY] * output_count + for position, arg in enumerate(_child_subtrees(generator)[1:]): + slot = position % output_count + per_output[slot] = per_output[slot].merge(_trace_expression(arg, traces)) + return per_output + if len(whole.leaves) > 1: + # e.g. inline(array(struct(a, b))): which output reads which input is + # not recoverable without modelling the generator + reason = f"columns produced by {generator_class} cannot be traced" + return [_AttrTrace(frozenset(), reason)] * output_count + # One input feeds every output (explode, posexplode, json_tuple, ...) + return [whole] * output_count + + +def _child_subtrees(expression: List[dict]) -> List[List[dict]]: + """The flattened subtrees of the root's children, in order.""" + children = [] + index = 1 + for _ in range(expression[0]["num-children"]): + end = _subtree_end(expression, index) + children.append(expression[index:end]) + index = end + return children def _trace_expression( diff --git a/src/datacustomcode/run.py b/src/datacustomcode/run.py index 16a8ecc..0552adf 100644 --- a/src/datacustomcode/run.py +++ b/src/datacustomcode/run.py @@ -144,10 +144,10 @@ def prepare_run( ) -> RunContext: """Resolve config and import dependencies for an entrypoint, without running it. - This is the setup half of :func:`run_entrypoint`. Callers that need to adjust - the resolved global ``config`` before the entrypoint runs (e.g. lineage - generation swapping in capturing readers/writers) call this, make their - changes, then call :func:`execute_entrypoint`. + This is the setup half of :func:`run_entrypoint`. Callers that need to act + between setup and running (e.g. lineage generation starting its capture) call + this, then :func:`execute_entrypoint`. Such callers can pass no dependencies + here and call :func:`import_dependencies` themselves at the right moment. Args: entrypoint: The entrypoint script to run. @@ -207,7 +207,17 @@ def prepare_run( config.streaming_source = _read_streaming_source(config_json) _update_config_options(profile, sf_cli_org) + import_dependencies(dependencies) + return RunContext(package_type=package_type, config_json=config_json) + + +def import_dependencies(dependencies: List[str]) -> None: + """Import each dependency module (or ``module.attribute``), as ``run`` does. + + Args: + dependencies: The dependencies to import. + """ for dependency in dependencies: try: importlib.import_module(dependency) @@ -222,8 +232,6 @@ def prepare_run( except (ModuleNotFoundError, AttributeError) as inner_exc: raise inner_exc from exc - return RunContext(package_type=package_type, config_json=config_json) - def execute_entrypoint( entrypoint: str, context: RunContext, test_file: Optional[str] = None diff --git a/tests/lineage/test_generate.py b/tests/lineage/test_generate.py index 46665ad..2a34e19 100644 --- a/tests/lineage/test_generate.py +++ b/tests/lineage/test_generate.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import sys import textwrap import pandas as pd @@ -16,13 +17,15 @@ ) from datacustomcode.config import config from datacustomcode.io.reader.base import BaseDataCloudReader -from datacustomcode.io.writer.base import BaseDataCloudWriter +from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode from datacustomcode.lineage import capture_lineage, generate_lineage from datacustomcode.lineage.capture import LineageCollector from datacustomcode.scan import SDK_CONFIG_DIR, SDK_CONFIG_FILE pytestmark = pytest.mark.spark +ORIGINAL_NEW = _BaseClient.__dict__["__new__"] + TABLES = { "orders__dll": pd.DataFrame( {"id": [1, 2], "first": ["a", "b"], "last": ["x", "y"], "amount": [1.0, 2.0]} @@ -203,30 +206,94 @@ def test_rejects_streaming_scripts(tmp_path, config_file): def test_script_error_propagates_and_restores_state(tmp_path, config_file): + config.load(config_file) + outside = Client() + reader, writer = outside._reader, outside._writer entrypoint = make_package(tmp_path, "raise RuntimeError('boom')\n") - sentinel = object.__new__(Client) - Client._instance = sentinel with pytest.raises(RuntimeError, match="boom"): generate_lineage(entrypoint, config_file) assert config.reader_config.type_config_name == InMemoryReader.CONFIG_NAME assert config.writer_config.type_config_name == RefusingWriter.CONFIG_NAME - assert Client._instance is sentinel + assert Client._instance is outside + assert (outside._reader, outside._writer) == (reader, writer) + assert "__new__" not in vars(Client) + assert _BaseClient.__dict__["__new__"] is ORIGINAL_NEW -def test_existing_client_is_not_reused_inside_capture(spark, config_file): +def test_existing_client_is_captured_and_restored(spark, config_file): config.load(config_file) outside = Client() + writer = outside._writer with capture_lineage(LineageCollector()): - inside = Client() - assert inside is not outside - assert type(inside._writer).__name__ == "LineageCaptureWriter" + assert Client() is outside + assert type(outside._writer).__name__ == "LineageCaptureWriter" + assert outside._writer._delegate is writer + assert outside._writer is writer assert Client() is outside +def test_retained_client_from_dependency_is_captured( + tmp_path, config_file, monkeypatch +): + # A dependency that builds its Client at import time keeps it in a module + # global; its writes must still be captured, never reach the real writer + (tmp_path / "lineage_dep_helpers.py").write_text( + "from datacustomcode.client import Client\n" + "from datacustomcode.io.writer.base import WriteMode\n" + "client = Client()\n" + "def save(df):\n" + " client.write_to_dlo('out__dll', df, WriteMode.APPEND)\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + monkeypatch.delitem(sys.modules, "lineage_dep_helpers", raising=False) + entrypoint = make_package( + tmp_path, + """ + import lineage_dep_helpers + lineage_dep_helpers.save( + client.read_dlo("orders__dll").select(F.col("id").alias("id__c")) + ) + """, + ) + + result = generate_lineage( + entrypoint, config_file, dependencies=["lineage_dep_helpers"] + ) + + assert result.data_lineage == {"out__dll.id__c": ["orders__dll.id"]} + assert VALIDATED == [("out__dll", ["id__c"])] + + +def test_client_built_with_reader_and_writer_is_captured(spark, config_file): + config.load(config_file) + reader, writer = InMemoryReader(spark), RefusingWriter(spark) + collector = LineageCollector() + + with capture_lineage(collector): + client = Client(reader=reader, writer=writer) + assert client._reader._delegate is reader + assert client._writer._delegate is writer + client.write_to_dlo( + "out__dll", + client.read_dlo("orders__dll").select("id"), + WriteMode.APPEND, + ) + + assert collector.result().data_lineage == {"out__dll.id": ["orders__dll.id"]} + assert (client._reader, client._writer) == (reader, writer) + + +def test_nested_capture_is_rejected(config_file): + with capture_lineage(LineageCollector()): + with pytest.raises(RuntimeError, match="already active"): + with capture_lineage(LineageCollector()): + pass + + def test_capture_falls_through_to_delegate_attributes(config_file): config.load(config_file) diff --git a/tests/lineage/test_tracer.py b/tests/lineage/test_tracer.py index dc3dc47..c0cc646 100644 --- a/tests/lineage/test_tracer.py +++ b/tests/lineage/test_tracer.py @@ -10,6 +10,7 @@ import pytest from datacustomcode.lineage.tracer import ( + AMBIGUOUS_SOURCE, UNTRACEABLE_SOURCE, LineageTraceError, SourceRegistry, @@ -223,3 +224,82 @@ class ConnectLike: with pytest.raises(LineageTraceError, match="Spark Connect"): registry.trace(ConnectLike()) # type: ignore[arg-type] + + +def test_derived_reader_column_does_not_rename_its_source(spark, registry): + # A reader that exposes both a column and a value computed from it: the + # computed column must not take over the source column's field identity + base = spark.createDataFrame(pd.DataFrame({"name": ["a", "b"]})) + source = base.select("name", F.upper("name").alias("normalized")) + registry.register(source, "source__dll") + + results = registry.trace( + source.select( + F.col("name").alias("name__c"), F.col("normalized").alias("norm__c") + ) + ) + + assert sources(results) == { + "name__c": {"source__dll.name"}, + "norm__c": {"source__dll.name"}, + } + + +def test_renamed_reader_column_names_the_field(spark, registry): + base = spark.createDataFrame(pd.DataFrame({"raw": ["a"]})) + registry.register(base.select(F.col("raw").alias("name")), "source__dll") + + traced = registry.trace(base.select(F.col("raw").alias("name__c"))) + + assert sources(traced) == {"name__c": {"source__dll.name"}} + + +def test_two_objects_over_one_relation_are_ambiguous(spark, registry): + base = spark.createDataFrame(pd.DataFrame({"id": [1], "other": [2]})) + registry.register(base, "first__dll") + registry.register(base, "second__dll") + + results = registry.trace(base.select("id")) + + assert results["id"].reason == AMBIGUOUS_SOURCE + assert results["id"].unresolved + + +def test_one_leaf_under_two_names_is_ambiguous(spark, registry): + base = spark.createDataFrame(pd.DataFrame({"id": [1]})) + registry.register(base.select("id", F.col("id").alias("id_copy")), "src__dll") + + assert registry.trace(base.select("id"))["id"].reason == AMBIGUOUS_SOURCE + + +def test_stack_outputs_trace_positionally(registry, orders): + stacked = orders.selectExpr("stack(1, first, last) AS (x, y)") + assert sources(registry.trace(stacked)) == { + "x": {"orders__dll.first"}, + "y": {"orders__dll.last"}, + } + + +def test_stack_with_several_rows(registry, orders): + stacked = orders.selectExpr( + "stack(2, first, amount, last, CAST(id AS DOUBLE)) AS (name, n)" + ) + assert sources(registry.trace(stacked)) == { + "name": {"orders__dll.first", "orders__dll.last"}, + "n": {"orders__dll.id", "orders__dll.amount"}, + } + + +def test_multi_output_generator_over_several_inputs_is_unresolved(registry, orders): + inlined = orders.selectExpr("inline(array(struct(first, last))) AS (x, y)") + results = registry.trace(inlined) + assert results["x"].unresolved and results["y"].unresolved + assert "Inline" in results["x"].reason + + +def test_posexplode_outputs_share_their_single_input(registry, orders): + exploded = orders.select(F.posexplode("tags").alias("pos", "tag")) + assert sources(registry.trace(exploded)) == { + "pos": {"orders__dll.tags"}, + "tag": {"orders__dll.tags"}, + } From 27dbe95690f047ac3e0b9efe74c74a39bed6236c Mon Sep 17 00:00:00 2001 From: Jesus Orosco Date: Fri, 9 Oct 2026 16:49:04 -0700 Subject: [PATCH 3/6] more bugfixes --- src/datacustomcode/client.py | 6 ++ src/datacustomcode/lineage/capture.py | 7 ++- src/datacustomcode/lineage/tracer.py | 91 ++++++++++++++++++++++++--- tests/lineage/test_generate.py | 32 ++++++++++ tests/lineage/test_tracer.py | 71 +++++++++++++++++++-- 5 files changed, 189 insertions(+), 18 deletions(-) diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 814ac5a..08a2419 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -26,6 +26,7 @@ Union, cast, ) +import weakref from datacustomcode.config import config from datacustomcode.einstein_predictions_config import spark_einstein_predictions_config @@ -349,6 +350,10 @@ class _BaseClient: # — and therefore one underlying connection — instead of opening two # containing differing state _shared_spark: ClassVar[Optional[SparkSession]] = None + # Every client built in this process, including ones no longer the + # singleton that a module still holds a reference to (lineage capture must + # reach those too) + _live_instances: ClassVar[weakref.WeakSet[_BaseClient]] = weakref.WeakSet() _reader: BaseDataCloudReader _writer: BaseDataCloudWriter _file: DefaultFindFilePath @@ -429,6 +434,7 @@ def __new__( DataCloudObjectType.DMO: set(), } cls._instance = instance + _BaseClient._live_instances.add(instance) elif reader is not None or writer is not None: raise ValueError("Cannot set reader or writer after client is initialized") return cast(_ClientT, cls._instance) diff --git a/src/datacustomcode/lineage/capture.py b/src/datacustomcode/lineage/capture.py index 2ee015c..9747b89 100644 --- a/src/datacustomcode/lineage/capture.py +++ b/src/datacustomcode/lineage/capture.py @@ -261,8 +261,9 @@ def restore(self) -> None: def capture_lineage(collector: LineageCollector) -> Iterator[LineageCollector]: """Capture lineage from every ``Client`` used inside the block. - Clients that already exist (including ones a dependency created and kept a - reference to) are wrapped on entry, and ``Client.__new__`` is hooked so + Every live client is wrapped on entry, including ones a module kept a + reference to after they stopped being the singleton (e.g. a dependency + imported by an earlier run), and ``Client.__new__`` is hooked so clients created inside the block are wrapped as they are built, however their reader and writer were supplied. Clients are unwrapped, the hook is removed, and the ``Client`` singletons are restored on exit. @@ -288,7 +289,7 @@ def capturing_new(cls: type, *args: Any, **kwargs: Any) -> _BaseClient: _active["session"] = session try: - for client in saved_clients: + for client in (*saved_clients, *_BaseClient._live_instances): if client is not None: session.wrap(client) _BaseClient.__new__ = staticmethod(capturing_new) # type: ignore[assignment] diff --git a/src/datacustomcode/lineage/tracer.py b/src/datacustomcode/lineage/tracer.py index 6eea9c4..7b43a40 100644 --- a/src/datacustomcode/lineage/tracer.py +++ b/src/datacustomcode/lineage/tracer.py @@ -291,25 +291,96 @@ def _trace_aliases(expression: List[dict], traces: Dict[int, _AttrTrace]) -> Non def _trace_generator( generator: List[dict], output_count: int, traces: Dict[int, _AttrTrace] ) -> List[_AttrTrace]: - """Trace each output of a Generate node's generator expression.""" + """Trace each output of a Generate node's generator expression. + + Outputs are traced individually where the generator's shape says which + input feeds which output. Otherwise they share the input's trace only when + that input is a source column passed through unchanged: a computed input + (e.g. ``map('key', name)``) can mix literals and columns, and the shared + trace would attribute the literal outputs to the columns. + """ whole = _trace_expression(generator, traces) - if output_count <= 1 or whole.opaque_reason is not None: + if output_count <= 1 or whole.opaque_reason is not None or not whole.leaves: return [whole] * output_count - generator_class = generator[0]["class"].rsplit(".", 1)[-1] + generator_class = _class_name(generator[0]) + children = _child_subtrees(generator) + per_output: Optional[List[_AttrTrace]] = None if generator_class == "Stack": # stack(n, e0, e1, ...) puts e_j in output j % output_count per_output = [_EMPTY] * output_count - for position, arg in enumerate(_child_subtrees(generator)[1:]): + for position, arg in enumerate(children[1:]): slot = position % output_count per_output[slot] = per_output[slot].merge(_trace_expression(arg, traces)) - return per_output - if len(whole.leaves) > 1: - # e.g. inline(array(struct(a, b))): which output reads which input is - # not recoverable without modelling the generator + elif generator_class == "JsonTuple": + # Every output is a value extracted from the one JSON input + per_output = [whole] * output_count + elif generator_class in ("Explode", "PosExplode", "Inline"): + with_position = generator_class == "PosExplode" + fields = _trace_collection( + children[0], + output_count - with_position, + generator_class == "Inline", + traces, + ) + if fields is not None: + per_output = [whole] * with_position + fields + if per_output is None: reason = f"columns produced by {generator_class} cannot be traced" return [_AttrTrace(frozenset(), reason)] * output_count - # One input feeds every output (explode, posexplode, json_tuple, ...) - return [whole] * output_count + return per_output + + +def _trace_collection( + collection: List[dict], + field_count: int, + of_structs: bool, + traces: Dict[int, _AttrTrace], +) -> Optional[List[_AttrTrace]]: + """Trace the fields a generator unpacks from each element of a collection. + + The fields are a map's key and value, or a struct's fields (``of_structs``). + Returns ``None`` when it is not known which input feeds which field. + """ + if _is_direct(collection, traces): + return [_trace_expression(collection, traces)] * field_count + root = _class_name(collection[0]) + if root == "CreateMap" and not of_structs and field_count == 2: + # map(k0, v0, k1, v1, ...) + fields = [_EMPTY, _EMPTY] + for position, entry in enumerate(_child_subtrees(collection)): + fields[position % 2] = fields[position % 2].merge( + _trace_expression(entry, traces) + ) + return fields + if root == "CreateArray" and of_structs: + fields = [_EMPTY] * field_count + for element in _child_subtrees(collection): + if _is_direct(element, traces): + element_fields = [_trace_expression(element, traces)] * field_count + elif _class_name(element[0]) == "CreateNamedStruct": + # named_struct(name0, value0, name1, value1, ...) + element_fields = [ + _trace_expression(value, traces) + for value in _child_subtrees(element)[1::2] + ] + else: + return None + fields = [f.merge(e) for f, e in zip(fields, element_fields, strict=True)] + return fields + return None + + +def _is_direct(expression: List[dict], traces: Dict[int, _AttrTrace]) -> bool: + """Whether a flattened expression is a source column passed through as-is.""" + return ( + len(expression) == 1 + and _class_name(expression[0]) == "AttributeReference" + and traces.get(_expr_id(expression[0]), _EMPTY).direct + ) + + +def _class_name(expr_node: dict) -> str: + return str(expr_node["class"]).rsplit(".", 1)[-1] def _child_subtrees(expression: List[dict]) -> List[List[dict]]: diff --git a/tests/lineage/test_generate.py b/tests/lineage/test_generate.py index 2a34e19..672bc90 100644 --- a/tests/lineage/test_generate.py +++ b/tests/lineage/test_generate.py @@ -268,6 +268,38 @@ def test_retained_client_from_dependency_is_captured( assert VALIDATED == [("out__dll", ["id__c"])] +def test_repeated_generation_captures_retained_client( + tmp_path, config_file, monkeypatch +): + # The helper module stays imported after the first run, keeping the client + # it built then; the second run must capture that client too + (tmp_path / "lineage_cached_helpers.py").write_text( + "from datacustomcode.client import Client\n" + "from datacustomcode.io.writer.base import WriteMode\n" + "client = Client()\n" + "def save(df):\n" + " client.write_to_dlo('out__dll', df, WriteMode.APPEND)\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + monkeypatch.delitem(sys.modules, "lineage_cached_helpers", raising=False) + entrypoint = make_package( + tmp_path, + """ + import lineage_cached_helpers + lineage_cached_helpers.save( + client.read_dlo("orders__dll").select(F.col("id").alias("id__c")) + ) + """, + ) + + for _ in range(2): + result = generate_lineage( + entrypoint, config_file, dependencies=["lineage_cached_helpers"] + ) + assert result.data_lineage == {"out__dll.id__c": ["orders__dll.id"]} + assert sys.modules["lineage_cached_helpers"].client is not Client._instance + + def test_client_built_with_reader_and_writer_is_captured(spark, config_file): config.load(config_file) reader, writer = InMemoryReader(spark), RefusingWriter(spark) diff --git a/tests/lineage/test_tracer.py b/tests/lineage/test_tracer.py index c0cc646..7099074 100644 --- a/tests/lineage/test_tracer.py +++ b/tests/lineage/test_tracer.py @@ -290,11 +290,72 @@ def test_stack_with_several_rows(registry, orders): } -def test_multi_output_generator_over_several_inputs_is_unresolved(registry, orders): - inlined = orders.selectExpr("inline(array(struct(first, last))) AS (x, y)") - results = registry.trace(inlined) - assert results["x"].unresolved and results["y"].unresolved - assert "Inline" in results["x"].reason +def test_inline_over_built_structs_traces_each_field(registry, orders): + inlined = orders.selectExpr( + "inline(array(named_struct('a', first, 'b', last), " + "named_struct('a', last, 'b', 'constant'))) AS (x, y)" + ) + assert sources(registry.trace(inlined)) == { + "x": {"orders__dll.first", "orders__dll.last"}, + "y": {"orders__dll.last"}, + } + + +def test_inline_literal_field_has_no_sources(spark, registry): + df = spark.createDataFrame(pd.DataFrame({"name": ["a"]})) + registry.register(df, "src__dll") + inlined = df.selectExpr("inline(array(struct(name, 'constant'))) AS (x, y)") + assert sources(registry.trace(inlined)) == {"x": {"src__dll.name"}, "y": set()} + + +def test_explode_of_built_map_traces_key_and_value(spark, registry): + df = spark.createDataFrame(pd.DataFrame({"name": ["a"]})) + registry.register(df, "src__dll") + assert sources(registry.trace(df.selectExpr("explode(map('key', name))"))) == { + "key": set(), + "value": {"src__dll.name"}, + } + exploded = df.selectExpr("posexplode(map('key', name)) AS (p, k, v)") + assert sources(registry.trace(exploded)) == { + "p": {"src__dll.name"}, + "k": set(), + "v": {"src__dll.name"}, + } + + +def test_inline_over_source_struct_column_shares_its_trace(spark, registry): + df = spark.createDataFrame([([("a", 1)],)], "items array>") + registry.register(df, "src__dll") + assert sources(registry.trace(df.selectExpr("inline(items)"))) == { + "s": {"src__dll.items"}, + "n": {"src__dll.items"}, + } + + +@pytest.mark.parametrize( + "build, generate", + [ + # The struct is built in an earlier projection, hidden behind a column + ( + "array(struct(name, 'constant')) AS items", + "inline(items) AS (x, y)", + ), + # The struct is built inline but wrapped in another function + ( + "name", + "inline(filter(array(struct(name, 'constant')), s -> true)) AS (x, y)", + ), + ("map('key', name) AS m", "explode(m) AS (k, v)"), + ], +) +def test_generator_over_computed_collection_is_unresolved( + spark, registry, build, generate +): + df = spark.createDataFrame(pd.DataFrame({"name": ["a"]})) + registry.register(df, "src__dll") + results = registry.trace(df.selectExpr(build).selectExpr(generate)) + assert all(r.unresolved for r in results.values()), results + assert "cannot be traced" in next(iter(results.values())).reason def test_posexplode_outputs_share_their_single_input(registry, orders): From eadd5db6a0757ed693b36d577c63aa4ee50bed53 Mon Sep 17 00:00:00 2001 From: Jesus Orosco Date: Mon, 12 Oct 2026 12:30:57 -0700 Subject: [PATCH 4/6] addressing PR comments --- README.md | 2 +- src/datacustomcode/client.py | 22 +-- src/datacustomcode/lineage/__init__.py | 98 +++++++--- src/datacustomcode/lineage/_worker.py | 97 ++++++++++ src/datacustomcode/lineage/capture.py | 160 +++++++--------- tests/lineage/lineage_test_io.py | 52 ++++++ tests/lineage/test_generate.py | 245 +++++++++++-------------- tests/test_client.py | 38 ++++ 8 files changed, 449 insertions(+), 265 deletions(-) create mode 100644 src/datacustomcode/lineage/_worker.py create mode 100644 tests/lineage/lineage_test_io.py diff --git a/README.md b/README.md index 56dcacc..169fb28 100644 --- a/README.md +++ b/README.md @@ -237,7 +237,7 @@ print(result.to_dict()["dataLineage"]) # "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. Reads use the reader's `default_row_limit` (1000 by default; set it in your `config.yaml`). +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`). - A target field's sources are the fields its *value* is computed from. Filter conditions and join keys aren't sources. - Fields computed only from literals are listed in `noSources`. diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 08a2419..49d1407 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -26,7 +26,6 @@ Union, cast, ) -import weakref from datacustomcode.config import config from datacustomcode.einstein_predictions_config import spark_einstein_predictions_config @@ -350,10 +349,6 @@ class _BaseClient: # — and therefore one underlying connection — instead of opening two # containing differing state _shared_spark: ClassVar[Optional[SparkSession]] = None - # Every client built in this process, including ones no longer the - # singleton that a module still holds a reference to (lineage capture must - # reach those too) - _live_instances: ClassVar[weakref.WeakSet[_BaseClient]] = weakref.WeakSet() _reader: BaseDataCloudReader _writer: BaseDataCloudWriter _file: DefaultFindFilePath @@ -381,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. @@ -409,9 +410,7 @@ 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 @@ -419,9 +418,7 @@ def __new__( 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 @@ -434,7 +431,6 @@ def __new__( DataCloudObjectType.DMO: set(), } cls._instance = instance - _BaseClient._live_instances.add(instance) elif reader is not None or writer is not None: raise ValueError("Cannot set reader or writer after client is initialized") return cast(_ClientT, cls._instance) diff --git a/src/datacustomcode/lineage/__init__.py b/src/datacustomcode/lineage/__init__.py index 66f80bc..cf96819 100644 --- a/src/datacustomcode/lineage/__init__.py +++ b/src/datacustomcode/lineage/__init__.py @@ -16,14 +16,26 @@ from __future__ import annotations -from typing import Optional, Sequence +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, - capture_lineage, + install_capture, ) from datacustomcode.lineage.tracer import LineageTraceError @@ -33,8 +45,8 @@ "LineageCollector", "LineageResult", "LineageTraceError", - "capture_lineage", "generate_lineage", + "install_capture", ] @@ -48,7 +60,10 @@ def generate_lineage( """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. + 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. @@ -62,26 +77,61 @@ def generate_lineage( 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``. """ - from datacustomcode.run import ( - execute_entrypoint, - import_dependencies, - prepare_run, - ) - - # Dependencies are imported inside the capture, 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" + 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 context.config_json.get("streamingSource"): - raise ValueError("Lineage generation does not support streaming scripts") + 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"]) + - collector = LineageCollector() - with capture_lineage(collector): - import_dependencies(list(dependencies)) - execute_entrypoint(entrypoint, context) - return collector.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}") diff --git a/src/datacustomcode/lineage/_worker.py b/src/datacustomcode/lineage/_worker.py new file mode 100644 index 0000000..62d15f7 --- /dev/null +++ b/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 ``. 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:])) diff --git a/src/datacustomcode/lineage/capture.py b/src/datacustomcode/lineage/capture.py index 9747b89..94c18cd 100644 --- a/src/datacustomcode/lineage/capture.py +++ b/src/datacustomcode/lineage/capture.py @@ -14,27 +14,27 @@ # limitations under the License. """Readers and writers that record lineage instead of writing data. -While :func:`capture_lineage` is active, the reader and writer of every client -(already existing or created inside the block) are wrapped in place by -:class:`LineageCaptureReader` and :class:`LineageCaptureWriter`. The wrappers -delegate to the reader/writer the client actually resolved, whether from config -or passed to ``Client(reader=..., writer=...)``, so the customer's reader choice -and credentials still apply: reads go through the delegate and are registered as -lineage sources, writes are validated through the delegate (when it supports -it), traced, and never written. +:func:`install_capture` replaces the configured reader and writer with forced +configs for :class:`LineageCaptureReader` and :class:`LineageCaptureWriter`, so +every client built afterwards uses them, even one passed its own +``Client(reader=..., writer=...)``. The wrappers delegate to the reader/writer +the config resolved, so the customer's reader choice and credentials still +apply: reads go through the delegate and are registered as lineage sources, +writes are validated through the delegate (when it supports it), traced, and +never written. + +Forcing config only reaches clients built after it, so capture runs in a fresh +process (see :mod:`datacustomcode.lineage._worker`), where no client exists yet. """ from __future__ import annotations -from contextlib import contextmanager from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, Any, Dict, - Iterator, List, - Optional, Union, ) @@ -47,10 +47,10 @@ from datacustomcode.lineage.tracer import SourceRegistry, TraceResult if TYPE_CHECKING: - from pyspark.sql import DataFrame as PySparkDataFrame + from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession from pyspark.sql.types import AtomicType, StructType - from datacustomcode.client import _BaseClient + from datacustomcode.config import AccessLayerObjectConfig @dataclass @@ -93,6 +93,22 @@ def to_dict(self) -> Dict[str, Any]: "pairCount": self.pair_count, } + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> LineageResult: + """Inverse of :meth:`to_dict`.""" + return cls( + data_lineage={ + entry["targetField"]: list(entry["sourceFields"]) + for entry in data["dataLineage"] + }, + unresolved={ + entry["targetField"]: entry["reason"] for entry in data["unresolved"] + }, + no_sources=list(data["noSources"]), + targets=list(data["targets"]), + sources=list(data["sources"]), + ) + class LineageCollector: """Accumulates sources read and target fields written during a run.""" @@ -145,13 +161,18 @@ def __getattr__(self, item: str) -> Any: class LineageCaptureReader(_Delegating, BaseDataCloudReader): - """Reads through a client's reader and registers lineage sources.""" + """Reads through the configured reader and registers lineage sources.""" + + CONFIG_NAME = "LineageCaptureReader" def __init__( - self, delegate: BaseDataCloudReader, collector: LineageCollector + self, + spark: SparkSession, + delegate_config: AccessLayerObjectConfig[BaseDataCloudReader], + collector: LineageCollector, ) -> None: - super().__init__(getattr(delegate, "spark", None)) # type: ignore[arg-type] - self._delegate = delegate + super().__init__(spark) + self._delegate = delegate_config.to_object(spark) self._collector = collector def read_dlo( @@ -184,11 +205,16 @@ def read_dmo( class LineageCaptureWriter(_Delegating, BaseDataCloudWriter): """Traces DataFrames handed to the writer; never writes or shows them.""" + CONFIG_NAME = "LineageCaptureWriter" + def __init__( - self, delegate: BaseDataCloudWriter, collector: LineageCollector + self, + spark: SparkSession, + delegate_config: AccessLayerObjectConfig[BaseDataCloudWriter], + collector: LineageCollector, ) -> None: - super().__init__(getattr(delegate, "spark", None)) # type: ignore[arg-type] - self._delegate = delegate + super().__init__(spark) + self._delegate = delegate_config.to_object(spark) self._collector = collector def _validate_dlo(self, name: str, dataframe: PySparkDataFrame) -> None: @@ -225,77 +251,27 @@ def auto_write_to_dmo(self, name: str, dataframe: PySparkDataFrame) -> None: self._collector.record_write(name, dataframe) -@dataclass -class _WrappedClient: - client: _BaseClient - reader: BaseDataCloudReader - writer: BaseDataCloudWriter - - -class _CaptureSession: - def __init__(self, collector: LineageCollector) -> None: - self.collector = collector - self._wrapped: Dict[int, _WrappedClient] = {} - - def wrap(self, client: _BaseClient) -> None: - """Route a client's reads and writes through capture (idempotent).""" - if id(client) in self._wrapped: - return - reader, writer = client._reader, client._writer - self._wrapped[id(client)] = _WrappedClient(client, reader, writer) - client._reader = LineageCaptureReader(reader, self.collector) - client._writer = LineageCaptureWriter(writer, self.collector) - - def restore(self) -> None: - for wrapped in self._wrapped.values(): - wrapped.client._reader = wrapped.reader - wrapped.client._writer = wrapped.writer - self._wrapped.clear() - - -# The capture in progress; guards against nested captures -_active: Dict[str, Optional[_CaptureSession]] = {"session": None} - - -@contextmanager -def capture_lineage(collector: LineageCollector) -> Iterator[LineageCollector]: - """Capture lineage from every ``Client`` used inside the block. - - Every live client is wrapped on entry, including ones a module kept a - reference to after they stopped being the singleton (e.g. a dependency - imported by an earlier run), and ``Client.__new__`` is hooked so - clients created inside the block are wrapped as they are built, however - their reader and writer were supplied. Clients are unwrapped, the hook is - removed, and the ``Client`` singletons are restored on exit. +def install_capture(collector: LineageCollector) -> None: + """Force every client built from now on to capture lineage into ``collector``. + + Wraps the current ``config.reader_config`` and ``config.writer_config`` (as + resolved by ``prepare_run``) in forced capture configs. Clients that already + exist are not affected, which is why capture runs in a fresh process. + + Raises: + ValueError: If no reader or writer is configured. """ - from datacustomcode.client import ( - Client, - StreamingClient, - _BaseClient, + from datacustomcode.config import AccessLayerObjectConfig, config + + if config.reader_config is None or config.writer_config is None: + raise ValueError("Lineage generation needs a reader_config and writer_config") + config.reader_config = AccessLayerObjectConfig( + type_config_name=LineageCaptureReader.CONFIG_NAME, + options={"delegate_config": config.reader_config, "collector": collector}, + force=True, + ) + config.writer_config = AccessLayerObjectConfig( + type_config_name=LineageCaptureWriter.CONFIG_NAME, + options={"delegate_config": config.writer_config, "collector": collector}, + force=True, ) - - if _active["session"] is not None: - raise RuntimeError("Lineage capture is already active") - - session = _CaptureSession(collector) - saved_clients = Client._instance, StreamingClient._instance - original_new = _BaseClient.__dict__["__new__"] - build = original_new.__func__ - - def capturing_new(cls: type, *args: Any, **kwargs: Any) -> _BaseClient: - client = build(cls, *args, **kwargs) - session.wrap(client) - return client # type: ignore[no-any-return] - - _active["session"] = session - try: - for client in (*saved_clients, *_BaseClient._live_instances): - if client is not None: - session.wrap(client) - _BaseClient.__new__ = staticmethod(capturing_new) # type: ignore[assignment] - yield collector - finally: - _BaseClient.__new__ = original_new # type: ignore[method-assign] - session.restore() - Client._instance, StreamingClient._instance = saved_clients - _active["session"] = None diff --git a/tests/lineage/lineage_test_io.py b/tests/lineage/lineage_test_io.py new file mode 100644 index 0000000..31913f9 --- /dev/null +++ b/tests/lineage/lineage_test_io.py @@ -0,0 +1,52 @@ +"""Reader/writer for the generate tests, importable by the lineage worker process. + +The worker re-imports this module (it is on ``sys.path`` via ``PYTHONPATH``), so +the writer reports validations through a file rather than shared memory. +""" + +from __future__ import annotations + +import json +import os + +import pandas as pd + +from datacustomcode.io.reader.base import BaseDataCloudReader +from datacustomcode.io.writer.base import BaseDataCloudWriter + +VALIDATED_FILE_ENV = "LINEAGE_TEST_VALIDATED_FILE" + +TABLES = { + "orders__dll": pd.DataFrame( + {"id": [1, 2], "first": ["a", "b"], "last": ["x", "y"], "amount": [1.0, 2.0]} + ), + "customers__dlm": pd.DataFrame({"id": [1, 2], "region": ["us", "eu"]}), +} + + +class InMemoryReader(BaseDataCloudReader): + CONFIG_NAME = "LineageTestInMemoryReader" + + def __init__(self, spark, dataspace=None): + super().__init__(spark) + + def read_dlo(self, name, schema=None): + return self.spark.createDataFrame(TABLES[name]) + + read_dmo = read_dlo + + +class RefusingWriter(BaseDataCloudWriter): + CONFIG_NAME = "LineageTestRefusingWriter" + + def __init__(self, spark, dataspace=None): + super().__init__(spark) + + def validate_dataframe_columns_against_dlo(self, dataframe, dlo_name): + with open(os.environ[VALIDATED_FILE_ENV], "a", encoding="utf-8") as f: + f.write(json.dumps([dlo_name, list(dataframe.columns)]) + "\n") + + def write_to_dlo(self, name, dataframe, write_mode): + raise AssertionError("lineage generation must never write") + + write_to_dmo = write_to_dlo diff --git a/tests/lineage/test_generate.py b/tests/lineage/test_generate.py index 672bc90..94af53f 100644 --- a/tests/lineage/test_generate.py +++ b/tests/lineage/test_generate.py @@ -3,76 +3,48 @@ from __future__ import annotations import json +from pathlib import Path import sys import textwrap -import pandas as pd import pytest import yaml -from datacustomcode.client import ( - Client, - StreamingClient, - _BaseClient, -) -from datacustomcode.config import config -from datacustomcode.io.reader.base import BaseDataCloudReader -from datacustomcode.io.writer.base import BaseDataCloudWriter, WriteMode -from datacustomcode.lineage import capture_lineage, generate_lineage -from datacustomcode.lineage.capture import LineageCollector +from datacustomcode.config import AccessLayerObjectConfig +from datacustomcode.lineage import LineageResult, generate_lineage +from datacustomcode.lineage.capture import LineageCaptureWriter, LineageCollector from datacustomcode.scan import SDK_CONFIG_DIR, SDK_CONFIG_FILE -pytestmark = pytest.mark.spark - -ORIGINAL_NEW = _BaseClient.__dict__["__new__"] - -TABLES = { - "orders__dll": pd.DataFrame( - {"id": [1, 2], "first": ["a", "b"], "last": ["x", "y"], "amount": [1.0, 2.0]} - ), - "customers__dlm": pd.DataFrame({"id": [1, 2], "region": ["us", "eu"]}), -} -VALIDATED: list[tuple[str, list[str]]] = [] - - -class InMemoryReader(BaseDataCloudReader): - CONFIG_NAME = "LineageTestInMemoryReader" - - def __init__(self, spark, dataspace=None): - super().__init__(spark) - - def read_dlo(self, name, schema=None): - return self.spark.createDataFrame(TABLES[name]) - - read_dmo = read_dlo - - -class RefusingWriter(BaseDataCloudWriter): - CONFIG_NAME = "LineageTestRefusingWriter" +# generate_lineage passes sys.path to its worker, so the worker can import this too +sys.path.insert(0, str(Path(__file__).parent)) +from lineage_test_io import ( + VALIDATED_FILE_ENV, + InMemoryReader, + RefusingWriter, +) - def __init__(self, spark, dataspace=None): - super().__init__(spark) +pytestmark = pytest.mark.spark - def validate_dataframe_columns_against_dlo(self, dataframe, dlo_name): - VALIDATED.append((dlo_name, list(dataframe.columns))) +PREAMBLE = ( + "from pyspark.sql import functions as F\n" + "import lineage_test_io\n" + "from datacustomcode.client import Client\n" + "from datacustomcode.io.writer.base import WriteMode\n" +) - def write_to_dlo(self, name, dataframe, write_mode): - raise AssertionError("lineage generation must never write") - write_to_dmo = write_to_dlo +@pytest.fixture(autouse=True) +def validated(tmp_path, monkeypatch): + """Reads back the (dlo, columns) validations the worker's writer recorded.""" + path = tmp_path / "validated.jsonl" + monkeypatch.setenv(VALIDATED_FILE_ENV, str(path)) + def read(): + if not path.exists(): + return [] + return [tuple(json.loads(line)) for line in path.read_text().splitlines()] -@pytest.fixture(autouse=True) -def isolated_client(spark): - saved = (config.reader_config, config.writer_config, config.spark_config) - saved_spark = _BaseClient._shared_spark - Client._instance = StreamingClient._instance = None - _BaseClient._shared_spark = spark - VALIDATED.clear() - yield - config.reader_config, config.writer_config, config.spark_config = saved - Client._instance = StreamingClient._instance = None - _BaseClient._shared_spark = saved_spark + return read @pytest.fixture @@ -89,16 +61,14 @@ def config_file(tmp_path): return str(path) -def make_package(tmp_path, body: str, config_json=None, package_type=None) -> str: +def make_package( + tmp_path, body: str, config_json=None, package_type=None, with_client=True +) -> str: payload = tmp_path / "payload" payload.mkdir() entrypoint = payload / "entrypoint.py" - entrypoint.write_text( - "from pyspark.sql import functions as F\n" - "from datacustomcode.client import Client\n" - "from datacustomcode.io.writer.base import WriteMode\n" - "client = Client()\n" + textwrap.dedent(body) - ) + preamble = PREAMBLE + ("client = Client()\n" if with_client else "") + entrypoint.write_text(preamble + textwrap.dedent(body)) (payload / "config.json").write_text( json.dumps(config_json or {"dataspace": "default"}) ) @@ -109,7 +79,7 @@ def make_package(tmp_path, body: str, config_json=None, package_type=None) -> st return str(entrypoint) -def test_generate_lineage_end_to_end(tmp_path, config_file): +def test_generate_lineage_end_to_end(tmp_path, config_file, validated): entrypoint = make_package( tmp_path, """ @@ -143,7 +113,9 @@ def test_generate_lineage_end_to_end(tmp_path, config_file): "sources": ["orders__dll"], "pairCount": 3, } - assert VALIDATED == [("out__dll", ["id__c", "name__c", "const__c", "lookup_id__c"])] + assert validated() == [ + ("out__dll", ["id__c", "name__c", "const__c", "lookup_id__c"]) + ] def test_repeat_writes_union_sources(tmp_path, config_file): @@ -170,7 +142,7 @@ def test_repeat_writes_union_sources(tmp_path, config_file): ] -def test_dmo_read_and_write(tmp_path, config_file): +def test_dmo_read_and_write(tmp_path, config_file, validated): entrypoint = make_package( tmp_path, """ @@ -184,7 +156,7 @@ def test_dmo_read_and_write(tmp_path, config_file): result = generate_lineage(entrypoint, config_file) assert result.data_lineage == {"out__dlm.region__c": ["customers__dlm.region"]} - assert VALIDATED == [] + assert validated() == [] def test_rejects_function_packages(tmp_path, config_file): @@ -205,43 +177,41 @@ def test_rejects_streaming_scripts(tmp_path, config_file): generate_lineage(entrypoint, config_file) -def test_script_error_propagates_and_restores_state(tmp_path, config_file): - config.load(config_file) - outside = Client() - reader, writer = outside._reader, outside._writer +def test_script_error_propagates(tmp_path, config_file): entrypoint = make_package(tmp_path, "raise RuntimeError('boom')\n") with pytest.raises(RuntimeError, match="boom"): generate_lineage(entrypoint, config_file) - assert config.reader_config.type_config_name == InMemoryReader.CONFIG_NAME - assert config.writer_config.type_config_name == RefusingWriter.CONFIG_NAME - assert Client._instance is outside - assert (outside._reader, outside._writer) == (reader, writer) - assert "__new__" not in vars(Client) - assert _BaseClient.__dict__["__new__"] is ORIGINAL_NEW +def test_non_builtin_error_is_reported_by_name(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + """ + class CustomerError(Exception): + pass + raise CustomerError("bad row") + """, + ) + + with pytest.raises(RuntimeError, match="CustomerError: bad row"): + generate_lineage(entrypoint, config_file) -def test_existing_client_is_captured_and_restored(spark, config_file): - config.load(config_file) - outside = Client() - writer = outside._writer - with capture_lineage(LineageCollector()): - assert Client() is outside - assert type(outside._writer).__name__ == "LineageCaptureWriter" - assert outside._writer._delegate is writer +def test_worker_exit_without_response_is_reported(tmp_path, config_file): + entrypoint = make_package(tmp_path, "import os\nos._exit(3)\n") - assert outside._writer is writer - assert Client() is outside + with pytest.raises(RuntimeError, match="exited with code 3"): + generate_lineage(entrypoint, config_file) def test_retained_client_from_dependency_is_captured( - tmp_path, config_file, monkeypatch + tmp_path, config_file, monkeypatch, validated ): # A dependency that builds its Client at import time keeps it in a module # global; its writes must still be captured, never reach the real writer (tmp_path / "lineage_dep_helpers.py").write_text( + "import lineage_test_io # noqa: F401 registers the test reader/writer\n" "from datacustomcode.client import Client\n" "from datacustomcode.io.writer.base import WriteMode\n" "client = Client()\n" @@ -249,7 +219,6 @@ def test_retained_client_from_dependency_is_captured( " client.write_to_dlo('out__dll', df, WriteMode.APPEND)\n" ) monkeypatch.syspath_prepend(str(tmp_path)) - monkeypatch.delitem(sys.modules, "lineage_dep_helpers", raising=False) entrypoint = make_package( tmp_path, """ @@ -265,71 +234,77 @@ def test_retained_client_from_dependency_is_captured( ) assert result.data_lineage == {"out__dll.id__c": ["orders__dll.id"]} - assert VALIDATED == [("out__dll", ["id__c"])] + assert validated() == [("out__dll", ["id__c"])] -def test_repeated_generation_captures_retained_client( - tmp_path, config_file, monkeypatch -): - # The helper module stays imported after the first run, keeping the client - # it built then; the second run must capture that client too - (tmp_path / "lineage_cached_helpers.py").write_text( - "from datacustomcode.client import Client\n" - "from datacustomcode.io.writer.base import WriteMode\n" - "client = Client()\n" - "def save(df):\n" - " client.write_to_dlo('out__dll', df, WriteMode.APPEND)\n" - ) - monkeypatch.syspath_prepend(str(tmp_path)) - monkeypatch.delitem(sys.modules, "lineage_cached_helpers", raising=False) +def test_repeated_generation_is_independent(tmp_path, config_file, validated): entrypoint = make_package( tmp_path, """ - import lineage_cached_helpers - lineage_cached_helpers.save( - client.read_dlo("orders__dll").select(F.col("id").alias("id__c")) + client.write_to_dlo( + "out__dll", + client.read_dlo("orders__dll").select(F.col("id").alias("id__c")), + WriteMode.APPEND, ) """, ) for _ in range(2): - result = generate_lineage( - entrypoint, config_file, dependencies=["lineage_cached_helpers"] - ) + result = generate_lineage(entrypoint, config_file) assert result.data_lineage == {"out__dll.id__c": ["orders__dll.id"]} - assert sys.modules["lineage_cached_helpers"].client is not Client._instance + assert validated() == [("out__dll", ["id__c"])] * 2 + +def test_client_built_with_explicit_reader_and_writer_is_captured( + tmp_path, config_file +): + # force=True overrides even a reader/writer passed straight to Client; the + # writer here would raise if it were used + entrypoint = make_package( + tmp_path, + """ + class ExplodingWriter(lineage_test_io.RefusingWriter): + CONFIG_NAME = "LineageTestExplodingWriter" -def test_client_built_with_reader_and_writer_is_captured(spark, config_file): - config.load(config_file) - reader, writer = InMemoryReader(spark), RefusingWriter(spark) - collector = LineageCollector() + def validate_dataframe_columns_against_dlo(self, dataframe, dlo_name): + raise AssertionError("explicit writer must not be used") - with capture_lineage(collector): - client = Client(reader=reader, writer=writer) - assert client._reader._delegate is reader - assert client._writer._delegate is writer + from pyspark.sql import SparkSession + spark = SparkSession.builder.getOrCreate() + client = Client( + reader=lineage_test_io.InMemoryReader(spark), + writer=ExplodingWriter(spark), + ) client.write_to_dlo( - "out__dll", - client.read_dlo("orders__dll").select("id"), - WriteMode.APPEND, + "out__dll", client.read_dlo("orders__dll").select("id"), WriteMode.APPEND ) + """, + with_client=False, + ) - assert collector.result().data_lineage == {"out__dll.id": ["orders__dll.id"]} - assert (client._reader, client._writer) == (reader, writer) + result = generate_lineage(entrypoint, config_file) + assert result.data_lineage == {"out__dll.id": ["orders__dll.id"]} -def test_nested_capture_is_rejected(config_file): - with capture_lineage(LineageCollector()): - with pytest.raises(RuntimeError, match="already active"): - with capture_lineage(LineageCollector()): - pass +def test_capture_falls_through_to_delegate_attributes(spark): + writer = LineageCaptureWriter( + spark, + AccessLayerObjectConfig(type_config_name=RefusingWriter.CONFIG_NAME), + LineageCollector(), + ) -def test_capture_falls_through_to_delegate_attributes(config_file): - config.load(config_file) + assert writer.validate_dataframe_columns_against_dlo is not None + assert isinstance(writer._delegate, RefusingWriter) + + +def test_result_round_trips_through_dict(): + result = LineageResult( + data_lineage={"out__dll.a__c": ["in__dll.a", "in__dll.b"]}, + unresolved={"out__dll.b__c": "untraceable source"}, + no_sources=["out__dll.c__c"], + targets=["out__dll"], + sources=["in__dll"], + ) - with capture_lineage(LineageCollector()): - writer = Client()._writer - assert writer.validate_dataframe_columns_against_dlo is not None - assert isinstance(writer._delegate, RefusingWriter) + assert LineageResult.from_dict(result.to_dict()).to_dict() == result.to_dict() diff --git a/tests/test_client.py b/tests/test_client.py index fdcdaa8..1293ea1 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -159,6 +159,44 @@ def test_initialization_with_config(self, mock_config, reset_client, mock_spark) assert client._reader is mock_reader assert client._writer is mock_writer + @patch("datacustomcode.client.config") + def test_forced_config_overrides_provided_reader_and_writer( + self, mock_config, reset_client, mock_spark + ): + """Forced configs replace a passed reader and writer, which needs a + session even though both were provided.""" + from datacustomcode.spark.default import DefaultSparkSessionProvider + + with patch.object( + DefaultSparkSessionProvider, "get_session" + ) as mock_get_session: + mock_get_session.return_value = mock_spark + + mock_reader = MagicMock(spec=BaseDataCloudReader) + mock_reader_config = MagicMock() + mock_reader_config.to_object.return_value = mock_reader + mock_reader_config.force = True + + mock_writer = MagicMock(spec=BaseDataCloudWriter) + mock_writer_config = MagicMock() + mock_writer_config.to_object.return_value = mock_writer + mock_writer_config.force = True + + mock_config.spark_provider_config = None + mock_config.reader_config = mock_reader_config + mock_config.writer_config = mock_writer_config + mock_config.spark_config = MagicMock(spec=SparkConfig) + + client = Client( + reader=MagicMock(spec=BaseDataCloudReader), + writer=MagicMock(spec=BaseDataCloudWriter), + ) + + mock_reader_config.to_object.assert_called_once_with(mock_spark) + mock_writer_config.to_object.assert_called_once_with(mock_spark) + assert client._reader is mock_reader + assert client._writer is mock_writer + def test_read_dlo(self, reset_client, mock_spark): reader = MagicMock(spec=BaseDataCloudReader) writer = MagicMock(spec=BaseDataCloudWriter) From 856d8636fc1c72fde4abdda7a74ea798cf97d817 Mon Sep 17 00:00:00 2001 From: Jesus Orosco Date: Thu, 15 Oct 2026 09:44:17 -0700 Subject: [PATCH 5/6] more UDF tracer things --- src/datacustomcode/lineage/capture.py | 18 +++++-- src/datacustomcode/lineage/tracer.py | 33 +++++++++++-- tests/lineage/test_generate.py | 51 ++++++++++++++++++++ tests/lineage/test_tracer.py | 68 ++++++++++++++++++++++++++- 4 files changed, 162 insertions(+), 8 deletions(-) diff --git a/src/datacustomcode/lineage/capture.py b/src/datacustomcode/lineage/capture.py index 94c18cd..404205f 100644 --- a/src/datacustomcode/lineage/capture.py +++ b/src/datacustomcode/lineage/capture.py @@ -61,6 +61,8 @@ class LineageResult: data_lineage: Target field to its sorted source fields. unresolved: Target field to the reason its lineage could not be traced. no_sources: Target fields computed from no source field (e.g. literals). + udf_fields: Target fields in ``data_lineage`` or ``no_sources`` computed + by a user-defined function, whose sources may be incomplete. targets: Objects written. sources: Objects read. """ @@ -68,6 +70,7 @@ class LineageResult: data_lineage: Dict[str, List[str]] = field(default_factory=dict) unresolved: Dict[str, str] = field(default_factory=dict) no_sources: List[str] = field(default_factory=list) + udf_fields: List[str] = field(default_factory=list) targets: List[str] = field(default_factory=list) sources: List[str] = field(default_factory=list) @@ -88,6 +91,7 @@ def to_dict(self) -> Dict[str, Any]: for target in sorted(self.unresolved) ], "noSources": sorted(self.no_sources), + "udfFields": sorted(self.udf_fields), "targets": sorted(self.targets), "sources": sorted(self.sources), "pairCount": self.pair_count, @@ -105,6 +109,7 @@ def from_dict(cls, data: Dict[str, Any]) -> LineageResult: entry["targetField"]: entry["reason"] for entry in data["unresolved"] }, no_sources=list(data["noSources"]), + udf_fields=list(data["udfFields"]), targets=list(data["targets"]), sources=list(data["sources"]), ) @@ -126,12 +131,16 @@ def record_write(self, object_name: str, dataframe: PySparkDataFrame) -> None: if column == MERGE_RECORD_TYPE_COLUMN: continue # Repeat writes to the same target: a field reads from the union of - # every write, and stays unresolved if any write was unresolved + # every write, stays unresolved if any write was unresolved, and is + # UDF-computed if any write computed it with a UDF previous = fields.get(column) if previous is None or (traced.unresolved and not previous.unresolved): fields[column] = traced elif not previous.unresolved: - fields[column] = TraceResult(sources=previous.sources | traced.sources) + fields[column] = TraceResult( + sources=previous.sources | traced.sources, + udf=previous.udf or traced.udf, + ) def result(self) -> LineageResult: result = LineageResult( @@ -142,10 +151,13 @@ def result(self) -> LineageResult: target = f"{object_name}.{column}" if traced.reason is not None: result.unresolved[target] = traced.reason - elif traced.sources: + continue + if traced.sources: result.data_lineage[target] = sorted(traced.sources) else: result.no_sources.append(target) + if traced.udf: + result.udf_fields.append(target) return result diff --git a/src/datacustomcode/lineage/tracer.py b/src/datacustomcode/lineage/tracer.py index 7b43a40..48455e3 100644 --- a/src/datacustomcode/lineage/tracer.py +++ b/src/datacustomcode/lineage/tracer.py @@ -60,6 +60,19 @@ {"ScalarSubquery", "Exists", "ListQuery", "InSubquery", "LateralSubquery"} ) +# User-defined functions: their arguments are traced, but the body can read data +# the plan does not show (a captured lookup, a broadcast, a file) +_UDF_EXPRESSIONS = frozenset( + { + "PythonUDF", + "PythonUDAF", + "PythonUDTF", + "ScalaUDF", + "ScalaUDAF", + "ScalaAggregator", + } +) + class LineageTraceError(Exception): """Raised when a DataFrame's plan cannot be traced at all.""" @@ -76,12 +89,15 @@ class ColumnTrace: ``mapInPandas`` or subquery); ``leaves`` is then incomplete. direct: The column is its single leaf column passed through unchanged (possibly renamed), not computed from it. + udf: The value is computed by a user-defined function, which may read + more than ``leaves``. """ name: str leaves: FrozenSet[LeafColumn] = field(default_factory=frozenset) opaque_reason: Optional[str] = None direct: bool = False + udf: bool = False @dataclass(frozen=True) @@ -89,11 +105,14 @@ class _AttrTrace: leaves: FrozenSet[LeafColumn] opaque_reason: Optional[str] = None direct: bool = False + udf: bool = False def merge(self, other: _AttrTrace) -> _AttrTrace: # A merge combines inputs, so the result is never a direct pass-through return _AttrTrace( - self.leaves | other.leaves, self.opaque_reason or other.opaque_reason + self.leaves | other.leaves, + self.opaque_reason or other.opaque_reason, + udf=self.udf or other.udf, ) @@ -108,10 +127,13 @@ class TraceResult: sources: ``"Object.field"`` names the value is computed from. reason: Set when the column cannot be fully traced; ``sources`` must then not be used. + udf: The value is computed by a user-defined function, so ``sources`` + may be incomplete and needs review. """ sources: FrozenSet[str] = field(default_factory=frozenset) reason: Optional[str] = None + udf: bool = False @property def unresolved(self) -> bool: @@ -156,7 +178,8 @@ def trace(self, df: PySparkDataFrame) -> Dict[str, TraceResult]: results[column.name] = TraceResult(reason=UNTRACEABLE_SOURCE) else: results[column.name] = TraceResult( - sources=frozenset(self._fields[leaf] for leaf in column.leaves) + sources=frozenset(self._fields[leaf] for leaf in column.leaves), + udf=column.udf, ) return results @@ -214,7 +237,9 @@ def trace_columns(df: PySparkDataFrame) -> List[ColumnTrace]: for name, expr_id in zip(df.columns, root_ids, strict=False): trace = traces.get(expr_id, _EMPTY) columns.append( - ColumnTrace(name, trace.leaves, trace.opaque_reason, trace.direct) + ColumnTrace( + name, trace.leaves, trace.opaque_reason, trace.direct, trace.udf + ) ) return columns @@ -405,6 +430,8 @@ def _trace_expression( result = result.merge( _AttrTrace(frozenset(), f"{expr_class} inputs cannot be traced") ) + elif expr_class in _UDF_EXPRESSIONS: + result = result.merge(_AttrTrace(frozenset(), udf=True)) elif expr_class == "AttributeReference": expr_id = _expr_id(expr_node) if expr_id in traces: diff --git a/tests/lineage/test_generate.py b/tests/lineage/test_generate.py index 94af53f..c8ea650 100644 --- a/tests/lineage/test_generate.py +++ b/tests/lineage/test_generate.py @@ -109,6 +109,7 @@ def test_generate_lineage_end_to_end(tmp_path, config_file, validated): {"targetField": "out__dll.lookup_id__c", "reason": "untraceable source"} ], "noSources": ["out__dll.const__c"], + "udfFields": [], "targets": ["out__dll"], "sources": ["orders__dll"], "pairCount": 3, @@ -142,6 +143,55 @@ def test_repeat_writes_union_sources(tmp_path, config_file): ] +def test_udf_fields_are_reported(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + """ + orders = client.read_dlo("orders__dll") + shout = F.udf(lambda a: a.upper(), "string") + constant = F.udf(lambda: "x", "string") + df = orders.select( + F.col("id").alias("id__c"), + shout("first").alias("loud__c"), + constant().alias("const__c"), + ) + client.write_to_dlo("out__dll", df, WriteMode.APPEND) + """, + ) + + result = generate_lineage(entrypoint, config_file).to_dict() + + assert result["udfFields"] == ["out__dll.const__c", "out__dll.loud__c"] + assert {e["targetField"]: e["sourceFields"] for e in result["dataLineage"]} == { + "out__dll.id__c": ["orders__dll.id"], + "out__dll.loud__c": ["orders__dll.first"], + } + assert result["noSources"] == ["out__dll.const__c"] + + +def test_repeat_writes_keep_udf_flag(tmp_path, config_file): + entrypoint = make_package( + tmp_path, + """ + orders = client.read_dlo("orders__dll") + shout = F.udf(lambda a: a.upper(), "string") + client.write_to_dlo( + "out__dll", orders.select(shout("first").alias("v__c")), WriteMode.APPEND + ) + client.write_to_dlo( + "out__dll", orders.select(F.col("last").alias("v__c")), WriteMode.APPEND + ) + """, + ) + + result = generate_lineage(entrypoint, config_file) + + assert result.data_lineage == { + "out__dll.v__c": ["orders__dll.first", "orders__dll.last"] + } + assert result.udf_fields == ["out__dll.v__c"] + + def test_dmo_read_and_write(tmp_path, config_file, validated): entrypoint = make_package( tmp_path, @@ -303,6 +353,7 @@ def test_result_round_trips_through_dict(): data_lineage={"out__dll.a__c": ["in__dll.a", "in__dll.b"]}, unresolved={"out__dll.b__c": "untraceable source"}, no_sources=["out__dll.c__c"], + udf_fields=["out__dll.a__c", "out__dll.c__c"], targets=["out__dll"], sources=["in__dll"], ) diff --git a/tests/lineage/test_tracer.py b/tests/lineage/test_tracer.py index 7099074..73728ee 100644 --- a/tests/lineage/test_tracer.py +++ b/tests/lineage/test_tracer.py @@ -94,9 +94,73 @@ def test_using_join_key_is_coalesced(registry, orders, customers): def test_python_udf_inputs(registry, orders): shout = F.udf(lambda a, b: f"{a}{b}".upper(), StringType()) df = orders.select(shout("first", "last").alias("loud")) - assert sources(registry.trace(df)) == { - "loud": {"orders__dll.first", "orders__dll.last"} + results = registry.trace(df) + assert sources(results) == {"loud": {"orders__dll.first", "orders__dll.last"}} + assert results["loud"].udf + + +def test_udf_flag_marks_only_udf_columns(registry, orders): + shout = F.udf(lambda a: a.upper(), StringType()) + df = orders.select("id", F.upper("last").alias("up"), shout("first").alias("s")) + results = registry.trace(df) + assert {name for name, r in results.items() if r.udf} == {"s"} + + +def test_pandas_udf_is_flagged(registry, orders): + @F.pandas_udf("double") + def double_it(s: pd.Series) -> pd.Series: + return s * 2 + + results = registry.trace(orders.select(double_it("amount").alias("dbl"))) + assert results["dbl"] == TraceResult( + sources=frozenset({"orders__dll.amount"}), udf=True + ) + + +def test_grouped_agg_pandas_udf_is_flagged(registry, orders): + @F.pandas_udf("double") + def total(s: pd.Series) -> float: + return float(s.sum()) + + results = registry.trace(orders.groupBy("first").agg(total("amount").alias("t"))) + assert sources(results) == { + "first": {"orders__dll.first"}, + "t": {"orders__dll.amount"}, } + assert results["t"].udf + assert not results["first"].udf + + +def test_sql_registered_udf_is_flagged(spark, registry, orders): + spark.udf.register("lineage_shout", lambda a: a.upper(), StringType()) + orders.createOrReplaceTempView("orders_udf_view") + df = spark.sql("SELECT id, lineage_shout(first) AS s FROM orders_udf_view") + results = registry.trace(df) + assert results["s"] == TraceResult( + sources=frozenset({"orders__dll.first"}), udf=True + ) + assert not results["id"].udf + + +def test_udf_flag_survives_rename_and_later_expressions(registry, orders): + shout = F.udf(lambda a: a.upper(), StringType()) + df = ( + orders.select(shout("first").alias("s"), "last") + .withColumnRenamed("s", "renamed") + .withColumn("joined", F.concat("renamed", "last")) + ) + results = registry.trace(df) + assert results["renamed"].udf + assert results["joined"] == TraceResult( + sources=frozenset({"orders__dll.first", "orders__dll.last"}), udf=True + ) + assert not results["last"].udf + + +def test_udf_without_arguments_has_no_sources_but_is_flagged(registry, orders): + constant = F.udf(lambda: "x", StringType()) + results = registry.trace(orders.select(constant().alias("c"))) + assert results["c"] == TraceResult(udf=True) def test_aggregate(registry, orders): From 99c909f916fdd88c4ab6b2bd5c7a551004465381 Mon Sep 17 00:00:00 2001 From: Jesus Orosco Date: Thu, 15 Oct 2026 14:29:53 -0700 Subject: [PATCH 6/6] few more PR comments, and updated docs --- README.md | 6 ++-- src/datacustomcode/lineage/tracer.py | 42 +++++++++++++++++++--- tests/lineage/test_tracer.py | 52 ++++++++++++++++++++++++++-- 3 files changed, 90 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 169fb28..1b22d9c 100644 --- a/README.md +++ b/README.md @@ -239,8 +239,10 @@ print(result.to_dict()["dataLineage"]) 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`). -- A target field's sources are the fields its *value* is computed from. Filter conditions and join keys aren't sources. -- Fields computed only from literals are listed in `noSources`. +> **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. diff --git a/src/datacustomcode/lineage/tracer.py b/src/datacustomcode/lineage/tracer.py index 48455e3..3d147c8 100644 --- a/src/datacustomcode/lineage/tracer.py +++ b/src/datacustomcode/lineage/tracer.py @@ -22,8 +22,12 @@ re-aliases one side of a self-join to fresh ``exprId`` values, but both sides still sit on the same underlying RDD / local data. -Only the columns that feed a column's *value* are sources. Filter predicates, -join conditions and sort keys gate or order rows and are deliberately ignored. +Only the columns that feed a column's *value* are sources. Columns that only +decide which rows, groups or branch are used are deliberately ignored, wherever +they appear: filter predicates, join conditions, grouping and sort keys, window +partition/order keys, ``CASE WHEN``/``IF`` conditions and aggregate ``FILTER`` +clauses. A value computed from none of its inputs' values (``row_number()``, +``rank()``, ``count(*)``) therefore has no sources. """ from __future__ import annotations @@ -74,6 +78,24 @@ ) +# Expressions with children that only decide which rows, groups or branch are used +_ROW_GATING_EXPRESSIONS = frozenset( + {"WindowExpression", "AggregateExpression", "CaseWhen", "If"} +) + + +def _keep_value_children(expr_class: str, position: int, count: int) -> bool: + """Whether child ``position`` of a row-gating expression feeds its value.""" + if expr_class in ("WindowExpression", "AggregateExpression"): + # (function, window spec) / (function, optional FILTER predicate) + return position == 0 + if expr_class == "CaseWhen": + # (cond, value)* then an optional else value + return position % 2 == 1 or (count % 2 == 1 and position == count - 1) + # If: (predicate, true value, false value) + return position != 0 + + class LineageTraceError(Exception): """Raised when a DataFrame's plan cannot be traced at all.""" @@ -422,10 +444,20 @@ def _child_subtrees(expression: List[dict]) -> List[List[dict]]: def _trace_expression( expression: List[dict], traces: Dict[int, _AttrTrace] ) -> _AttrTrace: - """Union the traces of every attribute an expression (sub)tree references.""" + """Union the traces of every attribute an expression (sub)tree's value reads.""" result = _EMPTY - for expr_node in expression: - expr_class = expr_node["class"].rsplit(".", 1)[-1] + index = 0 + while index < len(expression): + expr_node = expression[index] + expr_class = _class_name(expr_node) + if expr_class in _ROW_GATING_EXPRESSIONS: + children = _child_subtrees(expression[index:]) + for position, child in enumerate(children): + if _keep_value_children(expr_class, position, len(children)): + result = result.merge(_trace_expression(child, traces)) + index = _subtree_end(expression, index) + continue + index += 1 if expr_class in _OPAQUE_EXPRESSIONS: result = result.merge( _AttrTrace(frozenset(), f"{expr_class} inputs cannot be traced") diff --git a/tests/lineage/test_tracer.py b/tests/lineage/test_tracer.py index 73728ee..30eb956 100644 --- a/tests/lineage/test_tracer.py +++ b/tests/lineage/test_tracer.py @@ -194,15 +194,61 @@ def test_union_by_name(registry, orders, customers): } -def test_window(registry, orders): - window = Window.partitionBy("first").orderBy("amount") +def test_window_keys_are_not_sources(registry, orders): + window = Window.partitionBy("first").orderBy("id") df = orders.select("id", F.sum("amount").over(window).alias("running")) assert sources(registry.trace(df)) == { "id": {"orders__dll.id"}, - "running": {"orders__dll.amount", "orders__dll.first"}, + "running": {"orders__dll.amount"}, } +def test_window_matches_group_by(registry, orders): + grouped = orders.groupBy("first").agg(F.sum("amount").alias("total")) + windowed = orders.select( + F.sum("amount").over(Window.partitionBy("first")).alias("total") + ) + assert sources(registry.trace(grouped))["total"] == {"orders__dll.amount"} + assert sources(registry.trace(windowed))["total"] == {"orders__dll.amount"} + + +def test_ranking_function_has_no_sources(registry, orders): + window = Window.partitionBy("first").orderBy("amount") + df = orders.select(F.row_number().over(window).alias("rn")) + assert sources(registry.trace(df)) == {"rn": set()} + + +def test_when_condition_is_not_a_source(registry, orders): + df = orders.select( + F.when(F.col("first") == "a", F.col("amount")) + .when(F.col("last") == "y", F.col("id")) + .otherwise(F.col("last")) + .alias("picked") + ) + assert sources(registry.trace(df)) == { + "picked": {"orders__dll.amount", "orders__dll.id", "orders__dll.last"} + } + + +def test_sql_if_condition_is_not_a_source(registry, orders): + df = orders.select(F.expr("if(first = 'a', amount, id)").alias("picked")) + assert sources(registry.trace(df)) == { + "picked": {"orders__dll.amount", "orders__dll.id"} + } + + +def test_conditional_count_has_no_sources(registry, orders): + df = orders.select(F.sum(F.when(F.col("first") == "a", 1)).alias("count_a")) + assert sources(registry.trace(df)) == {"count_a": set()} + + +def test_aggregate_filter_matches_dataframe_filter(registry, orders): + filtered = orders.filter(F.col("first") == "a").agg(F.sum("amount").alias("t")) + sql_filter = orders.agg(F.expr("sum(amount) FILTER (WHERE first = 'a')").alias("t")) + assert sources(registry.trace(filtered))["t"] == {"orders__dll.amount"} + assert sources(registry.trace(sql_filter))["t"] == {"orders__dll.amount"} + + def test_explode(registry, orders): df = orders.select("id", F.explode("tags").alias("tag")) assert sources(registry.trace(df)) == {