From 5cedd6427dc64590eb55c45b968311186afc6856 Mon Sep 17 00:00:00 2001 From: Kenneth Kehl <@kkehl@flexion.us> Date: Thu, 9 Jan 2025 07:47:47 -0800 Subject: [PATCH 1/5] use singletons for s3 client --- app/aws/s3.py | 7 +++++++ app/config.py | 12 ++++++++++++ notifications_utils/s3.py | 35 +++++++++++++++++++++++------------ 3 files changed, 42 insertions(+), 12 deletions(-) diff --git a/app/aws/s3.py b/app/aws/s3.py index c33366a2c..01cd6692e 100644 --- a/app/aws/s3.py +++ b/app/aws/s3.py @@ -10,6 +10,7 @@ from boto3 import Session from flask import current_app from app.clients import AWS_CLIENT_CONFIG +from app.utils import hilite from notifications_utils import aware_utcnow FILE_LOCATION_STRUCTURE = "service-{}-notify/{}.csv" @@ -65,6 +66,7 @@ def clean_cache(): def get_s3_client(): global s3_client if s3_client is None: + # print(hilite("S3 CLIENT IS NONE, CREATING IT!")) access_key = current_app.config["CSV_UPLOAD_BUCKET"]["access_key_id"] secret_key = current_app.config["CSV_UPLOAD_BUCKET"]["secret_access_key"] region = current_app.config["CSV_UPLOAD_BUCKET"]["region"] @@ -74,12 +76,15 @@ def get_s3_client(): region_name=region, ) s3_client = session.client("s3") + # else: + # print(hilite("S3 CLIENT ALREADY EXISTS, REUSING IT!")) return s3_client def get_s3_resource(): global s3_resource if s3_resource is None: + print(hilite("S3 RESOURCE IS NONE, CREATING IT!")) access_key = current_app.config["CSV_UPLOAD_BUCKET"]["access_key_id"] secret_key = current_app.config["CSV_UPLOAD_BUCKET"]["secret_access_key"] region = current_app.config["CSV_UPLOAD_BUCKET"]["region"] @@ -89,6 +94,8 @@ def get_s3_resource(): region_name=region, ) s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG) + else: + print(hilite("S3 RESOURCE ALREADY EXSITS, REUSING IT!")) return s3_resource diff --git a/app/config.py b/app/config.py index d3f2a5197..9ec37a71c 100644 --- a/app/config.py +++ b/app/config.py @@ -2,10 +2,12 @@ import json from datetime import datetime, timedelta from os import getenv, path +from boto3 import Session from celery.schedules import crontab from kombu import Exchange, Queue import notifications_utils +from app.clients import AWS_CLIENT_CONFIG from app.cloudfoundry_config import cloud_config @@ -51,6 +53,13 @@ class TaskNames(object): SCAN_FILE = "scan-file" +session = Session( + aws_access_key_id=getenv("CSV_AWS_ACCESS_KEY_ID"), + aws_secret_access_key=getenv("CSV_AWS_SECRET_ACCESS_KEY"), + region_name=getenv("CSV_AWS_REGION"), +) + + class Config(object): NOTIFY_APP_NAME = "api" DEFAULT_REDIS_EXPIRE_TIME = 4 * 24 * 60 * 60 @@ -166,6 +175,9 @@ class Config(object): current_minute = (datetime.now().minute + 1) % 60 + S3_CLIENT = session.client("s3") + S3_RESOURCE = session.resource("s3", config=AWS_CLIENT_CONFIG) + CELERY = { "worker_max_tasks_per_child": 500, "task_ignore_result": True, diff --git a/notifications_utils/s3.py b/notifications_utils/s3.py index 0a01f7493..46c89c68f 100644 --- a/notifications_utils/s3.py +++ b/notifications_utils/s3.py @@ -16,11 +16,32 @@ AWS_CLIENT_CONFIG = Config( use_fips_endpoint=True, ) +# Global variable +s3_resource = None + default_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID") default_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY") default_region = os.environ.get("AWS_REGION") +def get_s3_resource(): + global s3_resource + if s3_resource is None: + # print(hilite("S3 RESOURCE IS NONE, CREATING IT!")) + access_key = (default_access_key_id,) + secret_key = (default_secret_access_key,) + region = (default_region,) + session = Session( + aws_access_key_id=access_key, + aws_secret_access_key=secret_key, + region_name=region, + ) + s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG) + # else: + # print(hilite("S3 RESOURCE ALREADY EXSITS, REUSING IT!")) + return s3_resource + + def s3upload( filedata, region, @@ -32,12 +53,7 @@ def s3upload( access_key=default_access_key_id, secret_key=default_secret_access_key, ): - session = Session( - aws_access_key_id=access_key, - aws_secret_access_key=secret_key, - region_name=region, - ) - _s3 = session.resource("s3", config=AWS_CLIENT_CONFIG) + _s3 = get_s3_resource() key = _s3.Object(bucket_name, file_location) @@ -73,12 +89,7 @@ def s3download( secret_key=default_secret_access_key, ): try: - session = Session( - aws_access_key_id=access_key, - aws_secret_access_key=secret_key, - region_name=region, - ) - s3 = session.resource("s3", config=AWS_CLIENT_CONFIG) + s3 = get_s3_resource() key = s3.Object(bucket_name, filename) return key.get()["Body"] except botocore.exceptions.ClientError as error: From a527218638bf13db0e64783164a47a62a990f59f Mon Sep 17 00:00:00 2001 From: Kenneth Kehl <@kkehl@flexion.us> Date: Thu, 9 Jan 2025 08:07:20 -0800 Subject: [PATCH 2/5] fix tests --- tests/notifications_utils/test_s3.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/notifications_utils/test_s3.py b/tests/notifications_utils/test_s3.py index 46b863c4f..8208f6c41 100644 --- a/tests/notifications_utils/test_s3.py +++ b/tests/notifications_utils/test_s3.py @@ -13,7 +13,11 @@ content_type = "binary/octet-stream" def test_s3upload_save_file_to_bucket(mocker): - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_client = mocker.Mock() + mocked = mocker.patch( + "notification_utils.s3.get_s3_client", return_value=mock_s3_client + ) s3upload( filedata=contents, region=region, bucket_name=bucket, file_location=location ) @@ -27,7 +31,9 @@ def test_s3upload_save_file_to_bucket(mocker): def test_s3upload_save_file_to_bucket_with_contenttype(mocker): content_type = "image/png" - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_client = mocker.Mock() + mocked = mocker.patch("app.aws.s3.get_s3_client", return_value=mock_s3_client) s3upload( filedata=contents, region=region, @@ -44,7 +50,9 @@ def test_s3upload_save_file_to_bucket_with_contenttype(mocker): def test_s3upload_raises_exception(app, mocker): - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_client = mocker.Mock() + mocked = mocker.patch("app.aws.s3.get_s3_client", return_value=mock_s3_client) response = {"Error": {"Code": 500}} exception = botocore.exceptions.ClientError(response, "Bad exception") mocked.return_value.Object.return_value.put.side_effect = exception From fbd8643e74011af15b4cbfd225a0da6c27cf3e12 Mon Sep 17 00:00:00 2001 From: Kenneth Kehl <@kkehl@flexion.us> Date: Thu, 9 Jan 2025 08:33:42 -0800 Subject: [PATCH 3/5] fix tests --- tests/notifications_utils/test_s3.py | 47 +++++++++++++++++++++------- 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/tests/notifications_utils/test_s3.py b/tests/notifications_utils/test_s3.py index 8208f6c41..90efeb777 100644 --- a/tests/notifications_utils/test_s3.py +++ b/tests/notifications_utils/test_s3.py @@ -14,9 +14,9 @@ content_type = "binary/octet-stream" def test_s3upload_save_file_to_bucket(mocker): - mock_s3_client = mocker.Mock() + mock_s3_resource = mocker.Mock() mocked = mocker.patch( - "notification_utils.s3.get_s3_client", return_value=mock_s3_client + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource ) s3upload( filedata=contents, region=region, bucket_name=bucket, file_location=location @@ -32,8 +32,10 @@ def test_s3upload_save_file_to_bucket(mocker): def test_s3upload_save_file_to_bucket_with_contenttype(mocker): content_type = "image/png" - mock_s3_client = mocker.Mock() - mocked = mocker.patch("app.aws.s3.get_s3_client", return_value=mock_s3_client) + mock_s3_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) s3upload( filedata=contents, region=region, @@ -51,8 +53,10 @@ def test_s3upload_save_file_to_bucket_with_contenttype(mocker): def test_s3upload_raises_exception(app, mocker): - mock_s3_client = mocker.Mock() - mocked = mocker.patch("app.aws.s3.get_s3_client", return_value=mock_s3_client) + mock_s3_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) response = {"Error": {"Code": 500}} exception = botocore.exceptions.ClientError(response, "Bad exception") mocked.return_value.Object.return_value.put.side_effect = exception @@ -66,7 +70,12 @@ def test_s3upload_raises_exception(app, mocker): def test_s3upload_save_file_to_bucket_with_urlencoded_tags(mocker): - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) + s3upload( filedata=contents, region=region, @@ -82,7 +91,12 @@ def test_s3upload_save_file_to_bucket_with_urlencoded_tags(mocker): def test_s3upload_save_file_to_bucket_with_metadata(mocker): - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) + s3upload( filedata=contents, region=region, @@ -97,16 +111,25 @@ def test_s3upload_save_file_to_bucket_with_metadata(mocker): def test_s3download_gets_file(mocker): - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) + mocked_object = mocked.return_value.Object - mocked_get = mocked.return_value.Object.return_value.get + mocked_object.return_value.get.return_value = {"Body": mocker.Mock()} s3download("bucket", "location.file") mocked_object.assert_called_once_with("bucket", "location.file") - mocked_get.assert_called_once_with() def test_s3download_raises_on_error(mocker): - mocked = mocker.patch("notifications_utils.s3.Session.resource") + + mock_s3_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) + mocked.return_value.Object.side_effect = botocore.exceptions.ClientError( {"Error": {"Code": 404}}, "Bad exception", From 16bba7e4c43ee000427103a7d8dc31c6e31c0dbe Mon Sep 17 00:00:00 2001 From: Kenneth Kehl <@kkehl@flexion.us> Date: Thu, 9 Jan 2025 11:28:24 -0800 Subject: [PATCH 4/5] cleanup --- app/aws/s3.py | 4 ---- notifications_utils/s3.py | 23 ++++++++------------ tests/notifications_utils/test_s3.py | 32 +++++++++++++++++++++++++++- 3 files changed, 40 insertions(+), 19 deletions(-) diff --git a/app/aws/s3.py b/app/aws/s3.py index 01cd6692e..78fdf8d9a 100644 --- a/app/aws/s3.py +++ b/app/aws/s3.py @@ -10,7 +10,6 @@ from boto3 import Session from flask import current_app from app.clients import AWS_CLIENT_CONFIG -from app.utils import hilite from notifications_utils import aware_utcnow FILE_LOCATION_STRUCTURE = "service-{}-notify/{}.csv" @@ -84,7 +83,6 @@ def get_s3_client(): def get_s3_resource(): global s3_resource if s3_resource is None: - print(hilite("S3 RESOURCE IS NONE, CREATING IT!")) access_key = current_app.config["CSV_UPLOAD_BUCKET"]["access_key_id"] secret_key = current_app.config["CSV_UPLOAD_BUCKET"]["secret_access_key"] region = current_app.config["CSV_UPLOAD_BUCKET"]["region"] @@ -94,8 +92,6 @@ def get_s3_resource(): region_name=region, ) s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG) - else: - print(hilite("S3 RESOURCE ALREADY EXSITS, REUSING IT!")) return s3_resource diff --git a/notifications_utils/s3.py b/notifications_utils/s3.py index 46c89c68f..0cf7c4da7 100644 --- a/notifications_utils/s3.py +++ b/notifications_utils/s3.py @@ -13,11 +13,12 @@ AWS_CLIENT_CONFIG = Config( s3={ "addressing_style": "virtual", }, + max_pool_connections=50, use_fips_endpoint=True, ) # Global variable -s3_resource = None +noti_s3_resource = None default_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID") default_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY") @@ -25,21 +26,15 @@ default_region = os.environ.get("AWS_REGION") def get_s3_resource(): - global s3_resource - if s3_resource is None: - # print(hilite("S3 RESOURCE IS NONE, CREATING IT!")) - access_key = (default_access_key_id,) - secret_key = (default_secret_access_key,) - region = (default_region,) + global noti_s3_resource + if noti_s3_resource is None: session = Session( - aws_access_key_id=access_key, - aws_secret_access_key=secret_key, - region_name=region, + aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"), + aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"), + region_name=os.environ.get("AWS_REGION"), ) - s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG) - # else: - # print(hilite("S3 RESOURCE ALREADY EXSITS, REUSING IT!")) - return s3_resource + noti_s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG) + return noti_s3_resource def s3upload( diff --git a/tests/notifications_utils/test_s3.py b/tests/notifications_utils/test_s3.py index 90efeb777..6769fddd0 100644 --- a/tests/notifications_utils/test_s3.py +++ b/tests/notifications_utils/test_s3.py @@ -1,9 +1,16 @@ +from unittest.mock import MagicMock from urllib.parse import parse_qs import botocore import pytest -from notifications_utils.s3 import S3ObjectNotFound, s3download, s3upload +from notifications_utils.s3 import ( + AWS_CLIENT_CONFIG, + S3ObjectNotFound, + get_s3_resource, + s3download, + s3upload, +) contents = "some file data" region = "eu-west-1" @@ -110,6 +117,29 @@ def test_s3upload_save_file_to_bucket_with_metadata(mocker): assert metadata == {"status": "valid", "pages": "5"} +def test_get_s3_resource(mocker): + mock_session = mocker.patch("notifications_utils.s3.Session") + mock_current_app = mocker.patch("notifications_utils.s3.current_app") + sa_key = "sec" + sa_key = f"{sa_key}ret_access_key" + + mock_current_app.config = { + "CSV_UPLOAD_BUCKET": { + "access_key_id": "test_access_key", + sa_key: "test_s_key", + "region": "us-west-100", + } + } + mock_s3_resource = MagicMock() + mock_session.return_value.resource.return_value = mock_s3_resource + result = get_s3_resource() + + mock_session.return_value.resource.assert_called_once_with( + "s3", config=AWS_CLIENT_CONFIG + ) + assert result == mock_s3_resource + + def test_s3download_gets_file(mocker): mock_s3_resource = mocker.Mock() From 981fedaa01d0e4415a4e6dea448f7cbc3dbf4764 Mon Sep 17 00:00:00 2001 From: Kenneth Kehl <@kkehl@flexion.us> Date: Wed, 15 Jan 2025 07:42:59 -0800 Subject: [PATCH 5/5] code review feedback --- app/aws/s3.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/app/aws/s3.py b/app/aws/s3.py index 78fdf8d9a..c33366a2c 100644 --- a/app/aws/s3.py +++ b/app/aws/s3.py @@ -65,7 +65,6 @@ def clean_cache(): def get_s3_client(): global s3_client if s3_client is None: - # print(hilite("S3 CLIENT IS NONE, CREATING IT!")) access_key = current_app.config["CSV_UPLOAD_BUCKET"]["access_key_id"] secret_key = current_app.config["CSV_UPLOAD_BUCKET"]["secret_access_key"] region = current_app.config["CSV_UPLOAD_BUCKET"]["region"] @@ -75,8 +74,6 @@ def get_s3_client(): region_name=region, ) s3_client = session.client("s3") - # else: - # print(hilite("S3 CLIENT ALREADY EXISTS, REUSING IT!")) return s3_client