diff --git a/app/config.py b/app/config.py index 167294a6c..13d9daf9d 100644 --- a/app/config.py +++ b/app/config.py @@ -2,10 +2,12 @@ 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..0cf7c4da7 100644 --- a/notifications_utils/s3.py +++ b/notifications_utils/s3.py @@ -13,14 +13,30 @@ s3={ "addressing_style": "virtual", }, + max_pool_connections=50, use_fips_endpoint=True, ) +# Global variable +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") default_region = os.environ.get("AWS_REGION") +def get_s3_resource(): + global noti_s3_resource + if noti_s3_resource is None: + session = Session( + 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"), + ) + noti_s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG) + return noti_s3_resource + + def s3upload( filedata, region, @@ -32,12 +48,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 +84,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: diff --git a/tests/notifications_utils/test_s3.py b/tests/notifications_utils/test_s3.py index 46b863c4f..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" @@ -13,7 +20,11 @@ def test_s3upload_save_file_to_bucket(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, bucket_name=bucket, file_location=location ) @@ -27,7 +38,11 @@ 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_resource = mocker.Mock() + mocked = mocker.patch( + "notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource + ) s3upload( filedata=contents, region=region, @@ -44,7 +59,11 @@ 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_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 @@ -58,7 +77,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, @@ -74,7 +98,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, @@ -88,17 +117,49 @@ 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): - 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",