Add the presigned url functionality to celery (#2496)

This change:
- adds the presigned S3 url functionality to celery
- adds a db migration to update the reports table (the presigned url is really long so I had to make the url column bigger, and I added the language column while I was at it)
- adds unit tests for the presigned url function
- renames the last migration file where the filename didn't match the revision

Author: smcmurtry

app/aws/s3.py

@@ -12,6 +12,7 @@
 
 FILE_LOCATION_STRUCTURE = "service-{}-notify/{}.csv"
 REPORTS_FILE_LOCATION_STRUCTURE = "service-{}/{}.csv"
+THREE_DAYS_IN_SECONDS = 3 * 24 * 60 * 60
 
 
 def get_s3_file(bucket_name, file_location):
@@ -150,14 +151,41 @@ def get_report_location(service_id, report_id):
 
 
 def upload_report_to_s3(service_id: str, report_id: str, file_data: bytes) -> str:
-    location = get_report_location(service_id, report_id)
+    object_key = get_report_location(service_id, report_id)
     utils_s3upload(
         filedata=file_data,
         region=current_app.config["AWS_REGION"],
         bucket_name=current_app.config["REPORTS_BUCKET_NAME"],
-        file_location=location,
+        file_location=object_key,
+    )
+    url = generate_presigned_url(
+        bucket_name=current_app.config["REPORTS_BUCKET_NAME"],
+        object_key=object_key,
+        expiration=THREE_DAYS_IN_SECONDS,
     )
-    # todo: generate a presigned url
-    # https://docs.aws.amazon.com/AmazonS3/latest/userguide/ShareObjectPreSignedURL.html
-    url = f"https://{current_app.config['REPORTS_BUCKET_NAME']}.s3.{current_app.config["AWS_REGION"]}.amazonaws.com/{location}"
     return url
+
+
+def generate_presigned_url(bucket_name: str, object_key: str, expiration: int = 3600) -> str:
+    """
+    Generate a presigned URL to share an S3 object
+
+    :param bucket_name: string
+    :param object_key: string
+    :param expiration: Time in seconds for the presigned URL to remain valid
+    :return: Presigned URL as string. If error, returns None.
+
+    Docs: https://docs.aws.amazon.com/AmazonS3/latest/userguide/ShareObjectPreSignedURL.html
+    """
+    s3_client = client("s3", current_app.config["AWS_REGION"])
+    try:
+        response = s3_client.generate_presigned_url(
+            "get_object",
+            Params={"Bucket": bucket_name, "Key": object_key},
+            ExpiresIn=expiration,
+        )
+    except botocore.exceptions.ClientError as e:
+        current_app.logger.error(e)
+        return ""
+
+    return response

app/models.py

@@ -2731,8 +2731,9 @@ class Report(BaseModel):
         nullable=False,
     )
     job_id = db.Column(UUID(as_uuid=True), db.ForeignKey("jobs.id"), nullable=True)  # only set if report is for a bulk job
-    url = db.Column(db.String(255), nullable=True)  # url to download the report from s3
+    url = db.Column(db.String(800), nullable=True)  # url to download the report from s3
     status = db.Column(db.String(255), nullable=False)
