Merge pull request #241 from bento-platform/feat/discovery/config-model

feat: discovery module with config model

Author: davidlougheed

README.md

@@ -97,6 +97,15 @@ additional code to set up error handling and basic Bento service boilerplate.
 
 `db` contains common base classes for setting up database managers.
 
+### `discovery`
+
+`discovery` contains models and helper functions for the Bento Discovery Configuration specification, used
+in [Katsu](https://github.com/bento-platform/katsu).
+
+#### Guides
+
+* [Discovery configuration: structure and validation](./docs/discovery/discovery_config.md)
+
 ### `drs`
 
 `drs` provides utilities for fetching data and record metadata from 

bento_lib/__init__.py

@@ -2,6 +2,7 @@
 
 from . import apps
 from . import auth
+from . import discovery
 from . import drs
 from . import events
 from . import schemas
@@ -15,6 +16,7 @@
     "__version__",
     "apps",
     "auth",
+    "discovery",
     "drs",
     "events",
     "schemas",

bento_lib/_internal.py

@@ -1,7 +1,8 @@
 import logging
+from structlog.stdlib import BoundLogger, get_logger
 
 __all__ = ["internal_logger"]
 
 logging.basicConfig(level=logging.NOTSET)
 
-internal_logger = logging.getLogger("bento_lib")
+internal_logger: BoundLogger = get_logger("bento_lib")

bento_lib/discovery/__init__.py

@@ -0,0 +1,4 @@
+from . import helpers
+from . import models
+
+__all__ = ["helpers", "models"]

bento_lib/discovery/exceptions.py

@@ -0,0 +1,5 @@
+__all__ = ["DiscoveryValidationError"]
+
+
+class DiscoveryValidationError(Exception):
+    pass

bento_lib/discovery/helpers.py

@@ -0,0 +1,75 @@
+from pathlib import Path
+from structlog.stdlib import BoundLogger
+
+from bento_lib._internal import internal_logger
+from .exceptions import DiscoveryValidationError
+from .models.config import DiscoveryConfig
+
+__all__ = [
+    "load_discovery_config",
+]
+
+
+FIELD_DEF_NOT_FOUND = "field definition not found"
+FIELD_ALREADY_SEEN = "field already seen"
+
+
+def _load_discovery_config_json(config_path: Path | str) -> DiscoveryConfig:
+    with open(config_path, "r") as fh:
+        return DiscoveryConfig.model_validate_json(fh.read())
+
+
+def _validate_references_and_duplicates(cfg: DiscoveryConfig, logger: BoundLogger) -> None:
+    fields = cfg.fields
+
+    # validate overview and check for chart duplicates:
+    seen_chart_fields: set[str] = set()
+    for s_idx, section in enumerate(cfg.overview):
+        for c_idx, chart in enumerate(section.charts):
+            exc_path = (
+                f"overview > section {section.section_title} [{s_idx}] > {chart.field} {chart.chart_type} [{c_idx}]"
+            )
+            if chart.field not in fields:
+                logger.error(
+                    f"overview {FIELD_DEF_NOT_FOUND}", section=section.section_title, field=chart.field, chart_idx=c_idx
+                )
+                raise DiscoveryValidationError(f"{exc_path}: {FIELD_DEF_NOT_FOUND}")
+            if chart.field in seen_chart_fields:
+                logger.error(
+                    f"overview {FIELD_ALREADY_SEEN}", section=section.section_title, field=chart.field, chart_idx=c_idx
+                )
+                raise DiscoveryValidationError(f"{exc_path}: {FIELD_ALREADY_SEEN}")
+            seen_chart_fields.add(chart.field)
+
+    # validate search:
+    seen_search_fields: set[str] = set()
+    for s_idx, section in enumerate(cfg.search):
+        for f_idx, f in enumerate(section.fields):
+            exc_path = f"search > section {section.section_title} [{s_idx}] > {f} [{f_idx}]"
+            if f not in fields:
+                logger.error(f"search {FIELD_DEF_NOT_FOUND}", section=section.section_title, field=f)
+                raise DiscoveryValidationError(f"{exc_path}: {FIELD_DEF_NOT_FOUND}")
+            if f in seen_search_fields:
+                logger.error(f"search {FIELD_ALREADY_SEEN}", section=section.section_title, field=f)
+                raise DiscoveryValidationError(f"{exc_path}: {FIELD_ALREADY_SEEN}")
+            seen_search_fields.add(f)
+
+    # issue warnings if there are fields defined that the config doesn't reference:
+    referenced_fields = seen_chart_fields | seen_search_fields
+    for fi, f in enumerate(fields.keys()):
+        if f not in referenced_fields:
+            logger.warning("field not referenced", field=f, field_idx=fi)
+
+
+def load_discovery_config(config_path: Path | str, logger: BoundLogger | None = None) -> DiscoveryConfig:
+    # 1. load the config object (or raise a Pydantic validation error if the config is in the wrong format)
+    cfg = _load_discovery_config_json(config_path)
+
+    # 2. validate the config's internal references and overview chart/search field entries
+    #  a) make sure all fields in overview and search are defined
+    #  b) make sure fields are not listed more than once as a chart or as a search filter
+    #  c) issue warnings if any fields are defined that the config doesn't reference anywhere
+    _validate_references_and_duplicates(cfg, logger or internal_logger)
+
+    # now that we've validated references, return the config
+    return cfg

bento_lib/discovery/models/__init__.py

@@ -0,0 +1,8 @@
+from . import config, fields, overview, search
+
+__all__ = [
+    "config",
+    "fields",
+    "overview",
+    "search",
+]

bento_lib/discovery/models/config.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel
+
+from .fields import FieldDefinition
+from .overview import OverviewSection
+from .search import SearchSection
+
+__all__ = [
+    "DiscoveryConfigRules",
+    "DiscoveryConfig",
+]
+
+
+class DiscoveryConfigRules(BaseModel):
+    count_threshold: int
+    max_query_parameters: int
+
+
+class DiscoveryConfig(BaseModel):
+    overview: list[OverviewSection]
+    search: list[SearchSection]
+    fields: dict[str, FieldDefinition]
+    rules: DiscoveryConfigRules

bento_lib/discovery/models/fields.py

@@ -0,0 +1,77 @@
+from pydantic import AliasChoices, BaseModel, Field
+from typing import Literal
+
+__all__ = [
+    "BaseFieldDefinition",
+    # string
+    "StringFieldConfig",
+    "StringFieldDefinition",
+    # number
+    "BaseNumberFieldConfig",
+    "ManualBinsNumberFieldConfig",
+    "AutoBinsNumberFieldConfig",
+    "NumberFieldDefinition",
+    # date
+    "DateFieldConfig",
+    "DateFieldDefinition",
+    # sum type:
+    "FieldDefinition",
+]
+
+
+DataTypeField = Field(validation_alias=AliasChoices("data_type", "datatype"))
+
+
+class BaseFieldDefinition(BaseModel):
+    mapping: str
+    title: str  # TODO: make optional and pull from Bento schema if not set
+    description: str  # TODO: make optional and pull from Bento schema if not set
+    data_type: Literal["string", "number", "date"] = DataTypeField
+    # --- The below fields are currently valid, but need to be reworked for new search ---------------------------------
+    mapping_for_search_filter: str | None = None
+    group_by: str | None = None
+    group_by_value: str | None = None
+    value_mapping: str | None = None
+    # ------------------------------------------------------------------------------------------------------------------
+
+
+class StringFieldConfig(BaseModel):
+    enum: list[str] | None
+
+
+class StringFieldDefinition(BaseFieldDefinition):
+    data_type: Literal["string"] = DataTypeField
+    config: StringFieldConfig
+
+
+class BaseNumberFieldConfig(BaseModel):
+    units: str
+
+
+class ManualBinsNumberFieldConfig(BaseNumberFieldConfig):
+    bins: list[int | float]
+
+
+class AutoBinsNumberFieldConfig(BaseNumberFieldConfig):
+    bin_size: int
+    taper_left: int
+    taper_right: int
+    minimum: int
+    maximum: int
+
+
+class NumberFieldDefinition(BaseFieldDefinition):
+    data_type: Literal["number"] = DataTypeField
+    config: ManualBinsNumberFieldConfig | AutoBinsNumberFieldConfig
+
+
+class DateFieldConfig(BaseModel):
+    bin_by: Literal["month"]  # Currently only binning by month is implemented
+
+
+class DateFieldDefinition(BaseFieldDefinition):
+    data_type: Literal["date"] = DataTypeField
+    config: DateFieldConfig
+
+
+FieldDefinition = DateFieldDefinition | NumberFieldDefinition | StringFieldDefinition

bento_lib/discovery/models/overview.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+from typing import Literal
+
+__all__ = [
+    "OverviewChart",
+    "OverviewSection",
+]
+
+
+class OverviewChart(BaseModel):
+    field: str
+    chart_type: Literal["bar", "choropleth", "histogram", "pie"]
+
+
+class OverviewSection(BaseModel):
+    section_title: str
+    charts: list[OverviewChart]