Merge pull request #2047 from alphagov/store-s3-metadata

Store info about files as S3 metadata

app/main/s3_client.py

@@ -15,6 +15,17 @@ def get_s3_object(bucket_name, filename):
     return s3.Object(bucket_name, filename)
 
 
+def get_csv_location(service_id, upload_id):
+    return (
+        current_app.config['CSV_UPLOAD_BUCKET_NAME'],
+        FILE_LOCATION_STRUCTURE.format(service_id, upload_id),
+    )
+
+
+def get_csv_upload(service_id, upload_id):
+    return get_s3_object(*get_csv_location(service_id, upload_id))
+
+
 def delete_s3_object(filename):
     bucket_name = current_app.config['LOGO_UPLOAD_BUCKET_NAME']
     get_s3_object(bucket_name, filename).delete()
@@ -39,20 +50,20 @@ def get_temp_truncated_filename(filename, user_id):
 
 def s3upload(service_id, filedata, region):
     upload_id = str(uuid.uuid4())
-    upload_file_name = FILE_LOCATION_STRUCTURE.format(service_id, upload_id)
+    bucket_name, file_location = get_csv_location(service_id, upload_id)
-    utils_s3upload(filedata=filedata['data'],
+    utils_s3upload(
-                   region=region,
+        filedata=filedata['data'],
-                   bucket_name=current_app.config['CSV_UPLOAD_BUCKET_NAME'],
+        region=region,
-                   file_location=upload_file_name)
+        bucket_name=bucket_name,
+        file_location=file_location,
+    )
     return upload_id
 
 
 def s3download(service_id, upload_id):
     contents = ''
     try:
-        bucket_name = current_app.config['CSV_UPLOAD_BUCKET_NAME']
+        key = get_csv_upload(service_id, upload_id)
-        upload_file_name = FILE_LOCATION_STRUCTURE.format(service_id, upload_id)
-        key = get_s3_object(bucket_name, upload_file_name)
         contents = key.get()['Body'].read().decode('utf-8')
     except botocore.exceptions.ClientError as e:
         current_app.logger.error("Unable to download s3 file {}".format(
@@ -121,3 +132,16 @@ def delete_temp_file(filename):
         raise ValueError('Not a temp file: {}'.format(filename))
 
     delete_s3_object(filename)
+
+
+def set_metadata_on_csv_upload(service_id, upload_id, **kwargs):
+    get_csv_upload(
+        service_id, upload_id
+    ).copy_from(
+        CopySource='{}/{}'.format(*get_csv_location(service_id, upload_id)),
+        ServerSideEncryption='AES256',
+        Metadata={
+            key: str(value) for key, value in kwargs.items()
+        },
+        MetadataDirective='REPLACE',
+    )

app/main/views/send.py

@@ -38,7 +38,7 @@ from app.main.forms import (
     SetSenderForm,
     get_placeholder_form_instance,
 )
-from app.main.s3_client import s3download, s3upload
+from app.main.s3_client import s3download, s3upload, set_metadata_on_csv_upload
 from app.template_previews import TemplatePreview, get_page_count_for_letter
 from app.utils import (
     Spreadsheet,
@@ -549,6 +549,13 @@ def _check_messages(service_id, template_id, upload_id, preview_row, letters_as_
         session['file_uploads'][upload_id]['notification_count'] = len(recipients)
         session['file_uploads'][upload_id]['template_id'] = str(template_id)
         session['file_uploads'][upload_id]['valid'] = True
+        set_metadata_on_csv_upload(
+            service_id,
+            upload_id,
+            notification_count=len(recipients),
+            template_id=str(template_id),
+            valid=True,
+        )
     else:
         session['file_uploads'].pop(upload_id)
 

tests/app/main/test_s3_client.py

@@ -1,5 +1,5 @@
 from collections import namedtuple
-from unittest.mock import call
+from unittest.mock import Mock, call
 
 import pytest
 
@@ -10,6 +10,7 @@ from app.main.s3_client import (
     delete_temp_files_created_by,
     get_temp_truncated_filename,
     persist_logo,
+    set_metadata_on_csv_upload,
     upload_logo,
 )
 
@@ -94,3 +95,21 @@ def test_does_not_delete_non_temp_file(client, mocker, fake_uuid):
 
     assert mocked_delete_s3_object.called_with_args(filename)
     assert str(error.value) == 'Not a temp file: {}'.format(filename)
+
+
+def test_sets_metadata(client, mocker):
+    mocked_s3_object = Mock()
+    mocked_get_s3_object = mocker.patch(
+        'app.main.s3_client.get_csv_upload',
+        return_value=mocked_s3_object,
+    )
+
+    set_metadata_on_csv_upload('1234', '5678', foo='bar', baz=True)
+
+    mocked_get_s3_object.assert_called_once_with('1234', '5678')
+    mocked_s3_object.copy_from.assert_called_once_with(
+        CopySource='test-notifications-csv-upload/service-1234-notify/5678.csv',
+        Metadata={'baz': 'True', 'foo': 'bar'},
+        MetadataDirective='REPLACE',
+        ServerSideEncryption='AES256',
+    )

tests/app/main/views/test_send.py

@@ -565,6 +565,7 @@ def test_upload_valid_csv_shows_preview_and_table(
     mock_get_service_template_with_placeholders,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
+    mock_s3_set_metadata,
     fake_uuid,
     extra_args,
     expected_link_in_first_row,
@@ -592,6 +593,14 @@ def test_upload_valid_csv_shows_preview_and_table(
         **extra_args
     )
 
+    mock_s3_set_metadata.assert_called_once_with(
+        SERVICE_ONE_ID,
+        fake_uuid,
+        notification_count=3,
+        template_id=fake_uuid,
+        valid=True,
+    )
+
     assert page.h1.text.strip() == 'Preview of Two week reminder'
     assert page.select_one('.sms-message-recipient').text.strip() == expected_recipient
     assert page.select_one('.sms-message-wrapper').text.strip() == expected_message
@@ -702,6 +711,7 @@ def test_404_for_previewing_a_row_out_of_range(
     mock_get_service_template_with_placeholders,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
+    mock_s3_set_metadata,
     fake_uuid,
     row_index,
     expected_status,
@@ -1464,6 +1474,7 @@ def test_upload_csvfile_with_valid_phone_shows_all_numbers(
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
     mock_get_live_service,
+    mock_s3_set_metadata,
     service_one,
     fake_uuid,
     mock_s3_upload,
@@ -1489,6 +1500,14 @@ def test_upload_csvfile_with_valid_phone_shows_all_numbers(
         assert sess['file_uploads'][fake_uuid]['template_id'] == fake_uuid
         assert sess['file_uploads'][fake_uuid]['valid'] is True
 
+    mock_s3_set_metadata.assert_called_once_with(
+        SERVICE_ONE_ID,
+        fake_uuid,
+        notification_count=53,
+        template_id=fake_uuid,
+        valid=True,
+    )
+
     content = response.get_data(as_text=True)
     assert response.status_code == 200
     assert '07700 900701' in content
@@ -1543,6 +1562,7 @@ def test_test_message_can_only_be_sent_now(
     mock_s3_download,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
+    mock_s3_set_metadata,
     fake_uuid
 ):
     with logged_in_client.session_transaction() as session:
@@ -1957,6 +1977,7 @@ def test_check_messages_back_link(
     mock_has_permissions,
     mock_get_detailed_service_for_today,
     mock_s3_download,
+    mock_s3_set_metadata,
     fake_uuid,
     mocker,
     template_mock,
@@ -2138,6 +2159,7 @@ def test_check_messages_shows_trial_mode_error_for_letters(
     mock_has_permissions,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
+    mock_s3_set_metadata,
     fake_uuid,
     mocker,
     service_mock,
@@ -2265,6 +2287,7 @@ def test_generate_test_letter_doesnt_block_in_trial_mode(
     fake_uuid,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
+    mock_s3_set_metadata,
 ):
 
     mocker.patch('app.main.views.send.s3download', return_value="""
@@ -2712,6 +2735,7 @@ def test_reply_to_is_previewed_if_chosen(
     mocker,
     mock_get_service_email_template,
     mock_s3_download,
+    mock_s3_set_metadata,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
     get_default_reply_to_email_address,
@@ -2760,6 +2784,7 @@ def test_sms_sender_is_previewed(
     mocker,
     mock_get_service_template,
     mock_s3_download,
+    mock_s3_set_metadata,
     mock_get_users_by_service,
     mock_get_detailed_service_for_today,
     get_default_sms_sender,

tests/conftest.py

@@ -2027,6 +2027,11 @@ def mock_s3_download(mocker, content=None):
     return mocker.patch('app.main.views.send.s3download', side_effect=_download)
 
 
+@pytest.fixture(scope='function')
+def mock_s3_set_metadata(mocker, content=None):
+    return mocker.patch('app.main.views.send.set_metadata_on_csv_upload')
+
+
 @pytest.fixture(scope='function')
 def sample_invite(mocker, service_one, status='pending'):
     id_ = str(generate_uuid())