+    language = db.Column(db.String(2), nullable=True)
 
     def serialize(self):
         return {

migrations/versions/0476_add_created_by_ids.py renamed to migrations/versions/0477_add_created_by_ids.py

@@ -0,0 +1,28 @@
+"""
+
+Revision ID: 0478_update_reports_table
+Revises: 0477_add_created_by_ids
+Create Date: 2025-04-07 23:44:02.031724
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+revision = '0478_update_reports_table'
+down_revision = '0477_add_created_by_ids'
+
+
+def upgrade():
+    op.add_column('reports', sa.Column('language', sa.String(length=2), nullable=True))
+    op.alter_column('reports', 'url',
+               existing_type=sa.VARCHAR(length=255),
+               type_=sa.String(length=800),
+               existing_nullable=True)
+
+
+def downgrade():
+    op.alter_column('reports', 'url',
+               existing_type=sa.String(length=800),
+               type_=sa.VARCHAR(length=255),
+               existing_nullable=True)
+    op.drop_column('reports', 'language')

migrations/versions/0478_update_reports_table.py

@@ -4,12 +4,14 @@
 
 import pytest
 import pytz
+from botocore.exceptions import ClientError
 from flask import current_app
 from freezegun import freeze_time
 from tests.app.conftest import datetime_in_past
 
 from app.aws.s3 import (
     filter_s3_bucket_objects_within_date_range,
+    generate_presigned_url,
     get_list_of_files_by_suffix,
     get_s3_bucket_objects,
     get_s3_file,
@@ -247,20 +249,62 @@ def test_remove_jobs_from_s3(notify_api, mocker):
 
 def test_upload_report_to_s3(notify_api, mocker):
     utils_mock = mocker.patch("app.aws.s3.utils_s3upload")
+    presigned_url_mock = mocker.patch("app.aws.s3.generate_presigned_url")
     service_id = uuid.uuid4()
     report_id = uuid.uuid4()
     csv_data = b"foo"
     file_location = f"service-{service_id}/{report_id}.csv"
-    url = upload_report_to_s3(service_id, report_id, csv_data)
+    upload_report_to_s3(service_id, report_id, csv_data)
 
     utils_mock.assert_called_once_with(
         filedata=csv_data,
         region=notify_api.config["AWS_REGION"],
         bucket_name=current_app.config["REPORTS_BUCKET_NAME"],
         file_location=file_location,
     )
-    expected_url = (
-        f"https://{current_app.config['REPORTS_BUCKET_NAME']}.s3.{current_app.config["AWS_REGION"]}.amazonaws.com/{file_location}"
+    presigned_url_mock.assert_called_once_with(
+        bucket_name=current_app.config["REPORTS_BUCKET_NAME"],
+        object_key=file_location,
+        expiration=259200,  # 3 days
+    )
+
+
+def test_generate_presigned_url_success(notify_api, mocker):
+    s3_client_mock = mocker.patch("app.aws.s3.client")
+    s3_client_mock.return_value.generate_presigned_url.return_value = "https://example.com/presigned-url"
+
+    bucket_name = "test-bucket"
+    object_key = "test-object"
+    expiration = 3600
+
+    url = generate_presigned_url(bucket_name, object_key, expiration)
+
+    s3_client_mock.assert_called_once_with("s3", notify_api.config["AWS_REGION"])
+    s3_client_mock.return_value.generate_presigned_url.assert_called_once_with(
+        "get_object",
+        Params={"Bucket": bucket_name, "Key": object_key},
+        ExpiresIn=expiration,
     )
+    assert url == "https://example.com/presigned-url"
 
-    assert url == expected_url
+
+def test_generate_presigned_url_error(notify_api, mocker):
+    s3_client_mock = mocker.patch("app.aws.s3.client")
+    s3_client_mock.return_value.generate_presigned_url.side_effect = ClientError(
+        error_response={"Error": {"Code": "403", "Message": "Forbidden"}},
+        operation_name="GeneratePresignedUrl",
+    )
+
+    bucket_name = "test-bucket"
+    object_key = "test-object"
+    expiration = 3600
+
+    url = generate_presigned_url(bucket_name, object_key, expiration)
+
+    s3_client_mock.assert_called_once_with("s3", notify_api.config["AWS_REGION"])
+    s3_client_mock.return_value.generate_presigned_url.assert_called_once_with(
+        "get_object",
+        Params={"Bucket": bucket_name, "Key": object_key},
+        ExpiresIn=expiration,
+    )
+    assert not url

tests/app/aws/test_s3.py

@@ -64,7 +64,7 @@ def test_expire_api_key_should_update_the_api_key_and_create_history_record(noti
 
 def test_last_used_should_update_the_api_key_and_not_create_history_record(notify_api, sample_api_key):
     last_used = datetime.utcnow()
-    update_last_used_api_key(api_key_id=sample_api_key.id, last_used=last_used)
+    update_last_used_api_key(sample_api_key.id, last_used)
     all_api_keys = get_model_api_keys(service_id=sample_api_key.service_id)
     assert len(all_api_keys) == 1
     assert all_api_keys[0].last_used_timestamp == last_used