Repository: PrefectHQ/prefect

Issue:
Client in Server mode makes a call to the server before having to setup the attached_headers.
## Description

In server mode the `prefect.client.Client` make a call to the server to obtain the tenant-id. However, in our case the server is protected with BasicAuth. `Client.attach_headers` is here to help us! However since this call is performed inside the `__init__` we can't pass the required headers (attached_headers have not been set yet).

https://github.com/PrefectHQ/prefect/blob/9721bb103ba85eedab863258588671e1cba8fd58/src/prefect/client/client.py#L124

```python
class Client:
    def __init__(self, api_server: str = None, api_token: str = None):
        self._access_token = None
        self._refresh_token = None
        self._access_token_expires_at = pendulum.now()
        self._active_tenant_id = None
        self._attached_headers = {}  # type: Dict[str, str]
        self.logger = create_diagnostic_logger("Diagnostics")

        # store api server
        self.api_server = api_server or prefect.context.config.cloud.get("graphql")

        # store api token
        self._api_token = api_token or prefect.context.config.cloud.get(
            "auth_token", None
        )

        if prefect.config.backend == "cloud":
          ...
        else:
            # TODO: Separate put this functionality and clean up initial tenant access handling
            if not self._active_tenant_id:
                # This call is performed without self._attached_headers set up
                tenant_info = self.graphql({"query": {"tenant": {"id"}}})
                if tenant_info.data.tenant:
                    self._active_tenant_id = tenant_info.data.tenant[0].id
```

In our case this is pretty annoying. Since that our server is protected with BasicAuth this example will raise an error at the constructor step:

```python
client = Client() # Here will have requests.exceptions.HTTPError: 401 Client Error: Unauthorized 
headers = {'Authorization': "Basic {}".format(PREFECT__SERVER__BASIC_AUTH)}
client.attach_headers(headers)
```

It will happens for all the classes depending on a client:

```python
agent = DockerAgent(name="trololo") # Here will have requests.exceptions.HTTPError: 401 Client Error: Unauthorized 
headers = {'Authorization': "Basic {}".format(PREFECT__SERVER__BASIC_AUTH)}
agent.client.attach_headers(headers)
```

## Expected Behavior

Move the initialization of `_active_tenant_id` further in order to let it enjoy our `attached_headers`.

```python
client = Client() #No call to the server here and no error.
headers = {'Authorization': "Basic {}".format(PREFECT__SERVER__BASIC_AUTH)}
client.attach_headers(headers)
```

## Environment
```json
{
  "config_overrides": {
    "flows": {
      "checkpointing": true
    }
  },
  "env_vars": [
    "PREFECT__SERVER__BASIC_AUTH"
  ],
  "system_information": {
    "platform": "Linux-4.15.0-118-generic-x86_64-with-debian-stretch-sid",
    "prefect_backend": "server",
    "prefect_version": "0.13.10",
    "python_version": "3.7.2"
  }
}
```


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/PrefectHQ/prefect/pull/3530?src=pr&el=h1) Report
> Merging [#3530](https://codecov.io/gh/PrefectHQ/prefect/pull/3530?src=pr&el=desc) into [master](https://codecov.io/gh/PrefectHQ/prefect/commit/9721bb103ba85eedab863258588671e1cba8fd58?el=desc) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.


EmGarr: 
> Do we want to run `Client._init_tenant` in the initializer still but wrapped in a try/catch with a warning log so we preserve the old behavior?

I agree. I don't all the implication of the tenant. If there are other changes I won't mind to do them.

>  Is the first `graphql` call initializing the tenant going to be a significant issue for anything? If so, should it be a public method?
> Do we need to initialize the tenant before other calls too? What about `Client.post`?

From what I read the tenant_id is only used in two methods in client.py

1. `create_project`

https://github.com/EmGarr/prefect/blob/63c0ed4aec2e652ae3f98b3a96494d83c581f910/src/prefect/client/client.py#L924

2. `register_agent`

https://github.com/EmGarr/prefect/blob/63c0ed4aec2e652ae3f98b3a96494d83c581f910/src/prefect/client/client.py#L1595

Since it is a private attribute. I assume we could move the initialization of the tenant to this part and let the graphql_method be a public one (without the private) ?



joshmeek: @EmGarr I like the idea of changing it to a property the way you have it! Especially because we use it outside of the client. Seems like a much cleaner approach
madkinsz: > @EmGarr I like the idea of changing it to a property the way you have it! Especially because we use it outside of the client. Seems like a much cleaner approach

I agree! This would be great and address any concerns I have
EmGarr: @joshmeek  @madkinsz What about this commit ? If it is okay for you I'll rebase and make a clearer explanation for the commit message.
madkinsz: > What about this commit ? If it is okay for you I'll rebase and make a clearer explanation for the commit message.

@EmGarr that looks good to me
EmGarr: Rebase done
joshmeek: This looks good @EmGarr! Let me test this with a few different setups to make sure the previous tenant init at client creation isn't expected anywhere (I don't think it is but want to make sure)
madkinsz: > Rebase done

