Set headers on S3 object directly

developmentseed · Sep 4, 2024 · c85f7b6 · c85f7b6
1 parent cbc085b
commit c85f7b6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/src/asgi_s3_response_middleware/__init__.py b/src/asgi_s3_response_middleware/__init__.py
@@ -1,3 +1,3 @@
 from .middleware import S3ResponseMiddleware  # noqa: F401
 
-__version__ = "0.0.2"
+__version__ = "0.0.3"
diff --git a/src/asgi_s3_response_middleware/middleware.py b/src/asgi_s3_response_middleware/middleware.py
@@ -27,9 +27,11 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
             return await self.app(scope, receive, send)
 
         s3_key = None  # Persist key across header & body messages
+        s3_headers = {}  # Persist headers across header & body messages
 
         async def send_with_s3_response(message: Message):
             nonlocal s3_key
+            nonlocal s3_headers
 
             if message["type"] == "http.response.start":
                 headers = MutableHeaders(scope=message)
@@ -43,25 +45,23 @@ async def send_with_s3_response(message: Message):
                         self.size_threshold,
                     )
                     s3_key = self.key_generator()
+                    s3_headers = {
+                        f"Content{key.title()}": headers[f"content-{key}"]
+                        for key in [
+                            "type",
+                            "encoding",
+                            "language",
+                            "disposition",
+                        ]
+                        if f"content-{key}" in headers
+                    }
                     message["status"] = 303
                     headers["content-length"] = "0"
                     headers["location"] = self.s3_client.generate_presigned_url(
                         "get_object",
                         Params={
                             "Bucket": self.s3_bucket_name,
                             "Key": s3_key,
-                            **{
-                                f"ResponseContent{key.title()}": headers[
-                                    f"content-{key}"
-                                ]
-                                for key in [
-                                    "type",
-                                    "encoding",
-                                    "language",
-                                    "disposition",
-                                ]
-                                if f"content-{key}" in headers
-                            },
                         },
                         ExpiresIn=self.url_expiry,
                     )
@@ -78,6 +78,7 @@ async def send_with_s3_response(message: Message):
                         Bucket=self.s3_bucket_name,
                         Key=s3_key,
                         Body=message["body"],
+                        **s3_headers,
                     )
                     message = {"type": "http.response.body", "body": b""}