diff --git a/README.md b/README.md index a9f5ae8..1b22d9c 100644 --- a/README.md +++ b/README.md @@ -224,6 +224,29 @@ If none of these exist, `find_file_path` raises `FileNotFoundError` with the lis `$LIBRARY_PATH` is set automatically to the root of the package at runtime inside Data Cloud. +### Generating lineage + +`generate_lineage` runs a script entry point locally and derives the field-level `dataLineage` section of `config.json` from the Spark plans of the DataFrames it writes: + +```python +from datacustomcode.lineage import generate_lineage + +result = generate_lineage("payload/entrypoint.py", sf_cli_org="my-org") +print(result.to_dict()["dataLineage"]) +# [{"targetField": "Output__dll.name__c", +# "sourceFields": ["Input__dll.first__c", "Input__dll.last__c"]}, ...] +``` + +The entry point runs exactly as it would locally (same config, reader, and credentials), except that writes are traced and validated against the target DLO instead of being performed. It runs in a separate Python process, which inherits your working directory, environment, and `sys.path`; errors raised by the script are re-raised from `generate_lineage`. Reads use the reader's `default_row_limit` (1000 by default; set it in your `config.yaml`). + +> **Review the output before you use it.** The generated lineage is a starting point for review, not a guarantee. It reflects only the code paths and data your local run exercised, and it can't see inside UDFs. Accurate and complete `dataLineage` in your package is your responsibility as the package author. + +- A target field's sources are the fields its *value* is computed from. Fields that only decide which rows, groups or branch are used aren't sources, however the logic is written: filter conditions, join keys, `groupBy` keys, sort keys, window `partitionBy`/`orderBy` keys, `when()`/`if()` conditions, and aggregate `FILTER (WHERE …)` clauses. For example, `groupBy("region").agg(sum("amount"))` and `sum("amount").over(Window.partitionBy("region"))` both trace to `amount` only. +- Fields computed only from literals are listed in `noSources`. So are fields whose value doesn't come from any input field's value, such as `row_number()`, `rank()`, `count("*")`, or `sum(when(cond, 1))`. +- Fields that can't be traced are listed in `unresolved` with a reason and get no lineage. This covers data that didn't come from `read_dlo`/`read_dmo` (for example a `toPandas()` → `createDataFrame` round trip), `mapInPandas`, and subqueries. +- `pairCount` is the total number of (target, source) field pairs. +- Only batch script packages are supported. + ## CLI diff --git a/src/datacustomcode/client.py b/src/datacustomcode/client.py index 814ac5a..49d1407 100644 --- a/src/datacustomcode/client.py +++ b/src/datacustomcode/client.py @@ -376,7 +376,13 @@ def __new__( instance._spark_named_credential = spark_named_credential # Initialize Readers and Writers from config # and/or provided reader and writer - if reader is None or writer is None: + reader_forced = ( + config.reader_config is not None and config.reader_config.force + ) + writer_forced = ( + config.writer_config is not None and config.writer_config.force + ) + if reader is None or writer is None or reader_forced or writer_forced: # We need a spark because we will initialize readers and writers. # Reuse the process-wide session if one client already built it, # so a Client and a StreamingClient share a single connection. @@ -404,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 @@ -414,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 diff --git a/src/datacustomcode/lineage/__init__.py b/src/datacustomcode/lineage/__init__.py new file mode 100644 index 0000000..cf96819 --- /dev/null +++ b/src/datacustomcode/lineage/__init__.py @@ -0,0 +1,137 @@ +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Generate config.json ``dataLineage`` by tracing a script's Spark plans.""" + +from __future__ import annotations + +import builtins +import json +import os +import subprocess +import sys +import tempfile +from typing import ( + Any, + Dict, + NoReturn, + Optional, + Sequence, +) + +from datacustomcode.lineage.capture import ( + LineageCaptureReader, + LineageCaptureWriter, + LineageCollector, + LineageResult, + install_capture, +) +from datacustomcode.lineage.tracer import LineageTraceError + +__all__ = [ + "LineageCaptureReader", + "LineageCaptureWriter", + "LineageCollector", + "LineageResult", + "LineageTraceError", + "generate_lineage", + "install_capture", +] + + +def generate_lineage( + entrypoint: str, + config_file: Optional[str] = None, + dependencies: Sequence[str] = (), + profile: str = "default", + sf_cli_org: Optional[str] = None, +) -> LineageResult: + """Run a script entrypoint and return the field-level lineage it produces. + + The entrypoint runs exactly as with ``run`` (same config resolution, reader + and credentials), except that writes are traced instead of performed. It runs + in a fresh Python process, so no client or config from the caller's process + leaks into the capture. The process inherits this one's working directory, + environment, ``sys.path`` and stdout/stderr. + + Args: + entrypoint: The script entrypoint to run. + config_file: Optional SDK config file, as for ``run``. + dependencies: Modules to import before running, as for ``run``. + profile: The credentials profile to use. + sf_cli_org: Optional SF CLI org alias or username. + + Returns: + The captured :class:`LineageResult`. + + Raises: + ValueError: If the package is not a batch script. + Exception: Whatever the run raised, re-raised as the same built-in + exception type or :class:`LineageTraceError`, else ``RuntimeError``. + """ + with tempfile.TemporaryDirectory(prefix="datacustomcode-lineage-") as tmp: + request_file = os.path.join(tmp, "request.json") + response_file = os.path.join(tmp, "response.json") + with open(request_file, "w", encoding="utf-8") as f: + json.dump( + { + "args": { + "entrypoint": entrypoint, + "config_file": config_file, + "dependencies": list(dependencies), + "profile": profile, + "sf_cli_org": sf_cli_org, + }, + "response_file": response_file, + }, + f, + ) + # The worker sees the same modules this process does + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join(path for path in sys.path if path) + completed = subprocess.run( + [sys.executable, "-m", "datacustomcode.lineage._worker", request_file], + env=env, + check=False, + ) + if not os.path.exists(response_file): + raise RuntimeError( + f"Lineage generation exited with code {completed.returncode} " + "without a result" + ) + with open(response_file, encoding="utf-8") as f: + response = json.load(f) + if "error" in response: + _raise(response["error"]) + return LineageResult.from_dict(response["result"]) + + +def _raise(error: Dict[str, Any]) -> NoReturn: + """Re-raise an error reported by the worker as the closest local type.""" + message = error["message"] + if error["type"] == LineageTraceError.__name__: + raise LineageTraceError(message) + exc_type = getattr(builtins, error["type"], None) + if ( + error["module"] == "builtins" + and isinstance(exc_type, type) + and issubclass(exc_type, Exception) + ): + try: + exc = exc_type(message) + except TypeError: # e.g. UnicodeDecodeError takes more arguments + pass + else: + raise exc + raise RuntimeError(f"{error['module']}.{error['type']}: {message}") 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 new file mode 100644 index 0000000..404205f --- /dev/null +++ b/src/datacustomcode/lineage/capture.py @@ -0,0 +1,289 @@ +# 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. + +: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 dataclasses import dataclass, field +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Union, +) + +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 + + from datacustomcode.config import AccessLayerObjectConfig + + +@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). + 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. + """ + + 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) + + @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), + "udfFields": sorted(self.udf_fields), + "targets": sorted(self.targets), + "sources": sorted(self.sources), + "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"]), + udf_fields=list(data["udfFields"]), + targets=list(data["targets"]), + sources=list(data["sources"]), + ) + + +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, 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, + udf=previous.udf or traced.udf, + ) + + 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 + 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 + + +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, + delegate_config: AccessLayerObjectConfig[BaseDataCloudReader], + collector: LineageCollector, + ) -> None: + super().__init__(spark) + self._delegate = delegate_config.to_object(spark) + self._collector = collector + + def read_dlo( + self, + name: str, + schema: Union[AtomicType, StructType, str, None] = None, + ) -> PySparkDataFrame: + 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 + + def read_dmo( + self, + name: str, + schema: Union[AtomicType, StructType, str, None] = None, + ) -> PySparkDataFrame: + 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 + + +class LineageCaptureWriter(_Delegating, BaseDataCloudWriter): + """Traces DataFrames handed to the writer; never writes or shows them.""" + + CONFIG_NAME = "LineageCaptureWriter" + + def __init__( + self, + spark: SparkSession, + delegate_config: AccessLayerObjectConfig[BaseDataCloudWriter], + collector: LineageCollector, + ) -> None: + super().__init__(spark) + self._delegate = delegate_config.to_object(spark) + self._collector = collector + + 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) + + +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.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, + ) diff --git a/src/datacustomcode/lineage/tracer.py b/src/datacustomcode/lineage/tracer.py new file mode 100644 index 0000000..3d147c8 --- /dev/null +++ b/src/datacustomcode/lineage/tracer.py @@ -0,0 +1,566 @@ +# 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. 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 + +from dataclasses import dataclass, field +import json +from typing import ( + TYPE_CHECKING, + Any, + Dict, + FrozenSet, + Hashable, + Iterator, + List, + Optional, + Set, + Tuple, +) + +if TYPE_CHECKING: + from pyspark.sql import DataFrame as PySparkDataFrame + +LeafColumn = Tuple[Hashable, int] +"""``(leaf relation identity, column position)``.""" + +UNTRACEABLE_SOURCE = "untraceable source" +AMBIGUOUS_SOURCE = "ambiguous 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"} +) + +# 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", + } +) + + +# 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.""" + + +@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. + 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) +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, + udf=self.udf or other.udf, + ) + + +_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. + 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: + return self.reason is not None + + +class SourceRegistry: + """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 direct columns of ``df`` are fields of ``object_name``.""" + for column in trace_columns(df): + 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) + + 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 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: + results[column.name] = TraceResult( + sources=frozenset(self._fields[leaf] for leaf in column.leaves), + udf=column.udf, + ) + 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)}" + ) + 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, trace.udf + ) + ) + return columns + + +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)}), direct=True) + 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": + 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) + + 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"): + 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. + + 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 or not whole.leaves: + return [whole] * output_count + 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(children[1:]): + slot = position % output_count + per_output[slot] = per_output[slot].merge(_trace_expression(arg, traces)) + 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 + 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]]: + """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( + expression: List[dict], traces: Dict[int, _AttrTrace] +) -> _AttrTrace: + """Union the traces of every attribute an expression (sub)tree's value reads.""" + result = _EMPTY + 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") + ) + 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: + 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..0552adf 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 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. + 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) @@ -169,7 +207,17 @@ def run_entrypoint( 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) @@ -184,8 +232,19 @@ def run_entrypoint( except (ModuleNotFoundError, AttributeError) as inner_exc: raise inner_exc from exc + +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/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 new file mode 100644 index 0000000..c8ea650 --- /dev/null +++ b/tests/lineage/test_generate.py @@ -0,0 +1,361 @@ +"""End-to-end tests for generate_lineage with an in-memory reader (no org).""" + +from __future__ import annotations + +import json +from pathlib import Path +import sys +import textwrap + +import pytest +import yaml + +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 + +# 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, +) + +pytestmark = pytest.mark.spark + +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" +) + + +@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()] + + return read + + +@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, with_client=True +) -> str: + payload = tmp_path / "payload" + payload.mkdir() + entrypoint = payload / "entrypoint.py" + 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"}) + ) + 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, validated): + 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"], + "udfFields": [], + "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_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, + """ + 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(tmp_path, config_file): + entrypoint = make_package(tmp_path, "raise RuntimeError('boom')\n") + + with pytest.raises(RuntimeError, match="boom"): + generate_lineage(entrypoint, config_file) + + +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_worker_exit_without_response_is_reported(tmp_path, config_file): + entrypoint = make_package(tmp_path, "import os\nos._exit(3)\n") + + 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, 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" + "def save(df):\n" + " client.write_to_dlo('out__dll', df, WriteMode.APPEND)\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + 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_repeated_generation_is_independent(tmp_path, config_file, validated): + entrypoint = make_package( + tmp_path, + """ + 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) + assert result.data_lineage == {"out__dll.id__c": ["orders__dll.id"]} + 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 validate_dataframe_columns_against_dlo(self, dataframe, dlo_name): + raise AssertionError("explicit writer must not be used") + + 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 + ) + """, + with_client=False, + ) + + result = generate_lineage(entrypoint, config_file) + + assert result.data_lineage == {"out__dll.id": ["orders__dll.id"]} + + +def test_capture_falls_through_to_delegate_attributes(spark): + writer = LineageCaptureWriter( + spark, + AccessLayerObjectConfig(type_config_name=RefusingWriter.CONFIG_NAME), + LineageCollector(), + ) + + 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"], + udf_fields=["out__dll.a__c", "out__dll.c__c"], + targets=["out__dll"], + sources=["in__dll"], + ) + + assert LineageResult.from_dict(result.to_dict()).to_dict() == result.to_dict() diff --git a/tests/lineage/test_tracer.py b/tests/lineage/test_tracer.py new file mode 100644 index 0000000..30eb956 --- /dev/null +++ b/tests/lineage/test_tracer.py @@ -0,0 +1,476 @@ +"""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 ( + AMBIGUOUS_SOURCE, + 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")) + 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): + 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_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"}, + } + + +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)) == { + "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] + + +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_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): + exploded = orders.select(F.posexplode("tags").alias("pos", "tag")) + assert sources(registry.trace(exploded)) == { + "pos": {"orders__dll.tags"}, + "tag": {"orders__dll.tags"}, + } 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) 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__"