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"""