From 5da9eaea88fad4c669f7185cf0bc8d62a511941d Mon Sep 17 00:00:00 2001 From: Ben Zhang Date: Sat, 17 Aug 2024 22:11:45 -0700 Subject: [PATCH] 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 | 5 +++ pdm.lock | 48 +++++++++++++++++++- pyproject.toml | 1 + src/watcloud_utils/env.py | 59 +++++++++++++++++++++++++ src/watcloud_utils/fastapi.py | 82 +++++++++++++++++++++++++++++++++++ src/watcloud_utils/logging.py | 13 ++++++ src/watcloud_utils/sentry.py | 60 ++++++++++++------------- src/watcloud_utils/typer.py | 45 +++++++++++++++++++ 8 files changed, 279 insertions(+), 34 deletions(-) create mode 100644 src/watcloud_utils/env.py create mode 100644 src/watcloud_utils/fastapi.py create mode 100644 src/watcloud_utils/logging.py create mode 100644 src/watcloud_utils/typer.py diff --git a/README.md b/README.md index 4801a66..1004fb8 100644 --- a/README.md +++ b/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 +``` diff --git a/pdm.lock b/pdm.lock index 3e4af65..27114af 100644 --- a/pdm.lock +++ b/pdm.lock @@ -5,7 +5,7 @@ groups = ["default", "testing"] strategy = ["inherit_metadata"] lock_version = "4.5.0" -content_hash = "sha256:7ce8c6de2bc4c00032a81d7b941be8eb14a4b6bc550a517ef83e3d8d14b4e2dd" +content_hash = "sha256:2ce844c1270c5621c586c241d69c9bf4add7ea99bb683757ca437d20ad2e87c7" [[metadata.targets]] requires_python = ">=3.10" @@ -318,6 +318,52 @@ files = [ {file = "pytest-8.3.2.tar.gz", hash = "sha256:c132345d12ce551242c87269de812483f5bcc87cdbb4722e48487ba194f9fdce"}, ] +[[package]] +name = "pyyaml" +version = "6.0.2" +requires_python = ">=3.8" +summary = "YAML parser and emitter for Python" +groups = ["default"] +files = [ + {file = "PyYAML-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086"}, + {file = "PyYAML-6.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf"}, + {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8824b5a04a04a047e72eea5cec3bc266db09e35de6bdfe34c9436ac5ee27d237"}, + {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7c36280e6fb8385e520936c3cb3b8042851904eba0e58d277dca80a5cfed590b"}, + {file = "PyYAML-6.0.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ec031d5d2feb36d1d1a24380e4db6d43695f3748343d99434e6f5f9156aaa2ed"}, + {file = "PyYAML-6.0.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:936d68689298c36b53b29f23c6dbb74de12b4ac12ca6cfe0e047bedceea56180"}, + {file = "PyYAML-6.0.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:23502f431948090f597378482b4812b0caae32c22213aecf3b55325e049a6c68"}, + {file = "PyYAML-6.0.2-cp310-cp310-win32.whl", hash = "sha256:2e99c6826ffa974fe6e27cdb5ed0021786b03fc98e5ee3c5bfe1fd5015f42b99"}, + {file = "PyYAML-6.0.2-cp310-cp310-win_amd64.whl", hash = "sha256:a4d3091415f010369ae4ed1fc6b79def9416358877534caf6a0fdd2146c87a3e"}, + {file = "PyYAML-6.0.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:cc1c1159b3d456576af7a3e4d1ba7e6924cb39de8f67111c735f6fc832082774"}, + {file = "PyYAML-6.0.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:1e2120ef853f59c7419231f3bf4e7021f1b936f6ebd222406c3b60212205d2ee"}, + {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5d225db5a45f21e78dd9358e58a98702a0302f2659a3c6cd320564b75b86f47c"}, + {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5ac9328ec4831237bec75defaf839f7d4564be1e6b25ac710bd1a96321cc8317"}, + {file = "PyYAML-6.0.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3ad2a3decf9aaba3d29c8f537ac4b243e36bef957511b4766cb0057d32b0be85"}, + {file = "PyYAML-6.0.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:ff3824dc5261f50c9b0dfb3be22b4567a6f938ccce4587b38952d85fd9e9afe4"}, + {file = "PyYAML-6.0.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:797b4f722ffa07cc8d62053e4cff1486fa6dc094105d13fea7b1de7d8bf71c9e"}, + {file = "PyYAML-6.0.2-cp311-cp311-win32.whl", hash = "sha256:11d8f3dd2b9c1207dcaf2ee0bbbfd5991f571186ec9cc78427ba5bd32afae4b5"}, + {file = "PyYAML-6.0.2-cp311-cp311-win_amd64.whl", hash = "sha256:e10ce637b18caea04431ce14fabcf5c64a1c61ec9c56b071a4b7ca131ca52d44"}, + {file = "PyYAML-6.0.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:c70c95198c015b85feafc136515252a261a84561b7b1d51e3384e0655ddf25ab"}, + {file = "PyYAML-6.0.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ce826d6ef20b1bc864f0a68340c8b3287705cae2f8b4b1d932177dcc76721725"}, + {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f71ea527786de97d1a0cc0eacd1defc0985dcf6b3f17bb77dcfc8c34bec4dc5"}, + {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9b22676e8097e9e22e36d6b7bda33190d0d400f345f23d4065d48f4ca7ae0425"}, + {file = "PyYAML-6.0.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80bab7bfc629882493af4aa31a4cfa43a4c57c83813253626916b8c7ada83476"}, + {file = "PyYAML-6.0.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:0833f8694549e586547b576dcfaba4a6b55b9e96098b36cdc7ebefe667dfed48"}, + {file = "PyYAML-6.0.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8b9c7197f7cb2738065c481a0461e50ad02f18c78cd75775628afb4d7137fb3b"}, + {file = "PyYAML-6.0.2-cp312-cp312-win32.whl", hash = "sha256:ef6107725bd54b262d6dedcc2af448a266975032bc85ef0172c5f059da6325b4"}, + {file = "PyYAML-6.0.2-cp312-cp312-win_amd64.whl", hash = "sha256:7e7401d0de89a9a855c839bc697c079a4af81cf878373abd7dc625847d25cbd8"}, + {file = "PyYAML-6.0.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:efdca5630322a10774e8e98e1af481aad470dd62c3170801852d752aa7a783ba"}, + {file = "PyYAML-6.0.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:50187695423ffe49e2deacb8cd10510bc361faac997de9efef88badc3bb9e2d1"}, + {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0ffe8360bab4910ef1b9e87fb812d8bc0a308b0d0eef8c8f44e0254ab3b07133"}, + {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:17e311b6c678207928d649faa7cb0d7b4c26a0ba73d41e99c4fff6b6c3276484"}, + {file = "PyYAML-6.0.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:70b189594dbe54f75ab3a1acec5f1e3faa7e8cf2f1e08d9b561cb41b845f69d5"}, + {file = "PyYAML-6.0.2-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:41e4e3953a79407c794916fa277a82531dd93aad34e29c2a514c2c0c5fe971cc"}, + {file = "PyYAML-6.0.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:68ccc6023a3400877818152ad9a1033e3db8625d899c72eacb5a668902e4d652"}, + {file = "PyYAML-6.0.2-cp313-cp313-win32.whl", hash = "sha256:bc2fa7c6b47d6bc618dd7fb02ef6fdedb1090ec036abab80d4681424b84c1183"}, + {file = "PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563"}, + {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"}, +] + [[package]] name = "rich" version = "13.7.1" diff --git a/pyproject.toml b/pyproject.toml index 5bd5b1a..2505a5b 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/watcloud_utils/env.py b/src/watcloud_utils/env.py new file mode 100644 index 0000000..20f0b9f --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/watcloud_utils/fastapi.py b/src/watcloud_utils/fastapi.py new file mode 100644 index 0000000..b983374 --- /dev/null +++ b/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 diff --git a/src/watcloud_utils/logging.py b/src/watcloud_utils/logging.py new file mode 100644 index 0000000..283034b --- /dev/null +++ b/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) diff --git a/src/watcloud_utils/sentry.py b/src/watcloud_utils/sentry.py index 04944d0..cc36fac 100644 --- a/src/watcloud_utils/sentry.py +++ b/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 diff --git a/src/watcloud_utils/typer.py b/src/watcloud_utils/typer.py new file mode 100644 index 0000000..85d57d3 --- /dev/null +++ b/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