Pass kwargs to base client`s execute and execute_ws methods

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Added escaping of enum values which are Python keywords by appending `_` to them.
 - Fixed `enums_module_name` option not being passed to generators.
 - Added additional base clients supporting the Open Telemetry tracing. Added `opentelemetry_client` config option.
+- Changed generated client's methods to pass `**kwargs` to base client's `execute` and `execute_ws` methods (breaking change for custom base clients).
 
 
 ## 0.9.0 (2023-09-11)

EXAMPLE.md

@@ -180,7 +180,7 @@ Generated client class inherits from `AsyncBaseClient` and has async method for
 ```py
 # graphql_client/client.py
 
-from typing import AsyncIterator, Optional, Union
+from typing import Any, AsyncIterator, Dict, Optional, Union
 
 from .async_base_client import AsyncBaseClient
 from .base_model import UNSET, UnsetType, Upload
@@ -197,7 +197,9 @@ def gql(q: str) -> str:
 
 
 class Client(AsyncBaseClient):
-    async def create_user(self, user_data: UserCreateInput) -> CreateUser:
+    async def create_user(
+        self, user_data: UserCreateInput, **kwargs: Any
+    ) -> CreateUser:
         query = gql(
             """
             mutation CreateUser($userData: UserCreateInput!) {
@@ -207,12 +209,12 @@ class Client(AsyncBaseClient):
             }
             """
         )
-        variables: dict[str, object] = {"userData": user_data}
-        response = await self.execute(query=query, variables=variables)
+        variables: Dict[str, object] = {"userData": user_data}
+        response = await self.execute(query=query, variables=variables, **kwargs)
         data = self.get_data(response)
         return CreateUser.model_validate(data)
 
-    async def list_all_users(self) -> ListAllUsers:
+    async def list_all_users(self, **kwargs: Any) -> ListAllUsers:
         query = gql(
             """
             query ListAllUsers {
@@ -228,13 +230,13 @@ class Client(AsyncBaseClient):
             }
             """
         )
-        variables: dict[str, object] = {}
-        response = await self.execute(query=query, variables=variables)
+        variables: Dict[str, object] = {}
+        response = await self.execute(query=query, variables=variables, **kwargs)
         data = self.get_data(response)
         return ListAllUsers.model_validate(data)
 
     async def list_users_by_country(
-        self, country: Union[Optional[str], UnsetType] = UNSET
+        self, country: Union[Optional[str], UnsetType] = UNSET, **kwargs: Any
     ) -> ListUsersByCountry:
         query = gql(
             """
@@ -257,33 +259,33 @@ class Client(AsyncBaseClient):
             }
             """
         )
-        variables: dict[str, object] = {"country": country}
-        response = await self.execute(query=query, variables=variables)
+        variables: Dict[str, object] = {"country": country}
+        response = await self.execute(query=query, variables=variables, **kwargs)
         data = self.get_data(response)
         return ListUsersByCountry.model_validate(data)
 
-    async def get_users_counter(self) -> AsyncIterator[GetUsersCounter]:
+    async def get_users_counter(self, **kwargs: Any) -> AsyncIterator[GetUsersCounter]:
         query = gql(
             """
             subscription GetUsersCounter {
               usersCounter
             }
             """
         )
-        variables: dict[str, object] = {}
-        async for data in self.execute_ws(query=query, variables=variables):
+        variables: Dict[str, object] = {}
+        async for data in self.execute_ws(query=query, variables=variables, **kwargs):
             yield GetUsersCounter.model_validate(data)
 
-    async def upload_file(self, file: Upload) -> UploadFile:
+    async def upload_file(self, file: Upload, **kwargs: Any) -> UploadFile:
         query = gql(
             """
             mutation uploadFile($file: Upload!) {
               fileUpload(file: $file)
             }
             """
         )
