From 29a9dcd03e452848227675192dec1f3ab099594d Mon Sep 17 00:00:00 2001 From: "fabian.ade" Date: Mon, 15 Jun 2026 14:40:10 +0200 Subject: [PATCH 1/3] draft solution for fork pr handling --- .github/workflows/acceptance.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 16a3e4cc..2ede4006 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -37,6 +37,10 @@ jobs: labels: linux-ubuntu-latest permissions: id-token: write + # Fork PRs get no OIDC token / secrets from GitHub, so JFrog auth (and therefore + # dependency installation) cannot run. Skip CI for them; fork PRs are to be tested + # by the reviewer(s) / maintainer(s) before merging. + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -63,6 +67,8 @@ jobs: needs: [ not-a-fork, lint ] permissions: id-token: write + # See the note on `lint`: fork PRs cannot authenticate to JFrog, so skip CI for them. + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork }} strategy: matrix: python-version: [ '3.12' ] From a5a3e3f91d79455031a3d34e9fc1199cbab0434d Mon Sep 17 00:00:00 2001 From: "fabian.ade" Date: Wed, 29 Jul 2026 18:07:25 +0200 Subject: [PATCH 2/3] WIP implemented concept which allows us to define custom solvers for different scenarios --- docs/impulse/docs/data_model/ingestion.md | 107 +++++- .../analyze/query/solvers/__init__.py | 22 +- .../analyze/query/solvers/default_solver.py | 35 +- .../analyze/query/solvers/query_solver.py | 55 +++- .../analyze/query/solvers/registry.py | 139 ++++++++ .../analyze/query/solvers/solver_context.py | 48 +++ src/impulse_reporting/config/config_parser.py | 44 ++- src/impulse_reporting/core/report.py | 52 +-- .../integration/custom_solver_test.py | 304 ++++++++++++++++++ .../analyze/query/solvers/from_config_test.py | 82 +++++ .../analyze/query/solvers/registry_test.py | 153 +++++++++ .../integration/custom_solver_report_test.py | 272 ++++++++++++++++ .../config/solver_registry_config_test.py | 118 +++++++ 13 files changed, 1392 insertions(+), 39 deletions(-) create mode 100644 src/impulse_query_engine/analyze/query/solvers/registry.py create mode 100644 src/impulse_query_engine/analyze/query/solvers/solver_context.py create mode 100644 tests/impulse_query_engine/integration/custom_solver_test.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/solvers/from_config_test.py create mode 100644 tests/impulse_query_engine/unit/analyze/query/solvers/registry_test.py create mode 100644 tests/impulse_reporting/integration/custom_solver_report_test.py create mode 100644 tests/impulse_reporting/unit/config/solver_registry_config_test.py diff --git a/docs/impulse/docs/data_model/ingestion.md b/docs/impulse/docs/data_model/ingestion.md index 133a247d..02a452e0 100644 --- a/docs/impulse/docs/data_model/ingestion.md +++ b/docs/impulse/docs/data_model/ingestion.md @@ -233,7 +233,8 @@ computed-column joins, JSON-encoded values, multi-column composite keys that need pre-processing, etc. — you can implement a custom solver by subclassing [`QuerySolver`](../references/api/impulse_query_engine/analyze/query/solvers/query_solver.md) -(or one of the existing solvers) and registering it in your report config. +(or one of the existing solvers, usually `DefaultSolver`) and selecting it +by name in your report config. This is significantly more invested than the `SolverConfig` path: you take on responsibility for the four solver pipeline stages @@ -245,6 +246,110 @@ yourself heading down this path, it is usually worth first asking whether a one-time ETL job to produce the standard silver-layer shape would be cheaper. +#### Registering a custom solver + +Decorate your subclass with `@register_solver(name)`. The report config then +selects it by that name via the existing `query_engine.solver` field — the +same field used for the built-in `DefaultSolver`. + +```python +from impulse_query_engine.analyze.query.solvers import ( + register_solver, + SolverConfig, +) +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver + + +class CustomSolverConfig(SolverConfig): + raw_signal_table: str # required — validated when the config is parsed + position_signal_name: str = "POSITION" # optional, with a default + + +@register_solver("CustomSolver", CustomSolverConfig) +class CustomSolver(DefaultSolver): + """Reshapes a custom raw table into the Impulse silver schema.""" + ... +``` + +```json +{ + "query_engine": { + "solver": "CustomSolver", + "solver_config": { + "raw_signal_table": "my_catalog.my_schema.raw_signals", + "channels": { "column_name_mapping": { "signal_value": "value" } } + } + } +} +``` + +Passing `CustomSolverConfig` to `register_solver` makes Impulse validate the +`solver_config` block through that subclass, so its extra fields are checked +at config-load time: a missing required `raw_signal_table`, or a mistyped +key, raises a `ValidationError` up front rather than being silently dropped. +Inside the solver you read them with normal typed attribute access +(`self.config.raw_signal_table`). + +**Loading the solver.** `@register_solver` runs only when its module is +imported, so import your solver package once in the driver before building +the report — that import is what registers the name: + +```python +import sys +sys.path.append("/Workspace/Repos/me/my_solvers_root") # dir containing my_solvers/ + +import my_solvers.custom_solver # runs @register_solver("CustomSolver") + +from impulse_reporting.core.report import Report + +report = Report(name="my_report", spark=spark, workspace_client=ws, + config_path="report_config.json") +report.determine_report() +``` + +If the name is not registered when the config is parsed, Impulse raises an +error listing the known names — a quick signal that the import is missing or +misspelled. (For an installed wheel the `sys.path` line is unnecessary; the +`import` alone suffices. In a notebook you can also just define the solver in +a cell — running that cell registers it.) + +Selection is by registered name only: a report config can never, on its own, +cause Impulse to import an arbitrary class. Which solvers exist is governed +entirely by what the driver imports. + +#### Example: a column-redaction solver + +A solver override plus one config field gives a config-switched masking layer: + +```python +import pyspark.sql.functions as F +from impulse_query_engine.analyze.query.solvers import register_solver, SolverConfig +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver + + +class RedactConfig(SolverConfig): + redact_columns: list[str] = [] + + +@register_solver("RedactingSolver", RedactConfig) +class RedactingSolver(DefaultSolver): + def solve(self, query, channels_df, selections, dtypes): + df = super().solve(query, channels_df, selections, dtypes) + for col in self.config.redact_columns: + if col in df.columns: + df = df.withColumn(col, F.lit(None).cast(df.schema[col].dataType)) + return df +``` + +```json +{ + "query_engine": { + "solver": "RedactingSolver", + "solver_config": { "redact_columns": ["vin", "gps_lat", "gps_lon"] } + } +} +``` + The general rule: **`SolverConfig` for naming differences, custom solver for structural differences, ETL into the standard shape for everything else.** diff --git a/src/impulse_query_engine/analyze/query/solvers/__init__.py b/src/impulse_query_engine/analyze/query/solvers/__init__.py index 1b736776..4182d6b1 100644 --- a/src/impulse_query_engine/analyze/query/solvers/__init__.py +++ b/src/impulse_query_engine/analyze/query/solvers/__init__.py @@ -1,4 +1,24 @@ from .default_solver import DefaultSolver, TimeSeriesCache +from .query_solver import QuerySolver +from .registry import ( + SolverRegistration, + is_registered, + register_solver, + registered_names, + resolve_registration, +) from .solver_config import SolverConfig +from .solver_context import SolverBuildContext -__all__ = ["DefaultSolver", "TimeSeriesCache", "SolverConfig"] +__all__ = [ + "DefaultSolver", + "TimeSeriesCache", + "SolverConfig", + "QuerySolver", + "SolverBuildContext", + "SolverRegistration", + "register_solver", + "resolve_registration", + "is_registered", + "registered_names", +] diff --git a/src/impulse_query_engine/analyze/query/solvers/default_solver.py b/src/impulse_query_engine/analyze/query/solvers/default_solver.py index 5e2e5d7d..858db8e4 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -14,6 +14,7 @@ from impulse_query_engine.model.series.sample_series import SampleSeries from .query_solver import QuerySolver +from .registry import register_solver from .series_cache import SeriesCache from .solver_config import RawEncoder, SolverConfig from .utils.interval_encoder import IntervalEncoder @@ -22,6 +23,8 @@ if TYPE_CHECKING: from impulse_query_engine.measurement_db import MeasurementDB + from .solver_context import SolverBuildContext + class TimeSeriesCache(SeriesCache): def __init__(self, pdf, col_map: dict[str, str]): @@ -134,6 +137,7 @@ def load_blob(self, mid, cid, uses_alias: bool = False): return SampleSeries(s[self._ts_col], s[self._te_col], values) +@register_solver("DefaultSolver", aliases=("DeltaSolver", "KeyValueStoreSolver")) class DefaultSolver(QuerySolver): """ The default query-engine solver. Adapts to the shape of the silver layer. @@ -197,6 +201,21 @@ def __init__( self.raw_encoder: RawEncoder = raw_encoder self.channel_encoder: RleEncoder | IntervalEncoder = self._build_channel_encoder() + @classmethod + def from_config(cls, ctx: "SolverBuildContext") -> "DefaultSolver": + """Build a :class:`DefaultSolver` from a :class:`SolverBuildContext`. + + Overrides the base hook to wire the SparkSession and the raw-data flags + that this solver requires beyond ``solver_config``. + """ + return cls( + ctx.spark, + config=ctx.solver_config, + is_raw_data=ctx.is_raw_data, + drop_implausible_data=ctx.drop_implausible_data, + raw_encoder=ctx.raw_encoder, + ) + def _build_channel_encoder(self) -> RleEncoder | IntervalEncoder: """Construct the raw -> interval encoder selected by ``raw_encoder``. @@ -261,7 +280,7 @@ def filter_container_tags(self, spark, query) -> DataFrame: required_elements.extend(filt.required_tags()) required_elements = set(required_elements) - tags = query.db.container_tags(self.spark) + tags = self.load_container_tags(query.db, self.spark) tags = self._apply_column_mapping(tags, self.config.container_tags.column_name_mapping) if self.config.project_id is not None: @@ -329,7 +348,7 @@ def filter_container_metrics( if pre_filtered_containers_df is not None: metrics = pre_filtered_containers_df else: - metrics = query.db.container_metrics(self.spark) + metrics = self.load_container_metrics(query.db, self.spark) metrics = self._apply_column_mapping( metrics, self.config.container_metrics.column_name_mapping @@ -403,7 +422,7 @@ def filter_channel_tags(self, spark, db, container_df, selectors) -> DataFrame: for selector in selectors: required_tags.update(selector.required_tags()) - tbl = db.channel_tags(spark) + tbl = self.load_channel_tags(db, spark) tbl = self._apply_column_mapping(tbl, self.config.channel_tags.column_name_mapping) expr = self._build_expr(selectors) @@ -456,7 +475,7 @@ def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: if len(selectors) == 0: return self._empty_channel_match_df(spark, db) - channel_metrics_df = db.channel_metrics(spark) + channel_metrics_df = self.load_channel_metrics(db, spark) channel_metrics_df = self._apply_column_mapping( channel_metrics_df, self.config.channel_metrics.column_name_mapping ) @@ -537,7 +556,7 @@ def filter_aliased_channel_metrics( if len(selectors) == 0: return self._empty_channel_match_df(spark, db) - channel_mapping = db.channel_mapping(spark) + channel_mapping = self.load_channel_mapping(db, spark) channel_mapping = self._apply_column_mapping( channel_mapping, self.config.channel_mapping.column_name_mapping ) @@ -552,7 +571,7 @@ def filter_aliased_channel_metrics( resolved_mapping = channel_mapping.where(self._build_expr(selectors)) - channel_metrics = db.channel_metrics(spark) + channel_metrics = self.load_channel_metrics(db, spark) channel_metrics = self._apply_column_mapping( channel_metrics, self.config.channel_metrics.column_name_mapping ) @@ -850,7 +869,7 @@ def _compute_conversion_factors(self, spark, query, channels_df: DataFrame) -> D :meth:`_validate_unit_conversion_table` for the underlying check. """ - uc_table = query.db.unit_conversion(spark) + uc_table = self.load_unit_conversion(query.db, spark) uc_table = self._apply_column_mapping( uc_table, self.config.unit_conversion.column_name_mapping ) @@ -1008,7 +1027,7 @@ def _prepare_channels_join(self, query, channels_df) -> tuple[DataFrame, DataFra if col_name in channels_df.columns: channels_df = channels_df.drop(col_name) - q = query.db.channels(self.spark) + q = self.load_channels(query.db, self.spark) q = self._apply_column_mapping(q, self.config.channels.column_name_mapping) if self.is_raw_data: diff --git a/src/impulse_query_engine/analyze/query/solvers/query_solver.py b/src/impulse_query_engine/analyze/query/solvers/query_solver.py index bfa35049..556e5b5a 100644 --- a/src/impulse_query_engine/analyze/query/solvers/query_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/query_solver.py @@ -12,6 +12,8 @@ if TYPE_CHECKING: from impulse_query_engine.measurement_db import MeasurementDB + from .solver_context import SolverBuildContext + from impulse_query_engine.analyze.metadata.time_series_expression import ( TimeSeriesSelector, ) @@ -35,6 +37,57 @@ class QuerySolver(ABC): def __init__(self, config: SolverConfig = None): self.config = config or SolverConfig() + @classmethod + def from_config(cls, ctx: "SolverBuildContext") -> "QuerySolver": + """Build a solver from a :class:`SolverBuildContext`. + + This is the uniform instantiation hook used by the report factory so + that solvers with different constructor signatures can all be built the + same way. The base implementation passes only ``solver_config`` (the + common case for a :class:`QuerySolver` subclass); solvers that need more + construction inputs — such as :class:`DefaultSolver`, which needs the + SparkSession and the raw-data flags — override this classmethod. + + Parameters + ---------- + ctx : SolverBuildContext + The construction inputs resolved from the report configuration. + + Returns + ------- + QuerySolver + A ready-to-use solver instance. + """ + return cls(config=ctx.solver_config) + + def load_container_tags(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``container_tags`` (narrow/EAV) table.""" + return db.container_tags(spark) + + def load_container_metrics(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``container_metrics`` table.""" + return db.container_metrics(spark) + + def load_channel_tags(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``channel_tags`` (narrow/EAV) table.""" + return db.channel_tags(spark) + + def load_channel_metrics(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``channel_metrics`` table.""" + return db.channel_metrics(spark) + + def load_channel_mapping(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``channel_mapping`` (alias) table.""" + return db.channel_mapping(spark) + + def load_unit_conversion(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``unit_conversion`` table.""" + return db.unit_conversion(spark) + + def load_channels(self, db: "MeasurementDB", spark) -> DataFrame: + """Read the raw ``channels`` (time-series data) table.""" + return db.channels(spark) + @staticmethod def _apply_column_mapping(df: DataFrame, mapping: dict[str, str]) -> DataFrame: """Rename DataFrame columns according to a physical → internal mapping.""" @@ -155,7 +208,7 @@ def _empty_channel_match_df(self, spark, db: MeasurementDB) -> DataFrame: schemas stay compatible regardless of the physical id types. """ ref = self._apply_column_mapping( - db.channel_metrics(spark), self.config.channel_metrics.column_name_mapping + self.load_channel_metrics(db, spark), self.config.channel_metrics.column_name_mapping ) return spark.createDataFrame( [], diff --git a/src/impulse_query_engine/analyze/query/solvers/registry.py b/src/impulse_query_engine/analyze/query/solvers/registry.py new file mode 100644 index 00000000..6cbb1324 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/solvers/registry.py @@ -0,0 +1,139 @@ +"""Registry mapping a solver name to its solver class and config class. + +The report config selects a solver by **name** (e.g. ``"solver": "DefaultSolver"``). +This registry resolves that name to the concrete :class:`QuerySolver` subclass to +instantiate **and** the :class:`SolverConfig` subclass to build for its +``solver_config`` block — so a customer can ship a solver that reshapes their raw +tables into the Impulse silver schema, plus a config subclass carrying extra +fields, without editing Impulse core. + +Built-in solvers self-register at import time (see the bottom of +``default_solver.py``). Customer solvers register themselves the same way from +their own package; importing that package runs the registration side effect. + +Selection is by registered name only: there is no fully-qualified-class-path +import path, so a report config can never, on its own, cause Impulse to import +and execute an arbitrary class. Which solvers exist is governed entirely by +what the driver imports. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import TypeVar + +from .query_solver import QuerySolver +from .solver_config import SolverConfig + +_SolverT = TypeVar("_SolverT", bound=type[QuerySolver]) + + +@dataclass(frozen=True) +class SolverRegistration: + """A registered solver: the class to instantiate and its config class.""" + + solver_cls: type[QuerySolver] + config_cls: type[SolverConfig] + + +_REGISTRY: dict[str, SolverRegistration] = {} + + +def register_solver( + name: str, + config_cls: type[SolverConfig] = SolverConfig, + *, + aliases: tuple[str, ...] = (), + overwrite: bool = False, +) -> Callable[[_SolverT], _SolverT]: + """Class decorator registering the decorated solver under *name*. + + Usage:: + + @register_solver("MySolver", MySolverConfig) + class MySolver(DefaultSolver): + ... + + Parameters + ---------- + name : str + The name used in the report config's ``query_engine.solver``. + config_cls : type[SolverConfig], optional + The :class:`SolverConfig` subclass to build for this solver's + ``solver_config`` block. Defaults to the base :class:`SolverConfig` + (for solvers that need no extra config fields). + aliases : tuple[str, ...], optional + Additional names resolving to the same registration (e.g. deprecated + solver names kept for backward compatibility). + overwrite : bool, optional + Allow replacing an existing registration. Defaults to ``False``, which + raises on a conflicting duplicate. + + Returns + ------- + Callable + A decorator that registers and returns the solver class unchanged. + + Raises + ------ + TypeError + If the decorated object is not a :class:`QuerySolver` subclass. + ValueError + If a name/alias is already registered to a different class and + *overwrite* is ``False``. + """ + + def decorator(solver_cls: _SolverT) -> _SolverT: + if not (isinstance(solver_cls, type) and issubclass(solver_cls, QuerySolver)): + raise TypeError(f"register_solver expects a QuerySolver subclass, got {solver_cls!r}.") + registration = SolverRegistration(solver_cls, config_cls) + for key in (name, *aliases): + existing = _REGISTRY.get(key) + if existing is not None and existing.solver_cls is not solver_cls and not overwrite: + raise ValueError( + f"Solver {key!r} is already registered to {existing.solver_cls.__name__}; " + f"pass overwrite=True to replace it." + ) + _REGISTRY[key] = registration + return solver_cls + + return decorator + + +def is_registered(name: str) -> bool: + """Return whether *name* resolves to a registered solver.""" + return name in _REGISTRY + + +def registered_names() -> list[str]: + """Return the sorted list of registered solver names and aliases.""" + return sorted(_REGISTRY) + + +def resolve_registration(name: str) -> SolverRegistration: + """Resolve a solver *name* to its :class:`SolverRegistration`. + + Parameters + ---------- + name : str + A registered solver name or alias. + + Returns + ------- + SolverRegistration + The resolved registration. + + Raises + ------ + KeyError + If *name* is not registered. The message lists the known names so a + missing driver-side import is easy to spot. + """ + try: + return _REGISTRY[name] + except KeyError: + raise KeyError( + f"Unknown solver {name!r}. Registered solvers: {registered_names()}. " + f"Import the package that registers it before building the report." + ) from None diff --git a/src/impulse_query_engine/analyze/query/solvers/solver_context.py b/src/impulse_query_engine/analyze/query/solvers/solver_context.py new file mode 100644 index 00000000..10cf357b --- /dev/null +++ b/src/impulse_query_engine/analyze/query/solvers/solver_context.py @@ -0,0 +1,48 @@ +"""Build context passed to :meth:`QuerySolver.from_config`. + +The context carries everything a solver needs to construct itself, decoupled +from the reporting-layer configuration model (``impulse_reporting`` must not be +imported here — the query engine is the lower layer). It is a frozen dataclass +rather than a Pydantic model because it holds a live ``SparkSession``. + +It is also the single extension point for solver construction inputs: new +inputs are added as fields here rather than to every solver's ``from_config`` +signature. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from .solver_config import RawEncoder, SolverConfig + +if TYPE_CHECKING: + from pyspark.sql import SparkSession + + +@dataclass(frozen=True) +class SolverBuildContext: + """Inputs used to build a :class:`QuerySolver` via ``from_config``. + + Attributes + ---------- + spark : SparkSession + Spark session used for query execution. + solver_config : SolverConfig or None + Optional per-table column mapping / filter configuration. Custom + solvers registered with a ``SolverConfig`` subclass receive that + subclass instance here (validated at config-parse time). + is_raw_data : bool + Whether the input data is raw point data rather than RLE intervals. + drop_implausible_data : bool + Whether to drop data points marked as implausible before processing. + raw_encoder : RawEncoder + Which encoder converts RAW point data into intervals for solving. + """ + + spark: SparkSession + solver_config: SolverConfig | None = None + is_raw_data: bool = False + drop_implausible_data: bool = False + raw_encoder: RawEncoder = RawEncoder.RLE diff --git a/src/impulse_reporting/config/config_parser.py b/src/impulse_reporting/config/config_parser.py index 6bff951e..e538cdd9 100644 --- a/src/impulse_reporting/config/config_parser.py +++ b/src/impulse_reporting/config/config_parser.py @@ -5,6 +5,7 @@ from pydantic import AfterValidator, BaseModel, field_validator, model_validator +from impulse_query_engine.analyze.query.solvers.registry import resolve_registration from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder, SolverConfig @@ -84,15 +85,20 @@ class DataType(StrEnum): RLE = "RLE" -class Solvers(Enum): +class Solvers(StrEnum): """ - Enumeration of available solver types for the query engine. + Names of the built-in solver types for the query engine. ``DEFAULT_SOLVER`` is the single, unified solver. ``DELTA_SOLVER`` and ``KEY_VALUE_STORE_SOLVER`` are **deprecated aliases** kept so that existing report configs continue to deserialize; both now resolve to the same ``DefaultSolver``. They will be removed in a future release. + This is a :class:`~enum.StrEnum`: each member *is* its string value, so the + ``query_engine.solver`` field is a plain ``str`` (accepting any registered + solver name, including customer solvers) while existing comparisons against + these members — ``qe.solver == Solvers.DEFAULT_SOLVER`` — keep working. + Attributes ---------- DEFAULT_SOLVER : str @@ -381,13 +387,45 @@ class QueryEngine(BaseModel): - RAW channel data must contain 'container_id', 'channel_id', 'timestamp', 'value' columns """ - solver: Solvers = Solvers.DEFAULT_SOLVER + solver: str = Solvers.DEFAULT_SOLVER data_type: DataType = DataType.RLE drop_implausible_data: bool = False raw_encoder: RawEncoder | None = None solver_config: SolverConfig | None = None batch_size: int = 500 + @model_validator(mode="before") + @classmethod + def _validate_solver_config_for_solver(cls, data): + """Validate ``solver_config`` through the selected solver's config class. + + A custom solver registered with a :class:`SolverConfig` subclass (via + ``register_solver(name, MyConfig)``) can carry extra config fields. By + default Pydantic would parse ``solver_config`` as the base + :class:`SolverConfig` and silently drop those fields. Here we resolve + the registered ``config_cls`` for ``solver`` and re-validate the raw + ``solver_config`` dict through it, so extra/required fields are enforced + at parse time and the subclass instance is preserved on the field. + + The solver must be registered when the config is parsed — i.e. the + driver imported the customer's package before building the report. An + unknown name is rejected here (as a ``ValidationError`` listing the + registered names), which surfaces a missing import early instead of + silently accepting an unusable config. + """ + if not isinstance(data, dict): + return data + name = str(data.get("solver", Solvers.DEFAULT_SOLVER)) + try: + config_cls = resolve_registration(name).config_cls + except KeyError as exc: + # Re-raise as ValueError so Pydantic surfaces it as a ValidationError. + raise ValueError(str(exc)) from exc + raw_config = data.get("solver_config") + if isinstance(raw_config, dict): + data["solver_config"] = config_cls.model_validate(raw_config) + return data + @model_validator(mode="after") def validate_drop_implausible_data_requires_raw(self): """`drop_implausible_data=True` currently only takes effect with RAW data. diff --git a/src/impulse_reporting/core/report.py b/src/impulse_reporting/core/report.py index a0e4b6f2..e9e6da2e 100644 --- a/src/impulse_reporting/core/report.py +++ b/src/impulse_reporting/core/report.py @@ -9,14 +9,14 @@ TimeSeriesExpression, ) from impulse_query_engine.analyze.query.query_builder import QueryBuilder -from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver from impulse_query_engine.analyze.query.solvers.query_solver import QuerySolver +from impulse_query_engine.analyze.query.solvers.registry import resolve_registration +from impulse_query_engine.analyze.query.solvers.solver_context import SolverBuildContext from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig from impulse_reporting.aggregations.aggregation_types import AggregationType from impulse_reporting.channels.channel_types import ChannelType from impulse_reporting.config.config_parser import ( ImpulseConfig, - Solvers, DataType, ) from impulse_reporting.core.page import Page @@ -127,7 +127,7 @@ def __init__( ) self.solver = Report.create_solver(self.spark, self.config) - log_telemetry(self.ws, "solver", self.config.query_engine.solver.name) + log_telemetry(self.ws, "solver", self.config.query_engine.solver) log_telemetry(self.ws, "data_type", self.config.query_engine.data_type.value) @property @@ -321,27 +321,25 @@ def create_solver(spark: SparkSession, config: ImpulseConfig) -> QuerySolver: Raises ------ - ValueError - If the solver type is unknown. - """ - # DELTA_SOLVER and KEY_VALUE_STORE_SOLVER are deprecated aliases retained - # for backward compatibility with existing report configs; all three - # resolve to the unified DefaultSolver. - match config.query_engine.solver: - case Solvers.DEFAULT_SOLVER | Solvers.DELTA_SOLVER | Solvers.KEY_VALUE_STORE_SOLVER: - return DefaultSolver( - spark, - config=config.query_engine.solver_config, - is_raw_data=config.query_engine.data_type is DataType.RAW, - drop_implausible_data=config.query_engine.drop_implausible_data, - raw_encoder=config.query_engine.raw_encoder, - ) - case _: - raise ValueError( - f"Unknown query engine solver: {config.query_engine.solver}. " - f"Supported: {Solvers.DEFAULT_SOLVER} (DELTA_SOLVER and " - f"KEY_VALUE_STORE_SOLVER are deprecated aliases)." - ) + KeyError + If ``query_engine.solver`` names a solver that is not registered. + The message lists the registered names; import the package that + registers the solver before building the report. + """ + # The solver name is resolved through the solver registry: the built-in + # DefaultSolver self-registers under "DefaultSolver" (plus the deprecated + # "DeltaSolver"/"KeyValueStoreSolver" aliases), and customer solvers + # register themselves when their package is imported. ``from_config`` + # builds the resolved class uniformly regardless of its constructor. + qe = config.query_engine + ctx = SolverBuildContext( + spark=spark, + solver_config=qe.solver_config, + is_raw_data=qe.data_type is DataType.RAW, + drop_implausible_data=qe.drop_implausible_data, + raw_encoder=qe.raw_encoder, + ) + return resolve_registration(qe.solver).solver_cls.from_config(ctx) def get_sink_config(self) -> SinkConfig: """ @@ -1183,7 +1181,11 @@ def _container_detection_args(self): if not self._has_sink: return None detector = ContainerUpsertDetector(self.spark) - silver_containers = self.db.container_metrics(self.spark) + # Route through the solver's container-metrics seam (not db.container_metrics + # directly) so a custom solver that combines/reshapes container sources feeds + # the same container set into incremental change detection. The default + # implementation reads the single configured table, preserving prior behavior. + silver_containers = self.solver.load_container_metrics(self.db, self.spark) measurement_dim_table = self.sink.config.get_output_uri_measurement_dimensions_table() silver_col = "last_modified" diff --git a/tests/impulse_query_engine/integration/custom_solver_test.py b/tests/impulse_query_engine/integration/custom_solver_test.py new file mode 100644 index 00000000..6ab00e10 --- /dev/null +++ b/tests/impulse_query_engine/integration/custom_solver_test.py @@ -0,0 +1,304 @@ +# pylint: disable=missing-function-docstring +"""End-to-end integration test for a customer-registered custom solver. + +Demonstrates the pluggable-solver contract via a column-redaction example: a +customer subclasses ``DefaultSolver``, registers it with a ``SolverConfig`` +subclass carrying an extra ``redact_columns`` field, builds it via +``from_config``, and runs the same solve flow as ``default_solver_test``. + +Asserts that the pipeline still resolves all containers and that the configured +columns are masked in the result. +""" + +import pyspark.sql.functions as F +import pytest +from pyspark.sql import DataFrame, SparkSession + +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver +from impulse_query_engine.analyze.query.solvers.registry import register_solver +from impulse_query_engine.analyze.query.solvers.solver_config import ( + SolverConfig, + TableConfig, +) +from impulse_query_engine.analyze.query.solvers.solver_context import SolverBuildContext +from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig + + +class RedactConfig(SolverConfig): + """SolverConfig subclass adding a column-redaction switch.""" + + redact_columns: list[str] = [] + + +@register_solver("RedactingSolver", RedactConfig) +class RedactingSolver(DefaultSolver): + """DefaultSolver that nulls out configured columns in the solve output.""" + + def solve(self, query, channels_df, selections, dtypes) -> DataFrame: # noqa: D102 + df = super().solve(query, channels_df, selections, dtypes) + for col in self.config.redact_columns: + if col in df.columns: + df = df.withColumn(col, F.lit(None).cast(df.schema[col].dataType)) + return df + + +def _redact_cfg(redact_columns: list[str]) -> RedactConfig: + """RedactConfig wired for the KVS test data (mirrors default_solver_test._kvs_cfg).""" + return RedactConfig( + project_id="SAMPLE_PROJECT", + container_tags=TableConfig(column_name_mapping={"element_id": "key"}), + container_metrics=TableConfig(column_name_mapping={"project": "project_id"}), + redact_columns=redact_columns, + ) + + +class TestCustomSolverIntegration: + def test_registered_custom_solver_built_via_from_config_runs( + self, spark: SparkSession, key_value_store_db: MeasurementDB + ): + """A custom solver built through the from_config hook runs the full pipeline.""" + cfg = _redact_cfg(redact_columns=[]) + solver = RedactingSolver.from_config(SolverBuildContext(spark=spark, solver_config=cfg)) + assert isinstance(solver, RedactingSolver) + + query = key_value_store_db.query + eng_rpm = query.channel(channel_name="Engine RPM") + result = query.select(eng_rpm.mean().alias("rpm_mean")).solve(spark=spark, solver=solver) + + assert {row.container_id for row in result.collect()} == {1, 2, 3} + + def test_redaction_masks_configured_column( + self, spark: SparkSession, key_value_store_db: MeasurementDB + ): + """The configured column is nulled while container_id still resolves.""" + cfg = _redact_cfg(redact_columns=["rpm_mean"]) + solver = RedactingSolver.from_config(SolverBuildContext(spark=spark, solver_config=cfg)) + + query = key_value_store_db.query + eng_rpm = query.channel(channel_name="Engine RPM") + result = query.select(eng_rpm.mean().alias("rpm_mean")).solve(spark=spark, solver=solver) + + rows = result.collect() + assert {row.container_id for row in rows} == {1, 2, 3} + # Every rpm_mean value has been redacted to NULL. + assert all(row.rpm_mean is None for row in rows) + + def test_without_redaction_column_has_values( + self, spark: SparkSession, key_value_store_db: MeasurementDB + ): + """Sanity check: without redaction the same column carries real values.""" + cfg = _redact_cfg(redact_columns=[]) + solver = RedactingSolver.from_config(SolverBuildContext(spark=spark, solver_config=cfg)) + + query = key_value_store_db.query + eng_rpm = query.channel(channel_name="Engine RPM") + result = query.select(eng_rpm.mean().alias("rpm_mean")).solve(spark=spark, solver=solver) + + rows = result.collect() + assert any(row.rpm_mean is not None for row in rows) + + +# --------------------------------------------------------------------------- +# Custom solver that UNIONs two channel-data sources via the load_channels seam +# --------------------------------------------------------------------------- + +_UNION_SCHEMA = "spark_catalog.silver_union_channels" + + +class UnionChannelsConfig(SolverConfig): + """SolverConfig subclass naming a SECOND channel-data table to union in. + + ``extra_channels_table`` is validated at config-parse time like any other + field on the subclass. + """ + + extra_channels_table: str + + +@register_solver("UnionChannelsSolver", UnionChannelsConfig) +class UnionChannelsSolver(DefaultSolver): + """DefaultSolver that loads channel data from TWO configured tables. + + Overrides only :meth:`load_channels`: instead of reading the single + ``channels`` table, it ``unionByName``s the primary table with a second + table named in the solver config. Everything downstream (column mapping, + encoding, the solve join) is inherited unchanged. + """ + + def load_channels(self, db, spark) -> DataFrame: + primary = db.channels(spark) + extra = spark.read.table(self.config.extra_channels_table) + return primary.unionByName(extra) + + +@pytest.fixture +def union_channels_db(spark: SparkSession, mock_workspace_client) -> MeasurementDB: + """Seed a DB whose channel data for a channel is split across two tables. + + Container 1 / channel 5 has its intervals split: the primary ``channels`` + table holds the first half, ``channels_extra`` holds the second half. A + solver reading only the primary table sees fewer samples than one that + unions both. + """ + spark.sql(f"CREATE SCHEMA IF NOT EXISTS {_UNION_SCHEMA}") + for table in spark.sql(f"SHOW TABLES IN {_UNION_SCHEMA}").collect(): + spark.sql(f"DROP TABLE IF EXISTS {_UNION_SCHEMA}.{table.tableName} PURGE") + + container_metrics = spark.createDataFrame( + [(1, "SAMPLE_PROJECT")], + schema="container_id long, project_id string", + ) + channel_metrics = spark.createDataFrame( + [(1, 5, "Engine RPM")], + schema="container_id long, channel_id int, channel_name string", + ) + channels_schema = "container_id long, channel_id int, tstart long, tend long, value double" + # Primary source: two low-value intervals. + channels_primary = spark.createDataFrame( + [ + (1, 5, 0, 1_000_000, 500.0), + (1, 5, 1_000_000, 2_000_000, 700.0), + ], + schema=channels_schema, + ) + # Second source: two high-value intervals continuing the same channel. + channels_extra = spark.createDataFrame( + [ + (1, 5, 2_000_000, 3_000_000, 6000.0), + (1, 5, 3_000_000, 4_000_000, 7000.0), + ], + schema=channels_schema, + ) + + container_metrics.write.format("delta").mode("overwrite").saveAsTable( + f"{_UNION_SCHEMA}.container_metrics" + ) + channel_metrics.write.format("delta").mode("overwrite").saveAsTable( + f"{_UNION_SCHEMA}.channel_metrics" + ) + channels_primary.write.format("delta").mode("overwrite").saveAsTable( + f"{_UNION_SCHEMA}.channels" + ) + channels_extra.write.format("delta").mode("overwrite").saveAsTable( + f"{_UNION_SCHEMA}.channels_extra" + ) + + tables = { + "container_metrics": spark.read.table(f"{_UNION_SCHEMA}.container_metrics"), + "channel_metrics": spark.read.table(f"{_UNION_SCHEMA}.channel_metrics"), + "channels": spark.read.table(f"{_UNION_SCHEMA}.channels"), + } + cfg = MeasurementDBConfig.for_debug(tables) + db = MeasurementDB(cfg, ws=mock_workspace_client) + yield db + spark.sql(f"DROP SCHEMA IF EXISTS {_UNION_SCHEMA} CASCADE") + + +class TestUnionChannelsSolver: + """A custom solver can union two configured channel-data sources.""" + + def _run(self, spark, db, solver) -> dict: + query = db.query + rpm = query.channel(channel_name="Engine RPM") + result = query.select( + rpm.min().alias("rpm_min"), + rpm.max().alias("rpm_max"), + ).solve(spark=spark, solver=solver) + return {row.container_id: row for row in result.collect()}[1] + + def test_baseline_default_solver_sees_only_primary_table( + self, spark: SparkSession, union_channels_db: MeasurementDB + ): + """The stock DefaultSolver reads only the primary channels table.""" + solver = DefaultSolver(spark, config=SolverConfig()) + row = self._run(spark, union_channels_db, solver) + # Only the primary source's low values are present. + assert row.rpm_min == 500.0 + assert row.rpm_max == 700.0 + + def test_union_solver_combines_both_channel_sources( + self, spark: SparkSession, union_channels_db: MeasurementDB + ): + """The custom solver unions both tables, so the high values appear too.""" + cfg = UnionChannelsConfig(extra_channels_table=f"{_UNION_SCHEMA}.channels_extra") + solver = UnionChannelsSolver.from_config( + SolverBuildContext(spark=spark, solver_config=cfg) + ) + assert isinstance(solver.config, UnionChannelsConfig) + + row = self._run(spark, union_channels_db, solver) + # min from the primary source, max from the second source — both tables + # contributed to the solved series. + assert row.rpm_min == 500.0 + assert row.rpm_max == 7000.0 + + +# --------------------------------------------------------------------------- +# The seam generalizes: the SAME override pattern works for container_metrics +# --------------------------------------------------------------------------- + + +class UnionContainersConfig(SolverConfig): + """Names a second container_metrics table to union in.""" + + extra_container_metrics_table: str + + +@register_solver("UnionContainersSolver", UnionContainersConfig) +class UnionContainersSolver(DefaultSolver): + """DefaultSolver that unions two container_metrics tables via the seam.""" + + def load_container_metrics(self, db, spark) -> DataFrame: + primary = db.container_metrics(spark) + extra = spark.read.table(self.config.extra_container_metrics_table) + return primary.unionByName(extra) + + +class TestUnionContainersSolver: + """The load_* seam approach is not channels-specific — container_metrics too.""" + + def test_union_solver_adds_containers_from_second_table( + self, spark: SparkSession, union_channels_db: MeasurementDB + ): + """Unioning a second container table brings container 2 into the result.""" + # A second container_metrics table introducing container 2 (with its + # own channel data already present in the primary channels table). + extra_containers = spark.createDataFrame( + [(2, "SAMPLE_PROJECT")], + schema="container_id long, project_id string", + ) + extra_containers.write.format("delta").mode("overwrite").saveAsTable( + f"{_UNION_SCHEMA}.container_metrics_extra" + ) + # Give container 2 a channel + channel data so it solves to a row. + spark.createDataFrame( + [(2, 5, "Engine RPM")], + schema="container_id long, channel_id int, channel_name string", + ).write.format("delta").mode("append").saveAsTable(f"{_UNION_SCHEMA}.channel_metrics") + spark.createDataFrame( + [(2, 5, 0, 1_000_000, 900.0)], + schema="container_id long, channel_id int, tstart long, tend long, value double", + ).write.format("delta").mode("append").saveAsTable(f"{_UNION_SCHEMA}.channels") + + # Rebuild the DB handle so it sees the appended channel rows. + tables = { + "container_metrics": spark.read.table(f"{_UNION_SCHEMA}.container_metrics"), + "channel_metrics": spark.read.table(f"{_UNION_SCHEMA}.channel_metrics"), + "channels": spark.read.table(f"{_UNION_SCHEMA}.channels"), + } + db = MeasurementDB(MeasurementDBConfig.for_debug(tables), ws=union_channels_db.ws) + + cfg = UnionContainersConfig( + extra_container_metrics_table=f"{_UNION_SCHEMA}.container_metrics_extra" + ) + solver = UnionContainersSolver.from_config( + SolverBuildContext(spark=spark, solver_config=cfg) + ) + + query = db.query + rpm = query.channel(channel_name="Engine RPM") + result = query.select(rpm.mean().alias("rpm_mean")).solve(spark=spark, solver=solver) + + # Baseline primary container_metrics has only container 1; the union + # brings in container 2 from the second table. + assert {row.container_id for row in result.collect()} == {1, 2} diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/from_config_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/from_config_test.py new file mode 100644 index 00000000..18de1d3b --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/from_config_test.py @@ -0,0 +1,82 @@ +# pylint: disable=missing-function-docstring +"""Unit tests for the ``from_config`` uniform-instantiation hook. + +``QuerySolver.from_config`` builds a config-only solver; ``DefaultSolver`` +overrides it to wire the SparkSession and the raw-data flags. These tests +do not need a live SparkSession — ``DefaultSolver.__init__`` only stores the +handle — so a sentinel object stands in for spark. +""" + +from pyspark.sql import DataFrame + +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver +from impulse_query_engine.analyze.query.solvers.query_solver import QuerySolver +from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder, SolverConfig +from impulse_query_engine.analyze.query.solvers.solver_context import SolverBuildContext + + +class _ConfigOnlySolver(QuerySolver): + """A solver that keeps the base single-arg constructor.""" + + def filter_container_tags(self, spark, query) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_container_metrics( + self, spark, query, container_df, pre_filtered_containers_df=None + ) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_channel_tags(self, spark, db, container_df, selectors) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def solve(self, query, channels_df, selections, dtypes): # noqa: D102 + raise NotImplementedError + + +class TestQuerySolverFromConfig: + def test_base_hook_passes_config_only(self): + cfg = SolverConfig(project_id="P1") + ctx = SolverBuildContext(spark=object(), solver_config=cfg) + + solver = _ConfigOnlySolver.from_config(ctx) + + assert isinstance(solver, _ConfigOnlySolver) + assert solver.config is cfg + + def test_base_hook_defaults_config_when_none(self): + ctx = SolverBuildContext(spark=object(), solver_config=None) + solver = _ConfigOnlySolver.from_config(ctx) + assert isinstance(solver.config, SolverConfig) + + +class TestDefaultSolverFromConfig: + def test_wires_spark_config_and_flags(self): + sentinel_spark = object() + cfg = SolverConfig(project_id="P2") + ctx = SolverBuildContext( + spark=sentinel_spark, + solver_config=cfg, + is_raw_data=True, + drop_implausible_data=True, + raw_encoder=RawEncoder.INTERVAL, + ) + + solver = DefaultSolver.from_config(ctx) + + assert isinstance(solver, DefaultSolver) + assert solver.spark is sentinel_spark + assert solver.config is cfg + assert solver.is_raw_data is True + assert solver.drop_implausible_data is True + assert solver.raw_encoder is RawEncoder.INTERVAL + + def test_defaults_match_constructor_defaults(self): + ctx = SolverBuildContext(spark=object()) + solver = DefaultSolver.from_config(ctx) + assert solver.is_raw_data is False + assert solver.drop_implausible_data is False + assert solver.raw_encoder is RawEncoder.RLE + assert isinstance(solver.config, SolverConfig) diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/registry_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/registry_test.py new file mode 100644 index 00000000..e3a8cc39 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/registry_test.py @@ -0,0 +1,153 @@ +# pylint: disable=missing-function-docstring +"""Unit tests for the solver registry. + +Covers ``register_solver`` (name + aliases + config_cls + overwrite), +``resolve_registration`` (name/alias lookup and unknown-name error), +``is_registered`` and ``registered_names``. + +The registry is process-global, so each test registers under unique names +and restores the registry via the ``clean_registry`` fixture to avoid +cross-test contamination. +""" + +import pytest +from pyspark.sql import DataFrame + +from impulse_query_engine.analyze.query.solvers import registry +from impulse_query_engine.analyze.query.solvers.query_solver import QuerySolver +from impulse_query_engine.analyze.query.solvers.registry import ( + SolverRegistration, + is_registered, + register_solver, + registered_names, + resolve_registration, +) +from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig + + +class _StubSolver(QuerySolver): + """Minimal concrete QuerySolver for registration tests.""" + + def filter_container_tags(self, spark, query) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_container_metrics( + self, spark, query, container_df, pre_filtered_containers_df=None + ) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_channel_tags(self, spark, db, container_df, selectors) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def solve(self, query, channels_df, selections, dtypes): # noqa: D102 + raise NotImplementedError + + +class _OtherStubSolver(_StubSolver): + """A distinct subclass used for conflict/overwrite tests.""" + + +class _StubConfig(SolverConfig): + """A SolverConfig subclass carrying an extra field.""" + + extra_table: str = "default_table" + + +@pytest.fixture +def clean_registry(): + """Snapshot and restore the process-global registry around each test.""" + saved = dict(registry._REGISTRY) + try: + yield + finally: + registry._REGISTRY.clear() + registry._REGISTRY.update(saved) + + +class TestRegisterSolver: + def test_registers_under_name_with_default_config(self, clean_registry): + register_solver("StubA")(_StubSolver) + + reg = resolve_registration("StubA") + assert isinstance(reg, SolverRegistration) + assert reg.solver_cls is _StubSolver + assert reg.config_cls is SolverConfig + + def test_registers_with_custom_config_cls(self, clean_registry): + register_solver("StubB", _StubConfig)(_StubSolver) + + reg = resolve_registration("StubB") + assert reg.solver_cls is _StubSolver + assert reg.config_cls is _StubConfig + + def test_decorator_returns_class_unchanged(self, clean_registry): + returned = register_solver("StubC")(_StubSolver) + assert returned is _StubSolver + + def test_aliases_resolve_to_same_registration(self, clean_registry): + register_solver("StubD", aliases=("StubDeprecated", "StubLegacy"))(_StubSolver) + + primary = resolve_registration("StubD") + assert resolve_registration("StubDeprecated") == primary + assert resolve_registration("StubLegacy") == primary + + def test_duplicate_same_class_is_idempotent(self, clean_registry): + register_solver("StubE")(_StubSolver) + # Re-registering the SAME class under the same name must not raise + # (e.g. a module re-imported in a notebook). + register_solver("StubE")(_StubSolver) + assert resolve_registration("StubE").solver_cls is _StubSolver + + def test_conflicting_duplicate_raises(self, clean_registry): + register_solver("StubF")(_StubSolver) + with pytest.raises(ValueError, match="already registered"): + register_solver("StubF")(_OtherStubSolver) + + def test_overwrite_replaces_registration(self, clean_registry): + register_solver("StubG")(_StubSolver) + register_solver("StubG", overwrite=True)(_OtherStubSolver) + assert resolve_registration("StubG").solver_cls is _OtherStubSolver + + def test_non_subclass_raises_type_error(self, clean_registry): + class NotASolver: + pass + + with pytest.raises(TypeError): + register_solver("StubH")(NotASolver) + + +class TestResolveRegistration: + def test_unknown_name_raises_keyerror_listing_names(self, clean_registry): + register_solver("StubKnown")(_StubSolver) + with pytest.raises(KeyError) as exc: + resolve_registration("DoesNotExist") + # The error message should help the user by listing known names. + assert "StubKnown" in str(exc.value) + + +class TestIntrospection: + def test_is_registered(self, clean_registry): + assert not is_registered("StubI") + register_solver("StubI")(_StubSolver) + assert is_registered("StubI") + + def test_registered_names_sorted_and_includes_aliases(self, clean_registry): + register_solver("StubJ", aliases=("StubJAlias",))(_StubSolver) + names = registered_names() + assert "StubJ" in names + assert "StubJAlias" in names + assert names == sorted(names) + + +class TestBuiltinRegistration: + def test_default_solver_is_registered(self): + # DefaultSolver self-registers at import time under "DefaultSolver" + # with the deprecated aliases. + from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver + + assert resolve_registration("DefaultSolver").solver_cls is DefaultSolver + assert resolve_registration("DeltaSolver").solver_cls is DefaultSolver + assert resolve_registration("KeyValueStoreSolver").solver_cls is DefaultSolver diff --git a/tests/impulse_reporting/integration/custom_solver_report_test.py b/tests/impulse_reporting/integration/custom_solver_report_test.py new file mode 100644 index 00000000..509fa388 --- /dev/null +++ b/tests/impulse_reporting/integration/custom_solver_report_test.py @@ -0,0 +1,272 @@ +# pylint: disable=missing-function-docstring +"""End-to-end report test driven by a custom solver selected via config. + +Scenario (a split silver layer): container information is split +across TWO physical tables — a base ``container_metrics`` table and a separate +``vehicle_info`` table that carries the ``brand`` column. A customer ships a +solver that JOINs the two into the container silver frame and prefilters rows to +a configurable set of brands. Both the second table's location and the brand +allowlist live in the solver's own ``SolverConfig`` subclass, so no Impulse core +change is needed — the report config just names the solver and its extra config. + +This exercises the full pluggable-solver path end to end: + report config (JSON dict) → registry resolves "MultiTableBrandSolver" → + extended config validated as MultiTableBrandConfig → solver built via + from_config → Report.determine_report runs the pipeline → the brand prefilter + is reflected in the container dimension and the aggregation results. + +The custom solver is defined in this module; importing the module registers it +(the same mechanism a customer package would use). +""" + +import datetime +from unittest.mock import create_autospec + +import pyspark.sql.functions as F +import pytest +from databricks.sdk import WorkspaceClient +from pyspark.sql import DataFrame, SparkSession + +from impulse_query_engine.analyze.query.solvers.default_solver import DefaultSolver +from impulse_query_engine.analyze.query.solvers.registry import register_solver +from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig +from impulse_reporting.aggregations.histogram import HistogramDuration +from impulse_reporting.core.page import Page +from impulse_reporting.core.report import Report + +_SCHEMA = "spark_catalog.silver_custom_solver" + + +class MultiTableBrandConfig(SolverConfig): + """SolverConfig subclass carrying the second table + brand prefilter. + + ``vehicle_info_table`` is required (validated at config-parse time); + ``include_brands`` is an optional allowlist — empty means no brand filter. + """ + + vehicle_info_table: str + include_brands: list[str] = [] + + +@register_solver("MultiTableBrandSolver", MultiTableBrandConfig) +class MultiTableBrandSolver(DefaultSolver): + """DefaultSolver that combines two container tables and prefilters by brand. + + Overrides only the ``load_container_metrics`` seam: it joins the base + ``container_metrics`` table with a separate ``vehicle_info`` table (read + from ``self.config.vehicle_info_table``) and applies the configurable brand + allowlist. Because the seam is the single source of container rows, both the + solve pipeline *and* incremental container detection read the same combined, + brand-filtered set — the rest of the pipeline is inherited unchanged. + """ + + def load_container_metrics(self, db, spark) -> DataFrame: + base = db.container_metrics(spark) + vehicle_info = spark.read.table(self.config.vehicle_info_table) + combined = base.join(vehicle_info, on=self.config.container_id_col, how="inner") + + if self.config.include_brands: + combined = combined.where(F.col("brand").isin(self.config.include_brands)) + + return combined + + +@pytest.fixture(scope="module") +def setup_custom_solver_db(spark: SparkSession): + """Seed two container tables (base + vehicle_info) plus channel data. + + Four containers across three brands: + 1, 2 → Seat 3 → VW 4 → Audi + """ + spark.sql(f"CREATE SCHEMA IF NOT EXISTS {_SCHEMA}") + for table in spark.sql(f"SHOW TABLES IN {_SCHEMA}").collect(): + spark.sql(f"DROP TABLE IF EXISTS {_SCHEMA}.{table.tableName} PURGE") + + ts = datetime.datetime(2025, 7, 1, 12, 0, 0) + + # Table 1: base container_metrics (no brand column). + base_container_metrics = spark.createDataFrame( + [ + (1, "UUT_1", "SAMPLE_PROJECT", ts), + (2, "UUT_2", "SAMPLE_PROJECT", ts), + (3, "UUT_3", "SAMPLE_PROJECT", ts), + (4, "UUT_4", "SAMPLE_PROJECT", ts), + ], + schema="container_id long, uut_id string, project_id string, start_dt timestamp", + ) + + # Table 2: vehicle_info — the second table combined in by the custom solver. + vehicle_info = spark.createDataFrame( + [ + (1, "Seat", "Leon"), + (2, "Seat", "Ibiza"), + (3, "VW", "Golf"), + (4, "Audi", "A3"), + ], + schema="container_id long, brand string, model string", + ) + + # Wide channel_metrics: one selectable channel per container. + channel_metrics = spark.createDataFrame( + [(cid, 5, "Engine RPM") for cid in (1, 2, 3, 4)], + schema="container_id long, channel_id int, channel_name string", + ) + + # RLE channel data: three intervals per container, values within the bins. + channel_rows = [] + for cid in (1, 2, 3, 4): + channel_rows += [ + (cid, 5, 0, 1_000_000, 800.0), + (cid, 5, 1_000_000, 2_000_000, 1500.0), + (cid, 5, 2_000_000, 3_000_000, 3000.0), + ] + channels = spark.createDataFrame( + channel_rows, + schema="container_id long, channel_id int, tstart long, tend long, value double", + ) + + base_container_metrics.write.format("delta").mode("overwrite").saveAsTable( + f"{_SCHEMA}.base_container_metrics" + ) + vehicle_info.write.format("delta").mode("overwrite").saveAsTable(f"{_SCHEMA}.vehicle_info") + channel_metrics.write.format("delta").mode("overwrite").saveAsTable( + f"{_SCHEMA}.channel_metrics" + ) + channels.write.format("delta").mode("overwrite").saveAsTable(f"{_SCHEMA}.channels") + + yield + + spark.sql(f"DROP SCHEMA IF EXISTS {_SCHEMA} CASCADE") + + +def _config(include_brands: list[str]) -> dict: + """Build a report config dict selecting the custom solver by name.""" + return { + "source": { + "container_metrics_table": f"{_SCHEMA}.base_container_metrics", + "channel_metrics_table": f"{_SCHEMA}.channel_metrics", + "channels_uri": f"{_SCHEMA}.channels", + }, + "unity_sink": { + "catalog": "spark_catalog", + "schema": "gold", + "table_prefix": "custom_solver_eval", + }, + "query_engine": { + "solver": "MultiTableBrandSolver", + "solver_config": { + "vehicle_info_table": f"{_SCHEMA}.vehicle_info", + "include_brands": include_brands, + }, + }, + "measurement_dimensions": ["container_id", "brand"], + } + + +def _build_report(spark: SparkSession, include_brands: list[str]) -> Report: + report = Report( + name="custom_solver_report", + spark=spark, + workspace_client=create_autospec(WorkspaceClient), + config=_config(include_brands), + ) + query = report.get_db().query + rpm = query.channel(channel_name="Engine RPM") + + page = Page(page_number=1) + report.add_page(page) + page.add_aggregation( + HistogramDuration( + name="rpm_hist", + base_expr=rpm, + bins=[float(i) for i in range(0, 8000, 250)], + ) + ) + return report + + +class TestCustomSolverReportE2E: + def test_config_selects_custom_solver_with_extended_config( + self, spark: SparkSession, setup_custom_solver_db + ): + """The config-named solver is built and its extended config is validated/typed.""" + report = _build_report(spark, include_brands=["Seat"]) + + solver = report.get_solver() + assert isinstance(solver, MultiTableBrandSolver) + assert isinstance(solver.config, MultiTableBrandConfig) + assert solver.config.vehicle_info_table == f"{_SCHEMA}.vehicle_info" + assert solver.config.include_brands == ["Seat"] + + def test_brand_prefilter_limits_containers_and_combines_tables( + self, spark: SparkSession, setup_custom_solver_db + ): + """Combining vehicle_info + prefiltering to Seat yields only containers 1 and 2.""" + report = _build_report(spark, include_brands=["Seat"]) + + report.determine_report() + + # container_dimension_df comes straight from the custom filter_container_metrics. + dim_rows = report.container_dimension_df.collect() + assert {r.container_id for r in dim_rows} == {1, 2} + # The `brand` column proves the second table was joined in. + assert {r.brand for r in dim_rows} == {"Seat"} + + # Aggregation ran and produced non-empty results for the surviving containers. + assert ( + report.aggregation_dfs["HISTOGRAM"]["changed"].filter(F.col("hist_value") > 0).count() + > 0 + ) + + def test_different_brand_selection(self, spark: SparkSession, setup_custom_solver_db): + """A different brand allowlist selects the corresponding containers.""" + report = _build_report(spark, include_brands=["VW", "Audi"]) + report.determine_report() + + dim_rows = report.container_dimension_df.collect() + assert {r.container_id for r in dim_rows} == {3, 4} + assert {r.brand for r in dim_rows} == {"VW", "Audi"} + + def test_empty_brand_list_keeps_all_containers( + self, spark: SparkSession, setup_custom_solver_db + ): + """No brand allowlist means the prefilter is skipped — all four containers survive.""" + report = _build_report(spark, include_brands=[]) + report.determine_report() + + dim_rows = report.container_dimension_df.collect() + assert {r.container_id for r in dim_rows} == {1, 2, 3, 4} + + def test_persist_writes_gold_tables(self, spark: SparkSession, setup_custom_solver_db): + """The full determine + persist path writes gold facts for the filtered set.""" + report = _build_report(spark, include_brands=["Seat"]) + report.determine_report() + report.persist_results() + + assert spark.catalog.tableExists("spark_catalog.gold.custom_solver_eval_histogram_fact") + measurement_dim = spark.read.table( + "spark_catalog.gold.custom_solver_eval_measurement_dimension" + ) + assert {r.container_id for r in measurement_dim.collect()} == {1, 2} + + def test_incremental_detection_uses_custom_solver_container_source( + self, spark: SparkSession, setup_custom_solver_db + ): + """Incremental container detection reads through the solver's seam. + + The detection input must reflect the custom solver's combined + (base + vehicle_info) and brand-filtered container set — containers + {1, 2} for Seat — not the raw base ``container_metrics`` table, which + alone has all four containers. + """ + report = _build_report(spark, include_brands=["Seat"]) + + args = report._container_detection_args() + assert args is not None + _detector, silver_containers, *_ = args + + detected_ids = {r.container_id for r in silver_containers.collect()} + assert detected_ids == {1, 2} + # The joined-in `brand` column is present, proving the second table + # flowed into detection via the seam. + assert "brand" in silver_containers.columns diff --git a/tests/impulse_reporting/unit/config/solver_registry_config_test.py b/tests/impulse_reporting/unit/config/solver_registry_config_test.py new file mode 100644 index 00000000..117f1361 --- /dev/null +++ b/tests/impulse_reporting/unit/config/solver_registry_config_test.py @@ -0,0 +1,118 @@ +# pylint: disable=missing-function-docstring +"""Unit tests for solver selection + extended-config validation in QueryEngine. + +Verifies the key requirement: when a report config selects a custom solver by +name, its ``solver_config`` block is validated through the *registered* config +class (a SolverConfig subclass), so extra/required fields are enforced at parse +time and the subclass instance type is preserved. + +Also verifies the StrEnum migration keeps existing enum-based comparisons and +the deprecated-alias names working. +""" + +import pytest +from pydantic import ValidationError +from pyspark.sql import DataFrame + +from impulse_query_engine.analyze.query.solvers import registry +from impulse_query_engine.analyze.query.solvers.query_solver import QuerySolver +from impulse_query_engine.analyze.query.solvers.registry import register_solver +from impulse_query_engine.analyze.query.solvers.solver_config import SolverConfig +from impulse_reporting.config.config_parser import QueryEngine, Solvers + + +class _RegConfig(SolverConfig): + """SolverConfig subclass with one required and one optional extra field.""" + + raw_signal_table: str # required + gps_signal_name: str = "GPS_LAT" # optional with default + + +class _RegSolver(QuerySolver): + def filter_container_tags(self, spark, query) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_container_metrics( + self, spark, query, container_df, pre_filtered_containers_df=None + ) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_channel_tags(self, spark, db, container_df, selectors) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def filter_channel_metrics(self, spark, db, channel_df, selectors) -> DataFrame: # noqa: D102 + raise NotImplementedError + + def solve(self, query, channels_df, selections, dtypes): # noqa: D102 + raise NotImplementedError + + +@pytest.fixture +def registered_custom_solver(): + """Register a custom solver+config for the duration of a test.""" + saved = dict(registry._REGISTRY) + register_solver("RegSolver", _RegConfig)(_RegSolver) + try: + yield + finally: + registry._REGISTRY.clear() + registry._REGISTRY.update(saved) + + +class TestExtendedConfigValidation: + def test_extended_config_validated_and_typed(self, registered_custom_solver): + qe = QueryEngine.model_validate( + { + "solver": "RegSolver", + "solver_config": {"raw_signal_table": "cat.sch.raw"}, + } + ) + # The subclass instance type is preserved on the base-typed field. + assert isinstance(qe.solver_config, _RegConfig) + assert qe.solver_config.raw_signal_table == "cat.sch.raw" + assert qe.solver_config.gps_signal_name == "GPS_LAT" + + def test_missing_required_field_raises_at_parse_time(self, registered_custom_solver): + with pytest.raises(ValidationError, match="raw_signal_table"): + QueryEngine.model_validate({"solver": "RegSolver", "solver_config": {}}) + + def test_unknown_solver_name_raises_with_config(self): + # An unknown name is rejected at parse time (surfaced as ValidationError). + with pytest.raises(ValidationError): + QueryEngine.model_validate({"solver": "NopeSolver", "solver_config": {"foo": "bar"}}) + + def test_unknown_solver_name_raises_without_config(self): + # Fail-fast even when no solver_config block is present. + with pytest.raises(ValidationError): + QueryEngine.model_validate({"solver": "NopeSolver"}) + + +class TestDefaultAndAliasBehavior: + def test_default_solver_uses_base_config(self): + qe = QueryEngine.model_validate( + {"solver": "DefaultSolver", "solver_config": {"project_id": "P"}} + ) + assert type(qe.solver_config) is SolverConfig + assert qe.solver_config.project_id == "P" + + def test_default_when_solver_omitted(self): + qe = QueryEngine.model_validate({}) + assert qe.solver == Solvers.DEFAULT_SOLVER + + def test_no_solver_config_stays_none(self): + qe = QueryEngine.model_validate({"solver": "DefaultSolver"}) + assert qe.solver_config is None + + @pytest.mark.parametrize( + "name,enum_member", + [ + ("DefaultSolver", Solvers.DEFAULT_SOLVER), + ("DeltaSolver", Solvers.DELTA_SOLVER), + ("KeyValueStoreSolver", Solvers.KEY_VALUE_STORE_SOLVER), + ], + ) + def test_strenum_regression_equality_holds(self, name, enum_member): + # After widening solver -> str, a StrEnum member still compares equal + # to the stored string, so existing `== Solvers.X` assertions pass. + qe = QueryEngine.model_validate({"solver": name}) + assert qe.solver == enum_member From 1a4a74151465840d2e3a1070e308e3b26ba86146 Mon Sep 17 00:00:00 2001 From: "fabian.ade" Date: Wed, 29 Jul 2026 18:19:51 +0200 Subject: [PATCH 3/3] corrected github action back to original --- .github/workflows/acceptance.yml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/acceptance.yml b/.github/workflows/acceptance.yml index 2ede4006..16a3e4cc 100644 --- a/.github/workflows/acceptance.yml +++ b/.github/workflows/acceptance.yml @@ -37,10 +37,6 @@ jobs: labels: linux-ubuntu-latest permissions: id-token: write - # Fork PRs get no OIDC token / secrets from GitHub, so JFrog auth (and therefore - # dependency installation) cannot run. Skip CI for them; fork PRs are to be tested - # by the reviewer(s) / maintainer(s) before merging. - if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -67,8 +63,6 @@ jobs: needs: [ not-a-fork, lint ] permissions: id-token: write - # See the note on `lint`: fork PRs cannot authenticate to JFrog, so skip CI for them. - if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.head.repo.fork }} strategy: matrix: python-version: [ '3.12' ]