diff --git a/docs/impulse/docs/config/configuration.md b/docs/impulse/docs/config/configuration.md index d8199953..0f27ab13 100644 --- a/docs/impulse/docs/config/configuration.md +++ b/docs/impulse/docs/config/configuration.md @@ -142,11 +142,17 @@ Two independent filter families: | `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). | +| `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. + +For a source schema that needs reshaping beyond column renames, select your own read seam — see +[Custom measurement DB](../extensibility/custom_measurement_db.md). ## Solver column mappings and filters diff --git a/docs/impulse/docs/extensibility/_category_.json b/docs/impulse/docs/extensibility/_category_.json new file mode 100644 index 00000000..909e70f6 --- /dev/null +++ b/docs/impulse/docs/extensibility/_category_.json @@ -0,0 +1,11 @@ +{ + "label": "Extensibility", + "position": 5, + "link": { + "type": "generated-index", + "title": "Extensibility", + "description": "Extension points for adapting Impulse to data or behaviour that the built-in configuration cannot express.", + "slug": "/extensibility" + }, + "collapsed": true +} diff --git a/docs/impulse/docs/extensibility/custom_measurement_db.md b/docs/impulse/docs/extensibility/custom_measurement_db.md new file mode 100644 index 00000000..e8fcd9b3 --- /dev/null +++ b/docs/impulse/docs/extensibility/custom_measurement_db.md @@ -0,0 +1,140 @@ +--- +sidebar_position: 1 +title: Custom measurement DB +--- + +# Custom measurement DB + +The **MeasurementDB** is the layer between your stored tables and everything that reads them — the +solve, incremental container detection, and ad-hoc analysis all read through it. Impulse ships a +built-in MeasurementDB that reads the [silver model](../data_model/silver_layer_schema.md) directly. +When your raw tables need more than column renames to match that model — unions, joins, a grain +change, an EAV unpivot, or a row filter that must run *before* a reshape — you can register your own +MeasurementDB and select it by name in the report config. The stock solver then runs unchanged, and +every read path sees the same silver-shaped data. + +This is an additive, opt-in extension point: with nothing configured, the built-in MeasurementDB is +used and behaviour is unchanged. + +## When to use it + +| Your data differs from the silver model by… | Use | +|---|---| +| 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)) | + +## Select it in the report config + +Two optional keys under `query_engine` choose the implementation and pass it settings: + +```json +"query_engine": { + "measurement_db": "MyMeasurementDB", + "measurement_db_config": { "session_table": "catalog.raw.sessions" } +} +``` + +- `measurement_db` — the **registered name** of the implementation (default `"MeasurementDB"`, the built-in reader). +- `measurement_db_config` — keyword arguments for that implementation's config class, merged with the + `source` tables (e.g. extra source tables, prefilter keys). + +In Python the same keys go on `QueryEngine`: +`QueryEngine(measurement_db="MyMeasurementDB", measurement_db_config={...})`. + +## Implement one + +A custom MeasurementDB is a small subclass of the built-in `MeasurementDB`, plus a config subclass of +`MeasurementDBConfig`, registered under a name: + +```python +from impulse_query_engine.measurement_db import MeasurementDB, MeasurementDBConfig +from impulse_query_engine.measurement_db_registry import register_measurement_db + + +class MyDBConfig(MeasurementDBConfig): + def __init__(self, *, session_table: str, **base_kwargs): # required: no default + super().__init__(**base_kwargs) # keeps the source tables, table_locations, and pinning state + self.session_table = session_table + + def configured_table_uris(self) -> list[str]: + # Declare extra tables so they are pinned to one Delta snapshot per run (see below). + return super().configured_table_uris() + [self.session_table] + + +@register_measurement_db("MyMeasurementDB", MyDBConfig) +class MyMeasurementDB(MeasurementDB): + def container_metrics(self, spark): + raw = super().container_metrics(spark) # raw read, pinned + sessions = self._read_table(spark, self.config.session_table) # extra read, pinned + return ... # reshape into the silver container_metrics shape +``` + +Override one method per silver table you need to reshape; the rest fall back to the base read. Each +override follows the same shape: call `super().(spark)` for the raw, pinned read, then +join / union / reshape into the silver frame the solver expects. + +### Optional tables must also be set in `source` + +Overriding the read of an optional table only takes effect if you also set that table in the +report config's `source`. Otherwise the solver doesn't use the table, and your override is never +called: + +- `container_tags()` needs `source.container_tags_table` +- `channel_tags()` needs `source.channel_tags_table` +- `channel_mapping()` needs `source.channel_mapping_table` +- `unit_conversion()` needs `source.unit_conversion_table` +- `poi_channels()` needs `has_poi_channels()` to return `True`; the base implementation checks + `source.poi_channels_uri` + +Set each entry to the raw table your override reshapes; `super().
(spark)` reads it. For +channel tags, unit conversion and POI, a missing entry doesn't raise an error: the feature just +isn't used. + +## Rules + +- **Import the registering package before the report config is parsed.** Selection is by name only + (config data never triggers an import), so the package whose import runs `@register_measurement_db` + must be imported first — including in any tool that only *generates* or *edits* a config. +- **Subclass `MeasurementDBConfig` and call `super().__init__(**base_kwargs)`.** The base initializer + sets the source tables, `table_locations`, and the per-run pinning state the reads rely on. +- **Keep the config `__init__` free of side effects** — it is constructed once to validate the config + at parse time and once more to build the DB. +- **Keep the DB constructor signature `__init__(self, config, ws)`** — Impulse builds it as `db_cls(config, ws)`. +- **If you override the DB `__init__`, call `super().__init__(config, ws)`.** The base initializer + stores the config and verifies the workspace client; reading `db.ws` (which `Report` and `db.query` + do) raises a `TypeError` if it was skipped. Don't assign `self.ws` yourself. +- **Your DB owns correctness.** Impulse does not validate the frames a custom DB returns; it must + provide the columns, types, row grain, and time unit the solver expects. The built-in `MeasurementDB` + and the schemas in `impulse_query_engine/schema.py` are the reference. + +## Parse-time validation + +The config is checked when it is loaded (`ImpulseConfig.model_validate`), so these fail fast as a +validation error rather than at run time: + +- an unregistered `measurement_db` name (the message lists the registered names); +- a missing required argument of the config class; +- an unknown or misspelled key; +- a key present in both `source` and `measurement_db_config`. + +## Snapshot pinning (extra tables) + +At the start of each run Impulse pins every configured table to one Delta version so all lazy reads +observe a consistent snapshot. The built-in pinning iterates `configured_table_uris()`, so any extra +table your DB reads must be: + +1. **declared** — override `configured_table_uris()` to add its URI; and +2. **read through `self._read_table(spark, uri)`** (never `spark.read.table`), using the **same URI + string** you declared. + +Views and non-Delta tables cannot be time-travelled: they are skipped with a warning and read at the +latest version. + +## Incremental reporting + +Because every reader goes through the MeasurementDB, incremental container detection reads the same +reshaped `container_metrics` as the solve, so the silver↔gold join stays aligned by construction. If +your reshape renames the physical container-id column, the rename is applied in exactly one place and +both the solve and detection see it. See [incremental](../config/configuration.md#incremental-optional) +for how the mode is resolved. 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 dcaca28e..94bd00e2 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 @@ -208,6 +208,9 @@ scoping. Key sub-fields: When omitted, all default column names are used and no project/toolbox filtering is applied. +- `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. #### validate\_drop\_implausible\_data\_requires\_raw @@ -221,6 +224,32 @@ The filter is applied inside the RAW -> interval conversion path by the selected ``raw_encoder`` (``RleEncoder`` / ``IntervalEncoder``). +#### reject\_channels\_filter\_on\_is\_plausible\_in\_raw + +```python +def reject_channels_filter_on_is_plausible_in_raw() +``` + +Reject an is_plausible channels filter in RAW mode (delegates to SolverConfig). + +The shared invariant lives on ``SolverConfig`` so direct query-engine use +(``DefaultSolver``) enforces the same rule; here we surface it at config-parse time. + + +#### warn\_channels\_filters\_bridge\_in\_raw + +```python +def warn_channels_filters_bridge_in_raw() +``` + +Warn (not reject) on any non-plausibility channels filter in RAW mode. + +Such filters run before raw encoding and bridge intervals across dropped +samples. That is fine for whole-channel scoping but corrupts per-sample cleaning, +and the two cannot be told apart statically, so we warn. ``is_plausible`` is +hard-rejected above as the one unambiguously per-sample case. + + #### default\_raw\_encoder\_for\_raw\_data ```python @@ -294,6 +323,22 @@ Names that do not match any registered entity are rejected at the start of - `events` (`list of str, default=[]`): Names of events to fully recalculate. - `calculated_channels` (`list of str, default=[]`): Names of calculated channels to fully recalculate. +#### build\_measurement\_db\_config + +```python +def build_measurement_db_config( + source_dict: dict[str, Any], query_engine: QueryEngine +) -> tuple[type[AbstractMeasurementDB], MeasurementDBConfig] +``` + +Resolve the selected MeasurementDB implementation and build its config class from + +``measurement_db_config`` merged with the ``source`` tables in *source_dict*. + +Raises ``KeyError`` for an unregistered ``measurement_db`` and ``TypeError`` for a missing, +unknown or duplicated config key. + + ## ImpulseConfig ```python 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 3d4a8fb2..d9e7b8e1 100644 --- a/docs/impulse/docs/references/api/impulse_reporting/core/report.md +++ b/docs/impulse/docs/references/api/impulse_reporting/core/report.md @@ -52,14 +52,14 @@ Returns a unique identifier for the report. #### get\_db ```python -def get_db() -> MeasurementDB +def get_db() -> AbstractMeasurementDB ``` Get the measurement database associated with this report. **Returns**: -`MeasurementDB`: The measurement database instance. +`AbstractMeasurementDB`: The measurement database instance. #### get\_solver @@ -109,11 +109,13 @@ Load Impulse configuration from a dictionary. ```python def create_measurement_db(config: ImpulseConfig, - ws: WorkspaceClient) -> MeasurementDB + ws: WorkspaceClient) -> AbstractMeasurementDB ``` -Create a measurement database based on the provided configuration. +Create the measurement database selected by ``query_engine.measurement_db``. +Builds the registered config class from the ``source`` tables plus +``query_engine.measurement_db_config``, then the registered DB class. Maps the optional ``container_tags`` field from the Source config to the ``container_tags_table`` parameter expected by ``MeasurementDBConfig``. @@ -125,12 +127,12 @@ to the ``container_tags_table`` parameter expected by **Returns**: -`MeasurementDB`: The measurement database instance. +`AbstractMeasurementDB`: The measurement database instance. #### create\_query\_builder ```python -def create_query_builder(db: MeasurementDB, +def create_query_builder(db: AbstractMeasurementDB, config: ImpulseConfig) -> QueryBuilder ``` @@ -143,7 +145,7 @@ supports tag and metric filters, but tag filters require the narrow **Arguments**: -- `db` (`MeasurementDB`): The measurement database instance. +- `db` (`AbstractMeasurementDB`): The measurement database instance. - `config` (`ImpulseConfig`): The Impulse configuration. **Raises**: diff --git a/skills/impulse-config/SKILL.md b/skills/impulse-config/SKILL.md index 81568e03..0bad9a30 100644 --- a/skills/impulse-config/SKILL.md +++ b/skills/impulse-config/SKILL.md @@ -6,7 +6,8 @@ description: > "configure an Impulse report", set the source/sink tables, filter which containers are processed, choose RLE vs RAW, turn on incremental processing, run without writing (sinkless), remap column names, or scope by project. Covers source, unity_sink, container_filters, query_engine, solver_config, - incremental, full_recalculation, measurement_dimensions, and calculated_channels, all validated by Pydantic. + custom measurement DBs, incremental, full_recalculation, measurement_dimensions, and calculated_channels, + all validated by Pydantic. --- # Impulse — configuration @@ -103,6 +104,8 @@ Omit entirely for `DefaultSolver` + `data_type="RLE"`. | `drop_implausible_data` | `false` | Drop `channels` rows where `is_plausible = false`. **Requires `data_type="RAW"`** (RLE raises). | | `max_channels_per_batch` | `500` | Max distinct channels (`channel(...)` / `poi_channel(...)` selections) read per solve batch. Applies to events, aggregations, and calculated channels (counts their *input* channels). `batch_size` is a **deprecated alias** (warns; removed in a future release). | | `solver_config` | `null` | Per-table column mappings, per-table filters, and project scoping (below). | +| `measurement_db` | `"MeasurementDB"` | Registered name of the MeasurementDB (the read seam). The default reads the silver tables directly; a custom one reshapes a different source schema at read time. See [Custom measurement DB](#custom-measurement-db). | +| `measurement_db_config` | `null` | Keyword args for the selected MeasurementDB's config class, merged with `source` (e.g. extra tables, prefilter keys). | ## solver_config (optional) — adapt to an existing layout @@ -176,6 +179,27 @@ to `channel_mapping` so aliased selectors auto-convert during `solve()`. Constan an aliased selector must then be in the target unit. Direct `channel(...)` selectors are never converted — conversion is a property of the alias. +## Custom measurement DB + +`solver_config` only renames and filters columns. When the raw tables need *reshaping* to reach the +silver model (unions, joins, a grain change, an EAV unpivot), register a custom MeasurementDB and select +it by name — the stock solver then runs unchanged: + +```python +"query_engine": { + "measurement_db": "MyMeasurementDB", + "measurement_db_config": {"session_table": "catalog.raw.sessions"}, +} +``` + +- `measurement_db` is a **registered name**; the package that registers it must be imported before the + config is parsed. Validation runs at load time — an unregistered name, a missing required argument, an + unknown key, or a key already present in `source` all raise. +- `measurement_db_config` holds that implementation's own constructor kwargs, merged with `source`. +- Implementing one (subclass `MeasurementDB` + `MeasurementDBConfig`, decorate with + `@register_measurement_db`, and declare any extra tables in `configured_table_uris()` so they are + pinned per run) is covered in the Impulse docs' *Custom measurement DB* section. + ## incremental (optional) Reuses prior results for unchanged definitions and reprocesses only new/updated containers. See diff --git a/src/impulse_query_engine/measurement_db.py b/src/impulse_query_engine/measurement_db.py index 2b1145fc..95a38afc 100644 --- a/src/impulse_query_engine/measurement_db.py +++ b/src/impulse_query_engine/measurement_db.py @@ -1,11 +1,9 @@ import warnings -from databricks.sdk import WorkspaceClient from pyspark.sql import DataFrame -from impulse_query_engine import __version__ -from impulse_query_engine.telemetry import verify_workspace_client -from .analyze.query.query_builder import QueryBuilder +from .measurement_db_abstract import AbstractMeasurementDB +from .measurement_db_registry import register_measurement_db class MeasurementDBConfig: @@ -100,15 +98,8 @@ def configured_table_uris(self) -> list[str]: ] -class MeasurementDB: - def __init__(self, config: MeasurementDBConfig, ws: WorkspaceClient): - self.config = config - self.ws = verify_workspace_client(ws, "databricks-impulse", __version__) - - @property - def query(self): - return QueryBuilder(db=self) - +@register_measurement_db("MeasurementDB", MeasurementDBConfig) +class MeasurementDB(AbstractMeasurementDB): def _current_delta_version(self, spark, uri: str) -> int | None: """Latest committed Delta version of ``uri``, or ``None`` if unpinnable. diff --git a/src/impulse_query_engine/measurement_db_abstract.py b/src/impulse_query_engine/measurement_db_abstract.py new file mode 100644 index 00000000..f2852632 --- /dev/null +++ b/src/impulse_query_engine/measurement_db_abstract.py @@ -0,0 +1,77 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from databricks.sdk import WorkspaceClient +from pyspark.sql import DataFrame + +from impulse_query_engine import __version__ +from impulse_query_engine.telemetry import verify_workspace_client + +from .analyze.query.query_builder import QueryBuilder + +if TYPE_CHECKING: + from .measurement_db import MeasurementDBConfig + + +class AbstractMeasurementDB(ABC): + """Contract for a MeasurementDB implementation: returns silver-shaped frames for the solver. + + Implementations are selected by name via ``register_measurement_db`` and built as + ``db_cls(config, ws)``. The returned frames must match the silver model the solver expects. + """ + + def __init__(self, config: "MeasurementDBConfig", ws: WorkspaceClient): + self.config = config + self._ws = verify_workspace_client(ws, "databricks-impulse", __version__) + + @property + def ws(self) -> WorkspaceClient: + """The verified workspace client; ``Report`` and ``QueryBuilder`` read it from here. + + Raises ``TypeError`` if a subclass ``__init__`` didn't call ``super().__init__(config, ws)``, + which would otherwise lose the client verification without any error. + """ + try: + return self._ws + except AttributeError: + raise TypeError( + f"{type(self).__name__}.__init__ must call super().__init__(config, ws)" + ) from None + + @property + def query(self) -> QueryBuilder: + return QueryBuilder(db=self) + + @abstractmethod + def pin_versions(self, spark) -> None: + """Pin every read table to one Delta snapshot for the run (called by ``Report``).""" + + @abstractmethod + def container_tags(self, spark) -> DataFrame: ... + + @abstractmethod + def container_metrics(self, spark) -> DataFrame: ... + + @abstractmethod + def channel_tags(self, spark) -> DataFrame: ... + + @abstractmethod + def channel_metrics(self, spark) -> DataFrame: ... + + @abstractmethod + def channels(self, spark) -> DataFrame: ... + + @abstractmethod + def has_poi_channels(self) -> bool: ... + + @abstractmethod + def poi_channels(self, spark) -> DataFrame: ... + + @abstractmethod + def channel_mapping(self, spark) -> DataFrame: ... + + @abstractmethod + def unit_conversion(self, spark) -> DataFrame: ... + + @abstractmethod + def channel_uri(self) -> str | None: ... diff --git a/src/impulse_query_engine/measurement_db_registry.py b/src/impulse_query_engine/measurement_db_registry.py new file mode 100644 index 00000000..9b9a0e40 --- /dev/null +++ b/src/impulse_query_engine/measurement_db_registry.py @@ -0,0 +1,42 @@ +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from .measurement_db_abstract import AbstractMeasurementDB + +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] = {} + + +def register_measurement_db(name: str, config_cls: type["MeasurementDBConfig"]): + """Register an ``AbstractMeasurementDB`` subclass under *name*, built with *config_cls*. + + 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 + + +def resolve_measurement_db(name: str) -> MeasurementDatabaseRegistration: + """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 diff --git a/src/impulse_reporting/config/config_parser.py b/src/impulse_reporting/config/config_parser.py index 9bbde633..025c9f88 100644 --- a/src/impulse_reporting/config/config_parser.py +++ b/src/impulse_reporting/config/config_parser.py @@ -3,12 +3,14 @@ import warnings from datetime import datetime from enum import Enum, StrEnum -from typing import Annotated +from typing import Annotated, Any, Self import pydantic -from pydantic import AfterValidator, BaseModel, field_validator, model_validator +from pydantic import AfterValidator, BaseModel, field_validator, model_validator, ConfigDict from impulse_query_engine.analyze.query.solvers.solver_config import RawEncoder, SolverConfig +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 @@ -388,6 +390,11 @@ class QueryEngine(BaseModel): When omitted, all default column names are used and no project/toolbox filtering is applied. + measurement_db : str, default="MeasurementDB" + Registered name of the MeasurementDB implementation to use. + measurement_db_config : dict, optional + Keyword arguments for that implementation's config class, + merged with the ``source`` tables. Notes ----- @@ -403,11 +410,17 @@ class QueryEngine(BaseModel): - RAW channel data must contain 'container_id', 'channel_id', 'timestamp', 'value' columns """ + # added to prevent misspelled args being silently accepted, which could have side effects + # (default measurement db silently used) + model_config = ConfigDict(extra="forbid") + solver: Solvers = Solvers.DEFAULT_SOLVER data_type: DataType = DataType.RLE drop_implausible_data: bool = False raw_encoder: RawEncoder | None = None solver_config: SolverConfig | None = None + measurement_db: str = "MeasurementDB" + measurement_db_config: dict[str, Any] | None = None max_channels_per_batch: int = 500 max_containers_per_batch: int | None = None @@ -651,6 +664,24 @@ class FullRecalculation(BaseModel): calculated_channels: list[str] = [] +def build_measurement_db_config( + source_dict: dict[str, Any], query_engine: QueryEngine +) -> tuple[type[AbstractMeasurementDB], MeasurementDBConfig]: + """Resolve the selected MeasurementDB implementation and build its config class from + ``measurement_db_config`` merged with the ``source`` tables in *source_dict*. + + Raises ``KeyError`` for an unregistered ``measurement_db`` and ``TypeError`` for a missing, + unknown or duplicated config key. + """ + registration = resolve_measurement_db(query_engine.measurement_db) + db_config = registration.config_cls( + **source_dict, + **(query_engine.measurement_db_config or {}), + table_locations="unity_catalog", + ) + return registration.db_cls, db_config + + class ImpulseConfig(BaseModel): """ Main configuration model. @@ -764,3 +795,14 @@ def _normalize_measurement_dimensions(cls, value: list[str]) -> list[str]: seen.add(name) normalized.append(name) return normalized + + @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. + """ + try: + build_measurement_db_config(dict(self.source), self.query_engine) + except (KeyError, 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 1a6b6f9f..17c9e8e2 100644 --- a/src/impulse_reporting/core/report.py +++ b/src/impulse_reporting/core/report.py @@ -13,13 +13,14 @@ 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.measurement_db import MeasurementDB, MeasurementDBConfig +from impulse_query_engine.measurement_db import AbstractMeasurementDB 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, + build_measurement_db_config, ) from impulse_reporting.core.page import Page from impulse_reporting.core.report_utils import ( @@ -157,13 +158,13 @@ def get_id(self) -> int: """ return zlib.crc32(self.name.encode()) & 0x7FFFFFFF # Ensures positive 32-bit int - def get_db(self) -> MeasurementDB: + def get_db(self) -> AbstractMeasurementDB: """ Get the measurement database associated with this report. Returns ------- - MeasurementDB + AbstractMeasurementDB The measurement database instance. """ return self.db @@ -215,10 +216,12 @@ def load_config_from_dict(config_info: dict[str, Any]) -> ImpulseConfig: return ImpulseConfig.model_validate(config_info) @staticmethod - def create_measurement_db(config: ImpulseConfig, ws: WorkspaceClient) -> MeasurementDB: + def create_measurement_db(config: ImpulseConfig, ws: WorkspaceClient) -> AbstractMeasurementDB: """ - Create a measurement database based on the provided configuration. + Create the measurement database selected by ``query_engine.measurement_db``. + Builds the registered config class from the ``source`` tables plus + ``query_engine.measurement_db_config``, then the registered DB class. Maps the optional ``container_tags`` field from the Source config to the ``container_tags_table`` parameter expected by ``MeasurementDBConfig``. @@ -232,18 +235,18 @@ def create_measurement_db(config: ImpulseConfig, ws: WorkspaceClient) -> Measure Returns ------- - MeasurementDB + AbstractMeasurementDB The measurement database instance. """ source_dict = dict(config.source) # Map config field name to MeasurementDBConfig parameter name if "container_tags" in source_dict: source_dict["container_tags_table"] = source_dict.pop("container_tags") - measurement_db_config = MeasurementDBConfig(**source_dict, table_locations="unity_catalog") - return MeasurementDB(config=measurement_db_config, ws=ws) + db_cls, db_config = build_measurement_db_config(source_dict, config.query_engine) + return db_cls(db_config, ws) @staticmethod - def create_query_builder(db: MeasurementDB, config: ImpulseConfig) -> QueryBuilder: + def create_query_builder(db: AbstractMeasurementDB, config: ImpulseConfig) -> QueryBuilder: """ Create a query builder based on the provided configuration and set container filters. @@ -254,7 +257,7 @@ def create_query_builder(db: MeasurementDB, config: ImpulseConfig) -> QueryBuild Parameters ---------- - db : MeasurementDB + db : AbstractMeasurementDB The measurement database instance. config : ImpulseConfig The Impulse configuration. diff --git a/tests/impulse_query_engine/unit/measurement_db_registry_test.py b/tests/impulse_query_engine/unit/measurement_db_registry_test.py new file mode 100644 index 00000000..33d736b3 --- /dev/null +++ b/tests/impulse_query_engine/unit/measurement_db_registry_test.py @@ -0,0 +1,98 @@ +# pylint: disable=missing-function-docstring +"""Tests for the ``AbstractMeasurementDB`` contract and the name registry.""" + +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, + MeasurementDB, + MeasurementDBConfig, +) +from impulse_query_engine.measurement_db_registry import ( + register_measurement_db, + resolve_measurement_db, +) + + +@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 + # than as a confusing instantiation error scattered across the rest of the suite. + assert issubclass(MeasurementDB, AbstractMeasurementDB) + assert not MeasurementDB.__abstractmethods__ + + +def test_incomplete_implementation_cannot_be_instantiated(): + class Incomplete(AbstractMeasurementDB): + pass + + with pytest.raises(TypeError, match="abstract"): + Incomplete(MeasurementDBConfig(), ws=None) + + +def test_ws_is_the_verified_client(): + ws = create_autospec(WorkspaceClient) + db = MeasurementDB(MeasurementDBConfig(), ws) + + assert db.ws is ws + assert ws.config._product_info == ("databricks-impulse", __version__) # tagged as Impulse + ws.clusters.select_spark_version.assert_called_once() # the connectivity check ran + + +def test_subclass_skipping_super_init_fails_loudly_on_ws(): + class SkipsSuperInit(MeasurementDB): + def __init__(self, config, ws): + self.config = config # forgets super().__init__(config, ws) + + db = SkipsSuperInit(MeasurementDBConfig(), create_autospec(WorkspaceClient)) + + expected = r"SkipsSuperInit\.__init__ must call super\(\)\.__init__\(config, ws\)" + with pytest.raises(TypeError, match=expected): + _ = db.ws + with pytest.raises(TypeError, match=expected): # standalone path: QueryBuilder reads db.ws + _ = db.query + + +def test_subclass_cannot_assign_ws_itself(): + class AssignsWs(MeasurementDB): + def __init__(self, config, ws): + self.config = config + self.ws = ws # would bypass verify_workspace_client + + with pytest.raises(AttributeError, match="has no setter"): + AssignsWs(MeasurementDBConfig(), create_autospec(WorkspaceClient)) + + +def test_register_rejects_non_subclass(): + with pytest.raises(TypeError, match="must subclass AbstractMeasurementDB"): + register_measurement_db("Bad", MeasurementDBConfig)(object) + + +def test_register_returns_class_and_last_registration_wins(): + class First(MeasurementDB): ... + + class Second(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 + + +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", + ): + resolve_measurement_db("Missing") diff --git a/tests/impulse_query_engine/unit/measurement_db_versioning_test.py b/tests/impulse_query_engine/unit/measurement_db_versioning_test.py index da45275b..ed904ce6 100644 --- a/tests/impulse_query_engine/unit/measurement_db_versioning_test.py +++ b/tests/impulse_query_engine/unit/measurement_db_versioning_test.py @@ -54,6 +54,28 @@ def test_pin_versions_freezes_snapshot(spark, pin_schema): # noqa: F811 assert db.container_metrics(spark).count() == 10 +def test_pin_versions_pins_extra_table_declared_by_config_subclass( + spark, pin_schema +): # noqa: F811 + # A custom read seam declares its extra tables via ``configured_table_uris``; the inherited + # ``pin_versions`` then pins them together with the silver tables. + extra = f"{pin_schema}.sessions" + spark.range(3).toDF("session_id").write.format("delta").mode("overwrite").saveAsTable(extra) + + class ConfigWithExtra(MeasurementDBConfig): + def configured_table_uris(self) -> list[str]: + return super().configured_table_uris() + [extra] + + cfg = ConfigWithExtra(table_locations="unity_catalog") + db = _db(cfg) + db.pin_versions(spark) + assert cfg.pinned_versions == {extra: 0} + + spark.range(3, 5).toDF("session_id").write.format("delta").mode("append").saveAsTable(extra) + + assert db._read_table(spark, extra).count() == 3 + + def test_pin_versions_freezes_snapshot_path_mode(spark, tmp_path): # noqa: F811 # Path mode (``external_locations``): reads and version resolution go through # the filesystem path rather than a catalog name. diff --git a/tests/impulse_reporting/integration/custom_measurement_db_report_test.py b/tests/impulse_reporting/integration/custom_measurement_db_report_test.py new file mode 100644 index 00000000..1fcdcb17 --- /dev/null +++ b/tests/impulse_reporting/integration/custom_measurement_db_report_test.py @@ -0,0 +1,152 @@ +# pylint: disable=missing-function-docstring +"""End-to-end: a customer-provided MeasurementDB implementation, selected from the report config. + +The custom DB reports only the containers listed in an extra allow-list table. Its config class +requires that table and declares it for Delta version pinning. Each test covers one promise of the +extension point: the selected DB serves all reads, incremental detection reads through it, and the +declared extra table is pinned. +""" + +from unittest.mock import create_autospec + +import pytest +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 +from impulse_reporting.core.page import Page +from impulse_reporting.core.report import Report + +_ALLOWED = "spark_catalog.silver.e2e_allowed_containers" +_GOLD_MEASUREMENTS = "spark_catalog.gold.e2e_measurement_dimension" + + +class AllowListConfig(MeasurementDBConfig): + def __init__(self, *, allowed_containers_table: str, **base_kwargs): + super().__init__(**base_kwargs) + self.allowed_containers_table = allowed_containers_table + + def configured_table_uris(self) -> list[str]: + return super().configured_table_uris() + [self.allowed_containers_table] + + +class AllowListDB(MeasurementDB): + """Reports only the containers listed in ``allowed_containers_table``.""" + + def container_metrics(self, spark): + allowed = self._read_table(spark, self.config.allowed_containers_table) + return super().container_metrics(spark).join(allowed, "container_id", "left_semi") + + +@pytest.fixture(autouse=True) +def allow_list_db(monkeypatch): + monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) + register_measurement_db("AllowListDB", AllowListConfig)(AllowListDB) + + +@pytest.fixture +def set_allowed_container_ids(spark): + """Overwrites the allow-list table with the given container ids.""" + + def write(*container_ids: int): + spark.createDataFrame([(i,) for i in container_ids], "container_id int").write.format( + "delta" + ).mode("overwrite").saveAsTable(_ALLOWED) + + yield write + spark.sql(f"DROP TABLE IF EXISTS {_ALLOWED}") + + +def _create_report(spark, *, incremental: bool = False) -> 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"}, + "query_engine": { + "measurement_db": "AllowListDB", + "measurement_db_config": {"allowed_containers_table": _ALLOWED}, + }, + "incremental": {"enabled": incremental}, + } + report = Report( + name="e2e_custom_db", + 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 _gold_created_at_by_container_id(spark) -> dict[int, object]: + rows = spark.read.table(_GOLD_MEASUREMENTS).select("container_id", "_created_at").collect() + return {row.container_id: row._created_at for row in rows} + + +def test_report_reads_through_selected_db(spark, set_allowed_container_ids): + set_allowed_container_ids(1, 2) + report = _create_report(spark) + report.determine_report() + + assert type(report.get_db()) is AllowListDB + # determine_report splits results into {"changed", "unchanged"} definition buckets; exactly + # one holds this run's histogram. Pull it out and check which containers it was computed for. + [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_report_restricts_to_single_allowed_container(spark, set_allowed_container_ids): + set_allowed_container_ids(1) + report = _create_report(spark) + + # The measurement_db_config sticks: the allow-list table reaches the built DB's config... + assert report.get_db().config.allowed_containers_table == _ALLOWED + + # ...and with only container 1 allowed, that is the only container the histogram covers. + report.determine_report() + [histogram] = [df for df in report.aggregation_dfs["HISTOGRAM"].values() if df is not None] + assert histogram.select("container_id").distinct().collect() == [Row(container_id=1)] + + +def test_incremental_detection_reads_through_custom_db(spark, set_allowed_container_ids): + set_allowed_container_ids(1) + first = _create_report(spark) + first.determine_report() + first.persist_results() + before = _gold_created_at_by_container_id(spark) + + set_allowed_container_ids(1, 2) + second = _create_report(spark, incremental=True) + second.determine_report() + second.persist_results() + after = _gold_created_at_by_container_id(spark) + + assert set(before.keys()) == {1} + # The second run computes only the newly allowed container 2 (unchanged definition bucket). + [histogram] = [df for df in second.aggregation_dfs["HISTOGRAM"].values() if df is not None] + assert histogram.select("container_id").distinct().collect() == [Row(container_id=2)] + assert after[1] == before[1] + + +def test_declared_extra_table_is_pinned(spark, set_allowed_container_ids): + set_allowed_container_ids(1) + db = _create_report(spark).get_db() + db.pin_versions(spark) + + set_allowed_container_ids(1, 2, 3) + + assert {row.container_id for row in db.container_metrics(spark).collect()} == {1} diff --git a/tests/impulse_reporting/unit/config/config_parser_test.py b/tests/impulse_reporting/unit/config/config_parser_test.py index 7ed5b12d..b935f0d0 100644 --- a/tests/impulse_reporting/unit/config/config_parser_test.py +++ b/tests/impulse_reporting/unit/config/config_parser_test.py @@ -3,6 +3,9 @@ import pytest from pydantic import ValidationError +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.config.config_parser import ( CalculatedChannels, CastType, @@ -16,6 +19,7 @@ RawEncoder, Solvers, TagFilter, + build_measurement_db_config, is_valid_table_name, is_valid_unity_entity_name, ) @@ -128,6 +132,18 @@ def test_impulse_config_data_format_raw(): assert config.query_engine.data_type == DataType.RAW +def test_impulse_config_measurement_db_defaults(): + config = ImpulseConfig.model_validate(impulse_config_JSON) + assert config.query_engine.measurement_db == "MeasurementDB" + assert config.query_engine.measurement_db_config is None + + +class _SessionDBConfig(MeasurementDBConfig): + def __init__(self, *, session_table: str, **base_kwargs): + super().__init__(**base_kwargs) + self.session_table = session_table + + def test_impulse_config_drop_in_plausible_data_defaults_to_false(): config_json = impulse_config_JSON.copy() config = ImpulseConfig.model_validate(config_json) @@ -1325,3 +1341,171 @@ def test_calculated_channels_unknown_kpi_rejected(): """An unknown KPI name is rejected at validation with a helpful message.""" with pytest.raises(ValidationError, match="Unknown calculated-channel KPI"): CalculatedChannels(kpis=["bogus"]) + + +@pytest.fixture +def session_db_registered(monkeypatch): + monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) + register_measurement_db("SessionDB", _SessionDBConfig)(MeasurementDB) + + +def _add_custom_measurement_db_to_report_config(name: str, extra: dict | None = None) -> dict: + query_engine = {"measurement_db": name, "measurement_db_config": extra} + return {**impulse_config_JSON, "query_engine": query_engine} + + +def test_impulse_config_measurement_db_selector_parsed(session_db_registered): + config_json = _add_custom_measurement_db_to_report_config( + "SessionDB", {"session_table": "cat.raw.sessions"} + ) + config = ImpulseConfig.model_validate(config_json) + assert config.query_engine.measurement_db == "SessionDB" + assert config.query_engine.measurement_db_config == {"session_table": "cat.raw.sessions"} + + +@pytest.mark.parametrize( + ("custom_measurement_db_name", "extra", "error"), + [ + ("NotRegistered", None, "Unknown measurement DB 'NotRegistered'"), + ("SessionDB", None, "missing 1 required keyword-only argument: 'session_table'"), + ("SessionDB", {"session_table": "t", "typo": 1}, "unexpected keyword argument 'typo'"), + ("SessionDB", {"session_table": "t", "channels_uri": "x"}, "multiple values"), + ], +) +def test_impulse_config_measurement_db_config_rejected_at_parse_time( + session_db_registered, custom_measurement_db_name, extra, error +): + with pytest.raises(ValidationError, match=error): + ImpulseConfig.model_validate( + _add_custom_measurement_db_to_report_config(custom_measurement_db_name, extra) + ) + + +@pytest.mark.parametrize( + ("query_engine", "misspelled_key"), + [ + # Selector typo: would otherwise silently fall back to the built-in MeasurementDB. + ( + {"measurment_db": "SessionDB", "measurement_db_config": {"session_table": "t"}}, + "measurment_db", + ), + ({"measurment_db": "SessionDB"}, "measurment_db"), + # Settings-key typo: would otherwise silently drop the custom DB's settings. + ( + {"measurement_db": "SessionDB", "measurment_db_config": {"session_table": "t"}}, + "measurment_db_config", + ), + # Any other query_engine key: would otherwise silently fall back to its default. + ({"data_typ": "RAW"}, "data_typ"), + ], +) +def test_impulse_config_misspelled_query_engine_key_rejected( + session_db_registered, query_engine, misspelled_key +): + """``QueryEngine`` forbids unknown keys, so a typo fails at parse time and names the key.""" + with pytest.raises(ValidationError) as exc_info: + ImpulseConfig.model_validate({**impulse_config_JSON, "query_engine": query_engine}) + + assert [(err["loc"], err["type"]) for err in exc_info.value.errors()] == [ + (("query_engine", misspelled_key), "extra_forbidden") + ] + + +# --------------------------------------------------------------------------- +# build_measurement_db_config — the config_parser factory that resolves the +# selected implementation and builds its config object, merging the ``source`` +# tables with the custom ``measurement_db_config`` kwargs. A custom DB subclass +# relies on every one of those params reaching its config object. Parsing +# itself calls this factory (see ``ImpulseConfig._validate_measurement_db``), so +# these go through ``model_validate`` first and then inspect the built config. +# --------------------------------------------------------------------------- +class _AcmeDBConfig(MeasurementDBConfig): + """A representative custom config: one required extra table, two optional prefilter keys, + and an extra table declared for pinning. Richer than ``_SessionDBConfig`` so a test can + assert that *every* param a custom config uses is present on the built object.""" + + def __init__( + self, + *, + recording_session_table: str, # required: a missing value is reported at parse time + signal_name=(), # optional prefilter key + start_dt=None, # optional prefilter key + **base_kwargs, # source tables + table_locations, forwarded to the base + ): + super().__init__(**base_kwargs) + self.recording_session_table = recording_session_table + self.signal_name = list(signal_name) + self.start_dt = start_dt + + def configured_table_uris(self) -> list[str]: + return super().configured_table_uris() + [self.recording_session_table] + + +class _AcmeDB(MeasurementDB): ... + + +@pytest.fixture +def acme_db_registered(monkeypatch): + monkeypatch.setattr(registry, "_REGISTRY", dict(registry._REGISTRY)) + register_measurement_db("AcmeDB", _AcmeDBConfig)(_AcmeDB) + + +def test_build_measurement_db_config_defaults_to_builtin(): + """No selector -> the factory resolves the built-in and builds a base config from source.""" + config = ImpulseConfig.model_validate(impulse_config_JSON) + db_cls, db_config = build_measurement_db_config(dict(config.source), config.query_engine) + + assert db_cls is MeasurementDB + assert type(db_config) is MeasurementDBConfig + assert db_config.container_metrics_table == "impulse_demo.silver.container_metric" + assert db_config.channel_metrics_table == "impulse_demo.silver.channel_metric" + assert db_config.channels_uri == "impulse_demo.silver.channel_data" + assert db_config.table_locations == "unity_catalog" + + +def test_build_measurement_db_config_builds_custom_config_with_all_params(acme_db_registered): + """The factory resolves the registered class and builds its config with EVERY param present: + the custom ``measurement_db_config`` kwargs, the merged ``source`` tables, and + ``table_locations`` -- and declares the custom extra table for pinning.""" + acme_config = { + "recording_session_table": "acme.raw.recording_session", + "signal_name": ["ENGINE_SPEED", "VEHICLE_SPEED"], + "start_dt": "2026-04-01", + } + config = ImpulseConfig.model_validate( + _add_custom_measurement_db_to_report_config("AcmeDB", acme_config) + ) + db_cls, db_config = build_measurement_db_config(dict(config.source), config.query_engine) + + # Resolves to the registered custom DB + config class, not the built-in. + assert db_cls is _AcmeDB + assert type(db_config) is _AcmeDBConfig + + # Every custom param the config declares is present on the built object. + assert db_config.recording_session_table == "acme.raw.recording_session" + assert db_config.signal_name == ["ENGINE_SPEED", "VEHICLE_SPEED"] + assert db_config.start_dt == "2026-04-01" + + # The source tables are merged in (base kwargs) and the location is forced to unity_catalog. + assert db_config.container_metrics_table == "impulse_demo.silver.container_metric" + assert db_config.channel_metrics_table == "impulse_demo.silver.channel_metric" + assert db_config.channels_uri == "impulse_demo.silver.channel_data" + assert db_config.table_locations == "unity_catalog" + + # The custom extra table is declared alongside the source tables for Delta version pinning. + assert "acme.raw.recording_session" in db_config.configured_table_uris() + + +def test_build_measurement_db_config_omits_optional_custom_params(acme_db_registered): + """Optional custom params are genuinely optional: the config builds with only the required one, + and the optional params fall back to their declared defaults.""" + config = ImpulseConfig.model_validate( + _add_custom_measurement_db_to_report_config( + "AcmeDB", {"recording_session_table": "acme.raw.recording_session"} + ) + ) + _, db_config = build_measurement_db_config(dict(config.source), config.query_engine) + + assert db_config.recording_session_table == "acme.raw.recording_session" + assert db_config.signal_name == [] + assert db_config.start_dt is None diff --git a/tests/impulse_reporting/unit/core/report_test.py b/tests/impulse_reporting/unit/core/report_test.py index a3e9fb5b..696125d4 100644 --- a/tests/impulse_reporting/unit/core/report_test.py +++ b/tests/impulse_reporting/unit/core/report_test.py @@ -4,8 +4,11 @@ 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.measurement_db import MeasurementDB, MeasurementDBConfig +from impulse_query_engine.measurement_db_registry import register_measurement_db from impulse_reporting.aggregations.aggregation_types import AggregationType from impulse_reporting.aggregations.histogram import HistogramDuration from impulse_reporting.aggregations.histogram2d import Histogram2DDuration @@ -170,6 +173,25 @@ def test_report_init(): assert report.aggregation_dfs == {} +def test_report_init_fails_loudly_for_custom_db_skipping_super_init(monkeypatch): + """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\(\)"): + Report( + name="test_report", + spark=None, + workspace_client=create_autospec(WorkspaceClient), + config={**DUMMY_CONFIG, "query_engine": {"measurement_db": "SkipsSuperInitDB"}}, + ) + + def test_set_config(): """Test setting config with config path"""