Persist token acquired through external-browser auth type (#110)

## Changes

This cache is local to the Python SDK and keyed off the workspace host,
the OAuth client ID, and the list of scopes to authorize for. The cache
path is `~/.config/databricks-sdk-py/oauth`. Files saved to this
directory are masked 0600.

It does not persist refreshes that happen during a session.

## Tests

Reproduce by running `./examples/local_browser_oauth.py` multiple times.

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

Author: pietern

databricks/sdk/core.py

@@ -21,7 +21,7 @@
 
 from .azure import ARM_DATABRICKS_RESOURCE_ID, ENVIRONMENTS, AzureEnvironment
 from .oauth import (ClientCredentials, OAuthClient, OidcEndpoints, Refreshable,
-                    Token, TokenSource)
+                    Token, TokenCache, TokenSource)
 from .version import __version__
 
 __all__ = ['Config', 'DatabricksError']
@@ -140,10 +140,20 @@ def external_browser(cfg: 'Config') -> Optional[HeaderFactory]:
                                client_id=client_id,
                                redirect_url='http://localhost:8020',
                                client_secret=cfg.client_secret)
-    consent = oauth_client.initiate_consent()
-    if not consent:
-        return None
-    credentials = consent.launch_external_browser()
+
+    # Load cached credentials from disk if they exist.
+    # Note that these are local to the Python SDK and not reused by other SDKs.
+    token_cache = TokenCache(oauth_client)
+    credentials = token_cache.load()
+    if credentials:
+        # Force a refresh in case the loaded credentials are expired.
+        credentials.token()
+    else:
+        consent = oauth_client.initiate_consent()
+        if not consent:
+            return None
+        credentials = consent.launch_external_browser()
+    token_cache.save(credentials)
     return credentials(cfg)
 
 

databricks/sdk/oauth.py

@@ -1,7 +1,9 @@
 import base64
 import functools
 import hashlib
+import json
 import logging
+import os
 import secrets
 import threading
 import urllib.parse
@@ -10,7 +12,7 @@
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from http.server import BaseHTTPRequestHandler, HTTPServer
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 import requests
 import requests.auth
@@ -402,3 +404,41 @@ def refresh(self) -> Token:
                               params,
                               use_params=self.use_params,
                               use_header=self.use_header)
+
+
+class TokenCache():
+    BASE_PATH = "~/.config/databricks-sdk-py/oauth"
+
+    def __init__(self, client: OAuthClient) -> None:
+        self.client = client
+
+    @property
+    def filename(self) -> str:
+        # Include host, client_id, and scopes in the cache filename to make it unique.
+        hash = hashlib.sha256()
+        for chunk in [self.client.host, self.client.client_id, ",".join(self.client._scopes), ]:
+            hash.update(chunk.encode('utf-8'))
+        return os.path.expanduser(os.path.join(self.__class__.BASE_PATH, hash.hexdigest() + ".json"))
+
+    def load(self) -> Optional[RefreshableCredentials]:
+        """
+        Load credentials from cache file. Return None if the cache file does not exist or is invalid.
+        """
+        if not os.path.exists(self.filename):
+            return None
+
+        try:
+            with open(self.filename, 'r') as f:
+                raw = json.load(f)
+                return RefreshableCredentials.from_dict(self.client, raw)
+        except Exception:
+            return None
+
+    def save(self, credentials: RefreshableCredentials) -> None:
+        """
+        Save credentials to cache file.
+        """
+        os.makedirs(os.path.dirname(self.filename), exist_ok=True)
+        with open(self.filename, 'w') as f:
+            json.dump(credentials.as_dict(), f)
+        os.chmod(self.filename, 0o600)

tests/test_oauth.py

@@ -0,0 +1,29 @@
+from databricks.sdk.core import Config, OidcEndpoints
+from databricks.sdk.oauth import OAuthClient, TokenCache
+
+
+def test_token_cache_unique_filename_by_host(mocker):
+    mocker.patch.object(Config, "oidc_endpoints",
+                        OidcEndpoints("http://localhost:1234", "http://localhost:1234"))
+    common_args = dict(client_id="abc", redirect_url="http://localhost:8020")
+    c1 = OAuthClient(host="http://localhost:", **common_args)
+    c2 = OAuthClient(host="https://bar.cloud.databricks.com", **common_args)
+    assert TokenCache(c1).filename != TokenCache(c2).filename
+
+
+def test_token_cache_unique_filename_by_client_id(mocker):
+    mocker.patch.object(Config, "oidc_endpoints",
+                        OidcEndpoints("http://localhost:1234", "http://localhost:1234"))
+    common_args = dict(host="http://localhost:", redirect_url="http://localhost:8020")
+    c1 = OAuthClient(client_id="abc", **common_args)
+    c2 = OAuthClient(client_id="def", **common_args)
+    assert TokenCache(c1).filename != TokenCache(c2).filename
+
+
+def test_token_cache_unique_filename_by_scopes(mocker):
+    mocker.patch.object(Config, "oidc_endpoints",
+                        OidcEndpoints("http://localhost:1234", "http://localhost:1234"))
+    common_args = dict(host="http://localhost:", client_id="abc", redirect_url="http://localhost:8020")
+    c1 = OAuthClient(scopes=["foo"], **common_args)
+    c2 = OAuthClient(scopes=["bar"], **common_args)
+    assert TokenCache(c1).filename != TokenCache(c2).filename