Allow passing custom headers through client

CHANGELOG.md

@@ -10,7 +10,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 
 ### Enhancements
 
-- None
+- Allow passing of custom headers in `Client` calls - [#1255](https://github.com/PrefectHQ/prefect/pull/1255)
 
 ### Task Library
 

src/prefect/client/client.py

@@ -92,7 +92,9 @@ def local_token_path(self) -> str:
     # -------------------------------------------------------------------------
     # Utilities
 
-    def get(self, path: str, server: str = None, **params: BuiltIn) -> dict:
+    def get(
+        self, path: str, server: str = None, headers: dict = None, **params: BuiltIn
+    ) -> dict:
         """
         Convenience function for calling the Prefect API with token auth and GET request
 
@@ -101,18 +103,23 @@ def get(self, path: str, server: str = None, **params: BuiltIn) -> dict:
                 http://prefect-server/v1/auth/login, path would be 'auth/login'.
             - server (str, optional): the server to send the GET request to;
                 defaults to `self.graphql_server`
+            - headers (dict, optional): Headers to pass with the request
             - **params (dict): GET parameters
 
         Returns:
             - dict: Dictionary representation of the request made
         """
-        response = self._request(method="GET", path=path, params=params, server=server)
+        response = self._request(
+            method="GET", path=path, params=params, server=server, headers=headers
+        )
         if response.text:
             return response.json()
         else:
             return {}
 
-    def post(self, path: str, server: str = None, **params: BuiltIn) -> dict:
+    def post(
+        self, path: str, server: str = None, headers: dict = None, **params: BuiltIn
+    ) -> dict:
         """
         Convenience function for calling the Prefect API with token auth and POST request
 
@@ -121,12 +128,15 @@ def post(self, path: str, server: str = None, **params: BuiltIn) -> dict:
                 http://prefect-server/v1/auth/login, path would be 'auth/login'.
             - server (str, optional): the server to send the POST request to;
                 defaults to `self.graphql_server`
+            - headers(dict): headers to pass with the request
             - **params (dict): POST parameters
 
         Returns:
             - dict: Dictionary representation of the request made
         """
-        response = self._request(method="POST", path=path, params=params, server=server)
+        response = self._request(
+            method="POST", path=path, params=params, server=server, headers=headers
+        )
         if response.text:
             return response.json()
         else:
@@ -136,6 +146,7 @@ def graphql(
         self,
         query: Any,
         raise_on_error: bool = True,
+        headers: dict = None,
         **variables: Union[bool, dict, str, int]
     ) -> GraphQLResult:
         """
@@ -146,6 +157,8 @@ def graphql(
                 parsed by prefect.utilities.graphql.parse_graphql().
             - raise_on_error (bool): if True, a `ClientError` will be raised if the GraphQL
                 returns any `errors`.
+            - headers (dict): any additional headers that should be passed as part of the
+                request
             - **variables (kwarg): Variables to be filled into a query with the key being
                 equivalent to the variables that are accepted by the query
 
@@ -160,6 +173,7 @@ def graphql(
             query=parse_graphql(query),
             variables=json.dumps(variables),
             server=self.graphql_server,
+            headers=headers,
         )
 
         if raise_on_error and "errors" in result:
@@ -168,7 +182,12 @@ def graphql(
             return as_nested_dict(result, GraphQLResult)  # type: ignore
 
     def _request(
-        self, method: str, path: str, params: dict = None, server: str = None
+        self,
+        method: str,
+        path: str,
+        params: dict = None,
+        server: str = None,
+        headers: dict = None,
     ) -> "requests.models.Response":
         """
         Runs any specified request (GET, POST, DELETE) against the server
@@ -179,6 +198,7 @@ def _request(
             - params (dict, optional): Parameters used for the request
             - server (str, optional): The server to make requests against, base API
                 server is used if not specified
+            - headers (dict, optional): Headers to pass with the request
 
         Returns:
             - requests.models.Response: The response returned from the request
@@ -202,7 +222,8 @@ def _request(
 
         params = params or {}
 
-        headers = {"Authorization": "Bearer {}".format(self.token)}
+        headers = headers or {}
+        headers.update({"Authorization": "Bearer {}".format(self.token)})
         session = requests.Session()
         retries = Retry(
             total=6,

tests/client/test_client.py

@@ -129,6 +129,57 @@ def test_client_posts_raises_with_no_token(monkeypatch):
     assert "Client.login" in str(exc.value)
 
 
+def test_headers_are_passed_to_get(monkeypatch):
+    get = MagicMock()
+    session = MagicMock()
+    session.return_value.get = get
+    monkeypatch.setattr("requests.Session", session)
+    with set_temporary_config(
+        {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
+    ):
+        client = Client()
+    client.get("/foo/bar", headers={"x": "y", "Authorization": "z"})
+    assert get.called
+    assert get.call_args[1]["headers"] == {
+        "x": "y",
+        "Authorization": "Bearer secret_token",
+    }
+
+
+def test_headers_are_passed_to_post(monkeypatch):
+    post = MagicMock()
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
+    with set_temporary_config(
+        {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
+    ):
+        client = Client()
+    client.post("/foo/bar", headers={"x": "y", "Authorization": "z"})
+    assert post.called
+    assert post.call_args[1]["headers"] == {
+        "x": "y",
+        "Authorization": "Bearer secret_token",
+    }
+
+
+def test_headers_are_passed_to_graphql(monkeypatch):
+    post = MagicMock()
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
+    with set_temporary_config(
+        {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
+    ):
+        client = Client()
+    client.graphql("query {}", headers={"x": "y", "Authorization": "z"})
+    assert post.called
+    assert post.call_args[1]["headers"] == {
+        "x": "y",
+        "Authorization": "Bearer secret_token",
+    }
+
+
 def test_client_posts_to_graphql_server(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(success=True)))