Implement fastapi, sentry, typer, and logging utilities (#2)

This PR contains an initial set of utilities. Developed alongside
https://github.com/WATonomous/mailing-list-gateway.

README.md

@@ -2,3 +2,8 @@
 
 This repo contains a set of common tools and configurations used by the WATcloud team.
 
+## Getting started
+
+```bash
+pip install git+https://github.com/WATonomous/watcloud-utils.git
+```

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "prometheus-fastapi-instrumentator>=7.0.0",
     "sentry-sdk[fastapi]>=2.13.0",
     "typer>=0.12.4",
+    "PyYAML>=6.0.2",
 ]
 
 dynamic = ["version"]

src/watcloud_utils/env.py

@@ -0,0 +1,59 @@
+import json
+import os
+import logging
+from enum import Enum
+
+class Vars(str, Enum):
+    SENTRY_DSN = "SENTRY_DSN"
+    SENTRY_RELEASE = "SENTRY_RELEASE"
+    BUILD_INFO = "BUILD_INFO"
+    DEPLOYMENT_ENVIRONMENT = "DEPLOYMENT_ENVIRONMENT"
+
+# This cache is for warning only once if a variable is missing.
+var_cache = {}
+def getvar(key: Vars, warn_if_missing: bool = True, logger=logging.getLogger(__name__)) -> str:
+    if key in var_cache:
+        return var_cache[key]
+
+    if key == Vars.BUILD_INFO:
+        # DOCKER_METADATA_OUTPUT_JSON is generated by the build pipeline (e.g. docker/metadata-action).
+        # Example:
+        # {
+        #     "tags": ["ghcr.io/watonomous/repo-ingestion:main"],
+        #     "labels": {
+        #         "org.opencontainers.image.title": "repo-ingestion",
+        #         "org.opencontainers.image.description": "Simple server to receive file changes and open GitHub pull requests",
+        #         "org.opencontainers.image.url": "https://github.com/WATonomous/repo-ingestion",
+        #         "org.opencontainers.image.source": "https://github.com/WATonomous/repo-ingestion",
+        #         "org.opencontainers.image.version": "main",
+        #         "org.opencontainers.image.created": "2024-01-20T16:10:39.421Z",
+        #         "org.opencontainers.image.revision": "1d55b62b15c78251e0560af9e97927591e260a98",
+        #         "org.opencontainers.image.licenses": "",
+        #     },
+        # }
+        BUILD_INFO_ENV_VAR = "DOCKER_METADATA_OUTPUT_JSON"
+        val = os.getenv(BUILD_INFO_ENV_VAR)
+        if not val:
+            if warn_if_missing:
+                logger.warning(f"Environment variable {BUILD_INFO_ENV_VAR} not found.")
+            ret = None
+            var_cache[key] = ret
+            return ret
+
+        try:
+            ret = json.loads(val)
+        except json.JSONDecodeError:
+            if warn_if_missing:
+                logger.warning(f"Failed to parse environment variable {BUILD_INFO_ENV_VAR}.")
+            ret = {}
+
+        var_cache[key] = ret
+        return ret
+
+    # Generic environment variable
+    ret = os.getenv(key.value)
+    if not ret and warn_if_missing:
+        logger.warning(f"Environment variable {key.value} not found.")
+
+    var_cache[key] = ret
+    return ret

src/watcloud_utils/fastapi.py

