added support for persistence

Signed-off-by: Moe Derakhshani <[email protected]>

Author: moderakh

src/databricks/sql/__init__.py

@@ -44,6 +44,6 @@ def TimestampFromTicks(ticks):
     return Timestamp(*time.localtime(ticks)[:6])
 
 
-def connect(server_hostname, http_path, **kwargs):
+def connect(server_hostname, http_path, experimental_oauth_persistence=None, **kwargs):
     from .client import Connection
-    return Connection(server_hostname, http_path, **kwargs)
+    return Connection(server_hostname, http_path, experimental_oauth_persistence, **kwargs)

src/databricks/sql/auth/auth.py

@@ -24,6 +24,7 @@
 from enum import Enum
 from databricks.sql.auth.authenticators import CredentialsProvider, \
     AccessTokenAuthProvider, BasicAuthProvider, DatabricksOAuthProvider
+from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 
 
 class AuthType(Enum):
@@ -42,7 +43,9 @@ def __init__(self,
                  oauth_scopes: List[str] = None,
                  oauth_client_id: str = None,
                  use_cert_as_auth: str = None,
-                 tls_client_cert_file: str = None):
+                 tls_client_cert_file: str = None,
+                 oauth_persistence=None
+                 ):
         self.hostname = hostname
         self.username = username
         self.password = password
@@ -52,11 +55,12 @@ def __init__(self,
         self.oauth_client_id = oauth_client_id
         self.use_cert_as_auth = use_cert_as_auth
         self.tls_client_cert_file = tls_client_cert_file
+        self.oauth_persistence = oauth_persistence
 
 
 def get_auth_provider(cfg: ClientContext):
     if cfg.auth_type == AuthType.DATABRICKS_OAUTH.value:
-        return DatabricksOAuthProvider(cfg.hostname, cfg.oauth_client_id, cfg.oauth_scopes)
+        return DatabricksOAuthProvider(cfg.hostname, cfg.oauth_persistence, cfg.oauth_client_id, cfg.oauth_scopes)
     elif cfg.access_token is not None:
         return AccessTokenAuthProvider(cfg.access_token)
     elif cfg.username is not None and cfg.password is not None:
@@ -73,7 +77,7 @@ def get_auth_provider(cfg: ClientContext):
 OAUTH_CLIENT_ID = "databricks-cli"
 
 
-def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
+def get_python_sql_connector_auth_provider(hostname: str, oauth_persistence: OAuthPersistence, **kwargs):
     cfg = ClientContext(hostname=hostname,
                         auth_type=kwargs.get("auth_type"),
                         access_token=kwargs.get("access_token"),
@@ -82,7 +86,8 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
                         use_cert_as_auth=kwargs.get("_use_cert_as_auth"),
                         tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
                         oauth_scopes=OAUTH_SCOPES,
-                        oauth_client_id=OAUTH_CLIENT_ID)
+                        oauth_client_id=OAUTH_CLIENT_ID,
+                        oauth_persistence=oauth_persistence)
     return get_auth_provider(cfg)
 
 

src/databricks/sql/auth/authenticators.py

@@ -27,6 +27,9 @@
 
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
+from databricks.sql.experimental.oauth_persistence import OAuthToken
+
+
 class CredentialsProvider:
     def add_headers(self, request_headers):
         pass
@@ -59,18 +62,17 @@ def add_headers(self, request_headers):
 # Please must not depend on it in your applications.
 class DatabricksOAuthProvider(CredentialsProvider):
     SCOPE_DELIM = ' '
-    # TODO: moderakh the refresh_token is only kept in memory. not saved on disk
-    # hence if application restarts the user may need to re-authenticate
-    # I will add support for this outside of the scope of current PR.
-    def __init__(self, hostname, client_id, scopes):
+
+    def __init__(self, hostname, oauth_persistence, client_id, scopes):
         self._hostname = self._normalize_host_name(hostname=hostname)
         self._scopes_as_str = DatabricksOAuthProvider.SCOPE_DELIM.join(scopes)
-        access_token, refresh_token = get_tokens(hostname=self._hostname, client_id=client_id, scope=self._scopes_as_str)
-        self._access_token = access_token
-        self._refresh_token = refresh_token
+        self._oauth_persistence = oauth_persistence
+        self._client_id = client_id
+        self._get_tokens()
 
     def add_headers(self, request_headers):
         check_and_refresh_access_token(hostname=self._hostname,
+                                       client_id=self._client_id,
                                        access_token=self._access_token,
                                        refresh_token=self._refresh_token)
         request_headers['Authorization'] = f"Bearer {self._access_token}"
@@ -80,3 +82,35 @@ def _normalize_host_name(hostname):
         maybe_scheme = "https://" if not hostname.startswith("https://") else ""
         maybe_trailing_slash = "/" if not hostname.endswith("/") else ""
         return f"{maybe_scheme}{hostname}{maybe_trailing_slash}"
