OAuth implementation

.github/workflows/code-quality-checks.yml

@@ -1,5 +1,5 @@
 name: Code Quality Checks
-on: [push]
+on: [pull_request, push]
 jobs:
   run-unit-tests:
     runs-on: ubuntu-latest

pyproject.toml

@@ -13,6 +13,8 @@ python = "^3.7.1"
 thrift = "^0.13.0"
 pandas = "^1.3.0"
 pyarrow = "^9.0.0"
+requests=">2.18.1"
+oauthlib=">=3.1.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.2"

src/databricks/sql/__init__.py

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

src/databricks/sql/auth/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2022 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"), except
+# that the use of services to which certain application programming
+# interfaces (each, an "API") connect requires that the user first obtain
+# a license for the use of the APIs from Databricks, Inc. ("Databricks"),
+# by creating an account at www.databricks.com and agreeing to either (a)
+# the Community Edition Terms of Service, (b) the Databricks Terms of
+# Service, or (c) another written agreement between Licensee and Databricks
+# for the use of the APIs.
+#
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

src/databricks/sql/auth/auth.py

@@ -0,0 +1,96 @@
+# Copyright 2022 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"), except
+# that the use of services to which certain application programming
+# interfaces (each, an "API") connect requires that the user first obtain
+# a license for the use of the APIs from Databricks, Inc. ("Databricks"),
+# by creating an account at www.databricks.com and agreeing to either (a)
+# the Community Edition Terms of Service, (b) the Databricks Terms of
+# Service, or (c) another written agreement between Licensee and Databricks
+# for the use of the APIs.
+#
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+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):
+    DATABRICKS_OAUTH = "databricks-oauth"
+    # other supported types (access_token, user/pass) can be inferred
+    # we can add more types as needed later
+
+
+class ClientContext:
+    def __init__(self,
+                 hostname: str,
+                 username: str = None,
+                 password: str = None,
+                 access_token: str = None,
+                 auth_type: str = None,
+                 oauth_scopes: List[str] = None,
+                 oauth_client_id: str = None,
+                 oauth_redirect_port_range: List[int] = None,
+                 use_cert_as_auth: str = None,
+                 tls_client_cert_file: str = None,
+                 oauth_persistence=None
+                 ):
+        self.hostname = hostname
+        self.username = username
+        self.password = password
+        self.access_token = access_token
+        self.auth_type = auth_type
+        self.oauth_scopes = oauth_scopes
+        self.oauth_client_id = oauth_client_id
+        self.oauth_redirect_port_range = oauth_redirect_port_range
+        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_persistence, cfg.oauth_redirect_port_range, 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:
+        return BasicAuthProvider(cfg.username, cfg.password)
+    elif cfg.use_cert_as_auth and cfg.tls_client_cert_file:
+        # no op authenticator. authentication is performed using ssl certificate outside of headers
+        return CredentialsProvider()
+    else:
+        raise RuntimeError("No valid authentication settings!")
+
+
+OAUTH_SCOPES = ["sql", "offline_access"]
+# TODO: moderakh to be changed once registered on the service side
+OAUTH_CLIENT_ID = "databricks-sql-python"
+OAUTH_REDIRECT_PORT_RANGE = range(8020, 8025)
+
+def get_python_sql_connector_auth_provider(hostname: str, oauth_persistence: OAuthPersistence = None, **kwargs):
+    cfg = ClientContext(hostname=hostname,
+                        auth_type=kwargs.get("auth_type"),
+                        access_token=kwargs.get("access_token"),
+                        username=kwargs.get("_username"),
+                        password=kwargs.get("_password"),
+                        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_redirect_port_range=OAUTH_REDIRECT_PORT_RANGE,
+                        oauth_persistence=oauth_persistence)
+    return get_auth_provider(cfg)
+
+

src/databricks/sql/auth/authenticators.py