@@ -0,0 +1,82 @@
+import logging
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from prometheus_fastapi_instrumentator import Instrumentator
+
+from .env import Vars, getvar
+from .sentry import sentry_sdk, set_up_sentry
+
+
+class WATcloudFastAPI(FastAPI):
+    def __init__(
+        self,
+        *args,
+        cors_allow_origins=None,
+        logger=logging.getLogger(__name__),
+        expose_metrics=True,
+        expose_build_info=True,
+        expose_health=True,
+        health_fns=[],
+        expose_runtime_info=True,
+        initial_runtime_info={},
+        enable_sentry=True,
+        **kwargs,
+    ):
+        """
+        A FastAPI wrapper that adds various convenience features.
+
+        Args:
+            cors_allow_origins (List[str], optional): A list of origins to allow CORS requests from. Defaults to ['*'] if the DEPLOYMENT_ENVIRONMENT environment variable is nonempty, otherwise []. This parameter is useful for local development. In production, CORS should be handled by the reverse proxy.
+        """
+        super().__init__(*args, **kwargs)
+
+        if cors_allow_origins is None:
+            cors_allow_origins = (
+                ["*"] if getvar(Vars.DEPLOYMENT_ENVIRONMENT, logger=logger) else []
+            )
+
+        if cors_allow_origins:
+            logger.info(
+                f"Adding CORS middleware with allow_origins={cors_allow_origins}. This should only be used for local development. Please handle CORS at the reverse proxy in production."
+            )
+            self.add_middleware(
+                CORSMiddleware,
+                allow_origins=cors_allow_origins,
+                allow_credentials=True,
+                allow_methods=["*"],
+                allow_headers=["*"],
+            )
+
+        self.logger = logger
+
+        if expose_metrics:
+            Instrumentator().instrument(self).expose(self)
+
+        if expose_build_info:
+            self.add_api_route("/build-info", self.read_build_info, methods=["GET"])
+
+        if expose_health:
+            self.health_fns = health_fns
+            self.add_api_route("/health", self.read_health, methods=["GET"])
+
+        if expose_runtime_info:
+            self.runtime_info = initial_runtime_info
+            self.add_api_route("/runtime-info", self.read_runtime_info, methods=["GET"])
+
+        if enable_sentry:
+            sentry_has_dsn = set_up_sentry(logger=self.logger)
+            if self.runtime_info is not None:
+                self.runtime_info["sentry_has_dsn"] = sentry_has_dsn
+                self.runtime_info["sentry_sdk_version"] = sentry_sdk.VERSION
+
+    def read_build_info(self):
+        return getvar(Vars.BUILD_INFO, logger=self.logger) or {}
+
+    def read_health(self):
+        for fn in self.health_fns:
+            fn(self)
+        return {"status": "ok"}
+
+    def read_runtime_info(self):
+        return self.runtime_info

src/watcloud_utils/logging.py

@@ -0,0 +1,13 @@
+import logging
+import os
+
+logger = logging.getLogger()
+
+def set_up_logging():
+    logger.setLevel(os.environ.get("APP_LOG_LEVEL", "INFO"))
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)

src/watcloud_utils/sentry.py

@@ -1,10 +1,10 @@
-import json
 import logging
-import os
 
 import sentry_sdk
 from sentry_sdk.integrations.logging import LoggingIntegration
 
