diff --git a/CHANGELOG.md b/CHANGELOG.md index 93a2858ed5d4..198300bf2c07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/ - Add a `save`/`load` interface to Flows - [#1685](https://github.com/PrefectHQ/prefect/pull/1685) - Add option to specify `aws_session_token` for the `FargateTaskEnvironment` - [#1688](https://github.com/PrefectHQ/prefect/pull/1688) +- Add an informative version header to all Cloud client requests - [#1690](https://github.com/PrefectHQ/prefect/pull/1690) ### Task Library diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py index 6b592c21d1cc..4a02ea09c419 100644 --- a/src/prefect/client/client.py +++ b/src/prefect/client/client.py @@ -265,6 +265,7 @@ def _request( headers = headers or {} if token: headers["Authorization"] = "Bearer {}".format(token) + headers["X-PREFECT-CORE-VERSION"] = str(prefect.__version__) session = requests.Session() retries = Retry( diff --git a/tests/client/test_client.py b/tests/client/test_client.py index 73b5baad5689..9ec88974ad12 100644 --- a/tests/client/test_client.py +++ b/tests/client/test_client.py @@ -31,6 +31,36 @@ def test_client_posts_to_api_server(patch_post): assert post.call_args[0][0] == "http://my-cloud.foo/foo/bar" +def test_version_header(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") + assert get.call_args[1]["headers"]["X-PREFECT-CORE-VERSION"] == str( + prefect.__version__ + ) + + +def test_version_header_cant_be_overridden(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-PREFECT-CORE-VERSION": "-1",}) + assert get.call_args[1]["headers"]["X-PREFECT-CORE-VERSION"] == str( + prefect.__version__ + ) + + def test_client_posts_graphql_to_api_server(patch_post): post = patch_post(dict(data=dict(success=True))) diff --git a/tests/client/test_client_auth.py b/tests/client/test_client_auth.py index 0bb288b32170..c77f11460cc2 100644 --- a/tests/client/test_client_auth.py +++ b/tests/client/test_client_auth.py @@ -180,7 +180,10 @@ def test_login_uses_api_token(self, patch_post): ) client = Client(api_token="api") client.login_to_tenant(tenant_id=tenant_id) - assert post.call_args[1]["headers"] == dict(Authorization="Bearer api") + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer api", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } def test_login_uses_api_token_when_access_token_is_set(self, patch_post): tenant_id = str(uuid.uuid4()) @@ -200,7 +203,10 @@ def test_login_uses_api_token_when_access_token_is_set(self, patch_post): client._access_token = "access" client.login_to_tenant(tenant_id=tenant_id) assert client.get_auth_token() == "ACCESS_TOKEN" - assert post.call_args[1]["headers"] == dict(Authorization="Bearer api") + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer api", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } def test_graphql_uses_access_token_after_login(self, patch_post): tenant_id = str(uuid.uuid4()) @@ -219,12 +225,18 @@ def test_graphql_uses_access_token_after_login(self, patch_post): client = Client(api_token="api") client.graphql({}) assert client.get_auth_token() == "api" - assert post.call_args[1]["headers"] == dict(Authorization="Bearer api") + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer api", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } client.login_to_tenant(tenant_id=tenant_id) client.graphql({}) assert client.get_auth_token() == "ACCESS_TOKEN" - assert post.call_args[1]["headers"] == dict(Authorization="Bearer ACCESS_TOKEN") + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer ACCESS_TOKEN", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } def test_login_to_tenant_writes_tenant_and_reloads_it_when_token_is_reloaded( self, patch_post @@ -364,7 +376,10 @@ def test_refresh_token_passes_refresh_token_as_header(self, patch_post): client = Client() client._refresh_token = "refresh" client._refresh_access_token() - assert post.call_args[1]["headers"] == dict(Authorization="Bearer refresh") + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer refresh", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } def test_get_available_tenants(self, patch_post): tenants = [ @@ -466,6 +481,7 @@ def test_headers_are_passed_to_get(self, monkeypatch): assert get.call_args[1]["headers"] == { "x": "y", "Authorization": "Bearer secret_token", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), } def test_headers_are_passed_to_post(self, monkeypatch): @@ -482,6 +498,7 @@ def test_headers_are_passed_to_post(self, monkeypatch): assert post.call_args[1]["headers"] == { "x": "y", "Authorization": "Bearer secret_token", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), } def test_headers_are_passed_to_graphql(self, monkeypatch): @@ -498,6 +515,7 @@ def test_headers_are_passed_to_graphql(self, monkeypatch): assert post.call_args[1]["headers"] == { "x": "y", "Authorization": "Bearer secret_token", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), } def test_tokens_are_passed_to_get(self, monkeypatch): @@ -509,7 +527,10 @@ def test_tokens_are_passed_to_get(self, monkeypatch): client = Client() client.get("/foo/bar", token="secret_token") assert get.called - assert get.call_args[1]["headers"] == {"Authorization": "Bearer secret_token"} + assert get.call_args[1]["headers"] == { + "Authorization": "Bearer secret_token", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } def test_tokens_are_passed_to_post(self, monkeypatch): post = MagicMock() @@ -520,7 +541,10 @@ def test_tokens_are_passed_to_post(self, monkeypatch): client = Client() client.post("/foo/bar", token="secret_token") assert post.called - assert post.call_args[1]["headers"] == {"Authorization": "Bearer secret_token"} + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer secret_token", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + } def test_tokens_are_passed_to_graphql(self, monkeypatch): post = MagicMock() @@ -531,4 +555,7 @@ def test_tokens_are_passed_to_graphql(self, monkeypatch): client = Client() client.graphql("query {}", token="secret_token") assert post.called - assert post.call_args[1]["headers"] == {"Authorization": "Bearer secret_token"} + assert post.call_args[1]["headers"] == { + "Authorization": "Bearer secret_token", + "X-PREFECT-CORE-VERSION": str(prefect.__version__), + }