Guess Content-Type

Author: trispera

src/tesk_core/filer_s3.py

@@ -4,6 +4,7 @@
 import re
 import botocore
 import boto3
+import mimetypes
 from tesk_core.transput import Transput, Type
 
 class S3Transput(Transput):
@@ -50,10 +51,24 @@ def download_file(self):
         os.makedirs(basedir, exist_ok=True)
         return self.get_s3_file(self.path, self.file_path)
 
+    def get_content_type(self):
+        # Guess content type based on filename; fallback to binary stream
+        mime, encoding = mimetypes.guess_type(self.path)
+        if mime is None:
+            return 'application/octet-stream'
+        elif mime.startswith('text/') or mime in ('application/json', 'application/xml', 'application/javascript'):
+            mime = f'{mime}; charset=utf-8'
+        return mime
+
+
     def upload_file(self):
         logging.debug('Uploading s3 object: "%s" Target: %s', self.path,  self.bucket + "/" + self.file_path)
+        content_type = self.get_content_type()
+        logging.debug('Guessed Content-Type: %s for file: %s', content_type, self.path)
         try:
-            self.bucket_obj.upload_file(Filename=self.path, Key=self.file_path)
+            # Pass ContentType via ExtraArgs so the object is uploaded with the right MIME type
+            self.bucket_obj.upload_file(Filename=self.path, Key=self.file_path,
+                                       ExtraArgs={'ContentType': content_type})
         except (botocore.exceptions.ClientError,  OSError) as err:
             logging.error("File upload failed for '%s'", self.bucket + "/" + self.file_path)
             logging.error(err)

tests/test_s3_filer.py

@@ -111,6 +111,27 @@ def test_s3_upload_file( moto_boto, path, url, ftype, expected,fs, caplog):
         otherwise an exception will be raised.
         '''
         assert client.Object('tesk', 'folder/file.txt').load() == None
+        # Check the ContentType metadata is set correctly for text files
+        head = client.meta.client.head_object(Bucket=trans.bucket, Key=trans.file_path)
+        assert head['ContentType'] == 'text/plain'
+
+
+@pytest.mark.parametrize("filename, url, expected_content", [
+    ("file.txt", "s3://tesk/folder/file.txt", "text/plain"),
+    ("file.zip", "s3://tesk/folder/file.zip", "application/zip"),
+])
+def test_s3_upload_file_content_type(moto_boto, filename, url, expected_content, fs):
+    """
+    Ensure uploaded objects have correct Content-Type metadata based on file extension
+    """
+    fs.create_file(f"/home/user/filer_test/{filename}")
+    client = boto3.resource('s3', endpoint_url="http://s3.amazonaws.com")
+    trans = S3Transput(f"/home/user/filer_test/{filename}", url, "FILE")
+    trans.bucket_obj = client.Bucket(trans.bucket)
+    assert trans.upload_file() == 0
+    head = client.meta.client.head_object(Bucket=trans.bucket, Key=trans.file_path)
+    assert head['ContentType'] == expected_content
+
 
 
 
@@ -133,8 +154,8 @@ def test_s3_upload_directory(path, url, ftype, expected, moto_boto, caplog):
         Checking if the file was uploaded, if the object is found load() method will return None 
         otherwise an exception will be raised.
         '''
-        assert client.Object('tesk', 'folder1/folder2/test_filer.py').load() == None
-
+        assert client.Object('tesk', 'folder1/folder2/test_filer.py').load() == None        head = client.meta.client.head_object(Bucket=trans.bucket, Key='folder1/folder2/test_filer.py')
+        assert head['ContentType'].startswith('text/')
 def test_upload_directory_for_unknown_file_type(moto_boto, fs, monkeypatch, caplog):
     """
         Checking whether an exception is raised when the object type is neither file or directory