feat(beetmover): allow not failing on unknown mimetypes

This is needed to support work from mozilla-releng#1141 where we're uploading all artifacts from upstream tasks, and can't reasonably predict all possible file types and extensions.

This also fixes a minor bug in the `upload_to_s3` tests, where they implicitly rely on `setup_mimetypes()` being called elsewhere to pass (which means they fail mysteriously if run independently).

Author: bhearsum

beetmoverscript/src/beetmoverscript/gcloud.py

@@ -102,11 +102,14 @@ def setup_gcs_credentials(raw_creds):
     os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = fp.name
 
 
-async def upload_to_gcs(context, target_path, path, expiry=None):
+async def upload_to_gcs(context, target_path, path, expiry=None, fail_on_unknown_mimetype=True):
     product = get_product_name(context.task, context.config)
     mime_type = mimetypes.guess_type(path)[0]
     if not mime_type:
-        raise ScriptWorkerTaskException("Unable to discover valid mime-type for path ({}), mimetypes.guess_type() returned {}".format(path, mime_type))
+        if fail_on_unknown_mimetype:
+            raise ScriptWorkerTaskException("Unable to discover valid mime-type for path ({}), mimetypes.guess_type() returned {}".format(path, mime_type))
+        else:
+            mime_type = "application/octet-stream"
     bucket_name = get_bucket_name(context, product, "gcloud")
 
     bucket = Bucket(context.gcs_client, name=bucket_name)

beetmoverscript/src/beetmoverscript/script.py

@@ -650,7 +650,7 @@ def enrich_balrog_manifest(context, locale):
 
 
 # retry_upload {{{1
-async def retry_upload(context, destinations, path, expiry=None):
+async def retry_upload(context, destinations, path, expiry=None, fail_on_unknown_mimetype=True):
     """Manage upload of `path` to `destinations`."""
     cloud_uploads = {key: [] for key in context.config["clouds"]}
 
@@ -661,11 +661,17 @@ async def retry_upload(context, destinations, path, expiry=None):
         # S3 upload
         enabled = is_cloud_enabled(context.config, "aws", context.resource)
         if enabled is True or (enabled == "buildhub-only" and path.endswith("buildhub.json")):
-            cloud_uploads["aws"].append(asyncio.ensure_future(upload_to_s3(context=context, s3_key=dest, path=path)))
+            cloud_uploads["aws"].append(
+                asyncio.ensure_future(upload_to_s3(context=context, s3_key=dest, path=path, fail_on_unknown_mimetype=fail_on_unknown_mimetype))
+            )
 
         # GCS upload
         if is_cloud_enabled(context.config, "gcloud", context.resource):
-            cloud_uploads["gcloud"].append(asyncio.ensure_future(upload_to_gcs(context=context, target_path=dest, path=path, expiry=expiry)))
+            cloud_uploads["gcloud"].append(
+                asyncio.ensure_future(
+                    upload_to_gcs(context=context, target_path=dest, path=path, expiry=expiry, fail_on_unknown_mimetype=fail_on_unknown_mimetype)
+                )
+            )
 
     await await_and_raise_uploads(cloud_uploads, context.config["clouds"], context.resource)
 
@@ -685,11 +691,15 @@ async def put(context, url, headers, abs_filename, session=None):
 
 
 # upload_to_s3 {{{1
-async def upload_to_s3(context, s3_key, path):
+async def upload_to_s3(context, s3_key, path, fail_on_unknown_mimetype=True):
     product = get_product_name(context.task, context.config)
     mime_type = mimetypes.guess_type(path)[0]
     if not mime_type:
-        raise ScriptWorkerTaskException("Unable to discover valid mime-type for path ({}), mimetypes.guess_type() returned {}".format(path, mime_type))
+        if fail_on_unknown_mimetype:
+            raise ScriptWorkerTaskException("Unable to discover valid mime-type for path ({}), mimetypes.guess_type() returned {}".format(path, mime_type))
+        else:
+            mime_type = "application/octet-stream"
+
     api_kwargs = {
         "Bucket": get_bucket_name(context, product, "aws"),
         "Key": s3_key,
@@ -699,6 +709,7 @@ async def upload_to_s3(context, s3_key, path):
         "Content-Type": mime_type,
         "Cache-Control": "public, max-age=%d" % CACHE_CONTROL_MAXAGE,
     }
+
     creds = get_credentials(context, "aws")
     s3 = boto3.client("s3", aws_access_key_id=creds["id"], aws_secret_access_key=creds["key"])
     url = s3.generate_presigned_url("put_object", api_kwargs, ExpiresIn=1800, HttpMethod="PUT")

beetmoverscript/tests/test_gcloud.py