@EmGarr if you could update the title to match the new implementation that'd be great too! :)
EmGarr: @joshmeek @madkinsz do you have any news regarding this MR ?

joshmeek: @EmGarr in its current state this actually will fail when running against cloud. I think we should move the `if prefect.config.backend == "cloud"` part back to the init and leave the else block in the `_init_tenant`
EmGarr: @joshmeek Sorry I had work to do! Do you have any information regarding what does not work with the cloud ?
EmGarr: I guess this is a problem with the local agent and this line:

https://github.com/PrefectHQ/prefect/blob/99f5701e209f78366f85f39ec82b5fd0590c3c65/src/prefect/agent/local/agent.py#L241

I'll make an update.



joshmeek: @EmGarr when communicating with Cloud as the backend the API token needs to be used for all calls in order to infer authentication. With this setup the `_init_tenant` is only being called when `active_tenant_id` is retrieved and that is only being retrieved in `create_project` and `register_agent`. Therefore when running something like `prefect get flows` this fails because the client has not yet loaded the API token

Patch:
diff --git a/changes/issue3515.yaml b/changes/issue3515.yaml
new file mode 100644
index 0000000000..ebf332fe53
--- /dev/null
+++ b/changes/issue3515.yaml
@@ -0,0 +1,5 @@
+enhancement:
+  - "Allow for setting `Client` headers before loading tenant when running with Prefect Server - [#3515](https://github.com/PrefectHQ/prefect/issues/3515)"
+
+contributor:
+  - "[Emilien Garreau](https://github.com/EmGarr)"
diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py
index 719b1f4713..1f98ce653d 100644
--- a/src/prefect/agent/agent.py
+++ b/src/prefect/agent/agent.py
@@ -524,7 +524,7 @@ class Agent:
                 "input": {
                     "before": now.isoformat(),
                     "labels": list(self.labels),
-                    "tenant_id": self.client._active_tenant_id,
+                    "tenant_id": self.client.active_tenant_id,
                 }
             },
         )
diff --git a/src/prefect/cli/auth.py b/src/prefect/cli/auth.py
index 0b83fbe198..dbc95eb162 100644
--- a/src/prefect/cli/auth.py
+++ b/src/prefect/cli/auth.py
@@ -114,7 +114,7 @@ def logout():
     )
 
     client = Client()
-    tenant_id = client._active_tenant_id
+    tenant_id = client.active_tenant_id
 
     if not tenant_id:
         click.secho("No tenant currently active", fg="red")
@@ -135,7 +135,7 @@ def list_tenants():
     client = Client()
 
     tenants = client.get_available_tenants()
-    active_tenant_id = client._active_tenant_id
+    active_tenant_id = client.active_tenant_id
 
     output = []
     for item in tenants:
diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py
index 9e8f680977..66e19c8b0c 100644
--- a/src/prefect/client/client.py
+++ b/src/prefect/client/client.py
@@ -98,32 +98,13 @@ class Client:
         # store api server
         self.api_server = api_server or prefect.context.config.cloud.get("graphql")
 
-        # store api token
         self._api_token = api_token or prefect.context.config.cloud.get(
             "auth_token", None
         )
 