+from .env import Vars, getvar
+
 
 def set_up_sentry(
     trace_sample_rate=1.0,
@@ -22,46 +22,38 @@ def set_up_sentry(
         Default is logging.INFO.
     - event_level: int, optional. The level at which events should be captured.
         Default is logging.ERROR.
+
+    Returns: bool. True if Sentry was set up successfully, False otherwise.
     """
-    if not os.getenv("SENTRY_DSN"):
-        logger.info("No Sentry DSN found. Skipping Sentry setup.")
-        return
-
-    # BUILD_INFO is generated by the build pipeline (e.g. docker/metadata-action).
-    # Example:
-    # {
-    #     "tags": ["ghcr.io/watonomous/repo-ingestion:main"],
-    #     "labels": {
-    #         "org.opencontainers.image.title": "repo-ingestion",
-    #         "org.opencontainers.image.description": "Simple server to receive file changes and open GitHub pull requests",
-    #         "org.opencontainers.image.url": "https://github.com/WATonomous/repo-ingestion",
-    #         "org.opencontainers.image.source": "https://github.com/WATonomous/repo-ingestion",
-    #         "org.opencontainers.image.version": "main",
-    #         "org.opencontainers.image.created": "2024-01-20T16:10:39.421Z",
-    #         "org.opencontainers.image.revision": "1d55b62b15c78251e0560af9e97927591e260a98",
-    #         "org.opencontainers.image.licenses": "",
-    #     },
-    # }
-    try:
-        BUILD_INFO = json.loads(os.getenv("DOCKER_METADATA_OUTPUT_JSON", "{}"))
-    except json.JSONDecodeError:
-        logger.warning("Failed to parse DOCKER_METADATA_OUTPUT_JSON. Not using build info.")
-        BUILD_INFO = {}
+    SENTRY_DSN = getvar(Vars.SENTRY_DSN, logger=logger)
+    if not SENTRY_DSN:
+        logger.warning("SENTRY_DSN not found. Skipping Sentry setup.")
+        return False
+
+    BUILD_INFO = getvar(Vars.BUILD_INFO, logger=logger) or {}
+    DEPLOYMENT_ENVIRONMENT = (
+        getvar(Vars.DEPLOYMENT_ENVIRONMENT, logger=logger) or "unknown"
+    )
 
     build_labels = BUILD_INFO.get("labels", {})
     image_title = build_labels.get("org.opencontainers.image.title", "unknown_image")
-    image_version = build_labels.get("org.opencontainers.image.version", "unknown_version")
+    image_version = build_labels.get(
+        "org.opencontainers.image.version", "unknown_version"
+    )
     image_rev = build_labels.get("org.opencontainers.image.revision", "unknown_rev")
 
+    SENTRY_RELEASE = (
+        getvar("SENTRY_RELEASE", logger=logger)
+        or f"{image_title}:{image_version}@{image_rev}"
+    )
+
     sentry_config = {
-        "dsn": os.environ["SENTRY_DSN"],
-        "environment": os.getenv("DEPLOYMENT_ENVIRONMENT", "unknown"),
-        "release": os.getenv(
-            "SENTRY_RELEASE", f"{image_title}:{image_version}@{image_rev}"
-        ),
+        "dsn": SENTRY_DSN,
+        "environment": DEPLOYMENT_ENVIRONMENT,
+        "release": SENTRY_RELEASE,
     }
 
-    logger.info(f"Sentry DSN found. Setting up Sentry with config: {sentry_config}")
+    logger.info(f"Setting up Sentry with config: {sentry_config}")
     logger.info(f"Sentry SDK version: {sentry_sdk.VERSION}")
 
     sentry_logging = LoggingIntegration(
@@ -87,3 +79,5 @@ def sentry_traces_sampler(sampling_context):
         traces_sampler=sentry_traces_sampler,
         enable_tracing=True,
     )
+
+    return True

src/watcloud_utils/typer.py

@@ -0,0 +1,45 @@
+import json
+import sys
+from enum import Enum
+
+import typer
+import yaml
+
+
+class OutputFormat(str, Enum):
+    yaml = "yaml"
+    json = "json"
+    raw = "raw"
+
+
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        # Encode all iterables as lists
+        # Derived from
+        # - https://stackoverflow.com/a/8230505
+        if hasattr(obj, "__iter__"):
+            return list(obj)
+        # Serialize date
+        if hasattr(obj, "isoformat"):
+            return obj.isoformat()
+        return json.JSONEncoder.default(self, obj)
+
+
+def cli_print_retval(ret: dict | list, output_format: OutputFormat, **kwargs):
+    if output_format == OutputFormat.yaml:
+        print(yaml.dump(ret, default_flow_style=False))
+    elif output_format == OutputFormat.json:
+        print(json.dumps(ret, indent=2, cls=CustomJSONEncoder))
+    elif output_format == OutputFormat.raw:
+        sys.stdout.write(ret)
+    else:
+        raise ValueError(f"Unknown output format: {output_format}")
+
+
+app = typer.Typer(result_callback=cli_print_retval)
+
+
+@app.callback()
+# This function is used to add global CLI options
+def main(output_format: OutputFormat = OutputFormat.yaml):
+    pass