Fix get_workspace_client in GCP (#532)

## Changes
The current implementation of get_workspace_client copies the entire
config, critically reusing the cached header factory so as to use the
same auth mechanism when getting a workspace client. However, at least
in GCP, account-level OAuth tokens can't be used to authenticate to a
workspace (probably because the audience for the account-level and
workspace-level tokens is different).

This PR fixes this by only copying the exported fields and not copying
the header factory. Subsequent use of the config in WorkspaceClient will
trigger config resolution. For GCP, this means creating a new token
source using the correct host as the audience.

This ports databricks/databricks-sdk-go#803 to
the Python SDK.

## Tests
Manually ran this integration test in all non-UC (Azure, AWS, GCP) and
UC (AWS, GCP) account-level environments.

- [ ] `make test` run locally
- [ ] `make fmt` applied
- [ ] relevant integration tests applied

.codegen/__init__.py.tmpl

@@ -157,6 +157,7 @@ class AccountClient:
         config.host = config.environment.deployment_url(workspace.deployment_name)
         config.azure_workspace_resource_id = azure.get_azure_resource_id(workspace)
         config.account_id = None
+        config.init_auth()
         return WorkspaceClient(config=config)
 
     def __repr__(self):

databricks/sdk/config.py

@@ -87,6 +87,7 @@ def __init__(self,
                  product_version="0.0.0",
                  clock: Clock = None,
                  **kwargs):
+        self._header_factory = None
         self._inner = {}
         self._user_agent_other_info = []
         self._credentials_provider = credentials_provider if credentials_provider else DefaultCredentials()
@@ -100,7 +101,7 @@ def __init__(self,
             self._known_file_config_loader()
             self._fix_host_if_needed()
             self._validate()
-            self._init_auth()
+            self.init_auth()
             self._product = product
             self._product_version = product_version
         except ValueError as e:
@@ -436,7 +437,7 @@ def _validate(self):
         names = " and ".join(sorted(auths_used))
         raise ValueError(f'validate: more than one authorization method configured: {names}')
 
-    def _init_auth(self):
+    def init_auth(self):
         try:
             self._header_factory = self._credentials_provider(self)
             self.auth_type = self._credentials_provider.auth_type()

tests/integration/test_client.py

@@ -1,13 +1,10 @@
 import pytest
 
 
-def test_get_workspace_client(a):
-    if a.config.is_azure or a.config.is_gcp:
-        pytest.skip('Not available on Azure and GCP currently')
-    wss = list(a.workspaces.list())
-    if len(wss) == 0:
-        pytest.skip("no workspaces")
-    w = a.get_workspace_client(wss[0])
+def test_get_workspace_client(a, env_or_skip):
+    workspace_id = env_or_skip("TEST_WORKSPACE_ID")
+    ws = a.workspaces.get(workspace_id)
+    w = a.get_workspace_client(ws)
     assert w.current_user.me().active