+
+    def _get_tokens(self):
+        if self._oauth_persistence:
+            token = self._oauth_persistence.read()
+            if token:
+                self._access_token = token.get_access_token()
+                self._refresh_token = token.get_refresh_token()
+                self._update_token_if_expired()
+            else:
+                (access_token, refresh_token) = get_tokens(hostname=self._hostname,
+                                                           client_id=self._client_id,
+                                                           scope=self._scopes_as_str)
+                self._access_token = access_token
+                self._refresh_token = refresh_token
+                self._oauth_persistence.persist(OAuthToken(access_token, refresh_token))
+
+    def _update_token_if_expired(self):
+        (fresh_access_token, fresh_refresh_token, is_refreshed) = check_and_refresh_access_token(
+            hostname=self._hostname,
+            client_id=self._client_id,
+            access_token=self._access_token,
+            refresh_token=self._refresh_token)
+
+        if not is_refreshed:
+            return
+        else:
+            self._access_token = fresh_access_token
+            self._refresh_token = fresh_refresh_token
+
+            if self._oauth_persistence:
+                token = OAuthToken(self._access_token, self._refresh_token)
+                self._oauth_persistence.persist(token)

src/databricks/sql/auth/oauth.py

@@ -220,7 +220,7 @@ def get_tokens_from_response(oauth_response):
     return access_token, refresh_token
 
 
-def check_and_refresh_access_token(hostname, access_token, refresh_token):
+def check_and_refresh_access_token(hostname, client_id, access_token, refresh_token):
     now = datetime.now(tz=UTC)
     # If we can't decode an expiration time, this will be expired by default.
     expiration_time = now
@@ -246,7 +246,7 @@ def check_and_refresh_access_token(hostname, access_token, refresh_token):
 
     # Try to refresh using the refresh token
     logger.debug(f"Attempting to refresh OAuth access token that expired on {expiration_time}")
-    oauth_response = send_refresh_token_request(hostname, refresh_token)
+    oauth_response = send_refresh_token_request(hostname, client_id, refresh_token)
     fresh_access_token, fresh_refresh_token = get_tokens_from_response(oauth_response)
     return fresh_access_token, fresh_refresh_token, True
 

src/databricks/sql/client.py

@@ -10,7 +10,7 @@
 from databricks.sql.utils import ExecuteResponse, ParamEscaper
 from databricks.sql.types import Row
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
-
+from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 logger = logging.getLogger(__name__)
 
 DEFAULT_RESULT_BUFFER_SIZE_BYTES = 10485760
@@ -22,6 +22,7 @@ def __init__(
         self,
         server_hostname: str,
         http_path: str,
+        oauth_persistence: OAuthPersistence = None,
         http_headers: Optional[List[Tuple[str, str]]] = None,
         session_configuration: Dict[str, Any] = None,
         catalog: Optional[str] = None,
@@ -85,7 +86,7 @@ def __init__(
         self.port = kwargs.get("_port", 443)
         self.disable_pandas = kwargs.get("_disable_pandas", False)
 
-        auth_provider = get_python_sql_connector_auth_provider(server_hostname, **kwargs)
+        auth_provider = get_python_sql_connector_auth_provider(server_hostname, oauth_persistence, **kwargs)
 
         if not kwargs.get("_user_agent_entry"):
             useragent_header = "{}/{}".format(USER_AGENT_NAME, __version__)

src/databricks/sql/experimental/__init__.py

@@ -0,0 +1,56 @@
+import logging
+import json
+logger = logging.getLogger(__name__)
+
+
+class OAuthToken:
+    def __init__(self, access_token, refresh_token):
+        self._access_token = access_token
+        self._refresh_token = refresh_token
+
+    def get_access_token(self) -> str:
+        return self._access_token
+
+    def get_refresh_token(self) -> str:
+        return self._refresh_token
+
+
+class OAuthPersistence:
+    def persist(self, oauth_token: OAuthToken):
+        pass
+
+    def read(self) -> OAuthToken:
+        pass
+
+
+# Note this is only intended to be used for development
+class DevOnlyFilePersistence(OAuthPersistence):
+
+    def __init__(self, file_path):
+        self._file_path = file_path
+
+    def persist(self, token: OAuthToken):
+        logger.info(f"persisting token in {self._file_path}")
+
+        # Data to be written
+        dictionary = {
+            "refresh_token": token.get_refresh_token(),
+            "access_token": token.get_access_token()
+        }
+
+        # Serializing json
+        json_object = json.dumps(dictionary, indent=4)
+
+        with open(self._file_path, "w") as outfile:
+            outfile.write(json_object)
+
+    def read(self) -> OAuthToken:
+        # TODO: validate the
+        try:
+            with open(self._file_path, "r") as infile:
+                json_as_string = infile.read()
+
+                token_as_json = json.loads(json_as_string)
+                return OAuthToken(token_as_json['access_token'], token_as_json['refresh_token'])
+        except Exception as e:
+            return None