Enable Chalice multipart uploads

Author: syrusakbary

src/graphql_server/chalice/views.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from io import BytesIO
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 from chalice.app import Request, Response
@@ -23,6 +24,8 @@
 class ChaliceHTTPRequestAdapter(SyncHTTPRequestAdapter):
     def __init__(self, request: Request) -> None:
         self.request = request
+        self._post_data: Optional[dict[str, Union[str, bytes]]] = None
+        self._files: Optional[dict[str, Any]] = None
 
     @property
     def query_params(self) -> QueryParams:
@@ -42,11 +45,49 @@ def headers(self) -> Mapping[str, str]:
 
     @property
     def post_data(self) -> Mapping[str, Union[str, bytes]]:
-        raise NotImplementedError
+        if self._post_data is None:
+            self._parse_body()
+        return self._post_data or {}
 
     @property
     def files(self) -> Mapping[str, Any]:
-        raise NotImplementedError
+        if self._files is None:
+            self._parse_body()
+        return self._files or {}
+
+    def _parse_body(self) -> None:
+        self._post_data = {}
+        self._files = {}
+
+        content_type = self.content_type or ""
+
+        if "multipart/form-data" in content_type:
+            import cgi
+
+            fp = BytesIO(self.request.raw_body)
+            environ = {
+                "REQUEST_METHOD": "POST",
+                "CONTENT_TYPE": content_type,
+                "CONTENT_LENGTH": str(len(self.request.raw_body)),
+            }
+            fs = cgi.FieldStorage(fp=fp, environ=environ, keep_blank_values=True)
+            for key in fs.keys():
+                field = fs[key]
+                if isinstance(field, list):
+                    field = field[0]
+                if getattr(field, "filename", None):
+                    data = field.file.read()
+                    self._files[key] = BytesIO(data)
+                else:
+                    self._post_data[key] = field.value
+        elif "application/x-www-form-urlencoded" in content_type:
+            from urllib.parse import parse_qs
+
+            data = parse_qs(self.request.raw_body.decode())
+            self._post_data = {k: v[0] for k, v in data.items()}
+        else:
+            self._post_data = {}
+            self._files = {}
 
     @property
     def content_type(self) -> Optional[str]:
@@ -65,9 +106,11 @@ def __init__(
         graphiql: Optional[bool] = None,
         graphql_ide: Optional[GraphQL_IDE] = "graphiql",
         allow_queries_via_get: bool = True,
+        multipart_uploads_enabled: bool = False,
     ) -> None:
         self.allow_queries_via_get = allow_queries_via_get
         self.schema = schema
+        self.multipart_uploads_enabled = multipart_uploads_enabled
         if graphiql is not None:
             warnings.warn(
                 "The `graphiql` argument is deprecated in favor of `graphql_ide`",

src/tests/http/clients/chalice.py

@@ -7,6 +7,7 @@
 from typing_extensions import Literal
 
 from graphql import ExecutionResult
+from urllib3 import encode_multipart_formdata
 
 from chalice.app import Chalice
 from chalice.app import Request as ChaliceRequest
@@ -60,11 +61,14 @@ def __init__(
             graphiql=graphiql,
             graphql_ide=graphql_ide,
             allow_queries_via_get=allow_queries_via_get,
+            multipart_uploads_enabled=multipart_uploads_enabled,
         )
         view.result_override = result_override
 
         @self.app.route(
-            "/graphql", methods=["GET", "POST"], content_types=["application/json"]
+            "/graphql",
+            methods=["GET", "POST"],
+            content_types=["application/json", "multipart/form-data"],
         )
         def handle_graphql():
             assert self.app.current_request is not None
@@ -80,7 +84,7 @@ async def _graphql_request(
         headers: Optional[dict[str, str]] = None,
         extensions: Optional[dict[str, Any]] = None,
         **kwargs: Any,
-    ) -> Response:
+        ) -> Response:
         body = self._build_body(
             query=query,
             operation_name=operation_name,
@@ -90,27 +94,27 @@ async def _graphql_request(
             extensions=extensions,
         )
 
-        data: Union[dict[str, object], str, None] = None
-
-        if body and files:
-            body.update({name: (file, name) for name, file in files.items()})
-
         url = "/graphql"
+        headers = self._get_headers(method=method, headers=headers, files=files)
 
         if method == "get":
             body_encoded = urllib.parse.urlencode(body or {})
             url = f"{url}?{body_encoded}"
-        else:
-            if body:
-                data = body if files else dumps(body)
-            kwargs["body"] = data
+        elif body:
+            if files:
+                fields = {"operations": body["operations"], "map": body["map"]}
+                for filename, file in files.items():
+                    fields[filename] = (filename, file.read(), "text/plain")
+                data, content_type = encode_multipart_formdata(fields)
+                headers.update(
+                    {"Content-Type": content_type, "Content-Length": f"{len(data)}"}
+                )
+                kwargs["body"] = data
+            else:
+                kwargs["body"] = dumps(body)
 
         with Client(self.app) as client:
-            response = getattr(client.http, method)(
-                url,
-                headers=self._get_headers(method=method, headers=headers, files=files),
-                **kwargs,
-            )
+            response = getattr(client.http, method)(url, headers=headers, **kwargs)
 
         return Response(
             status_code=response.status_code,

src/tests/http/test_upload.py

@@ -10,23 +10,11 @@
 
 @pytest.fixture
 def http_client(http_client_class: type[HttpClient]) -> HttpClient:
-    with contextlib.suppress(ImportError):
-        from .clients.chalice import ChaliceHttpClient
-
-        if http_client_class is ChaliceHttpClient:
-            pytest.xfail(reason="Chalice does not support uploads")
-
     return http_client_class()
 
 
 @pytest.fixture
 def enabled_http_client(http_client_class: type[HttpClient]) -> HttpClient:
-    with contextlib.suppress(ImportError):
-        from .clients.chalice import ChaliceHttpClient
-
-        if http_client_class is ChaliceHttpClient:
-            pytest.xfail(reason="Chalice does not support uploads")
-
     return http_client_class(multipart_uploads_enabled=True)
 
 
@@ -46,6 +34,17 @@ async def test_multipart_uploads_are_disabled_by_default(http_client: HttpClient
     )
 
     assert response.status_code == 400
+
+    with contextlib.suppress(ImportError):
+        from .clients.chalice import ChaliceHttpClient
+
+        if isinstance(http_client, ChaliceHttpClient):
+            assert response.json == {
+                "Code": "BadRequestError",
+                "Message": "Unsupported content type",
+            }
+            return
+
     assert response.data == b"Unsupported content type"