@@ -0,0 +1,130 @@
+# Copyright 2022 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"), except
+# that the use of services to which certain application programming
+# interfaces (each, an "API") connect requires that the user first obtain
+# a license for the use of the APIs from Databricks, Inc. ("Databricks"),
+# by creating an account at www.databricks.com and agreeing to either (a)
+# the Community Edition Terms of Service, (b) the Databricks Terms of
+# Service, or (c) another written agreement between Licensee and Databricks
+# for the use of the APIs.
+#
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import logging
+from typing import Dict, List
+
+from databricks.sql.auth.oauth import OAuthManager
+import base64
+
+
+# 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, OAuthPersistence
+
+
+class CredentialsProvider:
+    def add_headers(self, request_headers: Dict[str, str]):
+        pass
+
+
+# Private API: this is an evolving interface and it will change in the future.
+# Please must not depend on it in your applications.
+class AccessTokenAuthProvider(CredentialsProvider):
+    def __init__(self, access_token: str):
+        self.__authorization_header_value = "Bearer {}".format(access_token)
+
+    def add_headers(self, request_headers: Dict[str, str]):
+        request_headers['Authorization'] = self.__authorization_header_value
+
+
+# Private API: this is an evolving interface and it will change in the future.
+# Please must not depend on it in your applications.
+class BasicAuthProvider(CredentialsProvider):
+    def __init__(self, username: str, password: str):
+        auth_credentials = f"{username}:{password}".encode("UTF-8")
+        auth_credentials_base64 = base64.standard_b64encode(auth_credentials).decode("UTF-8")
+
+        self.__authorization_header_value = f"Basic {auth_credentials_base64}"
+
+    def add_headers(self, request_headers: Dict[str, str]):
+        request_headers['Authorization'] = self.__authorization_header_value
+
+
+# Private API: this is an evolving interface and it will change in the future.
+# Please must not depend on it in your applications.
+class DatabricksOAuthProvider(CredentialsProvider):
+    SCOPE_DELIM = ' '
+
+    def __init__(self, hostname: str, oauth_persistence: OAuthPersistence, redirect_port_range: List[int], client_id: str, scopes: List[str]):
+        try:
+            self.oauth_manager = OAuthManager(port_range=redirect_port_range, client_id=client_id)
+            self._hostname = self._normalize_host_name(hostname=hostname)
+            self._scopes_as_str = DatabricksOAuthProvider.SCOPE_DELIM.join(scopes)
+            self._oauth_persistence = oauth_persistence
+            self._client_id = client_id
+            self._access_token = None
+            self._refresh_token = None
+            self._initial_get_token()
+        except Exception as e:
+            logging.error(f"unexpected error", e, exc_info=True)
+            raise e
+
+    def add_headers(self, request_headers: Dict[str, str]):
+        self._update_token_if_expired()
+        request_headers['Authorization'] = f"Bearer {self._access_token}"
+
+    @staticmethod
+    def _normalize_host_name(hostname: str):
+        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 _initial_get_token(self):
+        try:
+            if self._access_token is None or self._refresh_token is None:
+                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()
+
+            if self._access_token and self._refresh_token:
+                self._update_token_if_expired()
+            else:
+                (access_token, refresh_token) = self.oauth_manager.get_tokens(
+                    hostname=self._hostname,
+                    scope=self._scopes_as_str)
+                self._access_token = access_token
+                self._refresh_token = refresh_token
+                self._oauth_persistence.persist(OAuthToken(access_token, refresh_token))
+        except Exception as e:
+            logging.error(f"unexpected error in oauth initialization", e, exc_info=True)
+            raise e
+
+    def _update_token_if_expired(self):
+        try:
+            (fresh_access_token, fresh_refresh_token, is_refreshed) = self.oauth_manager.check_and_refresh_access_token(
+                hostname=self._hostname,
+                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)
+        except Exception as e:
+            logging.error(f"unexpected error in oauth token update", e, exc_info=True)
+            raise e