Merge pull request #19 from bento-platform/feat/collect-data-types

feat: collect data types from data services into an endpoint

bento_service_registry/authz_header.py

@@ -0,0 +1,23 @@
+from fastapi import Depends, Request
+from typing import Annotated, Literal
+
+
+__all__ = [
+    "OptionalAuthzHeader",
+    "get_authz_header",
+    "OptionalAuthzHeaderDependency",
+]
+
+
+HeaderAuthorizationType = Literal["Authorization"]
+HEADER_AUTHORIZATION: HeaderAuthorizationType = "Authorization"
+
+OptionalAuthzHeader = dict[HeaderAuthorizationType, str] | None
+
+
+def get_authz_header(request: Request) -> OptionalAuthzHeader:
+    authz_header: str | None = request.headers.get(HEADER_AUTHORIZATION)
+    return {HEADER_AUTHORIZATION: authz_header} if authz_header is not None else None
+
+
+OptionalAuthzHeaderDependency = Annotated[OptionalAuthzHeader, Depends(get_authz_header)]

bento_service_registry/bento_services_json.py

@@ -0,0 +1,63 @@
+import aiofiles
+import json
+from fastapi import Depends
+
+from typing import Annotated
+
+from .config import ConfigDependency
+from .types import BentoService
+
+
+__all__ = [
+    "BentoServicesByComposeID",
+    "BentoServicesByKind",
+    "get_bento_services_by_compose_id",
+    "BentoServicesByComposeIDDependency",
+    "get_bento_services_by_kind",
+    "BentoServicesByKindDependency",
+]
+
+
+BentoServicesByComposeID = dict[str, BentoService]
+BentoServicesByKind = dict[str, BentoService]
+
+
+async def get_bento_services_by_compose_id(config: ConfigDependency) -> BentoServicesByComposeID:
+    async with aiofiles.open(config.bento_services, "r") as fh:
+        bento_services_data: BentoServicesByComposeID = json.loads(await fh.read())
+
+    return {
+        sk: BentoService(
+            **sv,
+            url=sv["url_template"].format(
+                BENTO_URL=config.bento_url,
+                BENTO_PUBLIC_URL=config.bento_public_url,
+                BENTO_PORTAL_PUBLIC_URL=config.bento_portal_public_url,
+                **sv,
+            ),
+        )  # type: ignore
+        for sk, sv in bento_services_data.items()
+        # Filter out disabled entries and entries without service_kind, which may be external/'transparent'
+        # - e.g., the gateway.
+        if not sv.get("disabled") and sv.get("service_kind")
+    }
+
+
+BentoServicesByComposeIDDependency = Annotated[BentoServicesByComposeID, Depends(get_bento_services_by_compose_id)]
+
+
+async def get_bento_services_by_kind(
+    bento_services_by_compose_id: BentoServicesByComposeIDDependency
+) -> BentoServicesByKind:
+    services_by_kind: BentoServicesByKind = {}
+
+    for sv in bento_services_by_compose_id.values():
+        # Disabled entries are already filtered out by get_bento_services_by_compose_id
+        # Filter out entries without service_kind, which may be external/'transparent' - e.g., the gateway.
+        if sk := sv.get("service_kind"):
+            services_by_kind[sk] = sv
+
+    return services_by_kind
+
+
+BentoServicesByKindDependency = Annotated[BentoServicesByKind, Depends(get_bento_services_by_kind)]

bento_service_registry/data_types.py

@@ -0,0 +1,84 @@
+import aiohttp
+import asyncio
+import itertools
+import logging
+
+from fastapi import Depends, status
+from pydantic import ValidationError
+from typing import Annotated
+from urllib.parse import urljoin
+
+from .authz_header import OptionalAuthzHeader, OptionalAuthzHeaderDependency
+from .http_session import HTTPSessionDependency
+from .logger import LoggerDependency
+from .models import DataTypeWithServiceURL
+from .services import ServicesDependency
+
+__all__ = [
+    "DataTypesTuple",
+    "get_data_types",
+    "DataTypesDependency",
+]
+
+
+DataTypesTuple = tuple[DataTypeWithServiceURL, ...]
+
+
+async def get_data_types_from_service(
+    authz_header: OptionalAuthzHeader,
+    http_session: aiohttp.ClientSession,
+    logger: logging.Logger,
+    service: dict,
+) -> DataTypesTuple:
+    service_url: str | None = service.get("url")
+
+    if service_url is None:
+        logger.error(f"Encountered a service missing a URL: {service}")
+        return ()
+
+    service_url_norm: str = service_url.rstrip("/") + "/"
+
+    async with http_session.get(urljoin(service_url_norm, "data-types"), headers=authz_header) as res:
+        if res.status != status.HTTP_200_OK:
+            logger.error(
+                f"Got non-200 response from data type service ({service_url=}): {res.status=}; body={await res.json()}")
+            return ()
+
+        dts: list[DataTypeWithServiceURL] = []
+
+        for dt in await res.json():
+            try:
+                dts.append(DataTypeWithServiceURL.model_validate({**dt, "service_base_url": service_url_norm}))
+            except ValidationError as err:
+                logger.error(f"Recieved malformatted data type: {dt} ({err=}); skipping")
+                continue
+
+        return tuple(dts)
+
+
+async def get_data_types(
+    authz_header: OptionalAuthzHeaderDependency,
+    http_session: HTTPSessionDependency,
+    logger: LoggerDependency,
+    services_tuple: ServicesDependency,
+) -> DataTypesTuple:
+    logger.debug("Collecting data types from data services")
+
+    data_services = [s for s in services_tuple if s.get("bento", {}).get("dataService", False)]
+
+    logger.debug(f"Found {len(data_services)} data services")
+
+    data_types_from_services: tuple[DataTypeWithServiceURL, ...] = tuple(
+        itertools.chain(
+            *await asyncio.gather(
+                *(get_data_types_from_service(authz_header, http_session, logger, s) for s in data_services)
+            )
+        )
+    )
+
+    logger.debug(f"Obtained {len(data_types_from_services)} data types")
+
+    return data_types_from_services
+
+
+DataTypesDependency = Annotated[DataTypesTuple, Depends(get_data_types)]

bento_service_registry/models.py

@@ -0,0 +1,16 @@
+from pydantic import BaseModel, Field
+
+__all__ = [
+    "DataTypeWithServiceURL",
+]
+
+
+class DataTypeWithServiceURL(BaseModel):
+    label: str | None = None
+    queryable: bool
+    item_schema: dict = Field(..., alias="schema")
+    metadata_schema: dict
+    id: str
+    count: int | None
+    # Injected rather than from service:
+    service_base_url: str