-        if prefect.config.backend == "cloud":
-            if not self._api_token:
-                # if no api token was passed, attempt to load state from local storage
-                settings = self._load_local_settings()
-                self._api_token = settings.get("api_token")
-
-                if self._api_token:
-                    self._active_tenant_id = settings.get("active_tenant_id")
-                if self._active_tenant_id:
-                    try:
-                        self.login_to_tenant(tenant_id=self._active_tenant_id)
-                    except AuthorizationError:
-                        # if an authorization error is raised, then the token is invalid and should
-                        # be cleared
-                        self.logout_from_tenant()
-        else:
-            # TODO: Separate put this functionality and clean up initial tenant access handling
-            if not self._active_tenant_id:
-                tenant_info = self.graphql({"query": {"tenant": {"id"}}})
-                if tenant_info.data.tenant:
-                    self._active_tenant_id = tenant_info.data.tenant[0].id
+        # Initialize the tenant and api token if not yet set
+        if not self._api_token and prefect.config.backend == "cloud":
+            self._init_tenant()
 
     def create_tenant(self, name: str, slug: str = None) -> str:
         """
@@ -241,6 +222,46 @@ class Client:
         else:
             return {}
 
+    @property
+    def active_tenant_id(self) -> Optional[str]:
+        """
+        Return the tenant to contact the server.
+        If your tenant has not been set yet it will be initialized.
+        """
+        if self._active_tenant_id is None:
+            self._init_tenant()
+        return self._active_tenant_id
+
+    def _init_tenant(self) -> None:
+        """
+        Init the tenant to contact the server.
+
+        If your backend is set to cloud the tenant will be read from: $HOME/.prefect/settings.toml.
+
+        For the server backend it will try to retrieve the default tenant. If the server is
+        protected with auth like BasicAuth do not forget to `attach_headers` before any call.
+        """
+        if prefect.config.backend == "cloud":
+            # if no api token was passed, attempt to load state from local storage
+            settings = self._load_local_settings()
+
+            if not self._api_token:
+                self._api_token = settings.get("api_token")
+            if self._api_token:
+                self._active_tenant_id = settings.get("active_tenant_id")
+
+            if self._active_tenant_id:
+                try:
+                    self.login_to_tenant(tenant_id=self._active_tenant_id)
+                except AuthorizationError:
+                    # if an authorization error is raised, then the token is invalid and should
+                    # be cleared
+                    self.logout_from_tenant()
+        else:
+            tenant_info = self.graphql({"query": {"tenant": {"id"}}})
+            if tenant_info.data.tenant:
+                self._active_tenant_id = tenant_info.data.tenant[0].id
+
     def graphql(
         self,
         query: Any,
@@ -508,7 +529,7 @@ class Client:
             - str: the access token
         """
         if not self._access_token:
-            return self._api_token
+            return self._api_token  # type: ignore
 
         expiration = self._access_token_expires_at or pendulum.now()
         if self._refresh_token and pendulum.now().add(seconds=30) > expiration:
@@ -599,7 +620,7 @@ class Client:
             )  # type: ignore
             self._refresh_token = payload.data.switch_tenant.refresh_token  # type: ignore
 
-        self._active_tenant_id = tenant_id
+        self._active_tenant_id = tenant_id  # type: ignore
 
         # save the tenant setting
         settings = self._load_local_settings()
@@ -919,7 +940,7 @@ class Client:
                 input=dict(
                     name=project_name,
                     description=project_description,
-                    tenant_id=self._active_tenant_id,
+                    tenant_id=self.active_tenant_id,
                 )
             ),
         )  # type: Any
@@ -1597,7 +1618,7 @@ class Client:
                     type=agent_type,
                     name=name,
                     labels=labels or [],
-                    tenant_id=self._active_tenant_id,
+                    tenant_id=self.active_tenant_id,
                     agent_config_id=agent_config_id,
                 )
             ),


Test patch:
diff --git a/tests/client/test_client_auth.py b/tests/client/test_client_auth.py
index b4b944c669..a259fa5059 100644
--- a/tests/client/test_client_auth.py
+++ b/tests/client/test_client_auth.py
@@ -76,6 +76,7 @@ class TestClientConfig:
                     toml.dump(dict(api_token="FILE_TOKEN"), f)
 
                 client = Client()
+                client._init_tenant()
         assert client._api_token == "FILE_TOKEN"
 
     def test_client_token_priotizes_config_over_file(selfmonkeypatch, cloud_api):
@@ -121,12 +122,14 @@ class TestClientConfig:
         with tempfile.TemporaryDirectory() as tmp:
             with set_temporary_config({"home_dir": tmp}):
                 client = Client(api_token="a")
+                client._init_tenant()
                 client.save_api_token()
 
                 client = Client(api_token="b")
                 assert client._api_token == "b"
-
-                assert Client()._api_token == "a"
+                client_local_api = Client()
+                client_local_api._init_tenant()
+                assert client_local_api._api_token == "a"
 
 
 class TestTenantAuth:
@@ -256,11 +259,13 @@ class TestTenantAuth:
         assert client._active_tenant_id is None
         client.login_to_tenant(tenant_id=tenant_id)
         client.save_api_token()
-        assert client._active_tenant_id == tenant_id
+        assert client.active_tenant_id == tenant_id
 
         # new client loads the active tenant and token
-        assert Client()._active_tenant_id == tenant_id
-        assert Client()._api_token == "abc"
+        client_load_active_tenant = Client()
+        # The tenant is initialized by calling the property active_tenant_id
+        assert client_load_active_tenant.active_tenant_id == tenant_id
+        assert client_load_active_tenant._api_token == "abc"
 
     def test_login_to_client_doesnt_reload_active_tenant_when_token_isnt_loaded(
         self, patch_post, cloud_api
@@ -459,6 +464,7 @@ class TestTenantAuth:
 
         # this initialization will fail with the patched error
         client = Client()
+        client._init_tenant()
         settings = client._load_local_settings()
         assert "active_tenant_id" not in settings