@@ -147,21 +147,23 @@ def close(self):
 
 
 @pytest.mark.parametrize(
-    "path,expiry,exists,raise_class",
+    "path,expiry,exists,expected_mimetype,raise_class,fail_on_unknown_mimetype",
     (
         # Raise when can't find a mimetype
-        ("foo/nomimetype", None, False, ScriptWorkerTaskException),
+        ("foo/nomimetype", None, False, None, ScriptWorkerTaskException, True),
+        # Don't raise when can't find a mimetype
+        ("foo/nomimetype", None, False, "application/octet-stream", None, False),
         # No expiration given
-        ("foo/target.zip", None, False, None),
+        ("foo/target.zip", None, False, "application/zip", None, True),
         # With expiration
-        ("foo/target.zip", datetime.now().isoformat(), False, None),
+        ("foo/target.zip", datetime.now().isoformat(), False, "application/zip", None, True),
         # With existing file
-        ("foo/target.zip", None, True, None),
+        ("foo/target.zip", None, True, "application/zip", None, True),
     ),
-    ids=["no mimetype", "no expiry", "with expiry", "with existing file"],
+    ids=["no mimetype", "no mimetype no fail", "no expiry", "with expiry", "with existing file"],
 )
 @pytest.mark.asyncio
-async def test_upload_to_gcs(context, monkeypatch, path, expiry, exists, raise_class):
+async def test_upload_to_gcs(context, monkeypatch, path, expiry, exists, expected_mimetype, raise_class, fail_on_unknown_mimetype):
     context.gcs_client = FakeClient()
     blob = FakeClient.FakeBlob()
     blob._exists = exists
@@ -181,13 +183,15 @@ async def test_upload_to_gcs(context, monkeypatch, path, expiry, exists, raise_c
                 target_path="path/target",
                 path=path,
                 expiry=expiry,
+                fail_on_unknown_mimetype=fail_on_unknown_mimetype,
             )
     else:
         await beetmoverscript.gcloud.upload_to_gcs(
             context=context,
             target_path="path/target",
             path=path,
             expiry=expiry,
+            fail_on_unknown_mimetype=fail_on_unknown_mimetype,
         )
 
         if expiry:
@@ -199,6 +203,8 @@ async def test_upload_to_gcs(context, monkeypatch, path, expiry, exists, raise_c
         else:
             log_warn.assert_not_called()
 
+        assert blob.upload_from_filename.call_args[1].get("content_type") == expected_mimetype
+
 
 @pytest.mark.parametrize(
     "candidate_blobs,release_blobs,partner_match,raises",

beetmoverscript/tests/test_script.py

@@ -208,6 +208,7 @@ async def test_retry_upload(context, mocker):
 # upload_to_s3 {{{1
 @pytest.mark.asyncio
 async def test_upload_to_s3(context, mocker):
+    setup_mimetypes()
     context.release_props["appName"] = "fake"
     mocker.patch.object(beetmoverscript.script, "retry_async", new=noop_async)
     mocker.patch.object(beetmoverscript.script, "boto3")
@@ -216,13 +217,41 @@ async def test_upload_to_s3(context, mocker):
 
 @pytest.mark.asyncio
 async def test_upload_to_s3_raises(context, mocker):
+    setup_mimetypes()
     context.release_props["appName"] = "fake"
     mocker.patch.object(beetmoverscript.script, "retry_async", new=noop_async)
     mocker.patch.object(beetmoverscript.script, "boto3")
     with pytest.raises(ScriptWorkerTaskException):
         await beetmoverscript.script.upload_to_s3(context, "foo", "mime.invalid")
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "fail_on_unknown_mime_type",
+    (
+        pytest.param(
+            True,
+            id="fail_on_unknown_mime_type",
+        ),
+        pytest.param(
+            False,
+            id="dont_fail_on_unknown_mime_type",
+        ),
+    ),
+)
+async def test_upload_to_s3_fail_on_missing_mime_type(context, mocker, fail_on_unknown_mime_type):
+    setup_mimetypes()
+    context.release_props["appName"] = "fake"
+    mocked_retry_async = mocker.patch.object(beetmoverscript.script, "retry_async")
+    mocker.patch.object(beetmoverscript.script, "boto3")
+    if fail_on_unknown_mime_type:
+        with pytest.raises(ScriptWorkerTaskException):
+            await beetmoverscript.script.upload_to_s3(context, "foo", "mime.invalid", fail_on_unknown_mime_type)
+    else:
+        await beetmoverscript.script.upload_to_s3(context, "foo", "mime.invalid", fail_on_unknown_mime_type)
+        assert mocked_retry_async.call_args[1]["args"][2].get("Content-Type") == "application/octet-stream"
+
+
 @pytest.fixture
 def restore_buildhub_file():
     original_location = "tests/test_work_dir/cot/eSzfNqMZT_mSiQQXu8hyqg/public/build/buildhub.json"