Skip to content

Commit 4fa1553

Browse files
authored
Refactor transports (#557)
* Refactor requests transport to be similar to aiohttp and httpx * Removing _prepare_batch_request to avoid code duplication
1 parent b3789ef commit 4fa1553

File tree

3 files changed

+188
-200
lines changed

3 files changed

+188
-200
lines changed

gql/transport/aiohttp.py

Lines changed: 16 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ async def connect(self) -> None:
127127

128128
# Adding custom parameters passed from init
129129
if self.client_session_args:
130-
client_session_args.update(self.client_session_args) # type: ignore
130+
client_session_args.update(self.client_session_args)
131131

132132
log.debug("Connecting transport")
133133

@@ -164,36 +164,22 @@ async def close(self) -> None:
164164

165165
self.session = None
166166

167-
def _prepare_batch_request(
168-
self,
169-
reqs: List[GraphQLRequest],
170-
extra_args: Optional[Dict[str, Any]] = None,
171-
) -> Dict[str, Any]:
172-
173-
payload = [req.payload for req in reqs]
174-
175-
post_args = {"json": payload}
176-
177-
# Log the payload
178-
if log.isEnabledFor(logging.DEBUG):
179-
log.debug(">>> %s", self.json_serialize(post_args["json"]))
180-
181-
# Pass post_args to aiohttp post method
182-
if extra_args:
183-
post_args.update(extra_args)
184-
185-
return post_args
186-
187167
def _prepare_request(
188168
self,
189-
request: GraphQLRequest,
169+
request: Union[GraphQLRequest, List[GraphQLRequest]],
190170
extra_args: Optional[Dict[str, Any]] = None,
191171
upload_files: bool = False,
192172
) -> Dict[str, Any]:
193173

194-
payload = request.payload
174+
payload: Dict | List
175+
if isinstance(request, GraphQLRequest):
176+
payload = request.payload
177+
else:
178+
payload = [req.payload for req in request]
195179

196180
if upload_files:
181+
assert isinstance(payload, Dict)
182+
assert isinstance(request, GraphQLRequest)
197183
post_args = self._prepare_file_uploads(request, payload)
198184
else:
199185
post_args = {"json": payload}
@@ -379,15 +365,15 @@ async def execute(
379365
:returns: an ExecutionResult object.
380366
"""
381367

368+
if self.session is None:
369+
raise TransportClosed("Transport is not connected")
370+
382371
post_args = self._prepare_request(
383372
request,
384373
extra_args,
385374
upload_files,
386375
)
387376

388-
if self.session is None:
389-
raise TransportClosed("Transport is not connected")
390-
391377
try:
392378
async with self.session.post(self.url, ssl=self.ssl, **post_args) as resp:
393379
return await self._prepare_result(resp)
@@ -413,14 +399,14 @@ async def execute_batch(
413399
if an error occurred.
414400
"""
415401

416-
post_args = self._prepare_batch_request(
402+
if self.session is None:
403+
raise TransportClosed("Transport is not connected")
404+
405+
post_args = self._prepare_request(
417406
reqs,
418407
extra_args,
419408
)
420409

421-
if self.session is None:
422-
raise TransportClosed("Transport is not connected")
423-
424410
async with self.session.post(self.url, ssl=self.ssl, **post_args) as resp:
425411
return await self._prepare_batch_result(reqs, resp)
426412

gql/transport/httpx.py

Lines changed: 21 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -59,15 +59,22 @@ def __init__(
5959

6060
def _prepare_request(
6161
self,
62-
req: GraphQLRequest,
62+
request: Union[GraphQLRequest, List[GraphQLRequest]],
63+
*,
6364
extra_args: Optional[Dict[str, Any]] = None,
6465
upload_files: bool = False,
6566
) -> Dict[str, Any]:
6667

67-
payload = req.payload
68+
payload: Dict | List
69+
if isinstance(request, GraphQLRequest):
70+
payload = request.payload
71+
else:
72+
payload = [req.payload for req in request]
6873

6974
if upload_files:
70-
post_args = self._prepare_file_uploads(req, payload)
75+
assert isinstance(payload, Dict)
76+
assert isinstance(request, GraphQLRequest)
77+
post_args = self._prepare_file_uploads(request, payload)
7178
else:
7279
post_args = {"json": payload}
7380

@@ -81,26 +88,6 @@ def _prepare_request(
8188

8289
return post_args
8390

84-
def _prepare_batch_request(
85-
self,
86-
reqs: List[GraphQLRequest],
87-
extra_args: Optional[Dict[str, Any]] = None,
88-
) -> Dict[str, Any]:
89-
90-
payload = [req.payload for req in reqs]
91-
92-
post_args = {"json": payload}
93-
94-
# Log the payload
95-
if log.isEnabledFor(logging.DEBUG):
96-
log.debug(">>> %s", self.json_serialize(payload))
97-
98-
# Pass post_args to aiohttp post method
99-
if extra_args:
100-
post_args.update(extra_args)
101-
102-
return post_args
103-
10491
def _prepare_file_uploads(
10592
self,
10693
request: GraphQLRequest,
@@ -244,7 +231,7 @@ def connect(self):
244231

245232
self.client = httpx.Client(**self.kwargs)
246233

247-
def execute( # type: ignore
234+
def execute(
248235
self,
249236
request: GraphQLRequest,
250237
*,
@@ -269,8 +256,8 @@ def execute( # type: ignore
269256

270257
post_args = self._prepare_request(
271258
request,
272-
extra_args,
273-
upload_files,
259+
extra_args=extra_args,
260+
upload_files=upload_files,
274261
)
275262

276263
try:
@@ -292,7 +279,7 @@ def execute_batch(
292279
:code:`execute_batch` on a client or a session.
293280
294281
:param reqs: GraphQL requests as a list of GraphQLRequest objects.
295-
:param extra_args: additional arguments to send to the aiohttp post method
282+
:param extra_args: additional arguments to send to the httpx post method
296283
:return: A list of results of execution.
297284
For every result `data` is the result of executing the query,
298285
`errors` is null if no errors occurred, and is a non-empty array
@@ -302,9 +289,9 @@ def execute_batch(
302289
if not self.client:
303290
raise TransportClosed("Transport is not connected")
304291

305-
post_args = self._prepare_batch_request(
292+
post_args = self._prepare_request(
306293
reqs,
307-
extra_args,
294+
extra_args=extra_args,
308295
)
309296

310297
response = self.client.post(self.url, **post_args)
@@ -361,8 +348,8 @@ async def execute(
361348

362349
post_args = self._prepare_request(
363350
request,
364-
extra_args,
365-
upload_files,
351+
extra_args=extra_args,
352+
upload_files=upload_files,
366353
)
367354

368355
try:
@@ -384,7 +371,7 @@ async def execute_batch(
384371
:code:`execute_batch` on a client or a session.
385372
386373
:param reqs: GraphQL requests as a list of GraphQLRequest objects.
387-
:param extra_args: additional arguments to send to the aiohttp post method
374+
:param extra_args: additional arguments to send to the httpx post method
388375
:return: A list of results of execution.
389376
For every result `data` is the result of executing the query,
390377
`errors` is null if no errors occurred, and is a non-empty array
@@ -394,9 +381,9 @@ async def execute_batch(
394381
if not self.client:
395382
raise TransportClosed("Transport is not connected")
396383

397-
post_args = self._prepare_batch_request(
384+
post_args = self._prepare_request(
398385
reqs,
399-
extra_args,
386+
extra_args=extra_args,
400387
)
401388

402389
response = await self.client.post(self.url, **post_args)

0 commit comments

Comments
 (0)