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, access_token=None, **kwargs):
  from .client import Connection
 
  return Connection(server_hostname, http_path, access_token, **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 (
+ AuthProvider,
+ 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 AuthProvider()
+ 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,120 @@
+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 AuthProvider:
+ 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(AuthProvider):
+ 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(AuthProvider):
+ 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(AuthProvider):
+ 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(self._hostname)
+ 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(
+ self._hostname, 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(self._hostname, token)
+ except Exception as e:
+ logging.error(f"unexpected error in oauth token update", e, exc_info=True)
+ raise e