From abce798f5b119069479d40f7946d5b5eca9a3159 Mon Sep 17 00:00:00 2001 From: Kenneth Kehl <@kkehl@flexion.us> Date: Mon, 4 Mar 2024 10:32:19 -0800 Subject: [PATCH] remove celery from notifications-utils (notify-api-130) --- app/__init__.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++-- poetry.lock | 9 +++--- pyproject.toml | 2 +- 3 files changed, 82 insertions(+), 8 deletions(-) diff --git a/app/__init__.py b/app/__init__.py index 5a364f151..9184c4815 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -3,15 +3,16 @@ import secrets import string import time import uuid +from contextlib import contextmanager from time import monotonic -from celery import current_task +from celery import Celery, Task, current_task from flask import current_app, g, has_request_context, jsonify, make_response, request +from flask.ctx import has_app_context from flask_marshmallow import Marshmallow from flask_migrate import Migrate from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy from notifications_utils import logging, request_helper -from notifications_utils.celery import NotifyCelery from notifications_utils.clients.encryption.encryption_client import Encryption from notifications_utils.clients.redis.redis_client import RedisClient from notifications_utils.clients.zendesk.zendesk_client import ZendeskClient @@ -27,6 +28,25 @@ from app.clients.email.aws_ses_stub import AwsSesStubClient from app.clients.sms.aws_sns import AwsSnsClient +class NotifyCelery(Celery): + def init_app(self, app): + self.task_cls = make_task(app) + + # Configure Celery app with options from the main app config. + self.config_from_object(app.config["CELERY"]) + + def send_task(self, name, args=None, kwargs=None, **other_kwargs): + other_kwargs["headers"] = other_kwargs.get("headers") or {} + + if has_request_context() and hasattr(request, "request_id"): + other_kwargs["headers"]["notify_request_id"] = request.request_id + + elif has_app_context() and "request_id" in g: + other_kwargs["headers"]["notify_request_id"] = g.request_id + + return super().send_task(name, args, kwargs, **other_kwargs) + + class SQLAlchemy(_SQLAlchemy): """We need to subclass SQLAlchemy in order to override create_engine options""" @@ -366,3 +386,58 @@ def setup_sqlalchemy_events(app): @event.listens_for(db.engine, "checkin") def checkin(dbapi_connection, connection_record): # noqa pass + + +def make_task(app): + class NotifyTask(Task): + abstract = True + start = None + + @property + def queue_name(self): + delivery_info = self.request.delivery_info or {} + return delivery_info.get("routing_key", "none") + + @property + def request_id(self): + # Note that each header is a direct attribute of the + # task context (aka "request"). + return self.request.get("notify_request_id") + + @contextmanager + def app_context(self): + with app.app_context(): + # Add 'request_id' to 'g' so that it gets logged. + g.request_id = self.request_id + yield + + def on_success(self, retval, task_id, args, kwargs): # noqa + # enables request id tracing for these logs + with self.app_context(): + elapsed_time = time.monotonic() - self.start + + app.logger.info( + "Celery task {task_name} (queue: {queue_name}) took {time}".format( + task_name=self.name, + queue_name=self.queue_name, + time="{0:.4f}".format(elapsed_time), + ) + ) + + def on_failure(self, exc, task_id, args, kwargs, einfo): # noqa + # enables request id tracing for these logs + with self.app_context(): + app.logger.exception( + "Celery task {task_name} (queue: {queue_name}) failed".format( + task_name=self.name, + queue_name=self.queue_name, + ) + ) + + def __call__(self, *args, **kwargs): + # ensure task has flask context to access config, logger, etc + with self.app_context(): + self.start = time.monotonic() + return super().__call__(*args, **kwargs) + + return NotifyTask diff --git a/poetry.lock b/poetry.lock index 8f8b2f11f..2030fc693 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "aiohttp" @@ -2670,8 +2670,8 @@ werkzeug = "^3.0.1" [package.source] type = "git" url = "https://github.com/GSA/notifications-utils.git" -reference = "HEAD" -resolved_reference = "df8ece836cad303c5d161edf485cf9bbdc666688" +reference = "053fe30" +resolved_reference = "053fe304a1bbc1ae6559869b1fd0a2c3f83981ea" [[package]] name = "numpy" @@ -3479,7 +3479,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -4771,4 +4770,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9,<3.12" -content-hash = "ccfd855abf7309cf520ea5ca02e77fc9d8ab911819ea9690298fa0549a9d7c42" +content-hash = "1041afd9fe9642567640966a7f7b1129f8f084f169419d8f2c7db9071672eb57" diff --git a/pyproject.toml b/pyproject.toml index 970059f5a..a87fefe96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ marshmallow = "==3.20.2" marshmallow-sqlalchemy = "==0.30.0" newrelic = "*" notifications-python-client = "==9.0.0" -notifications-utils = {git = "https://github.com/GSA/notifications-utils.git"} +notifications-utils = {git = "https://github.com/GSA/notifications-utils.git", rev="053fe30"} oscrypto = "==1.3.0" packaging = "==23.2" poetry-dotenv-plugin = "==0.2.0"