-        variables: dict[str, object] = {"file": file}
-        response = await self.execute(query=query, variables=variables)
+        variables: Dict[str, object] = {"file": file}
+        response = await self.execute(query=query, variables=variables, **kwargs)
         data = self.get_data(response)
         return UploadFile.model_validate(data)
 ```

ariadne_codegen/client_generators/arguments.py

@@ -27,7 +27,14 @@
 from ..exceptions import ParsingError
 from ..plugins.manager import PluginManager
 from ..utils import process_name
-from .constants import ANY, INPUT_SCALARS_MAP, OPTIONAL, UNSET_NAME, UNSET_TYPE_NAME
+from .constants import (
+    ANY,
+    INPUT_SCALARS_MAP,
+    KWARGS_NAMES,
+    OPTIONAL,
+    UNSET_NAME,
+    UNSET_TYPE_NAME,
+)
 from .scalars import ScalarData
 
 
@@ -83,6 +90,7 @@ def generate(
         arguments = generate_arguments(
             args=required_args + optional_args,
             defaults=[generate_name(UNSET_NAME) for _ in optional_args],
+            kwarg=generate_arg(KWARGS_NAMES, annotation=generate_name(ANY)),
         )
 
         if self.plugin_manager:

ariadne_codegen/client_generators/client.py

@@ -33,6 +33,7 @@
     ANY,
     ASYNC_ITERATOR,
     DICT,
+    KWARGS_NAMES,
     LIST,
     MODEL_VALIDATE_METHOD,
     OPTIONAL,
@@ -290,10 +291,13 @@ def _generate_execute_call(self) -> ast.Call:
         return generate_call(
             func=generate_attribute(generate_name("self"), "execute"),
             keywords=[
-                generate_keyword("query", generate_name(self._operation_str_variable)),
                 generate_keyword(
-                    "variables", generate_name(self._variables_dict_variable)
+                    value=generate_name(self._operation_str_variable), arg="query"
                 ),
+                generate_keyword(
+                    value=generate_name(self._variables_dict_variable), arg="variables"
+                ),
+                generate_keyword(value=generate_name(KWARGS_NAMES)),
             ],
         )
 
@@ -325,13 +329,13 @@ def _generate_async_generator_loop(
                 func=generate_attribute(value=generate_name("self"), attr="execute_ws"),
                 keywords=[
                     generate_keyword(
-                        arg="query",
-                        value=generate_name(self._operation_str_variable),
+                        value=generate_name(self._operation_str_variable), arg="query"
                     ),
                     generate_keyword(
-                        arg="variables",
                         value=generate_name(self._variables_dict_variable),
+                        arg="variables",
                     ),
+                    generate_keyword(value=generate_name(KWARGS_NAMES)),
                 ],
             ),
             body=[self._generate_yield_parsed_obj(return_type)],

ariadne_codegen/client_generators/constants.py

@@ -68,6 +68,7 @@
 SKIP_DIRECTIVE_NAME = "skip"
 INCLUDE_DIRECTIVE_NAME = "include"
 
+KWARGS_NAMES = "kwargs"
 
 DEFAULT_ASYNC_BASE_CLIENT_PATH = (
     Path(__file__).parent / "dependencies" / "async_base_client.py"

ariadne_codegen/client_generators/dependencies/async_base_client.py

@@ -84,7 +84,7 @@ async def __aexit__(
         await self.http_client.aclose()
 
     async def execute(
-        self, query: str, variables: Optional[Dict[str, Any]] = None
+        self, query: str, variables: Optional[Dict[str, Any]] = None, **kwargs: Any
     ) -> httpx.Response:
         processed_variables, files, files_map = self._process_variables(variables)
 
@@ -94,9 +94,12 @@ async def execute(
                 variables=processed_variables,
                 files=files,
                 files_map=files_map,
+                **kwargs,
             )
 
-        return await self._execute_json(query=query, variables=processed_variables)
+        return await self._execute_json(
+            query=query, variables=processed_variables, **kwargs
+        )
 
     def get_data(self, response: httpx.Response) -> Dict[str, Any]:
         if not response.is_success:
@@ -123,14 +126,20 @@ def get_data(self, response: httpx.Response) -> Dict[str, Any]:
         return cast(Dict[str, Any], data)
 
     async def execute_ws(
-        self, query: str, variables: Optional[Dict[str, Any]] = None
+        self, query: str, variables: Optional[Dict[str, Any]] = None, **kwargs: Any
     ) -> AsyncIterator[Dict[str, Any]]:
+        headers = self.ws_headers.copy()
+        headers.update(kwargs.get("extra_headers", {}))
+
+        merged_kwargs: Dict[str, Any] = {"origin": self.ws_origin}
+        merged_kwargs.update(kwargs)
+        merged_kwargs["extra_headers"] = headers
+
         operation_id = str(uuid4())
         async with ws_connect(
             self.ws_url,
             subprotocols=[Subprotocol(GRAPHQL_TRANSPORT_WS)],
-            origin=self.ws_origin,
-            extra_headers=self.ws_headers,
+            **merged_kwargs,
         ) as websocket:
             await self._send_connection_init(websocket)
             await self._send_subscribe(
@@ -220,6 +229,7 @@ async def _execute_multipart(
         variables: Dict[str, Any],
         files: Dict[str, Tuple[str, IO[bytes], str]],
         files_map: Dict[str, List[str]],
+        **kwargs: Any,
     ) -> httpx.Response:
         data = {
             "operations": json.dumps(
@@ -228,19 +238,25 @@ async def _execute_multipart(
             "map": json.dumps(files_map, default=to_jsonable_python),
         }
 
-        return await self.http_client.post(url=self.url, data=data, files=files)
+        return await self.http_client.post(
+            url=self.url, data=data, files=files, **kwargs
+        )
 
     async def _execute_json(
-        self,
-        query: str,
-        variables: Dict[str, Any],
+        self, query: str, variables: Dict[str, Any], **kwargs: Any
     ) -> httpx.Response:
+        headers: Dict[str, str] = {"Content-Type": "application/json"}
+        headers.update(kwargs.get("headers", {}))
+
+        merged_kwargs: Dict[str, Any] = kwargs.copy()
+        merged_kwargs["headers"] = headers
+
         return await self.http_client.post(
             url=self.url,
             content=json.dumps(
                 {"query": query, "variables": variables}, default=to_jsonable_python
             ),
-            headers={"Content-Type": "application/json"},
+            **merged_kwargs,
         )
 
     async def _send_connection_init(self, websocket: WebSocketClientProtocol) -> None: