diff --git a/docs/impulse/docs/config/configuration.md b/docs/impulse/docs/config/configuration.md index 3ad82f76..99926673 100644 --- a/docs/impulse/docs/config/configuration.md +++ b/docs/impulse/docs/config/configuration.md @@ -135,21 +135,23 @@ Two independent filter families: | Field | Type | Default | Description | |-------------------------|----------------|--------------------------|-----------------------------------------------------------------------------------------------------------------------------| -| `solver` | `str` | `"DefaultSolver"` | `"DefaultSolver"` adapts to the silver layer: it selects channels from a narrow EAV `channel_tags` table when `source.channel_tags_table` is set and otherwise from columns on `channel_metrics`; it filters containers via a narrow EAV `container_tags` table or, when `source.container_tags_table` is omitted, a wide-only `container_metrics`. `"DeltaSolver"` and `"KeyValueStoreSolver"` are **deprecated aliases** that resolve to `DefaultSolver`. | +| `solver` | `str` | `"DefaultSolver"` | Registered name of the solver to use. `"DefaultSolver"` adapts to the silver layer: it selects channels from a narrow EAV `channel_tags` table when `source.channel_tags_table` is set and otherwise from columns on `channel_metrics`; it filters containers via a narrow EAV `container_tags` table or, when `source.container_tags_table` is omitted, a wide-only `container_metrics`. `"DeltaSolver"` and `"KeyValueStoreSolver"` are **deprecated aliases** that resolve to `DefaultSolver`. A custom solver is selected by the name it registered; see [Custom solver](../extensibility/custom_solver.md). | | `data_type` | `str` | `"RLE"` | `"RLE"` (intervals `[tstart, tend)`) or `"RAW"` (raw timestamps; converted to intervals at query time by the configured `raw_encoder`). | | `raw_encoder` | `str` | `null` (resolves to `"RLE"` when `data_type = "RAW"`) | How RAW point data is converted to intervals. `"RLE"` (the default) run-length encodes on the fly to reduce memory consumption: consecutive samples with the same value collapse into a single `[tstart, tend)` interval per run. `"INTERVAL"` keeps every original sample, only deriving `tend` from the next sample's timestamp and dropping exact duplicate points — use it when downstream analysis needs all original timestamps. Only takes effect with `data_type = "RAW"`; ignored for `"RLE"` input. See [How Impulse interprets intervals](../data_model/silver_layer_schema.md#raw-format) for validity semantics and what counts as a duplicate point. | | `drop_implausible_data` | `bool` | `false` | When `true`, drops `channels` rows where `is_plausible = false`. Requires `data_type = "RAW"`; combining with `"RLE"` raises a validation error. | | `max_channels_per_batch` | `int` | `500` | Caps how many distinct channels are read in one solve batch, to bound per-solve memory. Counts channel selections (each distinct `channel(...)` / `poi_channel(...)` in your expressions). Applies to events, aggregations, and calculated channels; for calculated channels it counts the *input* channels they read, not the number of calculated channels. The cap is a packing target, not a hard limit: an expression that on its own reads more channels than the cap still runs, in a batch of its own. Must be `>= 1`. The old name `batch_size` is a **deprecated alias** (still accepted, warns, removed in a future release); setting both raises a validation error. | | `max_containers_per_batch` | `int` | `null` | Caps how many containers are solved at once, to bound per-solve memory. `null` disables the cap; when set it must be `>= 1`. On an **incremental** run, `Report.run()` commits at most this many upserted containers per iteration and loops until the population is drained (see [incremental](#incremental-optional)). On a **full** run, and when a changed definition recomputes over all containers, the population is split into chunks of about this size and solved chunk by chunk. Chunk sizes are approximate (the cap is a memory heuristic, not an exact bound). | -| `solver_config` | `SolverConfig` | `null` | Per-table column mappings, per-table equality filters, and project scoping. Set `project_id` to scope reads by project — it is applied to `container_tags` (if configured), `container_metrics`, and `channel_mapping` (if configured), so it works in both narrow EAV and wide-only data models. Omit it when you don't need project scoping. See [Solver column mappings and filters](#solver-column-mappings-and-filters). | +| `solver_config` | `SolverConfig` | `null` | Per-table column mappings, per-table equality filters, and project scoping. Set `project_id` to scope reads by project — it is applied to `container_tags` (if configured), `container_metrics`, and `channel_mapping` (if configured), so it works in both narrow EAV and wide-only data models. It is validated with the config class of the selected `solver`: with `DefaultSolver` you can omit it when you need none of these, while a [custom solver](../extensibility/custom_solver.md) can add its own settings here and make them required. See [Solver column mappings and filters](#solver-column-mappings-and-filters). | | `measurement_db` | `str` | `"MeasurementDB"` | Registered name of the MeasurementDB implementation to use. The default is the built-in reader; a custom implementation can reshape a different source schema at read time. See [Custom measurement DB](../extensibility/custom_measurement_db.md). | | `measurement_db_config` | `dict` | `null` | Keyword arguments for that implementation's config class, merged with the `source` tables (e.g. extra source tables, prefilter keys). | If `query_engine` is omitted, the default is `DefaultSolver` with `data_type = "RLE"`. -Unknown keys under `query_engine` are rejected when the config is loaded: a misspelled key such as -`"measurment_db"` fails with `Extra inputs are not permitted` instead of being silently ignored. +Unknown keys at the top level of the config, under `query_engine` and directly under +`query_engine.solver_config` are rejected when the config is loaded: a misspelled key such as +`"measurment_db"`, or `"solver"` placed outside `query_engine`, fails with `Extra inputs are not +permitted` instead of being silently ignored. For a source schema that needs reshaping beyond column renames, select your own read seam — see [Custom measurement DB](../extensibility/custom_measurement_db.md). diff --git a/docs/impulse/docs/data_model/index.md b/docs/impulse/docs/data_model/index.md index 7dc4373d..5822ee16 100644 --- a/docs/impulse/docs/data_model/index.md +++ b/docs/impulse/docs/data_model/index.md @@ -15,7 +15,8 @@ tables are **fully optional** add-ons, used only when configured — see the The rest of this page documents the full shape. Landing your data in it during ingest is the simplest path (see the [Ingestion guide](ingestion.md)); if you can't reshape, a [`SolverConfig`](ingestion.md#adapting-to-existing-data-layouts) -can remap column names, or you can implement a custom solver. +can remap column names, or a [custom MeasurementDB](../extensibility/custom_measurement_db.md) +can reshape your tables when they are read. ::: diff --git a/docs/impulse/docs/data_model/ingestion.md b/docs/impulse/docs/data_model/ingestion.md index 133a247d..a7b3068e 100644 --- a/docs/impulse/docs/data_model/ingestion.md +++ b/docs/impulse/docs/data_model/ingestion.md @@ -223,28 +223,30 @@ key names (`container_id`, `channel_id`) themselves are fixed constants. Their column *types* are not fixed, though — `container_id` in particular may be a `long`, `int`, or `string`; the engine adopts whatever type your tables use, as long as it is consistent across them. -For different relationships or composite keys, see custom solvers below. +For different relationships or composite keys, see the next section. -### Custom solvers +### Layouts `SolverConfig` can't express For physical layouts that do not match the silver-layer relationships at all — no EAV tag tables, alias lookup tables instead of `channel_tags`, 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. - -This is significantly more invested than the `SolverConfig` path: you take -on responsibility for the four solver pipeline stages -(`filter_container_tags`, `filter_container_metrics`, `filter_channel_tags`, -`filter_channel_metrics`) and the `solve` method. Some advanced deployments -do this — e.g. when the customer's silver layer pre-dates Impulse and -synthesises Impulse-shaped views via SQL CTEs at query time. If you find -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. - -The general rule: **`SolverConfig` for naming differences, custom solver -for structural differences, ETL into the standard shape for everything -else.** +that need pre-processing, etc. — implement a +[custom MeasurementDB](../extensibility/custom_measurement_db.md) that reshapes +your tables into the silver model when they are read. Every reader (the solve, +incremental detection, ad-hoc queries) then sees the same silver-shaped data, +and the stock solver runs unchanged. + +A [custom solver](../extensibility/custom_solver.md) is for the other kind of +difference: when the engine should do something different with silver-shaped +data, such as custom candidate selection, extra pruning, or post-processing the +solve output. Both are registered under a name and selected in the report +config (`query_engine.measurement_db` / `query_engine.solver`). + +Either path is significantly more invested than `SolverConfig`: you take on +responsibility for returning the shapes the engine expects. If you find +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. + +The general rule: **`SolverConfig` for naming differences, a custom +MeasurementDB for structural differences, a custom solver for engine +behaviour, ETL into the standard shape for everything else.** diff --git a/docs/impulse/docs/extensibility/custom_measurement_db.md b/docs/impulse/docs/extensibility/custom_measurement_db.md index e8fcd9b3..192f318c 100644 --- a/docs/impulse/docs/extensibility/custom_measurement_db.md +++ b/docs/impulse/docs/extensibility/custom_measurement_db.md @@ -22,7 +22,7 @@ used and behaviour is unchanged. |---|---| | Column names, or narrow-EAV vs. wide layout | `DefaultSolver` + [`solver_config`](../config/configuration.md#solver-column-mappings-and-filters) — no code | | Different tables, unions, joins, grain changes, EAV unpivots, or a pre-reshape row filter | **A custom MeasurementDB** (this page) | -| Query planning / execution the config cannot express | Implementation of a custom solver (see [Custom solvers](../data_model/ingestion.md#custom-solvers)) | +| Query planning / execution the config cannot express | Implementation of a custom solver (see [Custom solver](custom_solver.md)) | ## Select it in the report config diff --git a/docs/impulse/docs/extensibility/custom_solver.md b/docs/impulse/docs/extensibility/custom_solver.md new file mode 100644 index 00000000..96b8973c --- /dev/null +++ b/docs/impulse/docs/extensibility/custom_solver.md @@ -0,0 +1,117 @@ +--- +sidebar_position: 2 +title: Custom solver +--- + +# Custom solver + +The **solver** turns a query into results. It runs a fixed pipeline over the +[silver model](../data_model/silver_layer_schema.md): filter containers by tags and metrics, find the +candidate channels for each selection, pick the best candidate, and evaluate the TSAL expressions +(`solve`). Impulse ships `DefaultSolver`, which adapts to narrow or wide silver layouts through +[`solver_config`](../config/configuration.md#solver-column-mappings-and-filters). When you need the +engine to *behave* differently, you can register your own solver and select it by name in the report +config. + +With nothing configured, `DefaultSolver` is used and behaviour is unchanged. + +## Custom solver or custom MeasurementDB? + +Impulse has two extension points with different jobs. The **MeasurementDB** decides *where and how* +the silver tables are read; the **solver** decides *what the engine does* with them. + +| Your problem | Use | +|---|---| +| Column names differ, or narrow-EAV vs. wide layout | `DefaultSolver` + `solver_config`, no code | +| Your tables have a different shape: other tables, unions, joins, grain changes, EAV unpivots, or row filters that must run before a reshape | A [custom MeasurementDB](custom_measurement_db.md), or producing the silver shape upstream (ETL job or views) | +| The engine should do something different with silver-shaped data: custom candidate selection, extra pruning, post-processing the solve output, a different time-series encoding or storage | **A custom solver** (this page) | + +Rule of thumb: if the problem is **how the data looks**, change how it's read; if it's **what the +engine does with silver data**, change the solver. + +Prefer not to reshape data inside a solver. The solver isn't the only reader of the silver tables: +incremental reporting detects new and updated containers by running only the solver's container +stages (`filter_container_tags`, `filter_container_metrics`), outside `solve()`, and anything that +queries the MeasurementDB directly doesn't go through the solver at all. A reshape in a solver has to +be repeated in every stage that reads a table, and the other readers still see the original tables. +Reshape in a [custom MeasurementDB](custom_measurement_db.md) instead: every reader goes through it, +so they all see the same data. + +## Implement one + +A custom solver is a subclass of `DefaultSolver`, plus an optional config class for its own +settings, registered under a name: + +```python +import pyspark.sql.functions as F + +from impulse_query_engine.analyze.query.solvers import DefaultSolver, SolverConfig, register_solver + + +class ExcludeContainersConfig(SolverConfig): + excluded_container_ids: list[int] # required: no default + + +@register_solver("ExcludeContainersSolver", ExcludeContainersConfig) +class ExcludeContainersSolver(DefaultSolver): + def solve(self, query, channels_df, selections, dtypes): + solved = super().solve(query, channels_df, selections, dtypes) + return solved.filter( + ~F.col(self.config.container_id_col).isin(self.config.excluded_container_ids) + ) +``` + +- **Override only the stages you want to change** and call `super()` for the rest. Subclass + `DefaultSolver`. Subclassing `QuerySolver` directly is possible but means implementing all + stages, including the alias stages (`filter_aliased_channel_metrics`, + `resolve_channel_selections`), and overriding `from_config` if you need RAW data. +- **Put your settings on a `SolverConfig` subclass.** It keeps the per-table `column_name_mapping`, + `filters` and `project_id`, and `self.config` is an instance of your class. +- **`register_solver(name, config_cls=SolverConfig, *, aliases=())`** registers the class under the + name and each alias. Re-registering a name replaces the previous entry. +- **Impulse builds the solver with `from_config(ctx)`.** `DefaultSolver` subclasses inherit a version + that passes the Spark session, `solver_config` and the RAW-data settings (`data_type`, + `drop_implausible_data`, `raw_encoder`). Override `from_config` only if your constructor takes + different arguments. + +## Select it in the report config + +```json +"query_engine": { + "solver": "ExcludeContainersSolver", + "solver_config": { "excluded_container_ids": [3] } +} +``` + +- `solver` is the registered name. The default is `"DefaultSolver"`; `"DeltaSolver"` and + `"KeyValueStoreSolver"` are deprecated aliases for it. +- `solver_config` is validated with the config class the selected solver registered. +- In Python, the same keys go on `QueryEngine`: + `QueryEngine(solver="ExcludeContainersSolver", solver_config={"excluded_container_ids": [3]})`. + +For ad-hoc queries you don't need the registry: pass an instance, +`query.solve(spark, solver=ExcludeContainersSolver(spark, config=ExcludeContainersConfig(excluded_container_ids=[3])))`. + +## Parse-time validation + +The config is checked when it is loaded, so these fail as a validation error rather than at run time: + +- an unregistered `solver` name (the message lists the registered names); +- a missing required setting of your config class, or a setting with the wrong type, including when + `solver_config` is left out entirely; +- an unknown key in `solver_config`: a misspelled setting, or your settings without `"solver"` + selecting your solver. `SolverConfig` forbids unknown keys and your config class inherits that, so + neither is silently dropped. + +## Rules + +- **Import the package that registers the solver before the report config is parsed**, including in + tools that only generate or edit configs. Selection is by name only: config data never triggers an + import. +- **Keep each stage's output contract.** `filter_container_metrics` returns all container columns + needed downstream (container dimensions, event bounds), and `solve` returns one row per container: + `container_id` plus one column per selection. +- **Use the internal column names from `self.config`** (e.g. `self.config.container_id_col`), not + string literals: column mapping has already been applied when a stage runs. +- **Calculated channels need `solve_calculated_channels`**, which `DefaultSolver` implements. The + `QuerySolver` base raises, so subclass `DefaultSolver` if your reports use calculated channels. diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md index 9df87eef..f4526d22 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/default_solver.md @@ -49,6 +49,15 @@ silver layer. ``RawEncoder.INTERVAL`` only derives ``tend`` and drops exact duplicates. Only consulted when ``is_raw_data`` is ``True``. +#### from\_config + +```python +def from_config(cls, ctx: SolverBuildContext) -> DefaultSolver +``` + +Build the solver with the Spark session and RAW-data flags from the report config. + + #### filter\_container\_tags ```python diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/query_solver.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/query_solver.md index 10e4be49..f36dc488 100644 --- a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/query_solver.md +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/query_solver.md @@ -3,6 +3,21 @@ sidebar_label: query_solver title: impulse_query_engine.analyze.query.solvers.query_solver --- +## SolverBuildContext + +```python +class SolverBuildContext() +``` + +Everything a report passes to a solver when it builds one (see ``QuerySolver.from_config``). + +The fields copy the ``query_engine`` settings a solver needs rather than passing the report's +``QueryEngine`` config, because that model lives in ``impulse_reporting`` and the query engine +must not depend on the reporting layer (``config_parser`` imports this package, so the reverse +import would also be circular). When a solver needs another ``query_engine`` setting, add a +field here and set it in ``Report.create_solver``. + + ## QuerySolver ```python @@ -21,6 +36,18 @@ columns** needed for container dimensions and event bounds (e.g. not only ``container_id``. +#### from\_config + +```python +def from_config(cls, ctx: SolverBuildContext) -> QuerySolver +``` + +Build the solver from what a report provides (see ``register_solver``). + +The base implementation passes only ``ctx.solver_config``. Override it when the +constructor needs more, e.g. the Spark session or the RAW-data flags. + + #### scoped\_container\_metrics ```python diff --git a/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/solver_registry.md b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/solver_registry.md new file mode 100644 index 00000000..0ec4abf7 --- /dev/null +++ b/docs/impulse/docs/references/api/impulse_query_engine/analyze/query/solvers/solver_registry.md @@ -0,0 +1,36 @@ +--- +sidebar_label: solver_registry +title: impulse_query_engine.analyze.query.solvers.solver_registry +--- + +Name registry for query solvers, so a report config can select a solver by name. + +A solver registers under a name with ``@register_solver(...)`` and is selected with +``query_engine.solver`` in the report config. Selection is by name only: the package that +registers a solver must be imported before the report config is parsed. + + +#### register\_solver + +```python +def register_solver(name: str, + config_cls: type[SolverConfig] = SolverConfig, + *, + aliases: tuple[str, ...] = ()) +``` + +Register a ``QuerySolver`` subclass under *name* (and *aliases*), configured by *config_cls*. + +``config_cls`` validates ``query_engine.solver_config`` when this solver is selected. +Re-registering a name replaces the previous entry. + + +#### resolve\_solver + +```python +def resolve_solver(name: str) -> Registration[QuerySolver, SolverConfig] +``` + +Return the registration for *name*; raises ``KeyError`` listing the registered names. + + diff --git a/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md b/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md index 5c92a6b9..b69bbb0a 100644 --- a/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md +++ b/docs/impulse/docs/references/api/impulse_reporting/config/config_parser.md @@ -46,11 +46,13 @@ Validate if a string is a valid Unity Catalog entity name. ## Solvers ```python -class Solvers(Enum) +class Solvers(StrEnum) ``` -Enumeration of available solver types for the query engine. +Names of the built-in solvers. +``query_engine.solver`` takes any registered solver name (see +``register_solver``); these are the names Impulse registers itself. ``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 @@ -186,7 +188,11 @@ Configuration for the query engine solver. **Arguments**: -- `solver` (`Solvers, default=Solvers.DEFAULT_SOLVER`): The solver type to use for query execution. +- `solver` (`str, default="DefaultSolver"`): Registered name of the solver to use for query execution. Impulse +registers ``"DefaultSolver"`` (``"DeltaSolver"`` and +``"KeyValueStoreSolver"`` are deprecated aliases); a custom solver +registered with ``register_solver`` is selected by its name. An +unregistered name fails when the config is loaded. - `raw_encoder` (`RawEncoder, optional, default=None`): Encoder used to convert RAW point data into intervals. ``RLE`` collapses consecutive equal-valued samples into runs; ``INTERVAL`` only derives ``tend`` and drops exact duplicates. Only takes effect @@ -207,7 +213,9 @@ scoping. Key sub-fields: ``column_name_mapping`` and ``filters`` dicts. When omitted, all default column names are used and no -project/toolbox filtering is applied. +project/toolbox filtering is applied. It is validated with the +config class the selected solver registered, so a custom solver can +add its own settings to a ``SolverConfig`` subclass. - `measurement_db` (`str, default="MeasurementDB"`): Registered name of the MeasurementDB implementation to use. - `measurement_db_config` (`dict`): Keyword arguments for that implementation's config class, merged with the ``source`` tables. diff --git a/docs/impulse/docs/references/api/impulse_reporting/core/report.md b/docs/impulse/docs/references/api/impulse_reporting/core/report.md index 3171a8d0..cc6c9fed 100644 --- a/docs/impulse/docs/references/api/impulse_reporting/core/report.md +++ b/docs/impulse/docs/references/api/impulse_reporting/core/report.md @@ -179,7 +179,11 @@ Create a sink based on the provided configuration. def create_solver(spark: SparkSession, config: ImpulseConfig) -> QuerySolver ``` -Create a query solver based on the provided configuration. +Create the query solver selected by ``query_engine.solver``. + +Resolves the registered solver name (``"DeltaSolver"`` and +``"KeyValueStoreSolver"`` are deprecated aliases of ``"DefaultSolver"``) +and builds it with its ``from_config`` hook. **Arguments**: @@ -188,11 +192,12 @@ Create a query solver based on the provided configuration. **Raises**: -- `ValueError`: If the solver type is unknown. +- `KeyError`: If the solver name is not registered (normally already reported when the +config is parsed). **Returns**: -`QuerySolver`: An instance of the appropriate query solver based on the configuration. +`QuerySolver`: An instance of the registered solver. #### get\_sink\_config diff --git a/docs/impulse/docs/references/api/sidebar.json b/docs/impulse/docs/references/api/sidebar.json index 3f08df0f..f3648437 100644 --- a/docs/impulse/docs/references/api/sidebar.json +++ b/docs/impulse/docs/references/api/sidebar.json @@ -34,7 +34,8 @@ "items": [ "references/api/impulse_query_engine/analyze/query/solvers/default_solver", "references/api/impulse_query_engine/analyze/query/solvers/query_solver", - "references/api/impulse_query_engine/analyze/query/solvers/solver_config" + "references/api/impulse_query_engine/analyze/query/solvers/solver_config", + "references/api/impulse_query_engine/analyze/query/solvers/solver_registry" ], "label": "impulse_query_engine.analyze.query.solvers", "type": "category" diff --git a/docs/impulse/docs/references/query_engine/query_solvers.md b/docs/impulse/docs/references/query_engine/query_solvers.md index c71cd297..6e858fd9 100644 --- a/docs/impulse/docs/references/query_engine/query_solvers.md +++ b/docs/impulse/docs/references/query_engine/query_solvers.md @@ -109,9 +109,10 @@ the columns each table is expected to carry. Solver tuning lives under the `query_engine` section of your report config: - [`query_engine.solver`](../../config/configuration.md#query_engine-optional) - — the solver to use. Defaults to `"DefaultSolver"`. The values + — the registered name of the solver to use. Defaults to `"DefaultSolver"`. The values `"DeltaSolver"` and `"KeyValueStoreSolver"` are **deprecated aliases** - retained for backward compatibility; both resolve to `DefaultSolver`. + retained for backward compatibility; both resolve to `DefaultSolver`. A custom solver is + selected by the name it registered; see [Custom solver](../../extensibility/custom_solver.md). - [`query_engine.data_type` and `raw_encoder`](../../config/configuration.md#query_engine-optional) — declare whether `channels` holds pre-encoded `[tstart, tend)` intervals (`"RLE"`, the default) or raw point samples (`"RAW"`). For raw data, diff --git a/docs/impulse/pydoc-markdown.yml b/docs/impulse/pydoc-markdown.yml index f52dda2b..795cf488 100644 --- a/docs/impulse/pydoc-markdown.yml +++ b/docs/impulse/pydoc-markdown.yml @@ -24,6 +24,7 @@ loaders: - impulse_query_engine.analyze.query.solvers.default_solver - impulse_query_engine.analyze.query.solvers.solver_config - impulse_query_engine.analyze.query.solvers.query_solver + - impulse_query_engine.analyze.query.solvers.solver_registry - impulse_query_engine.model.series.sample_series - impulse_query_engine.model.series.intervals - impulse_query_engine.model.series.points_in_time diff --git a/skills/impulse-analyze/SKILL.md b/skills/impulse-analyze/SKILL.md index 9b0aef93..2be087d4 100644 --- a/skills/impulse-analyze/SKILL.md +++ b/skills/impulse-analyze/SKILL.md @@ -139,6 +139,10 @@ DefaultSolver(spark, config=None, is_raw_data=False, drop_implausible_data=False With **Option B**, pass `report.get_solver()` instead of constructing one, so the solver honors your report config. +A custom solver (a `DefaultSolver` subclass) is passed the same way: `solver=MySolver(spark, config=MyConfig(...))`. +No registration is needed for ad-hoc queries; `register_solver` only matters for selecting it by name in a +report config (see `impulse-config`). + ## What comes back The result DataFrame is `[container_id, ]`: **`container_id` is always the first column** diff --git a/skills/impulse-config/SKILL.md b/skills/impulse-config/SKILL.md index ddd5a32a..4086370f 100644 --- a/skills/impulse-config/SKILL.md +++ b/skills/impulse-config/SKILL.md @@ -98,7 +98,7 @@ Omit entirely for `DefaultSolver` + `data_type="RLE"`. | Field | Default | Description | |-------------------------|--------------------|---------------------------------------------------------------------------------------------------| -| `solver` | `"DefaultSolver"` | The solver. `"DeltaSolver"` / `"KeyValueStoreSolver"` are **deprecated aliases** for `DefaultSolver`. | +| `solver` | `"DefaultSolver"` | Registered name of the solver. `"DeltaSolver"` / `"KeyValueStoreSolver"` are **deprecated aliases** for `DefaultSolver`. A custom solver (registered with `register_solver`, imported before the config is parsed) is selected by its name; its own settings go in `solver_config`, validated by the config class it registered. | | `data_type` | `"RLE"` | `"RLE"` for interval-encoded `[tstart, tend)` samples; `"RAW"` for `(timestamp, value)` samples. | | `raw_encoder` | `null` (→ `"RLE"`) | How RAW point data becomes intervals. `"RLE"` collapses equal-valued runs (default); `"INTERVAL"` keeps every sample, only deriving `tend` + dropping exact duplicates. Only used when `data_type="RAW"`. | | `drop_implausible_data` | `false` | Drop `channels` rows where `is_plausible = false`. **Requires `data_type="RAW"`** (RLE raises). | diff --git a/skills/impulse-data-model/SKILL.md b/skills/impulse-data-model/SKILL.md index a52e2127..f788bc16 100644 --- a/skills/impulse-data-model/SKILL.md +++ b/skills/impulse-data-model/SKILL.md @@ -133,10 +133,12 @@ schemas and `impulse-config` for how `measurement_dimensions` picks the containe ## Adapting to an existing layout Reshaping into the silver shape at ingest is recommended. If the data already lives in Delta with -different names and rewriting is impractical, there are two escape hatches. The rule of thumb: +different names and rewriting is impractical, there are escape hatches. The rule of thumb: - **`SolverConfig` for naming differences** — same tables and relationships, different column names. -- **Custom solver for structural differences** — no EAV tags, alias lookups, composite keys. +- **Custom MeasurementDB for structural differences** — no EAV tags, alias lookups, composite keys: + reshape your tables when they're read. +- **Custom solver for engine behaviour** — candidate selection, pruning, post-processing the solve output. - **ETL into the standard shape** for everything else. ### Column-name remapping with SolverConfig @@ -163,10 +165,15 @@ The underlying tables must still follow the silver-layer relationships (EAV tag per-`(container_id, channel_id)` channel rows). The full `SolverConfig` schema — per-table `filters`, `project_id` scoping, `channel_mapping.join_keys`, and unit conversion — is in `impulse-config`. -### Custom solver +### Custom MeasurementDB or custom solver -For layouts that don't match the relationships at all, subclass `QuerySolver` (from -`impulse_query_engine.analyze.query.solvers`) and register it in config. You take on the solver -pipeline stages (`filter_container_tags`, `filter_container_metrics`, `filter_channel_tags`, -`filter_channel_metrics`, `solve`). This is a large investment — first check whether a one-time ETL -into the standard shape is cheaper. +For layouts that don't match the relationships at all, register a custom MeasurementDB (subclass +`MeasurementDB` and `MeasurementDBConfig`, decorate with `@register_measurement_db`) that reshapes your +tables into the silver model when they're read, and select it with `query_engine.measurement_db`. Every +reader then sees the same data, and the stock solver runs unchanged. + +For different engine behaviour on silver-shaped data, subclass `DefaultSolver`, decorate it with +`@register_solver` (from `impulse_query_engine.analyze.query.solvers`), and select it with +`query_engine.solver`; its own settings go on a `SolverConfig` subclass. Both are large investments — +first check whether a one-time ETL into the standard shape is cheaper. Config details are in +`impulse-config`. diff --git a/src/impulse_query_engine/analyze/query/solvers/__init__.py b/src/impulse_query_engine/analyze/query/solvers/__init__.py index 1b736776..2d22af61 100644 --- a/src/impulse_query_engine/analyze/query/solvers/__init__.py +++ b/src/impulse_query_engine/analyze/query/solvers/__init__.py @@ -1,4 +1,14 @@ from .default_solver import DefaultSolver, TimeSeriesCache +from .query_solver import QuerySolver, SolverBuildContext from .solver_config import SolverConfig +from .solver_registry import register_solver, resolve_solver -__all__ = ["DefaultSolver", "TimeSeriesCache", "SolverConfig"] +__all__ = [ + "DefaultSolver", + "TimeSeriesCache", + "SolverConfig", + "QuerySolver", + "SolverBuildContext", + "register_solver", + "resolve_solver", +] 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 f9e6d7b0..7cb204e8 100644 --- a/src/impulse_query_engine/analyze/query/solvers/default_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/default_solver.py @@ -19,9 +19,10 @@ from impulse_query_engine.model.series.points_in_time_series import PointsInTimeSeries from impulse_query_engine.model.series.sample_series import SampleSeries -from .query_solver import QuerySolver +from .query_solver import QuerySolver, SolverBuildContext from .series_cache import SeriesCache from .solver_config import RawEncoder, SolverConfig +from .solver_registry import register_solver from .utils.interval_encoder import IntervalEncoder from .utils.rle_encoder import RleEncoder from .utils.window_bounds import with_window_bounds @@ -210,6 +211,7 @@ def _container_meta(self, cols) -> dict: return {c: row[c] for c in cols if c in self.mdf.columns} +@register_solver("DefaultSolver", SolverConfig, aliases=("DeltaSolver", "KeyValueStoreSolver")) class DefaultSolver(QuerySolver): """ The default query-engine solver. Adapts to the shape of the silver layer. @@ -276,6 +278,17 @@ def __init__( self.config.reject_implausible_channels_filter_in_raw(is_raw=self.is_raw_data) self.channel_encoder: RleEncoder | IntervalEncoder = self._build_channel_encoder() + @classmethod + def from_config(cls, ctx: SolverBuildContext) -> DefaultSolver: + """Build the solver with the Spark session and RAW-data flags from the report 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``. 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 871f8206..b69911f8 100644 --- a/src/impulse_query_engine/analyze/query/solvers/query_solver.py +++ b/src/impulse_query_engine/analyze/query/solvers/query_solver.py @@ -2,11 +2,12 @@ import abc from abc import ABC +from dataclasses import dataclass from typing import TYPE_CHECKING import pyspark.sql.functions as F import pyspark.sql.types as T -from pyspark.sql import DataFrame +from pyspark.sql import DataFrame, SparkSession from pyspark.sql.column import Column if TYPE_CHECKING: @@ -16,7 +17,25 @@ TimeSeriesSelector, ) -from .solver_config import SolverConfig +from .solver_config import RawEncoder, SolverConfig + + +@dataclass(frozen=True, slots=True) +class SolverBuildContext: + """Everything a report passes to a solver when it builds one (see ``QuerySolver.from_config``). + + The fields copy the ``query_engine`` settings a solver needs rather than passing the report's + ``QueryEngine`` config, because that model lives in ``impulse_reporting`` and the query engine + must not depend on the reporting layer (``config_parser`` imports this package, so the reverse + import would also be circular). When a solver needs another ``query_engine`` setting, add a + field here and set it in ``Report.create_solver``. + """ + + spark: SparkSession + solver_config: SolverConfig | None = None + is_raw_data: bool = False + drop_implausible_data: bool = False + raw_encoder: RawEncoder = RawEncoder.RLE class QuerySolver(ABC): @@ -35,6 +54,15 @@ class QuerySolver(ABC): def __init__(self, config: SolverConfig = None): self.config = config or SolverConfig() + @classmethod + def from_config(cls, ctx: SolverBuildContext) -> QuerySolver: + """Build the solver from what a report provides (see ``register_solver``). + + The base implementation passes only ``ctx.solver_config``. Override it when the + constructor needs more, e.g. the Spark session or the RAW-data flags. + """ + return cls(config=ctx.solver_config) + @staticmethod def _apply_column_mapping(df: DataFrame, mapping: dict[str, str]) -> DataFrame: """Rename DataFrame columns according to a physical → internal mapping.""" diff --git a/src/impulse_query_engine/analyze/query/solvers/solver_config.py b/src/impulse_query_engine/analyze/query/solvers/solver_config.py index 559289ba..b276d50e 100644 --- a/src/impulse_query_engine/analyze/query/solvers/solver_config.py +++ b/src/impulse_query_engine/analyze/query/solvers/solver_config.py @@ -18,7 +18,7 @@ from enum import StrEnum from typing import Literal -from pydantic import BaseModel, model_validator +from pydantic import BaseModel, ConfigDict, model_validator class RawEncoder(StrEnum): @@ -154,6 +154,10 @@ class SolverConfig(BaseModel): their own unit. """ + # reject misspelled settings instead of silently dropping them; inherited by the config + # classes of custom solvers (a dropped custom setting would silently change results) + model_config = ConfigDict(extra="forbid") + project_id: str | None = None channel_time_unit: Literal["s", "ms", "us", "ns"] | None = None channel_time_origin: Literal["epoch", "container_start"] = "epoch" diff --git a/src/impulse_query_engine/analyze/query/solvers/solver_registry.py b/src/impulse_query_engine/analyze/query/solvers/solver_registry.py new file mode 100644 index 00000000..d99d8768 --- /dev/null +++ b/src/impulse_query_engine/analyze/query/solvers/solver_registry.py @@ -0,0 +1,29 @@ +"""Name registry for query solvers, so a report config can select a solver by name. + +A solver registers under a name with ``@register_solver(...)`` and is selected with +``query_engine.solver`` in the report config. Selection is by name only: the package that +registers a solver must be imported before the report config is parsed. +""" + +from impulse_query_engine.registry import Registration, Registry + +from .query_solver import QuerySolver +from .solver_config import SolverConfig + +_REGISTRY: Registry[QuerySolver, SolverConfig] = Registry("solver", QuerySolver, SolverConfig) + + +def register_solver( + name: str, config_cls: type[SolverConfig] = SolverConfig, *, aliases: tuple[str, ...] = () +): + """Register a ``QuerySolver`` subclass under *name* (and *aliases*), configured by *config_cls*. + + ``config_cls`` validates ``query_engine.solver_config`` when this solver is selected. + Re-registering a name replaces the previous entry. + """ + return _REGISTRY.register(name, config_cls, aliases=aliases) + + +def resolve_solver(name: str) -> Registration[QuerySolver, SolverConfig]: + """Return the registration for *name*; raises ``KeyError`` listing the registered names.""" + return _REGISTRY.resolve(name) diff --git a/src/impulse_query_engine/measurement_db_registry.py b/src/impulse_query_engine/measurement_db_registry.py index 9b9a0e40..69cecff5 100644 --- a/src/impulse_query_engine/measurement_db_registry.py +++ b/src/impulse_query_engine/measurement_db_registry.py @@ -1,19 +1,16 @@ -from dataclasses import dataclass from typing import TYPE_CHECKING from .measurement_db_abstract import AbstractMeasurementDB +from .registry import Registration, Registry if TYPE_CHECKING: from .measurement_db import MeasurementDBConfig - -@dataclass(frozen=True, slots=True) -class MeasurementDatabaseRegistration: - db_cls: type[AbstractMeasurementDB] - config_cls: type["MeasurementDBConfig"] - - -_REGISTRY: dict[str, MeasurementDatabaseRegistration] = {} +# No config_base: importing MeasurementDBConfig here would be circular (measurement_db.py +# imports this module to register the built-in DB). +_REGISTRY: Registry[AbstractMeasurementDB, "MeasurementDBConfig"] = Registry( + "measurement DB", AbstractMeasurementDB +) def register_measurement_db(name: str, config_cls: type["MeasurementDBConfig"]): @@ -21,22 +18,11 @@ def register_measurement_db(name: str, config_cls: type["MeasurementDBConfig"]): Re-registering a name replaces the previous entry. """ - - def decorator[T: type[AbstractMeasurementDB]](db_cls: T) -> T: - if not issubclass(db_cls, AbstractMeasurementDB): - raise TypeError(f"{db_cls!r} must subclass AbstractMeasurementDB") - _REGISTRY[name] = MeasurementDatabaseRegistration(db_cls, config_cls) - return db_cls - - return decorator + return _REGISTRY.register(name, config_cls) -def resolve_measurement_db(name: str) -> MeasurementDatabaseRegistration: +def resolve_measurement_db( + name: str, +) -> Registration[AbstractMeasurementDB, "MeasurementDBConfig"]: """Return the registration for *name*; raises ``KeyError`` listing the registered names.""" - try: - return _REGISTRY[name] - except KeyError: - raise KeyError( - f"Unknown measurement DB {name!r}; registered: {sorted(_REGISTRY)}. " - "Import the package that registers it before parsing the report config." - ) from None + return _REGISTRY.resolve(name) diff --git a/src/impulse_query_engine/registry.py b/src/impulse_query_engine/registry.py new file mode 100644 index 00000000..0c98e747 --- /dev/null +++ b/src/impulse_query_engine/registry.py @@ -0,0 +1,69 @@ +"""Name registries, so a report config can select an implementation by name. + +An implementation registers under a name with a registry's ``register`` decorator and is selected +by that name in the report config. Selection is by name only: the package that registers an +implementation must be imported before the report config is parsed. +""" + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class Registration[I, C]: + impl_cls: type[I] + config_cls: type[C] + + +class Registry[I, C]: + """Maps names to an implementation class (a subclass of *base*) and its config class. + + *kind* names the implementations in error messages. When *config_base* is given, a config + class must subclass it. + """ + + def __init__(self, kind: str, base: type[I], config_base: type[C] | None = None): + self.kind = kind + self.base = base + self.config_base = config_base + self._entries: dict[str, Registration[I, C]] = {} + + def register(self, name: str, config_cls: type[C], *, aliases: tuple[str, ...] = ()): + """Register a subclass of ``base`` under *name* (and *aliases*), configured by *config_cls*. + + Re-registering a name replaces the previous entry. + """ + self._check_config_cls(config_cls) + + def decorator[T: type](impl_cls: T) -> T: + if not _is_subclass(impl_cls, self.base): + raise TypeError(f"{impl_cls!r} must subclass {self.base.__name__}") + registration = Registration(impl_cls, config_cls) + for registered_name in (name, *aliases): + self._entries[registered_name] = registration + return impl_cls + + return decorator + + def resolve(self, name: str) -> Registration[I, C]: + """Return the registration for *name*; raises ``KeyError`` listing the registered names.""" + try: + return self._entries[name] + except KeyError: + raise KeyError( + f"Unknown {self.kind} {name!r}; registered: {sorted(self._entries)}. " + "Import the package that registers it before parsing the report config." + ) from None + + def _check_config_cls(self, config_cls: type) -> None: + """Raise ``TypeError`` unless *config_cls* subclasses ``config_base``. + + A registry without a ``config_base`` (the measurement-DB registry) accepts any config class. + """ + if self.config_base is None: + return + if not _is_subclass(config_cls, self.config_base): + raise TypeError(f"{config_cls!r} must subclass {self.config_base.__name__}") + + +def _is_subclass(obj: object, base: type) -> bool: + return isinstance(obj, type) and issubclass(obj, base) diff --git a/src/impulse_reporting/config/config_parser.py b/src/impulse_reporting/config/config_parser.py index 79b27472..286181f9 100644 --- a/src/impulse_reporting/config/config_parser.py +++ b/src/impulse_reporting/config/config_parser.py @@ -1,14 +1,25 @@ import os import re import warnings +from collections.abc import Callable from datetime import datetime from enum import Enum, StrEnum from typing import Annotated, Any, Self import pydantic -from pydantic import AfterValidator, BaseModel, field_validator, model_validator, ConfigDict +from pydantic import ( + AfterValidator, + BaseModel, + ConfigDict, + Field, + SerializeAsAny, + ValidationInfo, + field_validator, + model_validator, +) from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder, SolverConfig +from impulse_query_engine.analyze.query.solvers.solver_registry import resolve_solver from impulse_query_engine.measurement_db import AbstractMeasurementDB, MeasurementDBConfig from impulse_query_engine.measurement_db_registry import resolve_measurement_db from impulse_reporting.channels.calculated_channel_kpis import DEFAULT_KPIS, KPI_BUILDERS @@ -90,10 +101,12 @@ 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 solvers. + ``query_engine.solver`` takes any registered solver name (see + ``register_solver``); these are the names Impulse registers itself. ``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 @@ -359,14 +372,31 @@ class ContainerFilters(BaseModel): _PYDANTIC_DIR = os.path.dirname(pydantic.__file__) +def _registered_in(resolve: Callable[[str], object]) -> AfterValidator: + """Field validator: the name must be registered (*resolve* raises ``KeyError`` otherwise).""" + + def check(name: str) -> str: + try: + resolve(name) + except KeyError as exc: + raise ValueError(exc.args[0]) from None + return name + + return AfterValidator(check) + + class QueryEngine(BaseModel): """ Configuration for the query engine solver. Parameters ---------- - solver : Solvers, default=Solvers.DEFAULT_SOLVER - The solver type to use for query execution. + solver : str, default="DefaultSolver" + Registered name of the solver to use for query execution. Impulse + registers ``"DefaultSolver"`` (``"DeltaSolver"`` and + ``"KeyValueStoreSolver"`` are deprecated aliases); a custom solver + registered with ``register_solver`` is selected by its name. An + unregistered name fails when the config is loaded. raw_encoder : RawEncoder, optional, default=None Encoder used to convert RAW point data into intervals. ``RLE`` collapses consecutive equal-valued samples into runs; ``INTERVAL`` @@ -389,7 +419,9 @@ class QueryEngine(BaseModel): ``column_name_mapping`` and ``filters`` dicts. When omitted, all default column names are used and no - project/toolbox filtering is applied. + project/toolbox filtering is applied. It is validated with the + config class the selected solver registered, so a custom solver can + add its own settings to a ``SolverConfig`` subclass. measurement_db : str, default="MeasurementDB" Registered name of the MeasurementDB implementation to use. measurement_db_config : dict, optional @@ -414,12 +446,14 @@ class QueryEngine(BaseModel): # (default measurement db silently used) model_config = ConfigDict(extra="forbid") - solver: Solvers = Solvers.DEFAULT_SOLVER + solver: Annotated[str, _registered_in(resolve_solver)] = Solvers.DEFAULT_SOLVER.value data_type: DataType = DataType.RLE drop_implausible_data: bool = False raw_encoder: RawEncoder | None = None - solver_config: SolverConfig | None = None - measurement_db: str = "MeasurementDB" + # validate_default: a custom config class may require settings even when solver_config is omitted; + # SerializeAsAny: dumps (and so config_hash) include the custom config class's settings + solver_config: SerializeAsAny[SolverConfig] | None = Field(default=None, validate_default=True) + measurement_db: Annotated[str, _registered_in(resolve_measurement_db)] = "MeasurementDB" measurement_db_config: dict[str, Any] | None = None max_channels_per_batch: int = 500 max_containers_per_batch: int | None = None @@ -455,6 +489,32 @@ def _resolve_deprecated_batch_size(cls, data): data["max_channels_per_batch"] = data.pop("batch_size") return data + @field_validator("solver_config", mode="before") + @classmethod + def _validate_with_registered_config_cls(cls, value: Any, info: ValidationInfo) -> Any: + """Validate ``solver_config`` with the config class the selected solver registered. + + A custom solver can register a ``SolverConfig`` subclass with its own settings, so the + solver receives a typed config and missing or invalid settings fail when the config is + loaded. + """ + if "solver" not in info.data: # an unregistered solver name is already reported + return value + name = info.data["solver"] + config_cls = resolve_solver(name).config_cls + if config_cls is SolverConfig: + return value # validated as the declared field type + try: + # from_attributes: a plain SolverConfig instance is re-validated as config_cls + return config_cls.model_validate({} if value is None else value, from_attributes=True) + except pydantic.ValidationError as exc: + # Flatten the nested errors; raised as-is, pydantic would collapse their locations. + details = "; ".join( + f"{'.'.join(map(str, err['loc']))}: {err['msg']}" if err["loc"] else err["msg"] + for err in exc.errors() + ) + raise ValueError(f"invalid settings for solver {name!r}: {details}") from None + @property def batch_size(self) -> int: """Deprecated alias for :attr:`max_channels_per_batch`.""" @@ -679,7 +739,7 @@ def build_measurement_db_config( **(query_engine.measurement_db_config or {}), table_locations="unity_catalog", ) - return registration.db_cls, db_config + return registration.impl_cls, db_config class ImpulseConfig(BaseModel): @@ -774,6 +834,10 @@ class ImpulseConfig(BaseModel): >>> config = ImpulseConfig.model_validate(config_data) """ + # reject misplaced or misspelled top-level keys, e.g. "solver" outside "query_engine", + # which would otherwise be dropped and the default solver silently used + model_config = ConfigDict(extra="forbid") + source: Source unity_sink: UnitySink | None = None container_filters: ContainerFilters | None = None @@ -798,11 +862,12 @@ def _normalize_measurement_dimensions(cls, value: list[str]) -> list[str]: @model_validator(mode="after") def _validate_measurement_db(self) -> Self: - """Build the selected MeasurementDB implementation's config class once, so an unregistered - ``measurement_db`` or a missing, unknown or duplicated key fails when the config is parsed. + """Build the selected MeasurementDB implementation's config class once, so a missing, + unknown or duplicated key fails when the config is parsed (an unregistered + ``measurement_db`` is already rejected on the field). """ try: build_measurement_db_config(dict(self.source), self.query_engine) - except (KeyError, TypeError) as exc: + except TypeError as exc: raise ValueError(f"query_engine.measurement_db: {exc.args[0]}") from exc return self diff --git a/src/impulse_reporting/core/report.py b/src/impulse_reporting/core/report.py index f8d9bdab..b27c5e08 100644 --- a/src/impulse_reporting/core/report.py +++ b/src/impulse_reporting/core/report.py @@ -12,8 +12,12 @@ 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 import ( + QuerySolver, + SolverBuildContext, + resolve_solver, +) +from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder from impulse_query_engine.measurement_db import AbstractMeasurementDB from impulse_reporting.aggregations.aggregation_types import AggregationType from impulse_reporting.channels.channel_types import ChannelType @@ -83,6 +87,18 @@ ) +def _solver_telemetry_name(solver: str) -> str: + """Telemetry value for the selected solver. + + Built-in solvers keep their member name (``DEFAULT_SOLVER``, ...) so existing telemetry is + unchanged; custom solvers are reported as ``"custom"`` so customer class names aren't sent. + """ + try: + return Solvers(solver).name + except ValueError: + return "custom" + + class Report: """Represents a report containing pages, events, and configurations for data processing and persistence.""" @@ -152,7 +168,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", _solver_telemetry_name(self.config.query_engine.solver)) log_telemetry(self.ws, "data_type", self.config.query_engine.data_type.value) @property @@ -333,7 +349,12 @@ def create_sink(config: ImpulseConfig) -> Sink: @staticmethod def create_solver(spark: SparkSession, config: ImpulseConfig) -> QuerySolver: """ - Create a query solver based on the provided configuration. + Create the query solver selected by ``query_engine.solver``. + + Resolves the registered solver name (``"DeltaSolver"`` and + ``"KeyValueStoreSolver"`` are deprecated aliases of ``"DefaultSolver"``) + and builds it with its ``from_config`` hook. + Parameters ---------- spark : SparkSession @@ -344,31 +365,26 @@ def create_solver(spark: SparkSession, config: ImpulseConfig) -> QuerySolver: Returns ------- QuerySolver - An instance of the appropriate query solver based on the configuration. + An instance of the registered solver. 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 the solver name is not registered (normally already reported when the + config is parsed). + """ + query_engine = config.query_engine + registration = resolve_solver(query_engine.solver) + return registration.impl_cls.from_config( + SolverBuildContext( + spark=spark, + solver_config=query_engine.solver_config, + is_raw_data=query_engine.data_type is DataType.RAW, + drop_implausible_data=query_engine.drop_implausible_data, + # None for RLE input, where it isn't consulted; pass the solver default + raw_encoder=query_engine.raw_encoder or RawEncoder.RLE, + ) + ) def get_sink_config(self) -> SinkConfig: """ diff --git a/tests/conftest.py b/tests/conftest.py index 74afafd9..c2f9d2cb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,6 +11,8 @@ from pyspark.sql import SparkSession import impulse_query_engine.schema as S +from impulse_query_engine import measurement_db_registry +from impulse_query_engine.analyze.query.solvers import solver_registry from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig @@ -69,6 +71,13 @@ def mock_workspace_client(): return create_autospec(WorkspaceClient) +@pytest.fixture +def isolated_registries(monkeypatch): + """Drop the solver and MeasurementDB registrations a test makes (the registries are global).""" + for registry in (solver_registry._REGISTRY, measurement_db_registry._REGISTRY): + monkeypatch.setattr(registry, "_entries", dict(registry._entries)) + + @pytest.fixture def basic_narrow_db(spark, mock_workspace_client) -> MeasurementDB: """Return a basic narrow MeasurementDB instance with preloaded data.""" diff --git a/tests/impulse_query_engine/unit/analyze/query/solvers/solver_registry_test.py b/tests/impulse_query_engine/unit/analyze/query/solvers/solver_registry_test.py new file mode 100644 index 00000000..117a95e6 --- /dev/null +++ b/tests/impulse_query_engine/unit/analyze/query/solvers/solver_registry_test.py @@ -0,0 +1,86 @@ +# pylint: disable=missing-function-docstring +"""Tests for the solver registry wiring and the ``from_config`` build hook. + +The registry mechanics (aliases, last registration wins, the unknown-name message) are covered +in ``registry_test.py``. +""" + +from unittest.mock import sentinel + +import pytest + +from impulse_query_engine.analyze.query.solvers import ( + DefaultSolver, + QuerySolver, + SolverBuildContext, + SolverConfig, + register_solver, + resolve_solver, +) +from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder + + +@pytest.mark.parametrize("name", ["DefaultSolver", "DeltaSolver", "KeyValueStoreSolver"]) +def test_default_solver_is_registered_under_its_name_and_aliases(name): + registration = resolve_solver(name) + + assert registration.impl_cls is DefaultSolver + assert registration.config_cls is SolverConfig + + +def test_register_rejects_non_solver_class(): + with pytest.raises(TypeError, match="must subclass QuerySolver"): + register_solver("Bad")(object) + + +def test_register_rejects_non_solver_config_class(): + with pytest.raises(TypeError, match="must subclass SolverConfig"): + register_solver("Bad", dict) + + +def test_unknown_name_error_names_the_solver_kind(): + with pytest.raises(KeyError, match=r"Unknown solver 'Missing'"): + resolve_solver("Missing") + + +def test_base_from_config_passes_only_the_solver_config(): + class ConfigOnlySolver(QuerySolver): + # QuerySolver's five pipeline stages are abstract; this test only builds the solver. + def filter_container_tags(self, spark, query): ... + + def filter_container_metrics( + self, spark, query, container_df, pre_filtered_containers_df=None + ): ... + + def filter_channel_tags(self, spark, db, container_df, selectors): ... + + def filter_channel_metrics(self, spark, db, channel_df, selectors): ... + + def solve(self, query, channels_df, selections, dtypes): ... + + config = SolverConfig(project_id="p1") + solver = ConfigOnlySolver.from_config( + SolverBuildContext(spark=sentinel.spark, solver_config=config, is_raw_data=True) + ) + + assert type(solver) is ConfigOnlySolver + assert solver.config is config + + +def test_default_solver_from_config_wires_spark_and_raw_flags(): + config = SolverConfig() + solver = DefaultSolver.from_config( + SolverBuildContext( + spark=sentinel.spark, + solver_config=config, + is_raw_data=True, + drop_implausible_data=True, + raw_encoder=RawEncoder.INTERVAL, + ) + ) + + assert solver.spark is sentinel.spark + assert solver.config is config + assert solver.is_raw_data is True + assert solver.drop_implausible_data is True + assert solver.raw_encoder is RawEncoder.INTERVAL diff --git a/tests/impulse_query_engine/unit/measurement_db_registry_test.py b/tests/impulse_query_engine/unit/measurement_db_registry_test.py index 33d736b3..9f731377 100644 --- a/tests/impulse_query_engine/unit/measurement_db_registry_test.py +++ b/tests/impulse_query_engine/unit/measurement_db_registry_test.py @@ -1,12 +1,11 @@ # pylint: disable=missing-function-docstring -"""Tests for the ``AbstractMeasurementDB`` contract and the name registry.""" +"""Tests for the ``AbstractMeasurementDB`` contract and the MeasurementDB registry wiring.""" from unittest.mock import create_autospec import pytest from databricks.sdk import WorkspaceClient -import impulse_query_engine.measurement_db_registry as registry from impulse_query_engine import __version__ from impulse_query_engine.measurement_db import ( AbstractMeasurementDB, @@ -19,11 +18,6 @@ ) -@pytest.fixture(autouse=True) -def isolated_registry(monkeypatch): - monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) - - def test_measurement_db_implements_the_contract(): # Guard: if a new abstract method is added to AbstractMeasurementDB without an # implementation on the built-in, this fails here with a clear, local signal rather @@ -78,21 +72,14 @@ def test_register_rejects_non_subclass(): register_measurement_db("Bad", MeasurementDBConfig)(object) -def test_register_returns_class_and_last_registration_wins(): - class First(MeasurementDB): ... - - class Second(MeasurementDB): ... +def test_builtin_measurement_db_is_registered(): + registration = resolve_measurement_db("MeasurementDB") - assert register_measurement_db("Custom", MeasurementDBConfig)(First) is First - register_measurement_db("Custom", MeasurementDBConfig)(Second) - assert resolve_measurement_db("Custom").db_cls is Second + assert registration.impl_cls is MeasurementDB + assert registration.config_cls is MeasurementDBConfig -def test_unknown_name_lists_registered_names(): - # The error names the unknown DB, lists the registered names, and tells the caller how to - # fix it (import the registering package before parsing) -- all three are user-facing guidance. - with pytest.raises( - KeyError, - match=r"'Missing'.*\['MeasurementDB'\].*Import the package that registers it", - ): +def test_unknown_name_error_names_the_measurement_db_kind(): + # The registry mechanics are covered in registry_test.py; this pins the user-facing kind. + with pytest.raises(KeyError, match=r"Unknown measurement DB 'Missing'"): resolve_measurement_db("Missing") diff --git a/tests/impulse_query_engine/unit/registry_test.py b/tests/impulse_query_engine/unit/registry_test.py new file mode 100644 index 00000000..c55a51f1 --- /dev/null +++ b/tests/impulse_query_engine/unit/registry_test.py @@ -0,0 +1,73 @@ +# pylint: disable=missing-function-docstring +"""Tests for the generic name registry behind ``register_solver`` and ``register_measurement_db``.""" + +import pytest + +from impulse_query_engine.registry import Registry + + +class Widget: ... + + +class WidgetConfig: ... + + +@pytest.fixture +def widgets() -> Registry[Widget, WidgetConfig]: + # A local registry, so these tests don't touch the global solver / MeasurementDB registries. + return Registry("widget", Widget, WidgetConfig) + + +def test_register_returns_class_and_last_registration_wins(widgets): + class First(Widget): ... + + class Second(Widget): ... + + assert widgets.register("Custom", WidgetConfig)(First) is First + widgets.register("Custom", WidgetConfig)(Second) + assert widgets.resolve("Custom").impl_cls is Second + + +def test_register_with_config_class_and_aliases(widgets): + class CustomConfig(WidgetConfig): ... + + class Custom(Widget): ... + + widgets.register("Custom", CustomConfig, aliases=("OldCustom",))(Custom) + + assert widgets.resolve("Custom").config_cls is CustomConfig + assert widgets.resolve("OldCustom") == widgets.resolve("Custom") + + +@pytest.mark.parametrize("impl", [object, Widget(), "Widget"]) +def test_register_rejects_non_subclass(widgets, impl): + with pytest.raises(TypeError, match="must subclass Widget"): + widgets.register("Bad", WidgetConfig)(impl) + + +@pytest.mark.parametrize("config_cls", [dict, WidgetConfig()]) +def test_register_rejects_non_subclass_config_class(widgets, config_cls): + with pytest.raises(TypeError, match="must subclass WidgetConfig"): + widgets.register("Bad", config_cls) + + +def test_config_class_is_unchecked_without_config_base(): + registry = Registry("widget", Widget) + + registry.register("Custom", dict)(Widget) + + assert registry.resolve("Custom").config_cls is dict + + +def test_unknown_name_lists_registered_names(widgets): + # The error names the unknown entry, lists the registered names, and tells the caller how to + # fix it (import the registering package before parsing) -- all three are user guidance. + widgets.register("B", WidgetConfig)(Widget) + widgets.register("A", WidgetConfig)(Widget) + + with pytest.raises( + KeyError, + match=r"Unknown widget 'Missing'; registered: \['A', 'B'\]\. " + r"Import the package that registers it", + ): + widgets.resolve("Missing") diff --git a/tests/impulse_reporting/integration/custom_measurement_db_report_test.py b/tests/impulse_reporting/integration/custom_measurement_db_report_test.py index 7e006e68..7fdd280c 100644 --- a/tests/impulse_reporting/integration/custom_measurement_db_report_test.py +++ b/tests/impulse_reporting/integration/custom_measurement_db_report_test.py @@ -13,7 +13,6 @@ from databricks.sdk import WorkspaceClient from pyspark.sql import Row -import impulse_query_engine.measurement_db_registry as registry from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig from impulse_query_engine.measurement_db_registry import register_measurement_db from impulse_reporting.aggregations.histogram import HistogramDuration @@ -42,8 +41,7 @@ def container_metrics(self, spark): @pytest.fixture(autouse=True) -def allow_list_db(monkeypatch): - monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) +def allow_list_db(isolated_registries): register_measurement_db("AllowListDB", AllowListConfig)(AllowListDB) 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..7a1aa16b --- /dev/null +++ b/tests/impulse_reporting/integration/custom_solver_report_test.py @@ -0,0 +1,99 @@ +# pylint: disable=missing-function-docstring +"""End-to-end: a customer-provided solver, selected from the report config. + +The custom solver drops the containers listed in its own config setting from the solve output. +Each test covers one promise of the extension point: the report builds the selected solver with its +typed config and solves through it, and registering a custom solver leaves the built-in +``DefaultSolver`` unchanged. +""" + +from unittest.mock import create_autospec + +import pyspark.sql.functions as F +import pytest +from databricks.sdk import WorkspaceClient +from pyspark.sql import Row + +from impulse_query_engine.analyze.query.solvers import DefaultSolver, SolverConfig, register_solver +from impulse_reporting.aggregations.histogram import HistogramDuration +from impulse_reporting.core.page import Page +from impulse_reporting.core.report import Report + + +class ExcludeContainersConfig(SolverConfig): + excluded_container_ids: list[int] + + +class ExcludeContainersSolver(DefaultSolver): + """Drops the containers listed in ``excluded_container_ids`` from the solve output.""" + + def solve(self, query, channels_df, selections, dtypes): + solved = super().solve(query, channels_df, selections, dtypes) + return solved.filter( + ~F.col(self.config.container_id_col).isin(self.config.excluded_container_ids) + ) + + +@pytest.fixture(autouse=True) +def exclude_containers_solver(isolated_registries): + register_solver("ExcludeContainersSolver", ExcludeContainersConfig)(ExcludeContainersSolver) + + +def _create_report(spark, query_engine: dict) -> Report: + config = { + "source": { + "container_metrics_table": "spark_catalog.silver.container_metrics", + "channel_metrics_table": "spark_catalog.silver.channel_metrics", + "channels_uri": "spark_catalog.silver.channels", + }, + "unity_sink": {"catalog": "spark_catalog", "schema": "gold", "table_prefix": "e2e_solver"}, + "query_engine": query_engine, + } + report = Report( + name="e2e_custom_solver", + spark=spark, + workspace_client=create_autospec(WorkspaceClient), + config=config, + ) + rpm = report.get_db().query.channel(channel_name="Engine RPM") + page = Page(page_number=1) + report.add_page(page) + page.add_aggregation( + HistogramDuration("rpm_hist", base_expr=rpm, bins=[float(b) for b in range(0, 8000, 500)]) + ) + return report + + +def test_report_solves_through_selected_solver(spark): + report = _create_report( + spark, + {"solver": "ExcludeContainersSolver", "solver_config": {"excluded_container_ids": [3]}}, + ) + report.run(persist_results=False) + + # The report built the registered solver, and its typed config setting reached it... + solver = report.get_solver() + assert type(solver) is ExcludeContainersSolver + assert solver.config.excluded_container_ids == [3] + + # ...and every solve went through it: container 3 is missing from the histogram. + # The run splits results into {"changed", "unchanged"} definition buckets; exactly one + # holds this run's histogram. + [histogram] = [df for df in report.aggregation_dfs["HISTOGRAM"].values() if df is not None] + assert histogram.select("container_id").distinct().orderBy("container_id").collect() == [ + Row(container_id=1), + Row(container_id=2), + ] + + +def test_default_solver_is_unchanged_by_the_registration(spark): + report = _create_report(spark, {"solver": "DefaultSolver"}) + report.run(persist_results=False) + + assert type(report.get_solver()) is DefaultSolver + [histogram] = [df for df in report.aggregation_dfs["HISTOGRAM"].values() if df is not None] + assert histogram.select("container_id").distinct().orderBy("container_id").collect() == [ + Row(container_id=1), + Row(container_id=2), + Row(container_id=3), + ] diff --git a/tests/impulse_reporting/unit/config/config_parser_test.py b/tests/impulse_reporting/unit/config/config_parser_test.py index b935f0d0..7b16c04a 100644 --- a/tests/impulse_reporting/unit/config/config_parser_test.py +++ b/tests/impulse_reporting/unit/config/config_parser_test.py @@ -3,7 +3,7 @@ import pytest from pydantic import ValidationError -import impulse_query_engine.measurement_db_registry as registry +from impulse_query_engine.analyze.query.solvers import DefaultSolver, SolverConfig, register_solver from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig from impulse_query_engine.measurement_db_registry import register_measurement_db from impulse_reporting.config.config_parser import ( @@ -16,6 +16,7 @@ IncrementalConfig, ImpulseConfig, MetricFilter, + QueryEngine, RawEncoder, Solvers, TagFilter, @@ -621,7 +622,7 @@ def test_impulse_config_key_value_store_solver_no_project_id_accepted(): def test_impulse_config_basic_narrow_solver_rejected(): - """BasicNarrowSolver has been removed; the string is no longer a valid enum value.""" + """BasicNarrowSolver has been removed; the name is no longer a registered solver.""" config_json = impulse_config_JSON.copy() config_json["query_engine"] = {"solver": "BasicNarrowSolver"} with pytest.raises(ValidationError): @@ -635,6 +636,136 @@ def test_impulse_config_solver_config_none_by_default(): assert config.query_engine.solver_config is None +class _ExcludeContainersConfig(SolverConfig): + excluded_container_ids: list[int] # required custom setting + + +class _ExcludeContainersSolver(DefaultSolver): ... + + +@pytest.fixture +def custom_solver_registered(isolated_registries): + register_solver("ExcludeContainersSolver", _ExcludeContainersConfig)(_ExcludeContainersSolver) + + +def _with_query_engine(query_engine: dict) -> dict: + return {**impulse_config_JSON, "query_engine": query_engine} + + +def test_impulse_config_custom_solver_gets_its_typed_solver_config(custom_solver_registered): + config = ImpulseConfig.model_validate( + _with_query_engine( + { + "solver": "ExcludeContainersSolver", + "solver_config": {"excluded_container_ids": [3], "project_id": "p1"}, + } + ) + ) + + assert config.query_engine.solver == "ExcludeContainersSolver" + assert type(config.query_engine.solver_config) is _ExcludeContainersConfig + assert config.query_engine.solver_config.excluded_container_ids == [3] + assert config.query_engine.solver_config.project_id == "p1" # base settings still apply + + +def test_impulse_config_custom_solver_config_survives_dump_and_reload(custom_solver_registered): + config = ImpulseConfig.model_validate( + _with_query_engine( + {"solver": "ExcludeContainersSolver", "solver_config": {"excluded_container_ids": [3]}} + ) + ) + + # The custom settings are part of the dump, and so of the config_hash in measurement_dimension... + assert '"excluded_container_ids":[3]' in config.model_dump_json() + # ...and a dumped config loads again. + reloaded = ImpulseConfig.model_validate(config.model_dump()) + assert reloaded.query_engine.solver_config == config.query_engine.solver_config + + +_INVALID_EXCLUDE_SETTINGS = r"invalid settings for solver 'ExcludeContainersSolver': " + + +@pytest.mark.parametrize( + ("query_engine", "key", "error"), + [ + ( + {"solver": "NotRegistered"}, + "solver", + r"Unknown solver 'NotRegistered'; registered: \[.*'DefaultSolver'.*\]" + r".*Import the package that registers it", + ), + # A required custom setting is reported whether solver_config is empty or omitted. + ( + {"solver": "ExcludeContainersSolver", "solver_config": {}}, + "solver_config", + _INVALID_EXCLUDE_SETTINGS + "excluded_container_ids: Field required", + ), + ( + {"solver": "ExcludeContainersSolver"}, + "solver_config", + _INVALID_EXCLUDE_SETTINGS + "excluded_container_ids: Field required", + ), + ( + {"solver": "ExcludeContainersSolver", "solver_config": {"excluded_container_ids": 3}}, + "solver_config", + _INVALID_EXCLUDE_SETTINGS + "excluded_container_ids: Input should be a valid list", + ), + # The custom config class inherits SolverConfig's extra="forbid": a typo isn't dropped. + ( + { + "solver": "ExcludeContainersSolver", + "solver_config": {"excluded_container_ids": [3], "exclude_container_id": [4]}, + }, + "solver_config", + _INVALID_EXCLUDE_SETTINGS + "exclude_container_id: Extra inputs are not permitted", + ), + # An error on the whole value has no field name to report. + ( + {"solver": "ExcludeContainersSolver", "solver_config": 42}, + "solver_config", + _INVALID_EXCLUDE_SETTINGS + "Input should be a valid dictionary", + ), + ], +) +def test_impulse_config_solver_rejected_at_parse_time( + custom_solver_registered, query_engine, key, error +): + with pytest.raises(ValidationError, match=error) as exc_info: + ImpulseConfig.model_validate(_with_query_engine(query_engine)) + + [reported] = exc_info.value.errors() # one error, on the key that is wrong + assert reported["loc"] == ("query_engine", key) + + +@pytest.mark.parametrize( + ("config", "loc"), + [ + # Custom settings without selecting the custom solver are checked against SolverConfig. + ( + _with_query_engine({"solver_config": {"excluded_container_ids": [3]}}), + ("query_engine", "solver_config", "excluded_container_ids"), + ), + # "solver" outside query_engine. + ({**impulse_config_JSON, "solver": "ExcludeContainersSolver"}, ("solver",)), + ], +) +def test_impulse_config_rejects_misplaced_solver_settings(custom_solver_registered, config, loc): + # Silently dropping these would run DefaultSolver without the custom settings. + with pytest.raises(ValidationError, match="Extra inputs are not permitted") as exc_info: + ImpulseConfig.model_validate(config) + + [reported] = exc_info.value.errors() + assert reported["loc"] == loc + + +def test_query_engine_revalidates_base_solver_config_for_custom_solver(custom_solver_registered): + # In Python a user may pass a plain SolverConfig; the custom solver's settings still get checked. + with pytest.raises( + ValidationError, match=_INVALID_EXCLUDE_SETTINGS + "excluded_container_ids: Field required" + ): + QueryEngine(solver="ExcludeContainersSolver", solver_config=SolverConfig()) + + def test_is_valid_table_name_valid_inputs(): """ Test that valid table names pass validation. @@ -1344,8 +1475,7 @@ def test_calculated_channels_unknown_kpi_rejected(): @pytest.fixture -def session_db_registered(monkeypatch): - monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) +def session_db_registered(isolated_registries): register_measurement_db("SessionDB", _SessionDBConfig)(MeasurementDB) @@ -1381,6 +1511,14 @@ def test_impulse_config_measurement_db_config_rejected_at_parse_time( ) +def test_impulse_config_unregistered_measurement_db_reported_on_its_key(): + with pytest.raises(ValidationError) as exc_info: + ImpulseConfig.model_validate(_add_custom_measurement_db_to_report_config("NotRegistered")) + + [reported] = exc_info.value.errors() + assert reported["loc"] == ("query_engine", "measurement_db") + + @pytest.mark.parametrize( ("query_engine", "misspelled_key"), [ @@ -1445,8 +1583,7 @@ class _AcmeDB(MeasurementDB): ... @pytest.fixture -def acme_db_registered(monkeypatch): - monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) +def acme_db_registered(isolated_registries): register_measurement_db("AcmeDB", _AcmeDBConfig)(_AcmeDB) diff --git a/tests/impulse_reporting/unit/core/report_test.py b/tests/impulse_reporting/unit/core/report_test.py index 696125d4..1d3a1042 100644 --- a/tests/impulse_reporting/unit/core/report_test.py +++ b/tests/impulse_reporting/unit/core/report_test.py @@ -4,9 +4,10 @@ import pytest from databricks.sdk import WorkspaceClient -import impulse_query_engine.measurement_db_registry as registry from impulse_query_engine.analyze.metadata.tag_expression import TagSelector from impulse_query_engine.analyze.metadata.time_series_expression import TimeSeriesSelector +from impulse_query_engine.analyze.query.solvers import DefaultSolver, SolverConfig, register_solver +from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig from impulse_query_engine.measurement_db_registry import register_measurement_db from impulse_reporting.aggregations.aggregation_types import AggregationType @@ -173,14 +174,13 @@ def test_report_init(): assert report.aggregation_dfs == {} -def test_report_init_fails_loudly_for_custom_db_skipping_super_init(monkeypatch): +def test_report_init_fails_loudly_for_custom_db_skipping_super_init(isolated_registries): """A custom DB that skips ``super().__init__`` fails when the Report reads ``db.ws``.""" class SkipsSuperInitDB(MeasurementDB): def __init__(self, config, ws): self.config = config # forgets super().__init__(config, ws) - monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) register_measurement_db("SkipsSuperInitDB", MeasurementDBConfig)(SkipsSuperInitDB) with pytest.raises(TypeError, match=r"SkipsSuperInitDB\.__init__ must call super\(\)"): @@ -192,6 +192,70 @@ def __init__(self, config, ws): ) +@pytest.fixture +def logged_telemetry(monkeypatch): + """Records the (key, value) pairs the Report sends to telemetry.""" + logged = [] + monkeypatch.setattr( + "impulse_reporting.core.report.log_telemetry", + lambda _ws, key, value: logged.append((key, value)), + ) + return logged + + +def test_report_builds_registered_custom_solver(isolated_registries, logged_telemetry): + class ExcludeContainersConfig(SolverConfig): + excluded_container_ids: list[int] + + class ExcludeContainersSolver(DefaultSolver): ... + + register_solver("ExcludeContainersSolver", ExcludeContainersConfig)(ExcludeContainersSolver) + + report = Report( + name="test_report", + spark=None, + workspace_client=create_autospec(WorkspaceClient), + config={ + **DUMMY_CONFIG, + "query_engine": { + "solver": "ExcludeContainersSolver", + "solver_config": {"excluded_container_ids": [3]}, + }, + }, + ) + + solver = report.get_solver() + assert type(solver) is ExcludeContainersSolver + assert solver.config.excluded_container_ids == [3] + # Custom solver names aren't sent; built-in names stay as they were (next test). + assert ("solver", "custom") in logged_telemetry + + +def test_report_telemetry_keeps_builtin_solver_names(logged_telemetry): + report = Report( + name="test_report", + spark=None, + workspace_client=create_autospec(WorkspaceClient), + config=DUMMY_KEY_VALUE_STORE_CONFIG, + ) + + assert type(report.get_solver()) is DefaultSolver + assert ("solver", "KEY_VALUE_STORE_SOLVER") in logged_telemetry + + +def test_report_passes_a_concrete_raw_encoder_for_rle_data(): + report = Report( + name="test_report", + spark=None, + workspace_client=create_autospec(WorkspaceClient), + config=DUMMY_CONFIG, + ) + + # query_engine.raw_encoder stays unset for RLE input; the solver still gets the RLE default. + assert report.config.query_engine.raw_encoder is None + assert report.get_solver().raw_encoder is RawEncoder.RLE + + def test_set_config(): """Test setting config with config path"""