From 08a878dfd196cdec69c4073d716321280ffcc5d5 Mon Sep 17 00:00:00 2001
From: David Lougheed <david.lougheed@gmail.com>
Date: Wed, 27 Nov 2024 14:40:26 -0500
Subject: [PATCH] feat: discovery module with config model

---
 README.md                              |  5 ++
 bento_lib/__init__.py                  | 13 ++++-
 bento_lib/discovery/__init__.py        |  4 ++
 bento_lib/discovery/helpers.py         |  0
 bento_lib/discovery/models/__init__.py |  0
 bento_lib/discovery/models/config.py   | 22 ++++++++
 bento_lib/discovery/models/fields.py   | 69 ++++++++++++++++++++++++++
 bento_lib/discovery/models/overview.py | 17 +++++++
 bento_lib/discovery/models/search.py   |  8 +++
 9 files changed, 137 insertions(+), 1 deletion(-)
 create mode 100644 bento_lib/discovery/__init__.py
 create mode 100644 bento_lib/discovery/helpers.py
 create mode 100644 bento_lib/discovery/models/__init__.py
 create mode 100644 bento_lib/discovery/models/config.py
 create mode 100644 bento_lib/discovery/models/fields.py
 create mode 100644 bento_lib/discovery/models/overview.py
 create mode 100644 bento_lib/discovery/models/search.py

diff --git a/README.md b/README.md
index 509683c..f60441a 100644
--- a/README.md
+++ b/README.md
@@ -97,6 +97,11 @@ 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).
+
 ### `drs`
 
 `drs` provides utilities for fetching data and record metadata from 
diff --git a/bento_lib/__init__.py b/bento_lib/__init__.py
index 7053863..cd7f919 100644
--- a/bento_lib/__init__.py
+++ b/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
@@ -12,5 +13,15 @@
 
 __version__ = metadata.version(__name__)
 __all__ = [
-    "__version__", "apps", "auth", "drs", "events", "schemas", "search", "service_info", "streaming", "workflows"
+    "__version__",
+    "apps",
+    "auth",
+    "discovery",
+    "drs",
+    "events",
+    "schemas",
+    "search",
+    "service_info",
+    "streaming",
+    "workflows",
 ]
diff --git a/bento_lib/discovery/__init__.py b/bento_lib/discovery/__init__.py
new file mode 100644
index 0000000..fa30133
--- /dev/null
+++ b/bento_lib/discovery/__init__.py
@@ -0,0 +1,4 @@
+from . import helpers
+from . import models
+
+__all__ = ["helpers", "models"]
diff --git a/bento_lib/discovery/helpers.py b/bento_lib/discovery/helpers.py
new file mode 100644
index 0000000..e69de29
diff --git a/bento_lib/discovery/models/__init__.py b/bento_lib/discovery/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/bento_lib/discovery/models/config.py b/bento_lib/discovery/models/config.py
new file mode 100644
index 0000000..87cf79a
--- /dev/null
+++ b/bento_lib/discovery/models/config.py
@@ -0,0 +1,22 @@
+from pydantic import BaseModel
+
+from .fields import DateFieldDefinition, NumberFieldDefinition, StringFieldDefinition
+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, DateFieldDefinition | NumberFieldDefinition | StringFieldDefinition]
+    rules: DiscoveryConfigRules
diff --git a/bento_lib/discovery/models/fields.py b/bento_lib/discovery/models/fields.py
new file mode 100644
index 0000000..59bb442
--- /dev/null
+++ b/bento_lib/discovery/models/fields.py
@@ -0,0 +1,69 @@
+from pydantic import AliasChoices, BaseModel, Field
+from typing import Literal
+
+__all__ = [
+    "BaseFieldDefinition",
+    # string
+    "StringFieldConfig",
+    "StringFieldDefinition",
+    # number
+    "BaseNumberFieldConfig",
+    "ManualBinsNumberFieldConfig",
+    "AutoBinsNumberFieldConfig",
+    "NumberFieldDefinition",
+    # date
+    "DateFieldConfig",
+    "DateFieldDefinition",
+]
+
+
+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
+    datatype: Literal["string", "number", "date"] = Field(validation_alias=AliasChoices("data_type", "datatype"))
+    # --- 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):
+    datatype: Literal["string"]
+    config: StringFieldConfig
+
+
+class BaseNumberFieldConfig(BaseModel):
+    units: str
+
+
+class ManualBinsNumberFieldConfig(BaseNumberFieldConfig):
+    bins: list[int]
+
+
+class AutoBinsNumberFieldConfig(BaseNumberFieldConfig):
+    bin_size: int
+    taper_left: int
+    taper_right: int
+    minimum: int
+    maximum: int
+
+
+class NumberFieldDefinition(BaseFieldDefinition):
+    datatype: Literal["number"]
+    config: ManualBinsNumberFieldConfig | AutoBinsNumberFieldConfig
+
+
+class DateFieldConfig(BaseModel):
+    bin_by: str  # TODO
+
+
+class DateFieldDefinition(BaseFieldDefinition):
+    datatype: Literal["date"]
+    config: DateFieldConfig
diff --git a/bento_lib/discovery/models/overview.py b/bento_lib/discovery/models/overview.py
new file mode 100644
index 0000000..7aa3549
--- /dev/null
+++ b/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]
diff --git a/bento_lib/discovery/models/search.py b/bento_lib/discovery/models/search.py
new file mode 100644
index 0000000..7501d09
--- /dev/null
+++ b/bento_lib/discovery/models/search.py
@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+__all__ = ["SearchSection"]
+
+
+class SearchSection(BaseModel):
+    section_title: str
+    fields: list[str]