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 1/2] 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"

From c7ff5e44c64c0d9fc15e11ac8ffda215e33c4f17 Mon Sep 17 00:00:00 2001
From: Kenneth Kehl <@kkehl@flexion.us>
Date: Thu, 14 Mar 2024 11:27:26 -0700
Subject: [PATCH 2/2] duh

---
 poetry.lock    | 20 ++++++++++----------
 pyproject.toml |  2 +-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 91c32b2b9..1dbab6175 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2657,11 +2657,11 @@ flask = "^2.3.2"
 flask-redis = "^0.4.0"
 geojson = "^3.0.1"
 govuk-bank-holidays = "^0.13"
-idna = "^3.4"
+idna = "^3.6"
 itsdangerous = "^2.1.2"
 jinja2 = "^3.1.3"
 jmespath = "^1.0.1"
-markupsafe = "^2.1.2"
+markupsafe = "^2.1.5"
 mistune = "==0.8.4"
 numpy = "^1.24.2"
 orderedset = "^2.0.3"
@@ -2671,7 +2671,7 @@ python-dateutil = "^2.8.2"
 python-json-logger = "^2.0.7"
 pytz = "^2023.3"
 pyyaml = "^6.0"
-redis = "^5.0.1"
+redis = "^5.0.3"
 regex = "^2023.10.3"
 requests = "^2.31.0"
 s3transfer = "^0.7.0"
@@ -2685,8 +2685,8 @@ werkzeug = "^3.0.1"
 [package.source]
 type = "git"
 url = "https://github.com/GSA/notifications-utils.git"
-reference = "053fe30"
-resolved_reference = "053fe304a1bbc1ae6559869b1fd0a2c3f83981ea"
+reference = "HEAD"
+resolved_reference = "e421fb1936d4bb92294a11fb78fc81b77d3577b2"
 
 [[package]]
 name = "numpy"
@@ -3650,17 +3650,17 @@ full = ["numpy"]
 
 [[package]]
 name = "redis"
-version = "5.0.2"
+version = "5.0.3"
 description = "Python client for Redis database and key-value store"
 optional = false
 python-versions = ">=3.7"
 files = [
-    {file = "redis-5.0.2-py3-none-any.whl", hash = "sha256:4caa8e1fcb6f3c0ef28dba99535101d80934b7d4cd541bbb47f4a3826ee472d1"},
-    {file = "redis-5.0.2.tar.gz", hash = "sha256:3f82cc80d350e93042c8e6e7a5d0596e4dd68715babffba79492733e1f367037"},
+    {file = "redis-5.0.3-py3-none-any.whl", hash = "sha256:5da9b8fe9e1254293756c16c008e8620b3d15fcc6dde6babde9541850e72a32d"},
+    {file = "redis-5.0.3.tar.gz", hash = "sha256:4973bae7444c0fbed64a06b87446f79361cb7e4ec1538c022d696ed7a5015580"},
 ]
 
 [package.dependencies]
-async-timeout = ">=4.0.3"
+async-timeout = {version = ">=4.0.3", markers = "python_full_version < \"3.11.3\""}
 
 [package.extras]
 hiredis = ["hiredis (>=1.0.0)"]
@@ -4785,4 +4785,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9,<3.12"
-content-hash = "dfa333e33a5f72123b0f8f37d76478d89569a5395c31b052eb867bb9ac40a1b7"
+content-hash = "ed0a2c9f32e010a33a24e243a58550c4849d10ee8205f6c459ecdde57d3509ac"
diff --git a/pyproject.toml b/pyproject.toml
index 3a02d9b5f..4b8ab01ef 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", rev="053fe30"}
+notifications-utils = {git = "https://github.com/GSA/notifications-utils.git"}
 oscrypto = "==1.3.0"
 packaging = "==23.2"
 poetry-dotenv-plugin = "==0.2.0"