From 62fadc3dc851c2423352bf6b246567dac87b022c Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Tue, 3 Mar 2026 21:07:04 +0200 Subject: [PATCH 01/17] Add Spark Declarative Pipelines (SDP) component integration --- .../integrations/libraries/spark/index.md | 16 + .../dagster-spark/dagster_spark/__init__.py | 3 + .../dagster_spark/components/__init__.py | 1 + .../spark_declarative_pipeline/__init__.py | 21 ++ .../spark_declarative_pipeline/component.py | 191 ++++++++++++ .../spark_declarative_pipeline/discovery.py | 293 ++++++++++++++++++ .../spark_declarative_pipeline/resource.py | 139 +++++++++ .../spark_declarative_pipeline/scaffolder.py | 18 ++ .../components/__init__.py | 1 + .../spark_declarative_pipeline/__init__.py | 1 + .../test_component.py | 122 ++++++++ .../test_discovery.py | 114 +++++++ .../test_resource.py | 116 +++++++ .../libraries/dagster-spark/setup.py | 5 + 14 files changed, 1041 insertions(+) create mode 100644 python_modules/libraries/dagster-spark/dagster_spark/components/__init__.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark_tests/components/__init__.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/__init__.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py create mode 100644 python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py diff --git a/docs/docs/integrations/libraries/spark/index.md b/docs/docs/integrations/libraries/spark/index.md index 973adb2328e89..77994186b4e5d 100644 --- a/docs/docs/integrations/libraries/spark/index.md +++ b/docs/docs/integrations/libraries/spark/index.md @@ -40,3 +40,19 @@ Existing Spark jobs can be used with Pipes without any modifications. In this ca Additionally, it's possible to send events to Dagster from the job by utilizing the `dagster_pipes` module. This requires minimal code changes on the job side. This approach also works for Spark jobs written in Java or Scala, although we don't have Pipes implementations for emitting events from those languages yet. + +## Spark Declarative Pipelines (SDP) Integration + +Dagster also provides native support for the new **Spark Declarative Pipelines (SDP)** framework (Spark 4.0+). For organizations using SDP to define datasets via decorators or SQL, Dagster offers a dedicated component to seamlessly orchestrate these pipelines without duplicating code. + +The `SparkDeclarativePipelineComponent` leverages the `spark-pipelines` CLI to automatically discover your datasets and dependencies using `spark-pipelines dry-run`. + +**Key benefits include:** +* **Auto-Discovery:** No need to manually define `AssetSpec`s. The component infers Materialized Views and Streaming Tables automatically at load time. +* **Incremental & Full Refresh:** Natively supports both `--refresh` and `--full-refresh` execution modes. +* **Real-time Observability:** Streams execution logs and events directly back to the Dagster UI during execution. +* **UI Clutter Reduction:** Pipeline-scoped intermediate datasets (Temporary Views) are automatically filtered out from the Dagster lineage unless explicitly overridden. + +You can quickly initialize a new SDP component in your Dagster project using the `dg` CLI: +```bash +dg scaffold component dagster_spark.SparkDeclarativePipelineComponent my_sdp_pipeline --pipeline-spec-path ./path/to/spark-pipeline.yml \ No newline at end of file diff --git a/python_modules/libraries/dagster-spark/dagster_spark/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark/__init__.py index 772d65b2cafb1..11b6e37f914fc 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/__init__.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/__init__.py @@ -1,5 +1,8 @@ from dagster_shared.libraries import DagsterLibraryRegistry +from dagster_spark.components.spark_declarative_pipeline import ( + SparkDeclarativePipelineComponent as SparkDeclarativePipelineComponent, +) from dagster_spark.configs import define_spark_config as define_spark_config from dagster_spark.ops import create_spark_op as create_spark_op from dagster_spark.resources import spark_resource as spark_resource diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark/components/__init__.py new file mode 100644 index 0000000000000..5f0aa210e567f --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/__init__.py @@ -0,0 +1 @@ +# Dagster Spark components diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py new file mode 100644 index 0000000000000..2b3ff65c854b1 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py @@ -0,0 +1,21 @@ +"""Spark Declarative Pipeline (SDP) component for Dagster.""" + +from dagster_spark.components.spark_declarative_pipeline.component import ( + SparkDeclarativePipelineComponent as SparkDeclarativePipelineComponent, +) +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + DiscoveredDataset as DiscoveredDataset, + DryRunDatasetNode as DryRunDatasetNode, + DryRunReport as DryRunReport, + SparkPipelinesDryRunError as SparkPipelinesDryRunError, + SparkPipelineState as SparkPipelineState, + discover_datasets_fn as discover_datasets_fn, + discover_datasets_via_dry_run as discover_datasets_via_dry_run, + parse_dry_run_output_to_datasets as parse_dry_run_output_to_datasets, +) +from dagster_spark.components.spark_declarative_pipeline.resource import ( + SparkPipelinesResource as SparkPipelinesResource, +) +from dagster_spark.components.spark_declarative_pipeline.scaffolder import ( + SparkDeclarativePipelineScaffolder as SparkDeclarativePipelineScaffolder, +) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py new file mode 100644 index 0000000000000..263996658641f --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -0,0 +1,191 @@ +"""Spark Declarative Pipeline Dagster component (state-backed, resolvable). + +SparkDeclarativePipelineComponent discovers datasets via spark-pipelines dry-run (or +source_only), persists SparkPipelineState with Dagster serdes, and builds a multi_asset +that runs spark-pipelines run and yields MaterializeResults. +""" + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Annotated, Any, Literal, Optional + +import dagster as dg +from dagster import AssetSpec, Definitions, deserialize_value, serialize_value +from dagster.components.component.state_backed_component import StateBackedComponent +from dagster.components.core.context import ComponentLoadContext +from dagster.components.resolved.core_models import OpSpec +from dagster.components.resolved.model import Resolver +from dagster.components.scaffold.scaffold import scaffold_with +from dagster.components.utils.defs_state import ( + DefsStateConfig, + DefsStateConfigArgs, + ResolvedDefsStateConfig, +) + +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + DiscoveredDataset, + DiscoveryMode, + SparkPipelineState, + discover_datasets_fn, +) +from dagster_spark.components.spark_declarative_pipeline.resource import SparkPipelinesResource +from dagster_spark.components.spark_declarative_pipeline.scaffolder import ( + SparkDeclarativePipelineScaffolder, +) + +ExecutionMode = Literal["incremental", "full_refresh"] + + +def _resolve_spark_pipelines_resource(_context: Any, value: Any) -> SparkPipelinesResource: + """Resolve YAML/config to SparkPipelinesResource. Used by Resolver for spark_pipelines field.""" + if isinstance(value, SparkPipelinesResource): + return value + if value is None: + return SparkPipelinesResource() + return SparkPipelinesResource(**value) if isinstance(value, dict) else SparkPipelinesResource() + + +@scaffold_with(SparkDeclarativePipelineScaffolder) +@dataclass +class SparkDeclarativePipelineComponent(StateBackedComponent, dg.Resolvable): + """State-backed component for Spark Declarative Pipelines (SDP). + + Discovers datasets via spark-pipelines dry-run (or source_only), caches state, + and builds a multi_asset that runs spark-pipelines run and yields MaterializeResults. + """ + + pipeline_spec_path: str + defs_state: ResolvedDefsStateConfig = field( + default_factory=DefsStateConfigArgs.local_filesystem + ) + spark_pipelines: Annotated[ + SparkPipelinesResource, + Resolver(_resolve_spark_pipelines_resource), + ] = field(default_factory=SparkPipelinesResource) + op: Optional[OpSpec] = None + execution_mode: ExecutionMode = "incremental" + discovery_mode: DiscoveryMode = "dry_run_only" + asset_attributes_by_dataset: dict[str, dict[str, Any]] = field(default_factory=dict) + + @property + def defs_state_config(self) -> DefsStateConfig: + """Resolved DefsStateConfig for where to read/write component state.""" + return DefsStateConfig.from_args( + self.defs_state, + default_key=f"SparkDeclarativePipelineComponent[{self.pipeline_spec_path}]", + ) + + def write_state_to_path(self, state_path: Path) -> None: + """Run discovery and write SparkPipelineState (datasets) to state_path using Dagster serdes. + + Args: + state_path: File path to write serialized state (parent used as working_dir for dry-run). + """ + working_dir = state_path.parent + datasets = discover_datasets_fn( + pipeline_spec_path=self.pipeline_spec_path, + discovery_mode=self.discovery_mode, + working_dir=working_dir, + ) + state = SparkPipelineState( + datasets=datasets, + pipeline_spec_path=self.pipeline_spec_path, + ) + state_path.parent.mkdir(parents=True, exist_ok=True) + state_path.write_text(serialize_value(state)) + + def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: + """Build an AssetSpec for a discovered dataset. Override to customize key/metadata/group. + + Args: + dataset: Discovered dataset from state. + + Returns: + AssetSpec with key from dataset.name and optional description/metadata/group/tags + from asset_attributes_by_dataset. + """ + attrs = self.asset_attributes_by_dataset.get(dataset.name, {}) + return AssetSpec( + key=[dataset.name], + description=attrs.get("description") + or f"Spark Declarative Pipeline dataset: {dataset.name}", + metadata=attrs.get("metadata"), + group_name=attrs.get("group_name"), + tags=attrs.get("tags"), + ) + + def build_defs_from_state( + self, + context: ComponentLoadContext, + state_path: Optional[Path], + ) -> Definitions: + """Build Definitions with a multi_asset that runs spark_pipelines.run_and_observe. + + Deserializes SparkPipelineState from state_path, filters out temporary_view datasets + unless listed in asset_attributes_by_dataset, and builds one multi_asset with + can_subset=True that yields MaterializeResults via run_and_observe. + + Args: + context: Component load context (path, etc.). + state_path: Path to serialized state file; if None or missing, returns empty Definitions. + + Returns: + Definitions containing the multi_asset and spark_pipelines resource. + """ + if state_path is None or not state_path.exists(): + return Definitions() + + state = deserialize_value(state_path.read_text(), SparkPipelineState) + datasets = state.datasets + + # Filter out temporary_view datasets unless explicitly overridden in asset_attributes_by_dataset + def include_dataset(ds: DiscoveredDataset) -> bool: + if ds.name in self.asset_attributes_by_dataset: + return True + dataset_type = (ds.attributes or {}).get("dataset_type") if ds.attributes else None + if dataset_type == "temporary_view": + return False + return True + + datasets = [ds for ds in datasets if include_dataset(ds)] + if not datasets: + return Definitions() + + asset_specs = [self.get_asset_spec(ds) for ds in datasets] + op_spec = self.op or OpSpec() + pipeline_spec_path = state.pipeline_spec_path + # Resolve path relative to component path + if not Path(pipeline_spec_path).is_absolute(): + resolved_spec_path = (context.path / pipeline_spec_path).resolve() + else: + resolved_spec_path = Path(pipeline_spec_path) + working_dir = context.path + resource = self.spark_pipelines + execution_mode = self.execution_mode + + @dg.multi_asset( + specs=asset_specs, + can_subset=True, + name=op_spec.name or "spark_declarative_pipeline", + op_tags=op_spec.tags, + backfill_policy=op_spec.backfill_policy, + pool=op_spec.pool, + ) + def _spark_pipeline_asset(context: dg.AssetExecutionContext) -> Any: + keys = ( + list(context.selected_asset_keys) + if context.selected_asset_keys + else [s.key for s in asset_specs] + ) + yield from resource.run_and_observe( + context=context, + pipeline_spec_path=resolved_spec_path, + working_dir=working_dir, + execution_mode=execution_mode, + asset_keys=keys, + ) + + return Definitions( + assets=[_spark_pipeline_asset], + resources={"spark_pipelines": resource}, + ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py new file mode 100644 index 0000000000000..5bf8d320990f3 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -0,0 +1,293 @@ +"""Dataset discovery for Spark Declarative Pipelines via dry-run and optional fallbacks. + +This module runs ``spark-pipelines dry-run`` to discover datasets, parses JSON or +structured text output, and supports discovery_mode fallbacks (e.g. source_only). +State types (DiscoveredDataset, SparkPipelineState) are frozen dataclasses with +whitelist_for_serdes for Dagster serialize_value/deserialize_value compatibility. +""" + +import json +import re +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from dagster_shared import check +from dagster_shared.serdes import whitelist_for_serdes + +DiscoveryMode = Literal["dry_run_only", "dry_run_with_fallback", "source_only"] + + +class SparkPipelinesDryRunError(Exception): + """Raised when ``spark-pipelines dry-run`` or ``spark-pipelines run`` fails. + + Attributes: + message: Error description. + stderr: Captured stderr or combined stdout/stderr from the process. + returncode: Process exit code (non-zero on failure). + """ + + def __init__(self, message: str, stderr: str | None = None, returncode: int | None = None): + super().__init__(message) + self.stderr = stderr + self.returncode = returncode + + +@dataclass(frozen=True) +class DryRunDatasetNode: + """A single dataset node as reported by spark-pipelines dry-run. + + Attributes: + name: Dataset identifier. + raw: Raw dict from CLI output (e.g. JSON object for this node). + """ + + name: str + raw: dict[str, Any] + + +@dataclass(frozen=True) +class DryRunReport: + """Structured report produced by spark-pipelines dry-run (JSON or parsed text). + + Attributes: + datasets: List of dataset nodes from the report. + raw: Optional raw JSON dict when parsed from JSON; None for text fallback. + """ + + datasets: list[DryRunDatasetNode] + raw: dict[str, Any] | None = None + + +@whitelist_for_serdes +@dataclass(frozen=True) +class DiscoveredDataset: + """A dataset discovered for a Spark Declarative Pipeline (from dry-run or source). + + Used in cached state and for building AssetSpecs. Compatible with Dagster + serialize_value/deserialize_value when used inside SparkPipelineState. + + Attributes: + name: Dataset name (used as asset key component). + attributes: Arbitrary metadata (e.g. dataset_type, schema info). + """ + + name: str + attributes: dict[str, Any] + + +@whitelist_for_serdes +@dataclass(frozen=True) +class SparkPipelineState: + """Cached state for a Spark Declarative Pipeline (discovered datasets). + + Persisted via Dagster serialize_value/deserialize_value. Used by + SparkDeclarativePipelineComponent.write_state_to_path and build_defs_from_state. + + Attributes: + datasets: List of discovered datasets. + pipeline_spec_path: Path to the pipeline spec file (relative or absolute). + """ + + datasets: list[DiscoveredDataset] + pipeline_spec_path: str + + +def discover_datasets_via_dry_run( + pipeline_spec_path: str | Path, + working_dir: str | Path | None = None, + extra_args: list[str] | None = None, +) -> str: + """Run `spark-pipelines dry-run` and return raw stdout. + + Args: + pipeline_spec_path: Path to the pipeline spec file. + working_dir: Optional working directory for the subprocess. + extra_args: Optional extra CLI arguments. + + Returns: + Raw stdout from the command. + + Raises: + SparkPipelinesDryRunError: If the command fails or times out. + """ + path_str = str(pipeline_spec_path) + cmd = ["spark-pipelines", "dry-run", path_str] + if extra_args: + cmd.extend(extra_args) + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + cwd=working_dir, + timeout=300, + ) + + if result.returncode != 0: + raise SparkPipelinesDryRunError( + f"spark-pipelines dry-run failed with return code {result.returncode}", + stderr=result.stderr, + returncode=result.returncode, + ) + + return result.stdout + + +def extract_report(stdout: str) -> DryRunReport | None: + """Extract a DryRunReport from dry-run stdout (JSON or structured text). + + Tries JSON first (e.g. --output json), then a simple text fallback. + + Args: + stdout: Raw stdout string from spark-pipelines dry-run. + + Returns: + A DryRunReport if parsing succeeded, otherwise None. + """ + report = _extract_report_json(stdout) + if report is not None: + return report + return _extract_report_text(stdout) + + +def _extract_report_json(stdout: str) -> DryRunReport | None: + """Try to parse a JSON object or array from stdout and map to DryRunReport.""" + stripped = stdout.strip() + if not stripped: + return None + + # Try to find a JSON object or array (allow surrounding text) + for pattern in ( + r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", + r"\[\s*\{[^]]*\}\s*\]", + ): + match = re.search(pattern, stdout, re.DOTALL) + if match: + try: + data = json.loads(match.group(0)) + return _json_to_report(data) + except (json.JSONDecodeError, TypeError): + continue + + # Try parsing the whole output as JSON + try: + data = json.loads(stripped) + return _json_to_report(data) + except (json.JSONDecodeError, TypeError): + pass + + return None + + +def _json_to_report(data: Any) -> DryRunReport | None: + """Map a JSON structure to DryRunReport. Placeholder: adapt to real CLI output shape.""" + if isinstance(data, dict): + # Common keys: "datasets", "nodes", "sources", etc. + nodes = data.get("datasets") or data.get("nodes") or data.get("sources") or [] + elif isinstance(data, list): + nodes = data + else: + return None + + if not isinstance(nodes, list): + return None + + dataset_nodes: list[DryRunDatasetNode] = [] + for item in nodes: + if isinstance(item, dict): + name = item.get("name") or item.get("id") or item.get("dataset") or str(item) + if isinstance(name, dict): + name = name.get("name") or name.get("id") or str(name) + dataset_nodes.append(DryRunDatasetNode(name=str(name), raw=item)) + elif isinstance(item, str): + dataset_nodes.append(DryRunDatasetNode(name=item, raw={"name": item})) + + return DryRunReport(datasets=dataset_nodes, raw=data if isinstance(data, dict) else None) + + +def _extract_report_text(stdout: str) -> DryRunReport | None: + """Fallback: parse structured text lines (e.g. 'dataset: name' or bullet list).""" + datasets: list[DryRunDatasetNode] = [] + for raw_line in stdout.splitlines(): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + # Match "dataset: foo" or "- foo" or " - foo" + for pattern in (r"dataset:\s*(.+)", r"^[-*]\s*(.+)", r"^\d+\.\s*(.+)"): + match = re.match(pattern, line, re.IGNORECASE) + if match: + name = match.group(1).strip() + if name: + datasets.append(DryRunDatasetNode(name=name, raw={"name": name})) + break + + if not datasets: + return None + return DryRunReport(datasets=datasets, raw=None) + + +def parse_dry_run_output_to_datasets(stdout: str) -> list[DiscoveredDataset]: + """Parse dry-run stdout into a list of DiscoveredDataset. + + Uses extract_report for JSON and text fallback; maps each node to DiscoveredDataset. + + Args: + stdout: Raw stdout string from spark-pipelines dry-run. + + Returns: + List of DiscoveredDataset (empty if report could not be parsed). + """ + report = extract_report(stdout) + if report is None: + return [] + + return [ + DiscoveredDataset(name=node.name, attributes=dict(node.raw)) for node in report.datasets + ] + + +def discover_datasets_fn( + pipeline_spec_path: str | Path, + discovery_mode: DiscoveryMode, + working_dir: str | Path | None = None, + source_only_datasets: list[DiscoveredDataset] | None = None, +) -> list[DiscoveredDataset]: + """Discover datasets for a Spark Declarative Pipeline based on discovery_mode. + + - dry_run_only: Run spark-pipelines dry-run and parse output; raise if it fails. + - dry_run_with_fallback: Same as above, but on failure or empty result use source_only. + - source_only: Do not run dry-run; use source_only_datasets if provided, else []. + + Args: + pipeline_spec_path: Path to the pipeline spec. + discovery_mode: One of dry_run_only, dry_run_with_fallback, source_only. + working_dir: Optional working directory for dry-run. + source_only_datasets: Optional list used when mode is source_only or as fallback. + + Returns: + List of DiscoveredDataset. + """ + check.inst_param(discovery_mode, "discovery_mode", str) + source_list = source_only_datasets or [] + + if discovery_mode == "source_only": + return list(source_list) + + try: + raw = discover_datasets_via_dry_run( + pipeline_spec_path, + working_dir=working_dir, + ) + datasets = parse_dry_run_output_to_datasets(raw) + except SparkPipelinesDryRunError: + if discovery_mode == "dry_run_only": + raise + return list(source_list) + + if discovery_mode == "dry_run_with_fallback" and not datasets: + return list(source_list) + + return datasets diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py new file mode 100644 index 0000000000000..73a6bde70c0b7 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -0,0 +1,139 @@ +"""Resource for running Spark Declarative Pipelines and discovering datasets. + +SparkPipelinesResource provides discover_datasets (via dry-run or source_only) and +run_and_observe (run spark-pipelines with log streaming and MaterializeResult yields). +""" + +import subprocess +from collections.abc import Iterator +from pathlib import Path +from typing import Any, Literal + +from dagster import AssetKey, ConfigurableResource, MaterializeResult + +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + DiscoveredDataset, + DiscoveryMode, + SparkPipelinesDryRunError, + discover_datasets_fn, + parse_dry_run_output_to_datasets, +) + +ExecutionMode = Literal["incremental", "full_refresh"] + + +class SparkPipelinesResource(ConfigurableResource): + """Dagster resource for Spark Declarative Pipelines: discovery and run. + + Use discover_datasets to get datasets from spark-pipelines dry-run (or source_only). + Use run_and_observe inside an asset to run the pipeline and yield MaterializeResults. + """ + + def discover_datasets( + self, + pipeline_spec_path: str | Path, + discovery_mode: DiscoveryMode = "dry_run_only", + working_dir: str | Path | None = None, + source_only_datasets: list[DiscoveredDataset] | None = None, + ) -> list[DiscoveredDataset]: + """Discover datasets for the given pipeline spec using the configured discovery_mode. + + Args: + pipeline_spec_path: Path to the pipeline spec file (YAML). + discovery_mode: One of dry_run_only, dry_run_with_fallback, source_only. + working_dir: Optional working directory for the dry-run subprocess. + source_only_datasets: Optional list used when mode is source_only or as fallback. + + Returns: + List of DiscoveredDataset. + """ + return discover_datasets_fn( + pipeline_spec_path=pipeline_spec_path, + discovery_mode=discovery_mode, + working_dir=working_dir, + source_only_datasets=source_only_datasets, + ) + + def run_and_observe( + self, + context: Any, + pipeline_spec_path: str | Path, + working_dir: str | Path | None = None, + execution_mode: ExecutionMode = "incremental", + extra_args: list[str] | None = None, + asset_keys: list[AssetKey] | None = None, + ) -> Iterator[MaterializeResult]: + """Run spark-pipelines run with log streaming; yield MaterializeResult per asset on success. + + Uses Popen to stream stdout/stderr line-by-line and logs each line via context.log.info. + Passes --full-refresh or --refresh based on execution_mode, then optional comma-separated + dataset list from asset_keys. Only yields MaterializeResults if the process exits with + returncode == 0; otherwise raises SparkPipelinesDryRunError with the captured log. + + Args: + context: Asset execution context (used for context.log.info). + pipeline_spec_path: Path to the pipeline spec file (YAML). + working_dir: Optional working directory for the subprocess. + execution_mode: "incremental" (--refresh) or "full_refresh" (--full-refresh). + extra_args: Optional extra CLI arguments appended to the command. + asset_keys: Optional list of asset keys to materialize (passed as dataset list). + + Yields: + MaterializeResult for each materialized asset on success. + + Raises: + SparkPipelinesDryRunError: If spark-pipelines run exits with non-zero return code. + """ + path_str = str(pipeline_spec_path) + cmd = ["spark-pipelines", "run", path_str] + if execution_mode == "full_refresh": + cmd.append("--full-refresh") + else: + cmd.append("--refresh") + if asset_keys: + datasets_str = ",".join(k.to_user_string() for k in asset_keys) + if datasets_str: + cmd.append(datasets_str) + if extra_args: + cmd.extend(extra_args) + + cwd = str(working_dir) if working_dir else None + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + cwd=cwd, + bufsize=1, + ) + log_lines: list[str] = [] + if process.stdout: + for raw_line in process.stdout: + line = raw_line.rstrip("\n\r") + if line: + log_lines.append(line) + if context is not None and hasattr(context, "log"): + context.log.info(line) + process.wait() + returncode = process.returncode + + if returncode != 0: + captured = "\n".join(log_lines) if log_lines else "(no output)" + raise SparkPipelinesDryRunError( + f"spark-pipelines run failed with return code {returncode}", + stderr=captured, + returncode=returncode, + ) + + if asset_keys: + for k in asset_keys: + yield MaterializeResult(asset_key=k) + return + + # Fallback: parse stdout for reported materialized keys (use captured log as stdout) + stdout_text = "\n".join(log_lines) + datasets = parse_dry_run_output_to_datasets(stdout_text) + for ds in datasets: + yield MaterializeResult(asset_key=AssetKey([ds.name])) + if not datasets: + yield MaterializeResult(asset_key=AssetKey(["spark_pipeline"])) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py new file mode 100644 index 0000000000000..57c8df3802eb3 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py @@ -0,0 +1,18 @@ +"""Scaffolder for Spark Declarative Pipeline component.""" + +from dagster.components.component.component_scaffolder import Scaffolder +from dagster.components.component_scaffolding import scaffold_component +from dagster.components.scaffold.scaffold import ScaffoldRequest + + +class SparkDeclarativePipelineScaffolder(Scaffolder): + """Scaffolds a Spark Declarative Pipeline component defs.yaml and pipeline spec path.""" + + def scaffold(self, request: ScaffoldRequest) -> None: + scaffold_component( + request, + { + "pipeline_spec_path": "pipeline.yaml", + "discovery_mode": "dry_run_only", + }, + ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/__init__.py new file mode 100644 index 0000000000000..95fa19f231c93 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/__init__.py @@ -0,0 +1 @@ +# Dagster Spark component tests diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/__init__.py new file mode 100644 index 0000000000000..37c4c7eb0ae7b --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/__init__.py @@ -0,0 +1 @@ +# Spark Declarative Pipeline component tests diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py new file mode 100644 index 0000000000000..1ece804d0c612 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py @@ -0,0 +1,122 @@ +"""Tests for SparkDeclarativePipelineComponent (build_defs_from_state, temporary_view filtering).""" + +import tempfile +from pathlib import Path +from unittest.mock import MagicMock + +import dagster as dg +from dagster_spark.components.spark_declarative_pipeline import ( + DiscoveredDataset, + SparkDeclarativePipelineComponent, + SparkPipelineState, +) + + +def test_build_defs_from_state_returns_valid_definitions_with_multi_asset() -> None: + """build_defs_from_state returns a valid Definitions object containing a multi_asset.""" + component = SparkDeclarativePipelineComponent( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + ) + datasets = [ + DiscoveredDataset(name="table_a", attributes={"dataset_type": "table"}), + DiscoveredDataset(name="table_b", attributes={"dataset_type": "table"}), + ] + state = SparkPipelineState( + datasets=datasets, + pipeline_spec_path="pipeline.yaml", + ) + with tempfile.TemporaryDirectory() as tmpdir: + state_path = Path(tmpdir) / "state" + state_path.write_text(dg.serialize_value(state)) + context = MagicMock() + context.path = Path(tmpdir) + context.project_root = Path(tmpdir) + + defs = component.build_defs_from_state(context, state_path) + + assert defs is not None + assert isinstance(defs, dg.Definitions) + all_assets = list(defs.get_all_asset_specs()) + assert len(all_assets) == 2 + keys = {a.key.to_user_string() for a in all_assets} + assert "table_a" in keys + assert "table_b" in keys + + +def test_build_defs_from_state_filters_temporary_views() -> None: + """Temporary view datasets are filtered out unless overridden in asset_attributes_by_dataset.""" + component = SparkDeclarativePipelineComponent( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + asset_attributes_by_dataset={}, # no overrides + ) + datasets = [ + DiscoveredDataset(name="table_a", attributes={"dataset_type": "table"}), + DiscoveredDataset( + name="temp_view_x", + attributes={"dataset_type": "temporary_view"}, + ), + ] + state = SparkPipelineState( + datasets=datasets, + pipeline_spec_path="pipeline.yaml", + ) + with tempfile.TemporaryDirectory() as tmpdir: + state_path = Path(tmpdir) / "state" + state_path.write_text(dg.serialize_value(state)) + context = MagicMock() + context.path = Path(tmpdir) + context.project_root = Path(tmpdir) + + defs = component.build_defs_from_state(context, state_path) + + all_assets = list(defs.get_all_asset_specs()) + assert len(all_assets) == 1 + assert all_assets[0].key.to_user_string() == "table_a" + + +def test_build_defs_from_state_includes_temporary_view_when_overridden() -> None: + """A temporary_view is included when it has an entry in asset_attributes_by_dataset.""" + component = SparkDeclarativePipelineComponent( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + asset_attributes_by_dataset={"temp_view_x": {"description": "Included view"}}, + ) + datasets = [ + DiscoveredDataset(name="table_a", attributes={"dataset_type": "table"}), + DiscoveredDataset( + name="temp_view_x", + attributes={"dataset_type": "temporary_view"}, + ), + ] + state = SparkPipelineState( + datasets=datasets, + pipeline_spec_path="pipeline.yaml", + ) + with tempfile.TemporaryDirectory() as tmpdir: + state_path = Path(tmpdir) / "state" + state_path.write_text(dg.serialize_value(state)) + context = MagicMock() + context.path = Path(tmpdir) + context.project_root = Path(tmpdir) + + defs = component.build_defs_from_state(context, state_path) + + all_assets = list(defs.get_all_asset_specs()) + assert len(all_assets) == 2 + keys = {a.key.to_user_string() for a in all_assets} + assert "table_a" in keys + assert "temp_view_x" in keys + + +def test_build_defs_from_state_returns_empty_when_no_state_path() -> None: + """build_defs_from_state returns empty Definitions when state_path is None.""" + component = SparkDeclarativePipelineComponent( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + ) + context = MagicMock() + defs = component.build_defs_from_state(context, None) + assert isinstance(defs, dg.Definitions) + assert len(list(defs.get_all_asset_specs())) == 0 diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py new file mode 100644 index 0000000000000..df725d32b06a2 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -0,0 +1,114 @@ +"""Tests for Spark Declarative Pipeline discovery (dry-run parsing and discover_datasets_fn).""" + +import json +from unittest.mock import MagicMock, patch + +import pytest +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + DiscoveredDataset, + SparkPipelinesDryRunError, + discover_datasets_fn, + discover_datasets_via_dry_run, + extract_report, + parse_dry_run_output_to_datasets, +) + + +def test_parse_dry_run_output_to_datasets_parses_json_report() -> None: + """discover_datasets_fn / parse correctly parses a mock JSON report into DiscoveredDataset records.""" + mock_report = { + "datasets": [ + {"name": "dataset_a", "type": "table"}, + {"name": "dataset_b", "id": "ds_b"}, + ], + } + stdout = json.dumps(mock_report) + datasets = parse_dry_run_output_to_datasets(stdout) + assert len(datasets) == 2 + assert datasets[0].name == "dataset_a" + assert datasets[0].attributes.get("type") == "table" + assert datasets[1].name == "dataset_b" + + +def test_extract_report_returns_dry_run_report() -> None: + """extract_report returns a DryRunReport from valid JSON stdout.""" + stdout = json.dumps({"datasets": [{"name": "foo"}]}) + report = extract_report(stdout) + assert report is not None + assert len(report.datasets) == 1 + assert report.datasets[0].name == "foo" + + +def test_discover_datasets_fn_dry_run_only_raises_on_failure() -> None: + """SparkPipelinesDryRunError is raised when dry-run fails in dry_run_only mode.""" + with patch( + "dagster_spark.components.spark_declarative_pipeline.discovery.subprocess.run" + ) as mock_run: + mock_run.return_value = MagicMock(returncode=1, stderr="error", stdout="") + with pytest.raises(SparkPipelinesDryRunError) as exc_info: + discover_datasets_fn( + pipeline_spec_path="/path/to/spec.yaml", + discovery_mode="dry_run_only", + ) + assert exc_info.value.returncode == 1 + assert "error" in (exc_info.value.stderr or "") + + +def test_discover_datasets_fn_dry_run_with_fallback_uses_source_on_failure() -> None: + """In dry_run_with_fallback mode, source_only_datasets are returned when dry-run fails.""" + fallback = [DiscoveredDataset(name="fallback_ds", attributes={})] + with patch( + "dagster_spark.components.spark_declarative_pipeline.discovery.subprocess.run" + ) as mock_run: + mock_run.return_value = MagicMock(returncode=1, stderr="err", stdout="") + result = discover_datasets_fn( + pipeline_spec_path="/path/to/spec.yaml", + discovery_mode="dry_run_with_fallback", + source_only_datasets=fallback, + ) + assert result == fallback + + +def test_discover_datasets_fn_source_only_returns_source_list() -> None: + """In source_only mode, discover_datasets_fn returns source_only_datasets without running dry-run.""" + source = [ + DiscoveredDataset(name="a", attributes={}), + DiscoveredDataset(name="b", attributes={}), + ] + result = discover_datasets_fn( + pipeline_spec_path="/any/path", + discovery_mode="source_only", + source_only_datasets=source, + ) + assert result == source + + +def test_discover_datasets_via_dry_run_raises_on_nonzero_exit() -> None: + """discover_datasets_via_dry_run raises SparkPipelinesDryRunError when returncode != 0.""" + with patch( + "dagster_spark.components.spark_declarative_pipeline.discovery.subprocess.run" + ) as mock_run: + mock_run.return_value = MagicMock( + returncode=2, + stderr="stderr output", + stdout="", + ) + with pytest.raises(SparkPipelinesDryRunError) as exc_info: + discover_datasets_via_dry_run("/path/to/spec.yaml") + assert exc_info.value.returncode == 2 + assert exc_info.value.stderr == "stderr output" + + +def test_discover_datasets_via_dry_run_returns_stdout_on_success() -> None: + """discover_datasets_via_dry_run returns stdout when subprocess succeeds.""" + with patch( + "dagster_spark.components.spark_declarative_pipeline.discovery.subprocess.run" + ) as mock_run: + mock_run.return_value = MagicMock( + returncode=0, + stderr="", + stdout='{"datasets":[{"name":"x"}]}', + ) + out = discover_datasets_via_dry_run("/path/to/spec.yaml") + assert "datasets" in out + assert "x" in out diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py new file mode 100644 index 0000000000000..0738944d84fcb --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -0,0 +1,116 @@ +"""Tests for SparkPipelinesResource run_and_observe (log streaming and MaterializeResult yields).""" + +from unittest.mock import MagicMock, patch + +import pytest +from dagster import AssetKey +from dagster._core.definitions.result import MaterializeResult +from dagster_spark.components.spark_declarative_pipeline.discovery import SparkPipelinesDryRunError +from dagster_spark.components.spark_declarative_pipeline.resource import SparkPipelinesResource + + +def test_run_and_observe_yields_materialize_results_for_selected_datasets() -> None: + """run_and_observe yields correct MaterializeResults for the selected asset keys on success.""" + mock_context = MagicMock() + asset_keys = [ + AssetKey(["dataset_a"]), + AssetKey(["dataset_b"]), + ] + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter(["line1\n", "line2\n"]) + proc.wait.return_value = 0 + proc.returncode = 0 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + results = list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/to/spec.yaml", + execution_mode="incremental", + asset_keys=asset_keys, + ) + ) + + assert len(results) == 2 + assert all(isinstance(r, MaterializeResult) for r in results) + assert results[0].asset_key == AssetKey(["dataset_a"]) + assert results[1].asset_key == AssetKey(["dataset_b"]) + + +def test_run_and_observe_streams_logs_via_context() -> None: + """run_and_observe streams stdout line-by-line and logs each line with context.log.info.""" + mock_context = MagicMock() + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter(["log line 1\n", "log line 2\n"]) + proc.wait.return_value = 0 + proc.returncode = 0 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], + ) + ) + + assert mock_context.log.info.call_count >= 2 + calls = [str(c) for c in mock_context.log.info.call_args_list] + assert any("log line 1" in c for c in calls) + assert any("log line 2" in c for c in calls) + + +def test_run_and_observe_raises_with_captured_log_on_nonzero_exit() -> None: + """run_and_observe raises SparkPipelinesDryRunError with captured log when process exits non-zero.""" + mock_context = MagicMock() + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter(["error line 1\n", "error line 2\n"]) + proc.wait.return_value = 1 + proc.returncode = 1 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + with pytest.raises(SparkPipelinesDryRunError) as exc_info: + list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], + ) + ) + assert exc_info.value.returncode == 1 + assert "error line" in (exc_info.value.stderr or "") + + +def test_run_and_observe_only_yields_on_success() -> None: + """MaterializeResults are only yielded when process returncode is 0.""" + mock_context = MagicMock() + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter([]) + proc.wait.return_value = 1 + proc.returncode = 1 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + with pytest.raises(SparkPipelinesDryRunError): + list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], + ) + ) diff --git a/python_modules/libraries/dagster-spark/setup.py b/python_modules/libraries/dagster-spark/setup.py index 4ec89e5853531..c6fb7b6b1f28b 100644 --- a/python_modules/libraries/dagster-spark/setup.py +++ b/python_modules/libraries/dagster-spark/setup.py @@ -36,4 +36,9 @@ def get_version() -> str: python_requires=">=3.10,<3.15", install_requires=[f"dagster{pin}"], zip_safe=False, + entry_points={ + "dagster_dg_cli.registry_modules": [ + "dagster_spark = dagster_spark", + ], + }, ) From 8b0219de6c46aa7a0e03f02d29160ee1c58d3c0e Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 19:09:15 +0200 Subject: [PATCH 02/17] doc From ddcd6a1cdf713dd8f9e0f182700a249d68120d44 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 20:12:28 +0200 Subject: [PATCH 03/17] improve SDP dataset discovery and resource injection --- .../spark_declarative_pipeline/__init__.py | 1 + .../spark_declarative_pipeline/component.py | 12 +- .../spark_declarative_pipeline/discovery.py | 223 ++++++++++++++++-- .../spark_declarative_pipeline/resource.py | 13 +- .../spark_declarative_pipeline/scaffolder.py | 2 +- .../test_discovery.py | 16 ++ .../test_resource.py | 27 +++ 7 files changed, 256 insertions(+), 38 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py index 2b3ff65c854b1..fde5cc356f361 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py @@ -7,6 +7,7 @@ DiscoveredDataset as DiscoveredDataset, DryRunDatasetNode as DryRunDatasetNode, DryRunReport as DryRunReport, + DuplicateDatasetNamesError as DuplicateDatasetNamesError, SparkPipelinesDryRunError as SparkPipelinesDryRunError, SparkPipelineState as SparkPipelineState, discover_datasets_fn as discover_datasets_fn, diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 263996658641f..2218819db9d24 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -106,7 +106,7 @@ def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: """ attrs = self.asset_attributes_by_dataset.get(dataset.name, {}) return AssetSpec( - key=[dataset.name], + key=dataset.name.split("."), description=attrs.get("description") or f"Spark Declarative Pipeline dataset: {dataset.name}", metadata=attrs.get("metadata"), @@ -160,7 +160,6 @@ def include_dataset(ds: DiscoveredDataset) -> bool: else: resolved_spec_path = Path(pipeline_spec_path) working_dir = context.path - resource = self.spark_pipelines execution_mode = self.execution_mode @dg.multi_asset( @@ -171,13 +170,16 @@ def include_dataset(ds: DiscoveredDataset) -> bool: backfill_policy=op_spec.backfill_policy, pool=op_spec.pool, ) - def _spark_pipeline_asset(context: dg.AssetExecutionContext) -> Any: + def _spark_pipeline_asset( + context: dg.AssetExecutionContext, + spark_pipelines: SparkPipelinesResource, + ) -> Any: keys = ( list(context.selected_asset_keys) if context.selected_asset_keys else [s.key for s in asset_specs] ) - yield from resource.run_and_observe( + yield from spark_pipelines.run_and_observe( context=context, pipeline_spec_path=resolved_spec_path, working_dir=working_dir, @@ -187,5 +189,5 @@ def _spark_pipeline_asset(context: dg.AssetExecutionContext) -> Any: return Definitions( assets=[_spark_pipeline_asset], - resources={"spark_pipelines": resource}, + resources={"spark_pipelines": self.spark_pipelines}, ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index 5bf8d320990f3..b31c91b4c282f 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -13,8 +13,9 @@ from pathlib import Path from typing import Any, Literal -from dagster_shared import check -from dagster_shared.serdes import whitelist_for_serdes +from dagster import get_dagster_logger +from dagster_shared import check # type: ignore[reportMissingImports] +from dagster_shared.serdes import whitelist_for_serdes # type: ignore[reportMissingImports] DiscoveryMode = Literal["dry_run_only", "dry_run_with_fallback", "source_only"] @@ -34,6 +35,18 @@ def __init__(self, message: str, stderr: str | None = None, returncode: int | No self.returncode = returncode +class DuplicateDatasetNamesError(ValueError): + """Raised when discovered datasets contain duplicate names after normalization. + + Attributes: + duplicate_names: List of dataset names that appear more than once. + """ + + def __init__(self, message: str, duplicate_names: list[str]): + super().__init__(message) + self.duplicate_names = duplicate_names + + @dataclass(frozen=True) class DryRunDatasetNode: """A single dataset node as reported by spark-pipelines dry-run. @@ -153,32 +166,60 @@ def extract_report(stdout: str) -> DryRunReport | None: return _extract_report_text(stdout) +# Optional log-level prefix before JSON in stdout (e.g. "INFO: " or "WARN: ") +_LOG_PREFIX = re.compile(r"^\s*(?:\[?\w+\]?:\s*)?", re.IGNORECASE) + + def _extract_report_json(stdout: str) -> DryRunReport | None: - """Try to parse a JSON object or array from stdout and map to DryRunReport.""" + """Try to parse a JSON object or array from stdout and map to DryRunReport. + + Iterates through lines and tries json.loads on stripped content (with optional + log-style prefix stripped) to avoid capturing arbitrary log text between + first '{' and last '}' which can crash json.loads. + """ stripped = stdout.strip() if not stripped: return None - # Try to find a JSON object or array (allow surrounding text) - for pattern in ( - r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", - r"\[\s*\{[^]]*\}\s*\]", - ): - match = re.search(pattern, stdout, re.DOTALL) - if match: - try: - data = json.loads(match.group(0)) - return _json_to_report(data) - except (json.JSONDecodeError, TypeError): - continue - - # Try parsing the whole output as JSON + # Try parsing the whole output as JSON first try: data = json.loads(stripped) - return _json_to_report(data) + report = _json_to_report(data) + if report is not None: + return report except (json.JSONDecodeError, TypeError): pass + # Line-by-line: try each line (and optional log prefix strip) to find valid JSON + for raw_line in stdout.splitlines(): + line = raw_line.strip() + if not line or ("{" not in line and "[" not in line): + continue + for candidate in (line, _LOG_PREFIX.sub("", line)): + if not candidate: + continue + try: + data = json.loads(candidate) + report = _json_to_report(data) + if report is not None: + return report + except (json.JSONDecodeError, TypeError): + pass + # Single line may have leading/trailing text; try substring from first { or [ + for start_char, end_char in (("{", "}"), ("[", "]")): + i = line.find(start_char) + if i == -1: + continue + j = line.rfind(end_char) + if j != -1 and j > i: + try: + data = json.loads(line[i : j + 1]) + report = _json_to_report(data) + if report is not None: + return report + except (json.JSONDecodeError, TypeError): + pass + return None @@ -229,6 +270,100 @@ def _extract_report_text(stdout: str) -> DryRunReport | None: return DryRunReport(datasets=datasets, raw=None) +# Regex for source-only discovery: Python decorator (optional dp. prefix, optional args) + function name +_PY_DATASET_DECORATOR = re.compile( + r"@(?:dp\.)?(materialized_view|table|streaming_table|temporary_view)(?:\([^)]*\))?\s*\n\s*def\s+(\w+)\s*\(", + re.MULTILINE, +) +# Regex for source-only discovery: SQL CREATE statement (capture supports catalog.schema.table) +_SQL_CREATE_DATASET = re.compile( + r"CREATE\s+(?:MATERIALIZED\s+VIEW|STREAMING\s+TABLE|TABLE|VIEW)\s+(?:IF\s+NOT\s+EXISTS\s+)?([\w.]+)", + re.IGNORECASE, +) + + +def discover_datasets_from_sources(pipeline_spec_path: Path) -> list[DiscoveredDataset]: + """Discover datasets by scanning source files under the pipeline spec directory. + + Best-effort fallback only: scans the directory (and subdirectories) of + pipeline_spec_path for .py and .sql files using regex. Does not evaluate + dynamic Python arguments (e.g. @dp.table(name="tbl")) or resolve complex + imports (e.g. @table). Prefer dry-run discovery when available. + + For Python: matches @(dp.)?(materialized_view|table|streaming_table|temporary_view) + with optional parentheses and arguments, then def my_dataset_name(. + For SQL: CREATE (MATERIALIZED VIEW|...) [IF NOT EXISTS] my_dataset_name. + + Args: + pipeline_spec_path: Path to the pipeline spec file (its parent is scanned). + + Returns: + List of DiscoveredDataset (deduplicated by name). Returns [] on any failure. + """ + logger = get_dagster_logger() + try: + root = Path(pipeline_spec_path).resolve() + if root.suffix: + root = root.parent + except Exception as e: + logger.warning( + "Failed to resolve pipeline spec path %s: %s", + pipeline_spec_path, + e, + exc_info=True, + ) + return [] + + seen: set[str] = set() + result: list[DiscoveredDataset] = [] + + try: + for ext in ("*.py", "*.sql"): + for path in root.rglob(ext): + try: + text = path.read_text(encoding="utf-8", errors="replace") + except Exception as e: + logger.warning( + "Failed to read source file %s: %s", + path, + e, + exc_info=True, + ) + continue + if ext == "*.py": + for m in _PY_DATASET_DECORATOR.finditer(text): + name = m.group(2) + if name not in seen: + seen.add(name) + result.append( + DiscoveredDataset( + name=name, + attributes={"dataset_type": m.group(1), "source": str(path)}, + ) + ) + else: + for m in _SQL_CREATE_DATASET.finditer(text): + name = m.group(1) + if name not in seen: + seen.add(name) + result.append( + DiscoveredDataset( + name=name, + attributes={"source": str(path)}, + ) + ) + except Exception as e: + logger.warning( + "Failed to discover datasets from sources under %s: %s", + root, + e, + exc_info=True, + ) + return [] + + return result + + def parse_dry_run_output_to_datasets(stdout: str) -> list[DiscoveredDataset]: """Parse dry-run stdout into a list of DiscoveredDataset. @@ -249,6 +384,31 @@ def parse_dry_run_output_to_datasets(stdout: str) -> list[DiscoveredDataset]: ] +def _validate_no_duplicate_dataset_names(datasets: list[DiscoveredDataset]) -> None: + """Raise DuplicateDatasetNamesError if any dataset names are duplicated after normalization. + + Normalization: strip and lowercased name for comparison. + """ + from collections import Counter + + normalized = [ds.name.strip().lower() for ds in datasets] + counts = Counter(normalized) + duplicate_normalized = [k for k, c in counts.items() if c > 1] + if not duplicate_normalized: + return + # Report original names for duplicated normalized keys (first occurrence each) + seen_orig: dict[str, str] = {} + for ds in datasets: + key = ds.name.strip().lower() + if key not in seen_orig: + seen_orig[key] = ds.name + duplicate_names = [seen_orig[k] for k in duplicate_normalized] + raise DuplicateDatasetNamesError( + f"Duplicate dataset names after normalization: {duplicate_names}", + duplicate_names=duplicate_names, + ) + + def discover_datasets_fn( pipeline_spec_path: str | Path, discovery_mode: DiscoveryMode, @@ -259,7 +419,9 @@ def discover_datasets_fn( - dry_run_only: Run spark-pipelines dry-run and parse output; raise if it fails. - dry_run_with_fallback: Same as above, but on failure or empty result use source_only. - - source_only: Do not run dry-run; use source_only_datasets if provided, else []. + - source_only: Do not run dry-run; use source_only_datasets if provided, else discover from sources. + + Duplicate dataset names (after normalization: strip + lowercase) are a hard error. Args: pipeline_spec_path: Path to the pipeline spec. @@ -269,12 +431,21 @@ def discover_datasets_fn( Returns: List of DiscoveredDataset. + + Raises: + DuplicateDatasetNamesError: If the discovered list contains duplicate names. """ check.inst_param(discovery_mode, "discovery_mode", str) - source_list = source_only_datasets or [] + source_list = source_only_datasets + + def _return(datasets: list[DiscoveredDataset]) -> list[DiscoveredDataset]: + _validate_no_duplicate_dataset_names(datasets) + return list(datasets) if discovery_mode == "source_only": - return list(source_list) + if source_list is not None: + return _return(source_list) + return _return(discover_datasets_from_sources(Path(pipeline_spec_path))) try: raw = discover_datasets_via_dry_run( @@ -285,9 +456,13 @@ def discover_datasets_fn( except SparkPipelinesDryRunError: if discovery_mode == "dry_run_only": raise - return list(source_list) + if source_list is not None: + return _return(source_list) + return _return(discover_datasets_from_sources(Path(pipeline_spec_path))) if discovery_mode == "dry_run_with_fallback" and not datasets: - return list(source_list) + if source_list is not None: + return _return(source_list) + return _return(discover_datasets_from_sources(Path(pipeline_spec_path))) - return datasets + return _return(datasets) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index 73a6bde70c0b7..3ddc6147a43a6 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -16,7 +16,6 @@ DiscoveryMode, SparkPipelinesDryRunError, discover_datasets_fn, - parse_dry_run_output_to_datasets, ) ExecutionMode = Literal["incremental", "full_refresh"] @@ -130,10 +129,8 @@ def run_and_observe( yield MaterializeResult(asset_key=k) return - # Fallback: parse stdout for reported materialized keys (use captured log as stdout) - stdout_text = "\n".join(log_lines) - datasets = parse_dry_run_output_to_datasets(stdout_text) - for ds in datasets: - yield MaterializeResult(asset_key=AssetKey([ds.name])) - if not datasets: - yield MaterializeResult(asset_key=AssetKey(["spark_pipeline"])) + # No specific asset keys requested; complete gracefully without yielding (avoids UnexpectedAssetMaterializationError) + if context is not None and hasattr(context, "log"): + context.log.info( + "No specific asset keys requested; spark-pipelines run completed successfully." + ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py index 57c8df3802eb3..bd51422c5eda1 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py @@ -12,7 +12,7 @@ def scaffold(self, request: ScaffoldRequest) -> None: scaffold_component( request, { - "pipeline_spec_path": "pipeline.yaml", + "pipeline_spec_path": "spark-pipeline.yml", "discovery_mode": "dry_run_only", }, ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py index df725d32b06a2..4728aee23f3e1 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -6,6 +6,7 @@ import pytest from dagster_spark.components.spark_declarative_pipeline.discovery import ( DiscoveredDataset, + DuplicateDatasetNamesError, SparkPipelinesDryRunError, discover_datasets_fn, discover_datasets_via_dry_run, @@ -99,6 +100,21 @@ def test_discover_datasets_via_dry_run_raises_on_nonzero_exit() -> None: assert exc_info.value.stderr == "stderr output" +def test_discover_datasets_fn_raises_on_duplicate_dataset_names() -> None: + """Duplicate dataset names (after normalization) raise DuplicateDatasetNamesError.""" + source = [ + DiscoveredDataset(name="Foo", attributes={}), + DiscoveredDataset(name="foo", attributes={}), + ] + with pytest.raises(DuplicateDatasetNamesError) as exc_info: + discover_datasets_fn( + pipeline_spec_path="/any/path", + discovery_mode="source_only", + source_only_datasets=source, + ) + assert "foo" in exc_info.value.duplicate_names or "Foo" in exc_info.value.duplicate_names + + def test_discover_datasets_via_dry_run_returns_stdout_on_success() -> None: """discover_datasets_via_dry_run returns stdout when subprocess succeeds.""" with patch( diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py index 0738944d84fcb..3156755b5b479 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -93,6 +93,33 @@ def test_run_and_observe_raises_with_captured_log_on_nonzero_exit() -> None: assert "error line" in (exc_info.value.stderr or "") +def test_run_and_observe_yields_nothing_when_no_asset_keys() -> None: + """When asset_keys is empty, run_and_observe yields nothing and completes gracefully (logs only).""" + mock_context = MagicMock() + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter(["log line\n"]) + proc.wait.return_value = 0 + proc.returncode = 0 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + results = list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=None, + ) + ) + + assert len(results) == 0 + mock_context.log.info.assert_any_call( + "No specific asset keys requested; spark-pipelines run completed successfully." + ) + + def test_run_and_observe_only_yields_on_success() -> None: """MaterializeResults are only yielded when process returncode is 0.""" mock_context = MagicMock() From 17ea15b19f6cdb87ff4c58beede94b78506c1d4b Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 20:31:43 +0200 Subject: [PATCH 04/17] add asset lineage and configurable resource --- .../spark_declarative_pipeline/component.py | 21 ++++++- .../spark_declarative_pipeline/discovery.py | 13 +++- .../spark_declarative_pipeline/resource.py | 21 ++++++- .../test_component.py | 20 +++++++ .../test_discovery.py | 21 +++++++ .../test_resource.py | 60 +++++++++++++++++++ 6 files changed, 149 insertions(+), 7 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 2218819db9d24..3eea08f788827 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -10,7 +10,7 @@ from typing import Annotated, Any, Literal, Optional import dagster as dg -from dagster import AssetSpec, Definitions, deserialize_value, serialize_value +from dagster import AssetKey, AssetSpec, Definitions, deserialize_value, serialize_value from dagster.components.component.state_backed_component import StateBackedComponent from dagster.components.core.context import ComponentLoadContext from dagster.components.resolved.core_models import OpSpec @@ -82,10 +82,13 @@ def write_state_to_path(self, state_path: Path) -> None: state_path: File path to write serialized state (parent used as working_dir for dry-run). """ working_dir = state_path.parent + resource = self.spark_pipelines datasets = discover_datasets_fn( pipeline_spec_path=self.pipeline_spec_path, discovery_mode=self.discovery_mode, working_dir=working_dir, + spark_pipelines_cmd=resource.spark_pipelines_cmd, + dry_run_extra_args=resource.dry_run_extra_args, ) state = SparkPipelineState( datasets=datasets, @@ -101,12 +104,24 @@ def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: dataset: Discovered dataset from state. Returns: - AssetSpec with key from dataset.name and optional description/metadata/group/tags - from asset_attributes_by_dataset. + AssetSpec with key from dataset.name (split by '.' for multi-segment), optional deps + from attributes (dependencies, upstream_dataset_names, or deps), and optional + description/metadata/group/tags from asset_attributes_by_dataset. """ attrs = self.asset_attributes_by_dataset.get(dataset.name, {}) + source_attrs = dataset.attributes or {} + dep_names = ( + source_attrs.get("dependencies") + or source_attrs.get("upstream_dataset_names") + or source_attrs.get("deps") + or [] + ) + if isinstance(dep_names, str): + dep_names = [dep_names] + deps = [AssetKey(d.split(".")) for d in dep_names] if dep_names else [] return AssetSpec( key=dataset.name.split("."), + deps=deps, description=attrs.get("description") or f"Spark Declarative Pipeline dataset: {dataset.name}", metadata=attrs.get("metadata"), diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index b31c91b4c282f..a609bc0ac9fe4 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -111,6 +111,7 @@ def discover_datasets_via_dry_run( pipeline_spec_path: str | Path, working_dir: str | Path | None = None, extra_args: list[str] | None = None, + spark_pipelines_cmd: str = "spark-pipelines", ) -> str: """Run `spark-pipelines dry-run` and return raw stdout. @@ -118,6 +119,7 @@ def discover_datasets_via_dry_run( pipeline_spec_path: Path to the pipeline spec file. working_dir: Optional working directory for the subprocess. extra_args: Optional extra CLI arguments. + spark_pipelines_cmd: Executable name or path for the spark-pipelines CLI. Returns: Raw stdout from the command. @@ -126,7 +128,7 @@ def discover_datasets_via_dry_run( SparkPipelinesDryRunError: If the command fails or times out. """ path_str = str(pipeline_spec_path) - cmd = ["spark-pipelines", "dry-run", path_str] + cmd = [spark_pipelines_cmd, "dry-run", path_str] if extra_args: cmd.extend(extra_args) @@ -141,7 +143,7 @@ def discover_datasets_via_dry_run( if result.returncode != 0: raise SparkPipelinesDryRunError( - f"spark-pipelines dry-run failed with return code {result.returncode}", + f"{spark_pipelines_cmd} dry-run failed with return code {result.returncode}", stderr=result.stderr, returncode=result.returncode, ) @@ -414,6 +416,8 @@ def discover_datasets_fn( discovery_mode: DiscoveryMode, working_dir: str | Path | None = None, source_only_datasets: list[DiscoveredDataset] | None = None, + spark_pipelines_cmd: str = "spark-pipelines", + dry_run_extra_args: list[str] | None = None, ) -> list[DiscoveredDataset]: """Discover datasets for a Spark Declarative Pipeline based on discovery_mode. @@ -428,6 +432,8 @@ def discover_datasets_fn( discovery_mode: One of dry_run_only, dry_run_with_fallback, source_only. working_dir: Optional working directory for dry-run. source_only_datasets: Optional list used when mode is source_only or as fallback. + spark_pipelines_cmd: Executable name or path for the spark-pipelines CLI. + dry_run_extra_args: Optional extra CLI arguments for dry-run. Returns: List of DiscoveredDataset. @@ -437,6 +443,7 @@ def discover_datasets_fn( """ check.inst_param(discovery_mode, "discovery_mode", str) source_list = source_only_datasets + extra_args = dry_run_extra_args or [] def _return(datasets: list[DiscoveredDataset]) -> list[DiscoveredDataset]: _validate_no_duplicate_dataset_names(datasets) @@ -451,6 +458,8 @@ def _return(datasets: list[DiscoveredDataset]) -> list[DiscoveredDataset]: raw = discover_datasets_via_dry_run( pipeline_spec_path, working_dir=working_dir, + extra_args=extra_args, + spark_pipelines_cmd=spark_pipelines_cmd, ) datasets = parse_dry_run_output_to_datasets(raw) except SparkPipelinesDryRunError: diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index 3ddc6147a43a6..113cb80d49363 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -10,6 +10,7 @@ from typing import Any, Literal from dagster import AssetKey, ConfigurableResource, MaterializeResult +from pydantic import Field from dagster_spark.components.spark_declarative_pipeline.discovery import ( DiscoveredDataset, @@ -28,6 +29,19 @@ class SparkPipelinesResource(ConfigurableResource): Use run_and_observe inside an asset to run the pipeline and yield MaterializeResults. """ + spark_pipelines_cmd: str = Field( + default="spark-pipelines", + description="Executable name or path for the spark-pipelines CLI.", + ) + dry_run_extra_args: list[str] = Field( + default_factory=list, + description="Extra CLI arguments appended to spark-pipelines dry-run.", + ) + run_extra_args: list[str] = Field( + default_factory=list, + description="Extra CLI arguments appended to spark-pipelines run (before any per-call extra_args).", + ) + def discover_datasets( self, pipeline_spec_path: str | Path, @@ -51,6 +65,8 @@ def discover_datasets( discovery_mode=discovery_mode, working_dir=working_dir, source_only_datasets=source_only_datasets, + spark_pipelines_cmd=self.spark_pipelines_cmd, + dry_run_extra_args=self.dry_run_extra_args, ) def run_and_observe( @@ -84,15 +100,16 @@ def run_and_observe( SparkPipelinesDryRunError: If spark-pipelines run exits with non-zero return code. """ path_str = str(pipeline_spec_path) - cmd = ["spark-pipelines", "run", path_str] + cmd = [self.spark_pipelines_cmd, "run", path_str] if execution_mode == "full_refresh": cmd.append("--full-refresh") else: cmd.append("--refresh") if asset_keys: - datasets_str = ",".join(k.to_user_string() for k in asset_keys) + datasets_str = ",".join(".".join(k.path) for k in asset_keys) if datasets_str: cmd.append(datasets_str) + cmd.extend(self.run_extra_args) if extra_args: cmd.extend(extra_args) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py index 1ece804d0c612..ab065561bbb6d 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py @@ -44,6 +44,26 @@ def test_build_defs_from_state_returns_valid_definitions_with_multi_asset() -> N assert "table_b" in keys +def test_get_asset_spec_includes_deps_from_attributes() -> None: + """get_asset_spec sets deps from dataset.attributes (dependencies, upstream_dataset_names, or deps).""" + component = SparkDeclarativePipelineComponent( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + ) + dataset = DiscoveredDataset( + name="catalog.schema.orders", + attributes={ + "upstream_dataset_names": ["catalog.schema.customers", "catalog.schema.products"], + }, + ) + spec = component.get_asset_spec(dataset) + assert list(spec.key.path) == ["catalog", "schema", "orders"] + dep_keys = [dep.asset_key for dep in spec.deps] + assert len(dep_keys) == 2 + assert list(dep_keys[0].path) == ["catalog", "schema", "customers"] + assert list(dep_keys[1].path) == ["catalog", "schema", "products"] + + def test_build_defs_from_state_filters_temporary_views() -> None: """Temporary view datasets are filtered out unless overridden in asset_attributes_by_dataset.""" component = SparkDeclarativePipelineComponent( diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py index 4728aee23f3e1..a3810798777d2 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -128,3 +128,24 @@ def test_discover_datasets_via_dry_run_returns_stdout_on_success() -> None: out = discover_datasets_via_dry_run("/path/to/spec.yaml") assert "datasets" in out assert "x" in out + + +def test_discover_datasets_via_dry_run_uses_custom_cmd_and_extra_args() -> None: + """discover_datasets_via_dry_run uses spark_pipelines_cmd and extra_args when provided.""" + with patch( + "dagster_spark.components.spark_declarative_pipeline.discovery.subprocess.run" + ) as mock_run: + mock_run.return_value = MagicMock( + returncode=0, + stderr="", + stdout='{"datasets":[{"name":"x"}]}', + ) + discover_datasets_via_dry_run( + "/path/to/spec.yaml", + spark_pipelines_cmd="/custom/spark-pipelines", + extra_args=["--output", "json"], + ) + call_cmd = mock_run.call_args[0][0] + assert call_cmd[0] == "/custom/spark-pipelines" + assert "--output" in call_cmd + assert "json" in call_cmd diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py index 3156755b5b479..ba191413b0284 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -120,6 +120,66 @@ def test_run_and_observe_yields_nothing_when_no_asset_keys() -> None: ) +def test_run_and_observe_passes_dot_notation_datasets_to_cli() -> None: + """run_and_observe passes dataset names as dot-separated (catalog.db.table) to the CLI, not slash.""" + mock_context = MagicMock() + asset_keys = [ + AssetKey(["my_catalog", "my_db", "orders"]), + ] + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter([]) + proc.wait.return_value = 0 + proc.returncode = 0 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=asset_keys, + ) + ) + + call_cmd = mock_popen.call_args[0][0] + datasets_arg = call_cmd[-1] + assert datasets_arg == "my_catalog.my_db.orders" + assert "/" not in datasets_arg + + +def test_run_and_observe_uses_configurable_cmd_and_run_extra_args() -> None: + """run_and_observe uses spark_pipelines_cmd and appends run_extra_args to the command.""" + mock_context = MagicMock() + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter([]) + proc.wait.return_value = 0 + proc.returncode = 0 + mock_popen.return_value = proc + + resource = SparkPipelinesResource( + spark_pipelines_cmd="/usr/local/bin/spark-pipelines", + run_extra_args=["--option", "value"], + ) + list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], + ) + ) + + call_cmd = mock_popen.call_args[0][0] + assert call_cmd[0] == "/usr/local/bin/spark-pipelines" + assert "--option" in call_cmd + assert "value" in call_cmd + + def test_run_and_observe_only_yields_on_success() -> None: """MaterializeResults are only yielded when process returncode is 0.""" mock_context = MagicMock() From e52bba0cc815244a7abae0b1406f6c32afe1d47f Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 20:53:12 +0200 Subject: [PATCH 05/17] implements sql regax, Extract Dependencies during Parsing --- .../spark_declarative_pipeline/component.py | 29 +++--- .../spark_declarative_pipeline/discovery.py | 95 ++++++++++++++++--- .../spark_declarative_pipeline/resource.py | 7 +- .../spark_declarative_pipeline/scaffolder.py | 2 +- .../test_component.py | 46 +++++---- .../test_discovery.py | 27 ++++-- .../test_resource.py | 32 +++++++ 7 files changed, 176 insertions(+), 62 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 3eea08f788827..4e849410b65b7 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -64,7 +64,7 @@ class SparkDeclarativePipelineComponent(StateBackedComponent, dg.Resolvable): ] = field(default_factory=SparkPipelinesResource) op: Optional[OpSpec] = None execution_mode: ExecutionMode = "incremental" - discovery_mode: DiscoveryMode = "dry_run_only" + discovery_mode: DiscoveryMode = "dry_run_with_fallback" asset_attributes_by_dataset: dict[str, dict[str, Any]] = field(default_factory=dict) @property @@ -104,29 +104,25 @@ def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: dataset: Discovered dataset from state. Returns: - AssetSpec with key from dataset.name (split by '.' for multi-segment), optional deps - from attributes (dependencies, upstream_dataset_names, or deps), and optional - description/metadata/group/tags from asset_attributes_by_dataset. + AssetSpec with key from dataset.name (split by '.'), deps from dataset.inferred_deps, + kinds and metadata including dataset_type and source_file, and optional + description/group/tags from asset_attributes_by_dataset. """ attrs = self.asset_attributes_by_dataset.get(dataset.name, {}) - source_attrs = dataset.attributes or {} - dep_names = ( - source_attrs.get("dependencies") - or source_attrs.get("upstream_dataset_names") - or source_attrs.get("deps") - or [] - ) - if isinstance(dep_names, str): - dep_names = [dep_names] - deps = [AssetKey(d.split(".")) for d in dep_names] if dep_names else [] + deps = [AssetKey(d.split(".")) for d in dataset.inferred_deps] + metadata: dict[str, Any] = dict(attrs.get("metadata") or {}) + metadata["dataset_type"] = dataset.dataset_type + if dataset.source_file is not None: + metadata["source_file"] = dataset.source_file return AssetSpec( key=dataset.name.split("."), deps=deps, description=attrs.get("description") or f"Spark Declarative Pipeline dataset: {dataset.name}", - metadata=attrs.get("metadata"), + metadata=metadata, group_name=attrs.get("group_name"), tags=attrs.get("tags"), + kinds={"spark", dataset.dataset_type}, ) def build_defs_from_state( @@ -157,8 +153,7 @@ def build_defs_from_state( def include_dataset(ds: DiscoveredDataset) -> bool: if ds.name in self.asset_attributes_by_dataset: return True - dataset_type = (ds.attributes or {}).get("dataset_type") if ds.attributes else None - if dataset_type == "temporary_view": + if ds.dataset_type == "temporary_view": return False return True diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index a609bc0ac9fe4..7ca24331cbb05 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -54,10 +54,12 @@ class DryRunDatasetNode: Attributes: name: Dataset identifier. raw: Raw dict from CLI output (e.g. JSON object for this node). + inferred_deps: Upstream dataset names extracted from the node (e.g. deps, dependencies). """ name: str raw: dict[str, Any] + inferred_deps: tuple[str, ...] = () @dataclass(frozen=True) @@ -83,11 +85,19 @@ class DiscoveredDataset: Attributes: name: Dataset name (used as asset key component). - attributes: Arbitrary metadata (e.g. dataset_type, schema info). + dataset_type: Type of dataset (e.g. table, materialized_view, temporary_view). + source_file: Path to source file where dataset was defined, if known. + source_line: Line number in source file, if known. + inferred_deps: Upstream dataset names (e.g. catalog.schema.table). + discovery_method: How the dataset was discovered (e.g. dry_run, source_fallback). """ name: str - attributes: dict[str, Any] + dataset_type: str + source_file: str | None + source_line: int | None + inferred_deps: list[str] + discovery_method: str @whitelist_for_serdes @@ -128,7 +138,7 @@ def discover_datasets_via_dry_run( SparkPipelinesDryRunError: If the command fails or times out. """ path_str = str(pipeline_spec_path) - cmd = [spark_pipelines_cmd, "dry-run", path_str] + cmd = [spark_pipelines_cmd, "dry-run", "--spec", path_str] if extra_args: cmd.extend(extra_args) @@ -225,6 +235,16 @@ def _extract_report_json(stdout: str) -> DryRunReport | None: return None +def _extract_deps_from_node(item: dict[str, Any]) -> tuple[str, ...]: + """Extract upstream dataset names from a JSON node (deps, dependencies, upstream_dataset_names).""" + deps = item.get("upstream_dataset_names") or item.get("dependencies") or item.get("deps") or [] + if isinstance(deps, str): + deps = [deps] + if not isinstance(deps, list): + return () + return tuple(str(d) for d in deps if isinstance(d, str)) + + def _json_to_report(data: Any) -> DryRunReport | None: """Map a JSON structure to DryRunReport. Placeholder: adapt to real CLI output shape.""" if isinstance(data, dict): @@ -244,9 +264,12 @@ def _json_to_report(data: Any) -> DryRunReport | None: name = item.get("name") or item.get("id") or item.get("dataset") or str(item) if isinstance(name, dict): name = name.get("name") or name.get("id") or str(name) - dataset_nodes.append(DryRunDatasetNode(name=str(name), raw=item)) + inferred_deps = _extract_deps_from_node(item) + dataset_nodes.append( + DryRunDatasetNode(name=str(name), raw=item, inferred_deps=inferred_deps) + ) elif isinstance(item, str): - dataset_nodes.append(DryRunDatasetNode(name=item, raw={"name": item})) + dataset_nodes.append(DryRunDatasetNode(name=item, raw={"name": item}, inferred_deps=())) return DryRunReport(datasets=dataset_nodes, raw=data if isinstance(data, dict) else None) @@ -264,7 +287,9 @@ def _extract_report_text(stdout: str) -> DryRunReport | None: if match: name = match.group(1).strip() if name: - datasets.append(DryRunDatasetNode(name=name, raw={"name": name})) + datasets.append( + DryRunDatasetNode(name=name, raw={"name": name}, inferred_deps=()) + ) break if not datasets: @@ -277,9 +302,9 @@ def _extract_report_text(stdout: str) -> DryRunReport | None: r"@(?:dp\.)?(materialized_view|table|streaming_table|temporary_view)(?:\([^)]*\))?\s*\n\s*def\s+(\w+)\s*\(", re.MULTILINE, ) -# Regex for source-only discovery: SQL CREATE statement (capture supports catalog.schema.table) +# Regex for source-only discovery: SQL CREATE statement (capture type and catalog.schema.table) _SQL_CREATE_DATASET = re.compile( - r"CREATE\s+(?:MATERIALIZED\s+VIEW|STREAMING\s+TABLE|TABLE|VIEW)\s+(?:IF\s+NOT\s+EXISTS\s+)?([\w.]+)", + r"CREATE\s+(MATERIALIZED\s+VIEW|STREAMING\s+TABLE|TABLE|VIEW)\s+(?:IF\s+NOT\s+EXISTS\s+)?([\w.]+)", re.IGNORECASE, ) @@ -340,18 +365,28 @@ def discover_datasets_from_sources(pipeline_spec_path: Path) -> list[DiscoveredD result.append( DiscoveredDataset( name=name, - attributes={"dataset_type": m.group(1), "source": str(path)}, + dataset_type=m.group(1), + source_file=str(path), + source_line=None, + inferred_deps=[], + discovery_method="source_fallback", ) ) else: for m in _SQL_CREATE_DATASET.finditer(text): - name = m.group(1) + raw_type = m.group(1) + dataset_type_sql = raw_type.lower().replace(" ", "_") + name = m.group(2) if name not in seen: seen.add(name) result.append( DiscoveredDataset( name=name, - attributes={"source": str(path)}, + dataset_type=dataset_type_sql, + source_file=str(path), + source_line=None, + inferred_deps=[], + discovery_method="source_fallback", ) ) except Exception as e: @@ -366,10 +401,19 @@ def discover_datasets_from_sources(pipeline_spec_path: Path) -> list[DiscoveredD return result +def _node_dataset_type(raw: dict[str, Any]) -> str: + """Extract dataset_type from a raw node (type, dataset_type, or default 'table').""" + t = raw.get("type") or raw.get("dataset_type") + if isinstance(t, str): + return t.lower().replace(" ", "_") + return "table" + + def parse_dry_run_output_to_datasets(stdout: str) -> list[DiscoveredDataset]: """Parse dry-run stdout into a list of DiscoveredDataset. - Uses extract_report for JSON and text fallback; maps each node to DiscoveredDataset. + Uses extract_report for JSON and text fallback; maps each node to DiscoveredDataset + with explicit dataset_type, source_file, source_line, inferred_deps, discovery_method. Args: stdout: Raw stdout string from spark-pipelines dry-run. @@ -381,9 +425,30 @@ def parse_dry_run_output_to_datasets(stdout: str) -> list[DiscoveredDataset]: if report is None: return [] - return [ - DiscoveredDataset(name=node.name, attributes=dict(node.raw)) for node in report.datasets - ] + result: list[DiscoveredDataset] = [] + for node in report.datasets: + raw = node.raw + source_file = raw.get("source_file") or raw.get("source") + if isinstance(source_file, str): + pass + else: + source_file = None + source_line = raw.get("source_line") + if isinstance(source_line, int): + pass + else: + source_line = None + result.append( + DiscoveredDataset( + name=node.name, + dataset_type=_node_dataset_type(raw), + source_file=source_file, + source_line=source_line, + inferred_deps=list(node.inferred_deps), + discovery_method="dry_run", + ) + ) + return result def _validate_no_duplicate_dataset_names(datasets: list[DiscoveredDataset]) -> None: diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index 113cb80d49363..a4761a768a324 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -100,9 +100,12 @@ def run_and_observe( SparkPipelinesDryRunError: If spark-pipelines run exits with non-zero return code. """ path_str = str(pipeline_spec_path) - cmd = [self.spark_pipelines_cmd, "run", path_str] + cmd = [self.spark_pipelines_cmd, "run", "--spec", path_str] if execution_mode == "full_refresh": - cmd.append("--full-refresh") + if asset_keys: + cmd.append("--full-refresh") + else: + cmd.append("--full-refresh-all") else: cmd.append("--refresh") if asset_keys: diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py index bd51422c5eda1..416121b864f0d 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/scaffolder.py @@ -13,6 +13,6 @@ def scaffold(self, request: ScaffoldRequest) -> None: request, { "pipeline_spec_path": "spark-pipeline.yml", - "discovery_mode": "dry_run_only", + "discovery_mode": "dry_run_with_fallback", }, ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py index ab065561bbb6d..7e51f840cbd46 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py @@ -12,6 +12,22 @@ ) +def _ds( + name: str, + dataset_type: str = "table", + inferred_deps: list[str] | None = None, +) -> DiscoveredDataset: + """Helper to build DiscoveredDataset with required typed fields.""" + return DiscoveredDataset( + name=name, + dataset_type=dataset_type, + source_file=None, + source_line=None, + inferred_deps=inferred_deps or [], + discovery_method="dry_run", + ) + + def test_build_defs_from_state_returns_valid_definitions_with_multi_asset() -> None: """build_defs_from_state returns a valid Definitions object containing a multi_asset.""" component = SparkDeclarativePipelineComponent( @@ -19,8 +35,8 @@ def test_build_defs_from_state_returns_valid_definitions_with_multi_asset() -> N discovery_mode="source_only", ) datasets = [ - DiscoveredDataset(name="table_a", attributes={"dataset_type": "table"}), - DiscoveredDataset(name="table_b", attributes={"dataset_type": "table"}), + _ds("table_a"), + _ds("table_b"), ] state = SparkPipelineState( datasets=datasets, @@ -44,17 +60,19 @@ def test_build_defs_from_state_returns_valid_definitions_with_multi_asset() -> N assert "table_b" in keys -def test_get_asset_spec_includes_deps_from_attributes() -> None: - """get_asset_spec sets deps from dataset.attributes (dependencies, upstream_dataset_names, or deps).""" +def test_get_asset_spec_includes_deps_from_inferred_deps() -> None: + """get_asset_spec sets deps from dataset.inferred_deps.""" component = SparkDeclarativePipelineComponent( pipeline_spec_path="pipeline.yaml", discovery_mode="source_only", ) dataset = DiscoveredDataset( name="catalog.schema.orders", - attributes={ - "upstream_dataset_names": ["catalog.schema.customers", "catalog.schema.products"], - }, + dataset_type="table", + source_file=None, + source_line=None, + inferred_deps=["catalog.schema.customers", "catalog.schema.products"], + discovery_method="dry_run", ) spec = component.get_asset_spec(dataset) assert list(spec.key.path) == ["catalog", "schema", "orders"] @@ -72,11 +90,8 @@ def test_build_defs_from_state_filters_temporary_views() -> None: asset_attributes_by_dataset={}, # no overrides ) datasets = [ - DiscoveredDataset(name="table_a", attributes={"dataset_type": "table"}), - DiscoveredDataset( - name="temp_view_x", - attributes={"dataset_type": "temporary_view"}, - ), + _ds("table_a"), + _ds("temp_view_x", dataset_type="temporary_view"), ] state = SparkPipelineState( datasets=datasets, @@ -104,11 +119,8 @@ def test_build_defs_from_state_includes_temporary_view_when_overridden() -> None asset_attributes_by_dataset={"temp_view_x": {"description": "Included view"}}, ) datasets = [ - DiscoveredDataset(name="table_a", attributes={"dataset_type": "table"}), - DiscoveredDataset( - name="temp_view_x", - attributes={"dataset_type": "temporary_view"}, - ), + _ds("table_a"), + _ds("temp_view_x", dataset_type="temporary_view"), ] state = SparkPipelineState( datasets=datasets, diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py index a3810798777d2..7029f042c6a71 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -27,7 +27,7 @@ def test_parse_dry_run_output_to_datasets_parses_json_report() -> None: datasets = parse_dry_run_output_to_datasets(stdout) assert len(datasets) == 2 assert datasets[0].name == "dataset_a" - assert datasets[0].attributes.get("type") == "table" + assert datasets[0].dataset_type == "table" assert datasets[1].name == "dataset_b" @@ -55,9 +55,20 @@ def test_discover_datasets_fn_dry_run_only_raises_on_failure() -> None: assert "error" in (exc_info.value.stderr or "") +def _discovered(name: str) -> DiscoveredDataset: + return DiscoveredDataset( + name=name, + dataset_type="table", + source_file=None, + source_line=None, + inferred_deps=[], + discovery_method="source_fallback", + ) + + def test_discover_datasets_fn_dry_run_with_fallback_uses_source_on_failure() -> None: """In dry_run_with_fallback mode, source_only_datasets are returned when dry-run fails.""" - fallback = [DiscoveredDataset(name="fallback_ds", attributes={})] + fallback = [_discovered("fallback_ds")] with patch( "dagster_spark.components.spark_declarative_pipeline.discovery.subprocess.run" ) as mock_run: @@ -72,10 +83,7 @@ def test_discover_datasets_fn_dry_run_with_fallback_uses_source_on_failure() -> def test_discover_datasets_fn_source_only_returns_source_list() -> None: """In source_only mode, discover_datasets_fn returns source_only_datasets without running dry-run.""" - source = [ - DiscoveredDataset(name="a", attributes={}), - DiscoveredDataset(name="b", attributes={}), - ] + source = [_discovered("a"), _discovered("b")] result = discover_datasets_fn( pipeline_spec_path="/any/path", discovery_mode="source_only", @@ -102,10 +110,7 @@ def test_discover_datasets_via_dry_run_raises_on_nonzero_exit() -> None: def test_discover_datasets_fn_raises_on_duplicate_dataset_names() -> None: """Duplicate dataset names (after normalization) raise DuplicateDatasetNamesError.""" - source = [ - DiscoveredDataset(name="Foo", attributes={}), - DiscoveredDataset(name="foo", attributes={}), - ] + source = [_discovered("Foo"), _discovered("foo")] with pytest.raises(DuplicateDatasetNamesError) as exc_info: discover_datasets_fn( pipeline_spec_path="/any/path", @@ -147,5 +152,7 @@ def test_discover_datasets_via_dry_run_uses_custom_cmd_and_extra_args() -> None: ) call_cmd = mock_run.call_args[0][0] assert call_cmd[0] == "/custom/spark-pipelines" + assert call_cmd[1] == "dry-run" + assert "--spec" in call_cmd assert "--output" in call_cmd assert "json" in call_cmd diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py index ba191413b0284..50a74ea4e7481 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -145,6 +145,8 @@ def test_run_and_observe_passes_dot_notation_datasets_to_cli() -> None: ) call_cmd = mock_popen.call_args[0][0] + assert call_cmd[1] == "run" + assert "--spec" in call_cmd datasets_arg = call_cmd[-1] assert datasets_arg == "my_catalog.my_db.orders" assert "/" not in datasets_arg @@ -176,10 +178,40 @@ def test_run_and_observe_uses_configurable_cmd_and_run_extra_args() -> None: call_cmd = mock_popen.call_args[0][0] assert call_cmd[0] == "/usr/local/bin/spark-pipelines" + assert call_cmd[1] == "run" + assert "--spec" in call_cmd assert "--option" in call_cmd assert "value" in call_cmd +def test_run_and_observe_full_refresh_no_asset_keys_uses_full_refresh_all() -> None: + """When execution_mode is full_refresh and asset_keys is empty, CLI gets --full-refresh-all.""" + mock_context = MagicMock() + with patch( + "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" + ) as mock_popen: + proc = MagicMock() + proc.stdout = iter([]) + proc.wait.return_value = 0 + proc.returncode = 0 + mock_popen.return_value = proc + + resource = SparkPipelinesResource() + list( + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + execution_mode="full_refresh", + asset_keys=None, + ) + ) + + call_cmd = mock_popen.call_args[0][0] + assert "--spec" in call_cmd + assert "--full-refresh-all" in call_cmd + assert "--full-refresh" not in call_cmd + + def test_run_and_observe_only_yields_on_success() -> None: """MaterializeResults are only yielded when process returncode is 0.""" mock_context = MagicMock() From 47fce6cebad0a246c56628b7dc3972f1de15f85b Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 21:18:05 +0200 Subject: [PATCH 06/17] --- .../spark_declarative_pipeline/__init__.py | 2 ++ .../spark_declarative_pipeline/component.py | 16 ++++++--- .../spark_declarative_pipeline/discovery.py | 34 +++++++++++++++++-- .../spark_declarative_pipeline/resource.py | 13 ++++--- .../test_resource.py | 10 +++--- 5 files changed, 59 insertions(+), 16 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py index fde5cc356f361..8b083c9ccc924 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py @@ -9,6 +9,8 @@ DryRunReport as DryRunReport, DuplicateDatasetNamesError as DuplicateDatasetNamesError, SparkPipelinesDryRunError as SparkPipelinesDryRunError, + SparkPipelinesError as SparkPipelinesError, + SparkPipelinesExecutionError as SparkPipelinesExecutionError, SparkPipelineState as SparkPipelineState, discover_datasets_fn as discover_datasets_fn, discover_datasets_via_dry_run as discover_datasets_via_dry_run, diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 4e849410b65b7..57b10870f1ece 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -184,11 +184,17 @@ def _spark_pipeline_asset( context: dg.AssetExecutionContext, spark_pipelines: SparkPipelinesResource, ) -> Any: - keys = ( - list(context.selected_asset_keys) - if context.selected_asset_keys - else [s.key for s in asset_specs] - ) + # When the entire graph is executed (not a subset), pass keys=None to allow + # --full-refresh-all and avoid OS argument length limits. + is_subset = getattr(context, "is_subset", True) + if is_subset: + keys = ( + list(context.selected_asset_keys) + if context.selected_asset_keys + else [s.key for s in asset_specs] + ) + else: + keys = None yield from spark_pipelines.run_and_observe( context=context, pipeline_spec_path=resolved_spec_path, diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index 7ca24331cbb05..eb82b939b120e 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -20,8 +20,27 @@ DiscoveryMode = Literal["dry_run_only", "dry_run_with_fallback", "source_only"] -class SparkPipelinesDryRunError(Exception): - """Raised when ``spark-pipelines dry-run`` or ``spark-pipelines run`` fails. +class SparkPipelinesError(Exception): + """Base exception for Spark Declarative Pipeline operations (dry-run and execution).""" + + +class SparkPipelinesDryRunError(SparkPipelinesError): + """Raised when ``spark-pipelines dry-run`` fails. + + Attributes: + message: Error description. + stderr: Captured stderr or combined stdout/stderr from the process. + returncode: Process exit code (non-zero on failure). + """ + + def __init__(self, message: str, stderr: str | None = None, returncode: int | None = None): + super().__init__(message) + self.stderr = stderr + self.returncode = returncode + + +class SparkPipelinesExecutionError(SparkPipelinesError): + """Raised when ``spark-pipelines run`` fails. Attributes: message: Error description. @@ -341,14 +360,20 @@ def discover_datasets_from_sources(pipeline_spec_path: Path) -> list[DiscoveredD ) return [] + if not root.exists(): + logger.warning("Pipeline spec directory does not exist: %s", root) + return [] + seen: set[str] = set() result: list[DiscoveredDataset] = [] + scanned_files = 0 try: for ext in ("*.py", "*.sql"): for path in root.rglob(ext): try: text = path.read_text(encoding="utf-8", errors="replace") + scanned_files += 1 except Exception as e: logger.warning( "Failed to read source file %s: %s", @@ -398,6 +423,11 @@ def discover_datasets_from_sources(pipeline_spec_path: Path) -> list[DiscoveredD ) return [] + logger.info( + "Scanned %d files for static dataset discovery, found %d datasets.", + scanned_files, + len(result), + ) return result diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index a4761a768a324..0d2a6a58ca674 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -4,6 +4,7 @@ run_and_observe (run spark-pipelines with log streaming and MaterializeResult yields). """ +import os import subprocess from collections.abc import Iterator from pathlib import Path @@ -15,7 +16,7 @@ from dagster_spark.components.spark_declarative_pipeline.discovery import ( DiscoveredDataset, DiscoveryMode, - SparkPipelinesDryRunError, + SparkPipelinesExecutionError, discover_datasets_fn, ) @@ -82,8 +83,7 @@ def run_and_observe( Uses Popen to stream stdout/stderr line-by-line and logs each line via context.log.info. Passes --full-refresh or --refresh based on execution_mode, then optional comma-separated - dataset list from asset_keys. Only yields MaterializeResults if the process exits with - returncode == 0; otherwise raises SparkPipelinesDryRunError with the captured log. + dataset list from asset_keys. returncode == 0; otherwise raises SparkPipelinesExecutionError with the captured log. Args: context: Asset execution context (used for context.log.info). @@ -97,7 +97,7 @@ def run_and_observe( MaterializeResult for each materialized asset on success. Raises: - SparkPipelinesDryRunError: If spark-pipelines run exits with non-zero return code. + SparkPipelinesExecutionError: If spark-pipelines run exits with non-zero return code. """ path_str = str(pipeline_spec_path) cmd = [self.spark_pipelines_cmd, "run", "--spec", path_str] @@ -117,12 +117,15 @@ def run_and_observe( cmd.extend(extra_args) cwd = str(working_dir) if working_dir else None + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=cwd, + env=env, bufsize=1, ) log_lines: list[str] = [] @@ -138,7 +141,7 @@ def run_and_observe( if returncode != 0: captured = "\n".join(log_lines) if log_lines else "(no output)" - raise SparkPipelinesDryRunError( + raise SparkPipelinesExecutionError( f"spark-pipelines run failed with return code {returncode}", stderr=captured, returncode=returncode, diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py index 50a74ea4e7481..0d5387f67b16b 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -5,7 +5,9 @@ import pytest from dagster import AssetKey from dagster._core.definitions.result import MaterializeResult -from dagster_spark.components.spark_declarative_pipeline.discovery import SparkPipelinesDryRunError +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + SparkPipelinesExecutionError, +) from dagster_spark.components.spark_declarative_pipeline.resource import SparkPipelinesResource @@ -69,7 +71,7 @@ def test_run_and_observe_streams_logs_via_context() -> None: def test_run_and_observe_raises_with_captured_log_on_nonzero_exit() -> None: - """run_and_observe raises SparkPipelinesDryRunError with captured log when process exits non-zero.""" + """run_and_observe raises SparkPipelinesExecutionError with captured log when process exits non-zero.""" mock_context = MagicMock() with patch( "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" @@ -81,7 +83,7 @@ def test_run_and_observe_raises_with_captured_log_on_nonzero_exit() -> None: mock_popen.return_value = proc resource = SparkPipelinesResource() - with pytest.raises(SparkPipelinesDryRunError) as exc_info: + with pytest.raises(SparkPipelinesExecutionError) as exc_info: list( resource.run_and_observe( context=mock_context, @@ -225,7 +227,7 @@ def test_run_and_observe_only_yields_on_success() -> None: mock_popen.return_value = proc resource = SparkPipelinesResource() - with pytest.raises(SparkPipelinesDryRunError): + with pytest.raises(SparkPipelinesExecutionError): list( resource.run_and_observe( context=mock_context, From fb4f15532695b88695c47bea7e0de05f5fdb31a4 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 21:36:03 +0200 Subject: [PATCH 07/17] add const and a limit to log lines --- .../components/spark_declarative_pipeline/component.py | 8 ++++---- .../components/spark_declarative_pipeline/discovery.py | 7 +++++-- .../components/spark_declarative_pipeline/resource.py | 3 ++- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 57b10870f1ece..06457a8d2cc5d 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -127,7 +127,7 @@ def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: def build_defs_from_state( self, - context: ComponentLoadContext, + load_context: ComponentLoadContext, state_path: Optional[Path], ) -> Definitions: """Build Definitions with a multi_asset that runs spark_pipelines.run_and_observe. @@ -137,7 +137,7 @@ def build_defs_from_state( can_subset=True that yields MaterializeResults via run_and_observe. Args: - context: Component load context (path, etc.). + load_context: Component load context (path, etc.). state_path: Path to serialized state file; if None or missing, returns empty Definitions. Returns: @@ -166,10 +166,10 @@ def include_dataset(ds: DiscoveredDataset) -> bool: pipeline_spec_path = state.pipeline_spec_path # Resolve path relative to component path if not Path(pipeline_spec_path).is_absolute(): - resolved_spec_path = (context.path / pipeline_spec_path).resolve() + resolved_spec_path = (load_context.path / Path(pipeline_spec_path)).resolve() else: resolved_spec_path = Path(pipeline_spec_path) - working_dir = context.path + working_dir = load_context.path execution_mode = self.execution_mode @dg.multi_asset( diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index eb82b939b120e..a29e0598ba6ee 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -17,6 +17,9 @@ from dagster_shared import check # type: ignore[reportMissingImports] from dagster_shared.serdes import whitelist_for_serdes # type: ignore[reportMissingImports] +# Guardrail to prevent daemon hangs. +DRY_RUN_TIMEOUT_SECONDS = 60 + DiscoveryMode = Literal["dry_run_only", "dry_run_with_fallback", "source_only"] @@ -167,7 +170,7 @@ def discover_datasets_via_dry_run( text=True, check=False, cwd=working_dir, - timeout=300, + timeout=DRY_RUN_TIMEOUT_SECONDS, ) if result.returncode != 0: @@ -349,7 +352,7 @@ def discover_datasets_from_sources(pipeline_spec_path: Path) -> list[DiscoveredD logger = get_dagster_logger() try: root = Path(pipeline_spec_path).resolve() - if root.suffix: + if root.is_file(): root = root.parent except Exception as e: logger.warning( diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index 0d2a6a58ca674..fecd023d0cf9f 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -6,6 +6,7 @@ import os import subprocess +from collections import deque from collections.abc import Iterator from pathlib import Path from typing import Any, Literal @@ -128,7 +129,7 @@ def run_and_observe( env=env, bufsize=1, ) - log_lines: list[str] = [] + log_lines: deque[str] = deque(maxlen=1000) if process.stdout: for raw_line in process.stdout: line = raw_line.rstrip("\n\r") From f83f33dfa17220977cb02b4630078b8e438903de Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 22:17:42 +0200 Subject: [PATCH 08/17] add remarks --- .../spark_declarative_pipeline/component.py | 25 ++++++++++--------- .../spark_declarative_pipeline/discovery.py | 9 ++++++- .../spark_declarative_pipeline/resource.py | 13 ++++++---- .../test_resource.py | 3 +++ 4 files changed, 32 insertions(+), 18 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 06457a8d2cc5d..115ea958e3618 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -7,7 +7,7 @@ from dataclasses import dataclass, field from pathlib import Path -from typing import Annotated, Any, Literal, Optional +from typing import Annotated, Any, Literal import dagster as dg from dagster import AssetKey, AssetSpec, Definitions, deserialize_value, serialize_value @@ -62,7 +62,7 @@ class SparkDeclarativePipelineComponent(StateBackedComponent, dg.Resolvable): SparkPipelinesResource, Resolver(_resolve_spark_pipelines_resource), ] = field(default_factory=SparkPipelinesResource) - op: Optional[OpSpec] = None + op: OpSpec | None = None execution_mode: ExecutionMode = "incremental" discovery_mode: DiscoveryMode = "dry_run_with_fallback" asset_attributes_by_dataset: dict[str, dict[str, Any]] = field(default_factory=dict) @@ -127,8 +127,8 @@ def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: def build_defs_from_state( self, - load_context: ComponentLoadContext, - state_path: Optional[Path], + context: ComponentLoadContext, + state_path: Path | None, ) -> Definitions: """Build Definitions with a multi_asset that runs spark_pipelines.run_and_observe. @@ -137,7 +137,7 @@ def build_defs_from_state( can_subset=True that yields MaterializeResults via run_and_observe. Args: - load_context: Component load context (path, etc.). + context: Component load context (path, etc.). state_path: Path to serialized state file; if None or missing, returns empty Definitions. Returns: @@ -166,10 +166,10 @@ def include_dataset(ds: DiscoveredDataset) -> bool: pipeline_spec_path = state.pipeline_spec_path # Resolve path relative to component path if not Path(pipeline_spec_path).is_absolute(): - resolved_spec_path = (load_context.path / Path(pipeline_spec_path)).resolve() + resolved_spec_path = (context.path / Path(pipeline_spec_path)).resolve() else: resolved_spec_path = Path(pipeline_spec_path) - working_dir = load_context.path + working_dir = context.path execution_mode = self.execution_mode @dg.multi_asset( @@ -181,22 +181,23 @@ def include_dataset(ds: DiscoveredDataset) -> bool: pool=op_spec.pool, ) def _spark_pipeline_asset( - context: dg.AssetExecutionContext, + _context: dg.AssetExecutionContext, spark_pipelines: SparkPipelinesResource, ) -> Any: # When the entire graph is executed (not a subset), pass keys=None to allow # --full-refresh-all and avoid OS argument length limits. - is_subset = getattr(context, "is_subset", True) + # Inner param _context avoids shadowing outer context (ComponentLoadContext); + is_subset = getattr(_context, "is_subset", True) if is_subset: keys = ( - list(context.selected_asset_keys) - if context.selected_asset_keys + list(_context.selected_asset_keys) + if _context.selected_asset_keys else [s.key for s in asset_specs] ) else: keys = None yield from spark_pipelines.run_and_observe( - context=context, + context=_context, pipeline_spec_path=resolved_spec_path, working_dir=working_dir, execution_mode=execution_mode, diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index a29e0598ba6ee..3cd6211ad91cd 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -239,7 +239,10 @@ def _extract_report_json(stdout: str) -> DryRunReport | None: return report except (json.JSONDecodeError, TypeError): pass - # Single line may have leading/trailing text; try substring from first { or [ + # Note: This is an aggressive, brute-force fallback to extract JSON embedded within + # verbose or irregular log streams. It may attempt to parse invalid substrings + # (e.g., standard log lines with arbitrary brackets), but relies on catching + # json.JSONDecodeError to safely ignore them. for start_char, end_char in (("{", "}"), ("[", "]")): i = line.find(start_char) if i == -1: @@ -548,6 +551,10 @@ def _return(datasets: list[DiscoveredDataset]) -> list[DiscoveredDataset]: return list(datasets) if discovery_mode == "source_only": + get_dagster_logger().warning( + "Using 'source_only' discovery mode. This is a best-effort static parsing " + "fallback and may miss datasets dynamically generated by complex imports or wrappers." + ) if source_list is not None: return _return(source_list) return _return(discover_datasets_from_sources(Path(pipeline_spec_path))) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index fecd023d0cf9f..6198823d26159 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -105,14 +105,17 @@ def run_and_observe( if execution_mode == "full_refresh": if asset_keys: cmd.append("--full-refresh") + datasets_str = ",".join(".".join(k.path) for k in asset_keys) + if datasets_str: + cmd.append(datasets_str) else: cmd.append("--full-refresh-all") else: - cmd.append("--refresh") - if asset_keys: - datasets_str = ",".join(".".join(k.path) for k in asset_keys) - if datasets_str: - cmd.append(datasets_str) + if asset_keys: + cmd.append("--refresh") + datasets_str = ",".join(".".join(k.path) for k in asset_keys) + if datasets_str: + cmd.append(datasets_str) cmd.extend(self.run_extra_args) if extra_args: cmd.extend(extra_args) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py index 0d5387f67b16b..9a762dd8691aa 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -120,6 +120,9 @@ def test_run_and_observe_yields_nothing_when_no_asset_keys() -> None: mock_context.log.info.assert_any_call( "No specific asset keys requested; spark-pipelines run completed successfully." ) + # Incremental run with no asset_keys must not append --refresh (bare run = full pipeline) + call_cmd = mock_popen.call_args[0][0] + assert "--refresh" not in call_cmd def test_run_and_observe_passes_dot_notation_datasets_to_cli() -> None: From f3537df77e855b384ac562e5f14c16ca68019031 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 22:35:02 +0200 Subject: [PATCH 09/17] Update python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../components/spark_declarative_pipeline/component.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index 115ea958e3618..d9068d82bb219 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -42,7 +42,12 @@ def _resolve_spark_pipelines_resource(_context: Any, value: Any) -> SparkPipelin return value if value is None: return SparkPipelinesResource() - return SparkPipelinesResource(**value) if isinstance(value, dict) else SparkPipelinesResource() + if isinstance(value, dict): + return SparkPipelinesResource(**value) + raise ValueError( + f"Cannot resolve spark_pipelines field: expected a SparkPipelinesResource, dict, or None, " + f"got {type(value).__name__!r}: {value!r}" + ) @scaffold_with(SparkDeclarativePipelineScaffolder) From 4ea0805b78978e9a9fe78f93912ba188255b6b09 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 22:35:23 +0200 Subject: [PATCH 10/17] Update python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../components/spark_declarative_pipeline/resource.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index 6198823d26159..5c5cd34d9fbe0 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -84,7 +84,8 @@ def run_and_observe( Uses Popen to stream stdout/stderr line-by-line and logs each line via context.log.info. Passes --full-refresh or --refresh based on execution_mode, then optional comma-separated - dataset list from asset_keys. returncode == 0; otherwise raises SparkPipelinesExecutionError with the captured log. + dataset list from asset_keys. Yields MaterializeResult per asset on success when + returncode == 0; otherwise raises SparkPipelinesExecutionError with the captured log. Args: context: Asset execution context (used for context.log.info). From 5ae6bd4980da0e47db6ed8194fed52ca359e484a Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Wed, 4 Mar 2026 22:35:36 +0200 Subject: [PATCH 11/17] Update python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../spark_declarative_pipeline/resource.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index 5c5cd34d9fbe0..ec79311d31ce8 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -134,14 +134,17 @@ def run_and_observe( bufsize=1, ) log_lines: deque[str] = deque(maxlen=1000) - if process.stdout: - for raw_line in process.stdout: - line = raw_line.rstrip("\n\r") - if line: - log_lines.append(line) - if context is not None and hasattr(context, "log"): - context.log.info(line) - process.wait() + try: + if process.stdout: + for raw_line in process.stdout: + line = raw_line.rstrip("\n\r") + if line: + log_lines.append(line) + if context is not None and hasattr(context, "log"): + context.log.info(line) + finally: + process.wait() + returncode = process.returncode returncode = process.returncode if returncode != 0: From 1ff897906acb7931f4619c8f004ac8f9bfa70b7c Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Thu, 5 Mar 2026 13:56:03 +0200 Subject: [PATCH 12/17] Update python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py Co-authored-by: graphite-app[bot] <96075541+graphite-app[bot]@users.noreply.github.com> --- .../components/spark_declarative_pipeline/resource.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index ec79311d31ce8..d272d09201e06 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -145,7 +145,6 @@ def run_and_observe( finally: process.wait() returncode = process.returncode - returncode = process.returncode if returncode != 0: captured = "\n".join(log_lines) if log_lines else "(no output)" From ff9d221626657111b38fb4b0991960f8ebfb9b2b Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Sat, 14 Mar 2026 14:26:56 +0200 Subject: [PATCH 13/17] fix remarks --- .../spark_declarative_pipeline/discovery.py | 85 ++++++--- .../test_component.py | 102 ++++++++++- .../test_discovery.py | 165 ++++++++++++++++++ 3 files changed, 330 insertions(+), 22 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index 3cd6211ad91cd..117ee0ded642e 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -2,19 +2,19 @@ This module runs ``spark-pipelines dry-run`` to discover datasets, parses JSON or structured text output, and supports discovery_mode fallbacks (e.g. source_only). -State types (DiscoveredDataset, SparkPipelineState) are frozen dataclasses with +State types (DiscoveredDataset, SparkPipelineState) use @record with whitelist_for_serdes for Dagster serialize_value/deserialize_value compatibility. """ import json import re import subprocess -from dataclasses import dataclass from pathlib import Path -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal from dagster import get_dagster_logger from dagster_shared import check # type: ignore[reportMissingImports] +from dagster_shared.record import record # type: ignore[reportMissingImports] from dagster_shared.serdes import whitelist_for_serdes # type: ignore[reportMissingImports] # Guardrail to prevent daemon hangs. @@ -69,7 +69,7 @@ def __init__(self, message: str, duplicate_names: list[str]): self.duplicate_names = duplicate_names -@dataclass(frozen=True) +@record class DryRunDatasetNode: """A single dataset node as reported by spark-pipelines dry-run. @@ -83,8 +83,18 @@ class DryRunDatasetNode: raw: dict[str, Any] inferred_deps: tuple[str, ...] = () + if TYPE_CHECKING: -@dataclass(frozen=True) + def __init__( + self, + *, + name: str, + raw: dict[str, Any], + inferred_deps: tuple[str, ...] = (), + ) -> None: ... + + +@record class DryRunReport: """Structured report produced by spark-pipelines dry-run (JSON or parsed text). @@ -96,9 +106,18 @@ class DryRunReport: datasets: list[DryRunDatasetNode] raw: dict[str, Any] | None = None + if TYPE_CHECKING: + + def __init__( + self, + *, + datasets: list[DryRunDatasetNode], + raw: dict[str, Any] | None = None, + ) -> None: ... + @whitelist_for_serdes -@dataclass(frozen=True) +@record class DiscoveredDataset: """A dataset discovered for a Spark Declarative Pipeline (from dry-run or source). @@ -121,9 +140,22 @@ class DiscoveredDataset: inferred_deps: list[str] discovery_method: str + if TYPE_CHECKING: + + def __init__( + self, + *, + name: str, + dataset_type: str, + source_file: str | None, + source_line: int | None, + inferred_deps: list[str], + discovery_method: str, + ) -> None: ... + @whitelist_for_serdes -@dataclass(frozen=True) +@record class SparkPipelineState: """Cached state for a Spark Declarative Pipeline (discovered datasets). @@ -138,6 +170,15 @@ class SparkPipelineState: datasets: list[DiscoveredDataset] pipeline_spec_path: str + if TYPE_CHECKING: + + def __init__( + self, + *, + datasets: list[DiscoveredDataset], + pipeline_spec_path: str, + ) -> None: ... + def discover_datasets_via_dry_run( pipeline_spec_path: str | Path, @@ -261,13 +302,16 @@ def _extract_report_json(stdout: str) -> DryRunReport | None: def _extract_deps_from_node(item: dict[str, Any]) -> tuple[str, ...]: - """Extract upstream dataset names from a JSON node (deps, dependencies, upstream_dataset_names).""" + """Extract upstream dataset names from a JSON node (deps, dependencies, upstream_dataset_names). + + Empty or whitespace-only names are excluded to avoid AssetKey path components that are empty. + """ deps = item.get("upstream_dataset_names") or item.get("dependencies") or item.get("deps") or [] if isinstance(deps, str): deps = [deps] if not isinstance(deps, list): return () - return tuple(str(d) for d in deps if isinstance(d, str)) + return tuple(s for d in deps if isinstance(d, str) and (s := str(d).strip()) != "") def _json_to_report(data: Any) -> DryRunReport | None: @@ -306,8 +350,13 @@ def _extract_report_text(stdout: str) -> DryRunReport | None: line = raw_line.strip() if not line or line.startswith("#"): continue - # Match "dataset: foo" or "- foo" or " - foo" - for pattern in (r"dataset:\s*(.+)", r"^[-*]\s*(.+)", r"^\d+\.\s*(.+)"): + # Match "dataset: " or "- " or " - " or "1. "; require valid dataset id (alphanumeric, underscores, dots) + _dataset_id_pattern = r"[a-zA-Z0-9_.]+" + for pattern in ( + rf"dataset:\s*({_dataset_id_pattern})\s*$", + rf"^[-*]\s*({_dataset_id_pattern})\s*$", + rf"^\d+\.\s*({_dataset_id_pattern})\s*$", + ): match = re.match(pattern, line, re.IGNORECASE) if match: name = match.group(1).strip() @@ -464,16 +513,10 @@ def parse_dry_run_output_to_datasets(stdout: str) -> list[DiscoveredDataset]: result: list[DiscoveredDataset] = [] for node in report.datasets: raw = node.raw - source_file = raw.get("source_file") or raw.get("source") - if isinstance(source_file, str): - pass - else: - source_file = None - source_line = raw.get("source_line") - if isinstance(source_line, int): - pass - else: - source_line = None + raw_sf = raw.get("source_file") or raw.get("source") + source_file = raw_sf if isinstance(raw_sf, str) else None + raw_sl = raw.get("source_line") + source_line = raw_sl if isinstance(raw_sl, int) else None result.append( DiscoveredDataset( name=node.name, diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py index 7e51f840cbd46..4c1fa1b48bc74 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py @@ -1,10 +1,16 @@ -"""Tests for SparkDeclarativePipelineComponent (build_defs_from_state, temporary_view filtering).""" +"""Tests for SparkDeclarativePipelineComponent (YAML load, lifecycle, build_defs_from_state, temporary_view filtering).""" +import asyncio import tempfile +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from unittest.mock import MagicMock import dagster as dg +from dagster import AssetKey +from dagster._utils.test.definitions import scoped_definitions_load_context +from dagster.components.testing.utils import create_defs_folder_sandbox from dagster_spark.components.spark_declarative_pipeline import ( DiscoveredDataset, SparkDeclarativePipelineComponent, @@ -12,6 +18,35 @@ ) +@contextmanager +def setup_spark_component( + pipeline_spec_path: str, + discovery_mode: str = "source_only", + defs_state_type: str = "LOCAL_FILESYSTEM", +) -> Iterator[tuple[SparkDeclarativePipelineComponent, dg.Definitions]]: + """Set up a components project with a Spark component; yield (component, defs) from load_component_and_build_defs.""" + typename = "dagster_spark.components.spark_declarative_pipeline.component.SparkDeclarativePipelineComponent" + defs_yaml_contents: dict = { + "type": typename, + "attributes": { + "pipeline_spec_path": pipeline_spec_path, + "discovery_mode": discovery_mode, + "defs_state": {"management_type": defs_state_type}, + }, + } + with create_defs_folder_sandbox() as sandbox: + defs_path = sandbox.scaffold_component( + component_cls=SparkDeclarativePipelineComponent, + defs_yaml_contents=defs_yaml_contents, + ) + with sandbox.load_component_and_build_defs(defs_path=defs_path) as ( + component, + defs, + ): + assert isinstance(component, SparkDeclarativePipelineComponent) + yield component, defs + + def _ds( name: str, dataset_type: str = "table", @@ -28,6 +63,71 @@ def _ds( ) +def test_basic_component_load_via_sandbox() -> None: + """YAML -> Resolver -> component instantiation via create_defs_folder_sandbox and load_component_and_build_defs.""" + with create_defs_folder_sandbox() as sandbox: + project_root = sandbox.project_root + spec_path = project_root / "spark-pipeline.yml" + spec_path.write_text("") + (project_root / "models.py").write_text("@dp.table\ndef my_table(): pass\n") + typename = "dagster_spark.components.spark_declarative_pipeline.component.SparkDeclarativePipelineComponent" + defs_path = sandbox.scaffold_component( + component_cls=SparkDeclarativePipelineComponent, + defs_yaml_contents={ + "type": typename, + "attributes": { + "pipeline_spec_path": str(spec_path), + "discovery_mode": "source_only", + "defs_state": {"management_type": "LOCAL_FILESYSTEM"}, + }, + }, + ) + with sandbox.load_component_and_build_defs(defs_path=defs_path) as ( + component, + defs, + ): + assert isinstance(component, SparkDeclarativePipelineComponent) + # No state yet, so no assets + assert len(defs.resolve_asset_graph().get_all_asset_keys()) == 0 + + +def test_component_load_with_defs_state_lifecycle() -> None: + """Full lifecycle: no state -> refresh -> build defs (using create_defs_folder_sandbox and scoped_definitions_load_context).""" + with create_defs_folder_sandbox() as sandbox: + project_root = sandbox.project_root + spec_path = project_root / "spark-pipeline.yml" + spec_path.write_text("") + (project_root / "models.py").write_text("@dp.table\ndef my_table(): pass\n") + typename = "dagster_spark.components.spark_declarative_pipeline.component.SparkDeclarativePipelineComponent" + defs_path = sandbox.scaffold_component( + component_cls=SparkDeclarativePipelineComponent, + defs_yaml_contents={ + "type": typename, + "attributes": { + "pipeline_spec_path": str(spec_path), + "discovery_mode": "source_only", + "defs_state": {"management_type": "LOCAL_FILESYSTEM"}, + }, + }, + ) + with sandbox.load_component_and_build_defs(defs_path=defs_path) as ( + component, + defs, + ): + assert isinstance(component, SparkDeclarativePipelineComponent) + assert len(defs.resolve_asset_graph().get_all_asset_keys()) == 0 + asyncio.run(component.refresh_state(sandbox.project_root)) + with ( + scoped_definitions_load_context(), + sandbox.load_component_and_build_defs(defs_path=defs_path) as ( + _component, + defs, + ), + ): + keys = defs.resolve_asset_graph().get_all_asset_keys() + assert keys == {AssetKey(["my_table"])} + + def test_build_defs_from_state_returns_valid_definitions_with_multi_asset() -> None: """build_defs_from_state returns a valid Definitions object containing a multi_asset.""" component = SparkDeclarativePipelineComponent( diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py index 7029f042c6a71..3ec00ec38b6b3 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -1,6 +1,8 @@ """Tests for Spark Declarative Pipeline discovery (dry-run parsing and discover_datasets_fn).""" import json +import tempfile +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -9,6 +11,7 @@ DuplicateDatasetNamesError, SparkPipelinesDryRunError, discover_datasets_fn, + discover_datasets_from_sources, discover_datasets_via_dry_run, extract_report, parse_dry_run_output_to_datasets, @@ -31,6 +34,23 @@ def test_parse_dry_run_output_to_datasets_parses_json_report() -> None: assert datasets[1].name == "dataset_b" +def test_parse_dry_run_output_to_datasets_filters_empty_inferred_deps() -> None: + """Empty or whitespace-only dependency names are excluded from inferred_deps to avoid empty AssetKey path components.""" + mock_report = { + "datasets": [ + { + "name": "dataset_a", + "type": "table", + "deps": ["valid_dep", "", " ", "\t", "another_valid"], + }, + ], + } + stdout = json.dumps(mock_report) + datasets = parse_dry_run_output_to_datasets(stdout) + assert len(datasets) == 1 + assert datasets[0].inferred_deps == ["valid_dep", "another_valid"] + + def test_extract_report_returns_dry_run_report() -> None: """extract_report returns a DryRunReport from valid JSON stdout.""" stdout = json.dumps({"datasets": [{"name": "foo"}]}) @@ -156,3 +176,148 @@ def test_discover_datasets_via_dry_run_uses_custom_cmd_and_extra_args() -> None: assert "--spec" in call_cmd assert "--output" in call_cmd assert "json" in call_cmd + + +# ---- Real-world fixture tests: noisy Spark stdout with JSON or text ---- + +NOISY_SPARK_STDOUT_WITH_JSON = """ +INFO: Building Spark session... +INFO: JVM started. +WARN: Some config key was deprecated. +{"datasets": [{"name": "catalog.schema.table_a", "type": "table"}, {"name": "catalog.schema.table_b", "type": "materialized_view"}]} +INFO: Session closed. +""" + +NOISY_SPARK_STDOUT_BULLETED_TEXT = """ +INFO: Building Spark session... +INFO: JVM started. +- catalog.schema.table_a +* catalog.schema.table_b +1. catalog.schema.table_c +INFO: Session closed. +- Starting JVM... +* Some other log line that is not a dataset +""" + +NOISY_SPARK_STDOUT_DATASET_PREFIX = """ +INFO: Log line +dataset: my_dataset +dataset: another.dataset.name +INFO: More logs +""" + + +def test_parse_dry_run_output_to_datasets_isolates_json_from_noisy_spark_stdout() -> None: + """parse_dry_run_output_to_datasets correctly extracts datasets from stdout that contains Spark INFO/WARN logs and a JSON block.""" + datasets = parse_dry_run_output_to_datasets(NOISY_SPARK_STDOUT_WITH_JSON) + names = [d.name for d in datasets] + assert "catalog.schema.table_a" in names + assert "catalog.schema.table_b" in names + assert len(datasets) == 2 + assert datasets[0].dataset_type == "table" + assert datasets[1].dataset_type == "materialized_view" + + +def test_parse_dry_run_output_to_datasets_isolates_bulleted_text_from_noisy_spark_stdout() -> None: + """parse_dry_run_output_to_datasets (text fallback) extracts only valid dataset ids from bulleted lines, ignoring log-like lines.""" + datasets = parse_dry_run_output_to_datasets(NOISY_SPARK_STDOUT_BULLETED_TEXT) + names = [d.name for d in datasets] + assert "catalog.schema.table_a" in names + assert "catalog.schema.table_b" in names + assert "catalog.schema.table_c" in names + # These should NOT be included (regex requires alphanumeric/underscore/dot only, no spaces) + assert "Starting JVM..." not in names + assert "Some other log line that is not a dataset" not in names + assert len(datasets) == 3 + + +def test_parse_dry_run_output_to_datasets_isolates_dataset_prefix_from_noisy_stdout() -> None: + """parse_dry_run_output_to_datasets (text fallback) extracts dataset: lines with valid identifiers.""" + datasets = parse_dry_run_output_to_datasets(NOISY_SPARK_STDOUT_DATASET_PREFIX) + names = [d.name for d in datasets] + assert "my_dataset" in names + assert "another.dataset.name" in names + assert len(datasets) == 2 + + +def test_extract_report_text_rejects_log_like_bullet_lines() -> None: + """Text fallback does not treat '- Starting JVM...' or similar as dataset names.""" + stdout = "- Starting JVM...\n* Some log message with spaces\n1. Not a valid id with spaces" + report = extract_report(stdout) + # No line is a valid dataset id (alphanumeric, underscores, dots only); report may be None or empty. + if report is not None: + for d in report.datasets: + assert all(c.isalnum() or c in "_." for c in d.name), ( + "Dataset names must be valid identifiers (no spaces)" + ) + report2 = extract_report("- Starting JVM...\n* Foo bar") + if report2 is not None: + names = [d.name for d in report2.datasets] + assert "Starting JVM..." not in names + assert "Foo bar" not in names + + +# ---- Unit tests for discover_datasets_from_sources ---- + + +def test_discover_datasets_from_sources_python_decorators() -> None: + """discover_datasets_from_sources finds @dp.table and @dp.materialized_view def names from .py files.""" + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "pipeline.py").write_text( + """ +import dp + +@dp.table +def my_table(): + pass + +@dp.materialized_view +def my_mv(): + pass +""" + ) + result = discover_datasets_from_sources(root / "pipeline.py") + names = [d.name for d in result] + assert "my_table" in names + assert "my_mv" in names + assert len(result) == 2 + + +def test_discover_datasets_from_sources_sql_create_table() -> None: + """discover_datasets_from_sources finds CREATE STREAMING TABLE / CREATE TABLE names from .sql files.""" + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "schema.sql").write_text( + """ +CREATE STREAMING TABLE IF NOT EXISTS catalog.schema.events; +CREATE TABLE catalog.schema.dim_customer; +""" + ) + result = discover_datasets_from_sources(root / "schema.sql") + names = [d.name for d in result] + assert "catalog.schema.events" in names + assert "catalog.schema.dim_customer" in names + assert len(result) == 2 + + +def test_discover_datasets_from_sources_mixed_py_and_sql() -> None: + """discover_datasets_from_sources finds datasets from both .py and .sql under the pipeline spec directory.""" + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "spec.yaml").write_text( + "" + ) # pipeline spec path must exist so parent is used as root + (root / "a.py").write_text( + """ +@dp.table +def py_table(): + pass +""" + ) + (root / "b.sql").write_text("CREATE TABLE sql_dataset;") + result = discover_datasets_from_sources(root / "spec.yaml") + names = [d.name for d in result] + assert "py_table" in names + assert "sql_dataset" in names + assert len(result) == 2 From c13610a9b41d8791d1929c12d8283e63f3e690af Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Sat, 14 Mar 2026 14:38:09 +0200 Subject: [PATCH 14/17] fix remarks --- .../spark_declarative_pipeline/component.py | 30 ++--- .../spark_declarative_pipeline/discovery.py | 4 +- .../spark_declarative_pipeline/resource.py | 26 ++--- .../test_discovery.py | 11 ++ .../test_resource.py | 108 +++++++----------- 5 files changed, 79 insertions(+), 100 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py index d9068d82bb219..694213b631d5d 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -100,7 +100,7 @@ def write_state_to_path(self, state_path: Path) -> None: pipeline_spec_path=self.pipeline_spec_path, ) state_path.parent.mkdir(parents=True, exist_ok=True) - state_path.write_text(serialize_value(state)) + state_path.write_text(serialize_value(state), encoding="utf-8") def get_asset_spec(self, dataset: DiscoveredDataset) -> AssetSpec: """Build an AssetSpec for a discovered dataset. Override to customize key/metadata/group. @@ -189,25 +189,25 @@ def _spark_pipeline_asset( _context: dg.AssetExecutionContext, spark_pipelines: SparkPipelinesResource, ) -> Any: - # When the entire graph is executed (not a subset), pass keys=None to allow - # --full-refresh-all and avoid OS argument length limits. - # Inner param _context avoids shadowing outer context (ComponentLoadContext); - is_subset = getattr(_context, "is_subset", True) - if is_subset: - keys = ( - list(_context.selected_asset_keys) - if _context.selected_asset_keys - else [s.key for s in asset_specs] - ) - else: - keys = None - yield from spark_pipelines.run_and_observe( + selected_keys = ( + list(_context.selected_asset_keys) + if _context.selected_asset_keys + else [s.key for s in asset_specs] + ) + is_subset = ( + len(_context.selected_asset_keys) < len(asset_specs) + if _context.selected_asset_keys + else False + ) + spark_pipelines.run_and_observe( context=_context, pipeline_spec_path=resolved_spec_path, working_dir=working_dir, execution_mode=execution_mode, - asset_keys=keys, + asset_keys=selected_keys if is_subset else None, ) + for key in selected_keys: + yield dg.MaterializeResult(asset_key=key) return Definitions( assets=[_spark_pipeline_asset], diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index 117ee0ded642e..1ca0b40d1991f 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -350,8 +350,8 @@ def _extract_report_text(stdout: str) -> DryRunReport | None: line = raw_line.strip() if not line or line.startswith("#"): continue - # Match "dataset: " or "- " or " - " or "1. "; require valid dataset id (alphanumeric, underscores, dots) - _dataset_id_pattern = r"[a-zA-Z0-9_.]+" + # Match "dataset: " or "- " or " - " or "1. "; require valid dataset id (alphanumeric, underscores, dots, hyphens) + _dataset_id_pattern = r"[a-zA-Z0-9_.-]+" for pattern in ( rf"dataset:\s*({_dataset_id_pattern})\s*$", rf"^[-*]\s*({_dataset_id_pattern})\s*$", diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py index d272d09201e06..e4e9e15be9eb4 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -1,17 +1,16 @@ """Resource for running Spark Declarative Pipelines and discovering datasets. SparkPipelinesResource provides discover_datasets (via dry-run or source_only) and -run_and_observe (run spark-pipelines with log streaming and MaterializeResult yields). +run_and_observe (run spark-pipelines with log streaming). The asset yields MaterializeResults. """ import os import subprocess from collections import deque -from collections.abc import Iterator from pathlib import Path from typing import Any, Literal -from dagster import AssetKey, ConfigurableResource, MaterializeResult +from dagster import AssetKey, ConfigurableResource from pydantic import Field from dagster_spark.components.spark_declarative_pipeline.discovery import ( @@ -79,13 +78,13 @@ def run_and_observe( execution_mode: ExecutionMode = "incremental", extra_args: list[str] | None = None, asset_keys: list[AssetKey] | None = None, - ) -> Iterator[MaterializeResult]: - """Run spark-pipelines run with log streaming; yield MaterializeResult per asset on success. + ) -> None: + """Run spark-pipelines run with log streaming; does not yield (asset yields MaterializeResults). Uses Popen to stream stdout/stderr line-by-line and logs each line via context.log.info. Passes --full-refresh or --refresh based on execution_mode, then optional comma-separated - dataset list from asset_keys. Yields MaterializeResult per asset on success when - returncode == 0; otherwise raises SparkPipelinesExecutionError with the captured log. + dataset list from asset_keys. The calling multi_asset must yield one MaterializeResult per + selected asset key after this returns. Args: context: Asset execution context (used for context.log.info). @@ -95,9 +94,6 @@ def run_and_observe( extra_args: Optional extra CLI arguments appended to the command. asset_keys: Optional list of asset keys to materialize (passed as dataset list). - Yields: - MaterializeResult for each materialized asset on success. - Raises: SparkPipelinesExecutionError: If spark-pipelines run exits with non-zero return code. """ @@ -154,13 +150,7 @@ def run_and_observe( returncode=returncode, ) - if asset_keys: - for k in asset_keys: - yield MaterializeResult(asset_key=k) - return - - # No specific asset keys requested; complete gracefully without yielding (avoids UnexpectedAssetMaterializationError) - if context is not None and hasattr(context, "log"): + if asset_keys is None and context is not None and hasattr(context, "log"): context.log.info( - "No specific asset keys requested; spark-pipelines run completed successfully." + "spark-pipelines run completed successfully (full graph; asset will yield results)." ) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py index 3ec00ec38b6b3..d15a3656595d2 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -240,6 +240,17 @@ def test_parse_dry_run_output_to_datasets_isolates_dataset_prefix_from_noisy_std assert len(datasets) == 2 +def test_extract_report_text_accepts_hyphenated_dataset_ids() -> None: + """Text fallback accepts dataset identifiers with hyphens (e.g. my-catalog.db.table).""" + stdout = "- my-catalog.schema.events\n* other-dataset" + report = extract_report(stdout) + assert report is not None + names = [d.name for d in report.datasets] + assert "my-catalog.schema.events" in names + assert "other-dataset" in names + assert len(report.datasets) == 2 + + def test_extract_report_text_rejects_log_like_bullet_lines() -> None: """Text fallback does not treat '- Starting JVM...' or similar as dataset names.""" stdout = "- Starting JVM...\n* Some log message with spaces\n1. Not a valid id with spaces" diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py index 9a762dd8691aa..89a3a3692f832 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -1,18 +1,17 @@ -"""Tests for SparkPipelinesResource run_and_observe (log streaming and MaterializeResult yields).""" +"""Tests for SparkPipelinesResource run_and_observe (log streaming; asset yields MaterializeResults).""" from unittest.mock import MagicMock, patch import pytest from dagster import AssetKey -from dagster._core.definitions.result import MaterializeResult from dagster_spark.components.spark_declarative_pipeline.discovery import ( SparkPipelinesExecutionError, ) from dagster_spark.components.spark_declarative_pipeline.resource import SparkPipelinesResource -def test_run_and_observe_yields_materialize_results_for_selected_datasets() -> None: - """run_and_observe yields correct MaterializeResults for the selected asset keys on success.""" +def test_run_and_observe_completes_successfully_with_asset_keys() -> None: + """run_and_observe runs the subprocess and returns None; the asset yields MaterializeResults.""" mock_context = MagicMock() asset_keys = [ AssetKey(["dataset_a"]), @@ -28,19 +27,14 @@ def test_run_and_observe_yields_materialize_results_for_selected_datasets() -> N mock_popen.return_value = proc resource = SparkPipelinesResource() - results = list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/to/spec.yaml", - execution_mode="incremental", - asset_keys=asset_keys, - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/to/spec.yaml", + execution_mode="incremental", + asset_keys=asset_keys, ) - assert len(results) == 2 - assert all(isinstance(r, MaterializeResult) for r in results) - assert results[0].asset_key == AssetKey(["dataset_a"]) - assert results[1].asset_key == AssetKey(["dataset_b"]) + mock_popen.assert_called_once() def test_run_and_observe_streams_logs_via_context() -> None: @@ -56,12 +50,10 @@ def test_run_and_observe_streams_logs_via_context() -> None: mock_popen.return_value = proc resource = SparkPipelinesResource() - list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - asset_keys=[AssetKey(["a"])], - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], ) assert mock_context.log.info.call_count >= 2 @@ -84,19 +76,17 @@ def test_run_and_observe_raises_with_captured_log_on_nonzero_exit() -> None: resource = SparkPipelinesResource() with pytest.raises(SparkPipelinesExecutionError) as exc_info: - list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - asset_keys=[AssetKey(["a"])], - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], ) assert exc_info.value.returncode == 1 assert "error line" in (exc_info.value.stderr or "") -def test_run_and_observe_yields_nothing_when_no_asset_keys() -> None: - """When asset_keys is empty, run_and_observe yields nothing and completes gracefully (logs only).""" +def test_run_and_observe_completes_when_asset_keys_none() -> None: + """When asset_keys is None (full graph), run_and_observe runs and logs; asset yields results.""" mock_context = MagicMock() with patch( "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" @@ -108,19 +98,15 @@ def test_run_and_observe_yields_nothing_when_no_asset_keys() -> None: mock_popen.return_value = proc resource = SparkPipelinesResource() - results = list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - asset_keys=None, - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=None, ) - assert len(results) == 0 mock_context.log.info.assert_any_call( - "No specific asset keys requested; spark-pipelines run completed successfully." + "spark-pipelines run completed successfully (full graph; asset will yield results)." ) - # Incremental run with no asset_keys must not append --refresh (bare run = full pipeline) call_cmd = mock_popen.call_args[0][0] assert "--refresh" not in call_cmd @@ -141,12 +127,10 @@ def test_run_and_observe_passes_dot_notation_datasets_to_cli() -> None: mock_popen.return_value = proc resource = SparkPipelinesResource() - list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - asset_keys=asset_keys, - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=asset_keys, ) call_cmd = mock_popen.call_args[0][0] @@ -173,12 +157,10 @@ def test_run_and_observe_uses_configurable_cmd_and_run_extra_args() -> None: spark_pipelines_cmd="/usr/local/bin/spark-pipelines", run_extra_args=["--option", "value"], ) - list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - asset_keys=[AssetKey(["a"])], - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], ) call_cmd = mock_popen.call_args[0][0] @@ -202,13 +184,11 @@ def test_run_and_observe_full_refresh_no_asset_keys_uses_full_refresh_all() -> N mock_popen.return_value = proc resource = SparkPipelinesResource() - list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - execution_mode="full_refresh", - asset_keys=None, - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + execution_mode="full_refresh", + asset_keys=None, ) call_cmd = mock_popen.call_args[0][0] @@ -217,8 +197,8 @@ def test_run_and_observe_full_refresh_no_asset_keys_uses_full_refresh_all() -> N assert "--full-refresh" not in call_cmd -def test_run_and_observe_only_yields_on_success() -> None: - """MaterializeResults are only yielded when process returncode is 0.""" +def test_run_and_observe_raises_on_nonzero_exit() -> None: + """run_and_observe raises SparkPipelinesExecutionError when process returncode is not 0.""" mock_context = MagicMock() with patch( "dagster_spark.components.spark_declarative_pipeline.resource.subprocess.Popen" @@ -231,10 +211,8 @@ def test_run_and_observe_only_yields_on_success() -> None: resource = SparkPipelinesResource() with pytest.raises(SparkPipelinesExecutionError): - list( - resource.run_and_observe( - context=mock_context, - pipeline_spec_path="/path/spec.yaml", - asset_keys=[AssetKey(["a"])], - ) + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=[AssetKey(["a"])], ) From c6dd2607bd4de024bc03ed0627e1872b25910585 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Sat, 14 Mar 2026 15:19:08 +0200 Subject: [PATCH 15/17] remove unnecessary pyright type ignores --- .../components/spark_declarative_pipeline/discovery.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py index 1ca0b40d1991f..b97d717e76872 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -13,9 +13,9 @@ from typing import TYPE_CHECKING, Any, Literal from dagster import get_dagster_logger -from dagster_shared import check # type: ignore[reportMissingImports] -from dagster_shared.record import record # type: ignore[reportMissingImports] -from dagster_shared.serdes import whitelist_for_serdes # type: ignore[reportMissingImports] +from dagster_shared import check +from dagster_shared.record import record +from dagster_shared.serdes import whitelist_for_serdes # Guardrail to prevent daemon hangs. DRY_RUN_TIMEOUT_SECONDS = 60 From 46da043a8b17db6fbec4688413763eb3a3f0d8c2 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Fri, 20 Mar 2026 13:22:06 +0200 Subject: [PATCH 16/17] create test integration file --- .../test_integration.py | 156 ++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py new file mode 100644 index 0000000000000..911d00b4a65d8 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py @@ -0,0 +1,156 @@ +import os +import shutil +import tempfile +from pathlib import Path + +import pytest +from dagster_spark.components.spark_declarative_pipeline.discovery import discover_datasets_fn + + +# Check if the spark-pipelines CLI is available on the system. +def _spark_pipelines_available_locally() -> bool: + if ( + shutil.which("spark-pipelines") is not None + or shutil.which("spark-pipelines.cmd") is not None + or shutil.which("spark-pipelines.bat") is not None + ): + return True + + try: + import pyspark # type: ignore + + pyspark_bin_dir = Path(pyspark.__file__).resolve().parent / "bin" + for candidate in ("spark-pipelines.cmd", "spark-pipelines.bat", "spark-pipelines"): + if (pyspark_bin_dir / candidate).exists(): + return True + except ModuleNotFoundError: + pass + + return False + + +HAS_SPARK_PIPELINES = _spark_pipelines_available_locally() + +# Skip this entire test module if Spark 4.0+ is not installed locally. +pytestmark = pytest.mark.skipif( + not HAS_SPARK_PIPELINES, + reason="spark-pipelines CLI not found on this machine", +) + + +def _find_spark_pipelines_script() -> str | None: + script = ( + shutil.which("spark-pipelines.cmd") + or shutil.which("spark-pipelines.bat") + or shutil.which("spark-pipelines") + ) + if script: + return script + + try: + import pyspark # type: ignore + + pyspark_bin_dir = Path(pyspark.__file__).resolve().parent / "bin" + candidate = pyspark_bin_dir / "spark-pipelines" + if candidate.exists(): + return str(candidate) + + for fallback in ("spark-pipelines.cmd", "spark-pipelines.bat"): + p = pyspark_bin_dir / fallback + if p.exists(): + return str(p) + except ModuleNotFoundError: + pass + + return None + + +def _detect_java_home() -> str | None: + """Best-effort JAVA_HOME detection for local testing.""" + if os.environ.get("JAVA_HOME"): + return os.environ["JAVA_HOME"] + + java_exe = shutil.which("java") + if java_exe: + java_path = Path(java_exe).resolve() + if java_path.parent.name.lower() == "bin": + return str(java_path.parent.parent) + + if os.name == "nt": + program_files = Path(os.environ.get("ProgramFiles", r"C:\Program Files")) + java_dir = program_files / "Java" + if java_dir.exists(): + candidates = [ + p + for p in java_dir.iterdir() + if p.is_dir() and p.name.lower().startswith(("jdk", "jre")) + ] + if candidates: + candidates.sort(key=lambda p: p.name, reverse=True) + return str(candidates[0]) + + return None + + +def test_real_spark_dry_run_integration() -> None: + """Integration test that invokes the real spark-pipelines CLI and tests the fallback mechanism.""" + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + storage_path = (root / "storage").as_uri() + + # 1. Scaffold a minimal SDP project + spec_path = root / "spark-pipeline.yml" + spec_path.write_text( + f"name: integration_test_pipeline\n" + f"storage: {storage_path}\n" + "catalog: spark_catalog\n" + "database: default\n" + "libraries:\n" + " - glob:\n" + " include: 'models.py'\n", + encoding="utf-8", + ) + + models_path = root / "models.py" + models_path.write_text( + "import pyspark.pipelines as dp\n" + "from pyspark.sql import SparkSession\n\n" + "@dp.table(name='real_integration_table')\n" + "def real_integration_table():\n" + " spark = SparkSession.builder.getOrCreate()\n" + " return spark.readStream.format('rate').load()\n\n" + "@dp.materialized_view(name='real_integration_mv')\n" + "def real_integration_mv():\n" + " spark = SparkSession.builder.getOrCreate()\n" + " return spark.range(0)\n", + encoding="utf-8", + ) + + # Ensure JAVA_HOME is set for the subprocess + java_home = _detect_java_home() + if java_home: + os.environ["JAVA_HOME"] = java_home + + # 2. Call the discovery function (which encapsulates the CLI call and fallback logic) + try: + cmd = _find_spark_pipelines_script() or "spark-pipelines" + datasets = discover_datasets_fn( + spark_pipelines_cmd=cmd, + pipeline_spec_path=spec_path, + discovery_mode="dry_run_with_fallback", + dry_run_extra_args=[], + ) + except Exception as e: + pytest.skip(f"Failed to execute discovery locally: {e}") + + # 3. Assertions + if len(datasets) == 0: + pytest.fail("Discovery found 0 datasets. Both CLI dry-run and source fallback failed.") + + assert len(datasets) == 2 + names = {ds.name for ds in datasets} + assert "real_integration_table" in names + assert "real_integration_mv" in names + + # Verify it actually used the fallback mechanism since the Spark CLI outputs no metadata + assert datasets[0].discovery_method == "source_fallback" From a64fb29403aade39f1144701776765add613eb33 Mon Sep 17 00:00:00 2001 From: michalcabir-ui Date: Mon, 23 Mar 2026 20:03:59 +0200 Subject: [PATCH 17/17] pyright fix --- .../components/spark_declarative_pipeline/test_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py index 911d00b4a65d8..07ce6596fcb1d 100644 --- a/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_integration.py @@ -17,7 +17,7 @@ def _spark_pipelines_available_locally() -> bool: return True try: - import pyspark # type: ignore + import pyspark pyspark_bin_dir = Path(pyspark.__file__).resolve().parent / "bin" for candidate in ("spark-pipelines.cmd", "spark-pipelines.bat", "spark-pipelines"): @@ -48,7 +48,7 @@ def _find_spark_pipelines_script() -> str | None: return script try: - import pyspark # type: ignore + import pyspark pyspark_bin_dir = Path(pyspark.__file__).resolve().parent / "bin" candidate = pyspark_bin_dir / "spark-pipelines"