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
@@ -154,4 +154,4 @@ jobs:
       #              black the code
       #----------------------------------------------
       - name: Mypy
-        run: poetry run mypy src
+        run: poetry run mypy --install-types --non-interactive src

pyproject.toml

@@ -13,10 +13,13 @@ 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"
 mypy = "^0.950"
+pylint = ">=2.12.0"
 black = "^22.3.0"
 
 [tool.poetry.urls]

src/databricks/sql/__init__.py

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

src/databricks/sql/auth/auth.py

@@ -0,0 +1,96 @@
+from enum import Enum
+from typing import List
+
+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:
+        assert cfg.oauth_redirect_port_range is not None
+        assert cfg.oauth_client_id is not None
+        assert cfg.oauth_scopes is not None
+
+        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!")
+
+
+PYSQL_OAUTH_SCOPES = ["sql", "offline_access"]
+PYSQL_OAUTH_CLIENT_ID = "databricks-sql-python"
+PYSQL_OAUTH_REDIRECT_PORT_RANGE = list(range(8020, 8025))
+
+
+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 get_python_sql_connector_auth_provider(hostname: str, **kwargs):
+    cfg = ClientContext(
+        hostname=normalize_host_name(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=PYSQL_OAUTH_SCOPES,
+        oauth_client_id=PYSQL_OAUTH_CLIENT_ID,
+        oauth_redirect_port_range=PYSQL_OAUTH_REDIRECT_PORT_RANGE,
+        oauth_persistence=kwargs.get("experimental_oauth_persistence"),
+    )
+    return get_auth_provider(cfg)

src/databricks/sql/auth/authenticators.py

@@ -0,0 +1,118 @@
+import base64
+import logging
+from typing import Dict, List
+
+from databricks.sql.auth.oauth import OAuthManager
+
+# 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 = 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}"
+
+    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.access_token
+                        self._refresh_token = token.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