diff --git a/docs/docs/integrations/libraries/spark/index.md b/docs/docs/integrations/libraries/spark/index.md index 9f27fff7cf334..5f3f2a542c155 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..8b083c9ccc924 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/__init__.py @@ -0,0 +1,24 @@ +"""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, + 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, + 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..694213b631d5d --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/component.py @@ -0,0 +1,215 @@ +"""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 + +import dagster as dg +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 +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() + 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) +@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: 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) + + @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 + 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, + pipeline_spec_path=self.pipeline_spec_path, + ) + state_path.parent.mkdir(parents=True, exist_ok=True) + 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. + + Args: + dataset: Discovered dataset from state. + + Returns: + 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, {}) + 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=metadata, + group_name=attrs.get("group_name"), + tags=attrs.get("tags"), + kinds={"spark", dataset.dataset_type}, + ) + + def build_defs_from_state( + self, + context: ComponentLoadContext, + state_path: Path | None, + ) -> 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 + if ds.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 / Path(pipeline_spec_path)).resolve() + else: + resolved_spec_path = Path(pipeline_spec_path) + working_dir = context.path + 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, + spark_pipelines: SparkPipelinesResource, + ) -> Any: + 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=selected_keys if is_subset else None, + ) + for key in selected_keys: + yield dg.MaterializeResult(asset_key=key) + + return Definitions( + assets=[_spark_pipeline_asset], + 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 new file mode 100644 index 0000000000000..b97d717e76872 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/discovery.py @@ -0,0 +1,625 @@ +"""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) use @record with +whitelist_for_serdes for Dagster serialize_value/deserialize_value compatibility. +""" + +import json +import re +import subprocess +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +from dagster import get_dagster_logger +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 + +DiscoveryMode = Literal["dry_run_only", "dry_run_with_fallback", "source_only"] + + +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. + 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 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 + + +@record +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). + inferred_deps: Upstream dataset names extracted from the node (e.g. deps, dependencies). + """ + + name: str + raw: dict[str, Any] + inferred_deps: tuple[str, ...] = () + + if TYPE_CHECKING: + + 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). + + 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 + + if TYPE_CHECKING: + + def __init__( + self, + *, + datasets: list[DryRunDatasetNode], + raw: dict[str, Any] | None = None, + ) -> None: ... + + +@whitelist_for_serdes +@record +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). + 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 + dataset_type: str + source_file: str | None + source_line: int | None + 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 +@record +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 + + 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, + 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. + + Args: + 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. + + Raises: + SparkPipelinesDryRunError: If the command fails or times out. + """ + path_str = str(pipeline_spec_path) + cmd = [spark_pipelines_cmd, "dry-run", "--spec", path_str] + if extra_args: + cmd.extend(extra_args) + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + cwd=working_dir, + timeout=DRY_RUN_TIMEOUT_SECONDS, + ) + + if result.returncode != 0: + raise SparkPipelinesDryRunError( + f"{spark_pipelines_cmd} 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) + + +# 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. + + 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 parsing the whole output as JSON first + try: + data = json.loads(stripped) + 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 + # 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: + 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 + + +def _extract_deps_from_node(item: dict[str, Any]) -> tuple[str, ...]: + """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(s for d in deps if isinstance(d, str) and (s := str(d).strip()) != "") + + +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) + 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}, inferred_deps=())) + + 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: " 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*$", + rf"^\d+\.\s*({_dataset_id_pattern})\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}, inferred_deps=()) + ) + break + + if not datasets: + return 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 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.]+)", + 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.is_file(): + 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 [] + + 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", + 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, + 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): + 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, + dataset_type=dataset_type_sql, + source_file=str(path), + source_line=None, + inferred_deps=[], + discovery_method="source_fallback", + ) + ) + except Exception as e: + logger.warning( + "Failed to discover datasets from sources under %s: %s", + root, + e, + exc_info=True, + ) + return [] + + logger.info( + "Scanned %d files for static dataset discovery, found %d datasets.", + scanned_files, + len(result), + ) + 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 + with explicit dataset_type, source_file, source_line, inferred_deps, discovery_method. + + 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 [] + + result: list[DiscoveredDataset] = [] + for node in report.datasets: + raw = node.raw + 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, + 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: + """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, + 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. + + - 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 discover from sources. + + Duplicate dataset names (after normalization: strip + lowercase) are a hard error. + + 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. + 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. + + Raises: + DuplicateDatasetNamesError: If the discovered list contains duplicate names. + """ + 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) + 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))) + + try: + 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: + if discovery_mode == "dry_run_only": + raise + 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: + if source_list is not None: + return _return(source_list) + return _return(discover_datasets_from_sources(Path(pipeline_spec_path))) + + 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 new file mode 100644 index 0000000000000..e4e9e15be9eb4 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark/components/spark_declarative_pipeline/resource.py @@ -0,0 +1,156 @@ +"""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). The asset yields MaterializeResults. +""" + +import os +import subprocess +from collections import deque +from pathlib import Path +from typing import Any, Literal + +from dagster import AssetKey, ConfigurableResource +from pydantic import Field + +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + DiscoveredDataset, + DiscoveryMode, + SparkPipelinesExecutionError, + discover_datasets_fn, +) + +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. + """ + + 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, + 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, + spark_pipelines_cmd=self.spark_pipelines_cmd, + dry_run_extra_args=self.dry_run_extra_args, + ) + + 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, + ) -> 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. 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). + 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). + + Raises: + 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] + 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: + 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) + + 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: deque[str] = deque(maxlen=1000) + 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 + + if returncode != 0: + captured = "\n".join(log_lines) if log_lines else "(no output)" + raise SparkPipelinesExecutionError( + f"spark-pipelines run failed with return code {returncode}", + stderr=captured, + returncode=returncode, + ) + + if asset_keys is None and context is not None and hasattr(context, "log"): + context.log.info( + "spark-pipelines run completed successfully (full graph; asset will yield results)." + ) 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..416121b864f0d --- /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": "spark-pipeline.yml", + "discovery_mode": "dry_run_with_fallback", + }, + ) 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..4c1fa1b48bc74 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_component.py @@ -0,0 +1,254 @@ +"""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, + SparkPipelineState, +) + + +@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", + 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_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( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + ) + datasets = [ + _ds("table_a"), + _ds("table_b"), + ] + 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_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", + 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"] + 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( + pipeline_spec_path="pipeline.yaml", + discovery_mode="source_only", + asset_attributes_by_dataset={}, # no overrides + ) + datasets = [ + _ds("table_a"), + _ds("temp_view_x", 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 = [ + _ds("table_a"), + _ds("temp_view_x", 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..d15a3656595d2 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_discovery.py @@ -0,0 +1,334 @@ +"""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 +from dagster_spark.components.spark_declarative_pipeline.discovery import ( + DiscoveredDataset, + DuplicateDatasetNamesError, + SparkPipelinesDryRunError, + discover_datasets_fn, + discover_datasets_from_sources, + 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].dataset_type == "table" + 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"}]}) + 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 _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 = [_discovered("fallback_ds")] + 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 = [_discovered("a"), _discovered("b")] + 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_fn_raises_on_duplicate_dataset_names() -> None: + """Duplicate dataset names (after normalization) raise DuplicateDatasetNamesError.""" + source = [_discovered("Foo"), _discovered("foo")] + 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( + "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 + + +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 call_cmd[1] == "dry-run" + 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_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" + 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 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..07ce6596fcb1d --- /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 + + 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 + + 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" 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..89a3a3692f832 --- /dev/null +++ b/python_modules/libraries/dagster-spark/dagster_spark_tests/components/spark_declarative_pipeline/test_resource.py @@ -0,0 +1,218 @@ +"""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_spark.components.spark_declarative_pipeline.discovery import ( + SparkPipelinesExecutionError, +) +from dagster_spark.components.spark_declarative_pipeline.resource import SparkPipelinesResource + + +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"]), + 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() + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/to/spec.yaml", + execution_mode="incremental", + asset_keys=asset_keys, + ) + + mock_popen.assert_called_once() + + +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() + 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 SparkPipelinesExecutionError 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(SparkPipelinesExecutionError) as exc_info: + 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_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" + ) 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() + resource.run_and_observe( + context=mock_context, + pipeline_spec_path="/path/spec.yaml", + asset_keys=None, + ) + + mock_context.log.info.assert_any_call( + "spark-pipelines run completed successfully (full graph; asset will yield results)." + ) + 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: + """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() + 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] + 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 + + +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"], + ) + 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 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() + 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_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" + ) 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(SparkPipelinesExecutionError): + 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/pyproject.toml b/python_modules/libraries/dagster-spark/pyproject.toml index aa22643ba7698..4fc4a6bd7cb62 100644 --- a/python_modules/libraries/dagster-spark/pyproject.toml +++ b/python_modules/libraries/dagster-spark/pyproject.toml @@ -23,3 +23,6 @@ dependencies = [ [project.urls] Homepage = "https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-spark" + +[project.entry-points."dagster_dg_cli.registry_modules"] +dagster_spark = "dagster_spark" \ No newline at end of file