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:
@EmGarr thanks for opening the issue! ~I think the best option here is to actually move the tenant check down to where it makes the actual requests to the API. That way it can be initialized no problem but won't actually attempt to check the tenant until it makes its first request.~

Edit: actually let me think on this a bit because my "solution" written will actually lead to a recursion error because it needs to make the initial query in order to retrieve the tenant but in order to retrieve the tenant it needs the headers. The better option here might be allowing for a way to set headers _on_ client initialization.

Resolving patch:
diff --git a/src/prefect/agent/agent.py b/src/prefect/agent/agent.py
--- a/src/prefect/agent/agent.py
+++ b/src/prefect/agent/agent.py
@@ -524,7 +524,7 @@ def query_flow_runs(self) -> list:
                 "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
--- 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
--- a/src/prefect/client/client.py
+++ b/src/prefect/client/client.py
@@ -98,32 +98,13 @@ def __init__(self, api_server: str = None, api_token: str = None):
         # 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 @@ def post(
         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 @@ def get_auth_token(self) -> str:
             - 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 @@ def login_to_tenant(self, tenant_slug: str = None, tenant_id: str = None) -> boo
             )  # 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 @@ def create_project(self, project_name: str, project_description: str = None) ->
                 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 @@ def register_agent(
                     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
--- a/tests/client/test_client_auth.py
+++ b/tests/client/test_client_auth.py
@@ -76,6 +76,7 @@ def test_client_token_initializes_from_file(selfmonkeypatch, cloud_api):
                     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 @@ def test_load_local_api_token_is_called_when_the_client_is_initialized_without_t
         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 @@ def test_login_to_tenant_writes_tenant_and_reloads_it_when_token_is_reloaded(
         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 @@ def test_client_clears_active_tenant_if_login_fails_on_initialization(
 
         # 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