Repository: databricks/dbt-databricks
Pull request #627 (2024-03-29 21:39:23)

Move credentials to their own file
### Description

Since the connections file is massive and complicated, starting to refactor to chunks

### Checklist

- [x] I have run this code in development and it appears to resolve the stated issue
- [x] This PR includes tests, or tests are not required/relevant for this PR
- [ ] I have updated the `CHANGELOG.md` and added information about my change to the "dbt-databricks next" section.



Patch:
diff --git a/dbt/adapters/databricks/__init__.py b/dbt/adapters/databricks/__init__.py
index fa126c5c..76392f7f 100644
--- a/dbt/adapters/databricks/__init__.py
+++ b/dbt/adapters/databricks/__init__.py
@@ -1,7 +1,7 @@
 from dbt.adapters.base import AdapterPlugin
 from dbt.adapters.databricks.column import DatabricksColumn  # noqa
 from dbt.adapters.databricks.connections import DatabricksConnectionManager  # noqa
-from dbt.adapters.databricks.connections import DatabricksCredentials
+from dbt.adapters.databricks.credentials import DatabricksCredentials
 from dbt.adapters.databricks.impl import DatabricksAdapter
 from dbt.adapters.databricks.relation import DatabricksRelation  # noqa
 from dbt.include import databricks
diff --git a/dbt/adapters/databricks/auth.py b/dbt/adapters/databricks/auth.py
index f5e0c348..b7b87d33 100644
--- a/dbt/adapters/databricks/auth.py
+++ b/dbt/adapters/databricks/auth.py
@@ -9,6 +9,8 @@ from databricks.sdk.core import HeaderFactory
 from databricks.sdk.oauth import ClientCredentials
 from databricks.sdk.oauth import Token
 from databricks.sdk.oauth import TokenSource
+from requests import PreparedRequest
+from requests.auth import AuthBase
 
 
 class token_auth(CredentialsProvider):
@@ -84,3 +86,21 @@ class m2m_auth(CredentialsProvider):
             return {"Authorization": f"{token.token_type} {token.access_token}"}
 
         return inner
+
+
+class BearerAuth(AuthBase):
+    """This mix-in is passed to our requests Session to explicitly
+    use the bearer authentication method.
+
+    Without this, a local .netrc file in the user's home directory
+    will override the auth headers provided by our header_factory.
+
+    More details in issue #337.
+    """
+
+    def __init__(self, headers: HeaderFactory):
+        self.headers = headers()
+
+    def __call__(self, r: PreparedRequest) -> PreparedRequest:
+        r.headers.update(**self.headers)
+        return r
diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py
index 43da876e..56e0015c 100644
--- a/dbt/adapters/databricks/connections.py
+++ b/dbt/adapters/databricks/connections.py
@@ -1,10 +1,7 @@
 import decimal
-import itertools
-import json
 import os
 import re
 import sys
-import threading
 import time
 import uuid
 import warnings
@@ -18,48 +15,37 @@ from typing import Callable
 from typing import cast
 from typing import Dict
 from typing import Hashable
-from typing import Iterable
 from typing import Iterator
 from typing import List
 from typing import Optional
 from typing import Sequence
 from typing import Tuple
-from typing import Union
 
 import databricks.sql as dbsql
-import keyring
 from agate import Table
-from databricks.sdk.core import CredentialsProvider
-from databricks.sdk.core import HeaderFactory
-from databricks.sdk.oauth import OAuthClient
-from databricks.sdk.oauth import SessionCredentials
 from databricks.sql.client import Connection as DatabricksSQLConnection
 from databricks.sql.client import Cursor as DatabricksSQLCursor
 from databricks.sql.exc import Error
 from dbt_common.clients import agate_helper
 from dbt_common.events.contextvars import get_node_info
 from dbt_common.events.functions import fire_event
-from dbt_common.exceptions import DbtConfigError
 from dbt_common.exceptions import DbtInternalError
 from dbt_common.exceptions import DbtRuntimeError
-from dbt_common.exceptions import DbtValidationError
 from dbt_common.utils import cast_to_str
-from requests import PreparedRequest
 from requests import Session
-from requests.auth import AuthBase
 
 from dbt.adapters.base.query_headers import MacroQueryStringSetter
 from dbt.adapters.contracts.connection import AdapterRequiredConfig
 from dbt.adapters.contracts.connection import AdapterResponse
 from dbt.adapters.contracts.connection import Connection
 from dbt.adapters.contracts.connection import ConnectionState
-from dbt.adapters.contracts.connection import Credentials
 from dbt.adapters.contracts.connection import DEFAULT_QUERY_COMMENT
 from dbt.adapters.contracts.connection import Identifier
 from dbt.adapters.contracts.connection import LazyHandle
 from dbt.adapters.databricks.__version__ import version as __version__
-from dbt.adapters.databricks.auth import m2m_auth
-from dbt.adapters.databricks.auth import token_auth
+from dbt.adapters.databricks.auth import BearerAuth
+from dbt.adapters.databricks.credentials import DatabricksCredentials
+from dbt.adapters.databricks.credentials import TCredentialProvider
 from dbt.adapters.databricks.logging import logger
 from dbt.adapters.databricks.utils import redact_credentials
 from dbt.adapters.events.types import ConnectionClosedInCleanup
@@ -77,20 +63,8 @@ st_refresh_regex = re.compile(
 )
 
 
-TCredentialProvider = Union[CredentialsProvider, SessionCredentials]
-
-
-CATALOG_KEY_IN_SESSION_PROPERTIES = "databricks.catalog"
 DBR_VERSION_REGEX = re.compile(r"([1-9][0-9]*)\.(x|0|[1-9][0-9]*)")
-DBT_DATABRICKS_INVOCATION_ENV = "DBT_DATABRICKS_INVOCATION_ENV"
-DBT_DATABRICKS_INVOCATION_ENV_REGEX = re.compile("^[A-z0-9\\-]+$")
-EXTRACT_CLUSTER_ID_FROM_HTTP_PATH_REGEX = re.compile(r"/?sql/protocolv1/o/\d+/(.*)")
-DBT_DATABRICKS_HTTP_SESSION_HEADERS = "DBT_DATABRICKS_HTTP_SESSION_HEADERS"
 
-REDIRECT_URL = "http://localhost:8020"
-CLIENT_ID = "dbt-databricks"
-SCOPES = ["all-apis", "offline_access"]
-MAX_NT_PASSWORD_SIZE = 1280
 
 # toggle for session managements that minimizes the number of sessions opened/closed
 USE_LONG_SESSIONS = os.getenv("DBT_DATABRICKS_LONG_SESSIONS", "True").upper() == "TRUE"
@@ -100,379 +74,6 @@ USE_LONG_SESSIONS = os.getenv("DBT_DATABRICKS_LONG_SESSIONS", "True").upper() ==
 DEFAULT_MAX_IDLE_TIME = 600
 
 
-class BearerAuth(AuthBase):
-    """This mix-in is passed to our requests Session to explicitly
-    use the bearer authentication method.
-
-    Without this, a local .netrc file in the user's home directory
-    will override the auth headers provided by our header_factory.
-
-    More details in issue #337.
-    """
-
-    def __init__(self, headers: HeaderFactory):
-        self.headers = headers()
-
-    def __call__(self, r: PreparedRequest) -> PreparedRequest:
-        r.headers.update(**self.headers)
-        return r
-
-
-@dataclass
-class DatabricksCredentials(Credentials):
-    database: Optional[str] = None  # type: ignore[assignment]
-    schema: Optional[str] = None  # type: ignore[assignment]
-    host: Optional[str] = None
-    http_path: Optional[str] = None
-    token: Optional[str] = None
-    client_id: Optional[str] = None
-    client_secret: Optional[str] = None
-    session_properties: Optional[Dict[str, Any]] = None
-    connection_parameters: Optional[Dict[str, Any]] = None
-    auth_type: Optional[str] = None
-
-    # Named compute resources specified in the profile. Used for
-    # creating a connection when a model specifies a compute resource.
-    compute: Optional[Dict[str, Any]] = None
-
-    connect_retries: int = 1
-    connect_timeout: Optional[int] = None
-    retry_all: bool = False
-    connect_max_idle: Optional[int] = None
-
-    _credentials_provider: Optional[Dict[str, Any]] = None
-    _lock = threading.Lock()  # to avoid concurrent auth
-
-    _ALIASES = {
-        "catalog": "database",
-        "target_catalog": "target_database",
-    }
-
-    @classmethod
-    def __pre_deserialize__(cls, data: Dict[Any, Any]) -> Dict[Any, Any]:
-        data = super().__pre_deserialize__(data)
-        if "database" not in data:
-            data["database"] = None
-        return data
-
-    def __post_init__(self) -> None:
-        if "." in (self.schema or ""):
-            raise DbtValidationError(
-                f"The schema should not contain '.': {self.schema}\n"
-                "If you are trying to set a catalog, please use `catalog` instead.\n"
-            )
-
-        session_properties = self.session_properties or {}
-        if CATALOG_KEY_IN_SESSION_PROPERTIES in session_properties:
-            if self.database is None:
-                self.database = session_properties[CATALOG_KEY_IN_SESSION_PROPERTIES]
-                del session_properties[CATALOG_KEY_IN_SESSION_PROPERTIES]
-            else:
-                raise DbtValidationError(
-                    f"Got duplicate keys: (`{CATALOG_KEY_IN_SESSION_PROPERTIES}` "
-                    'in session_properties) all map to "database"'
-                )
-        self.session_properties = session_properties
-
-        if self.database is not None:
-            database = self.database.strip()
-            if not database:
-                raise DbtValidationError(f"Invalid catalog name : `{self.database}`.")
-            self.database = database
-        else:
-            self.database = "hive_metastore"
-
-        connection_parameters = self.connection_parameters or {}
-        for key in (
-            "server_hostname",
-            "http_path",
-            "access_token",
-            "client_id",
-            "client_secret",
-            "session_configuration",
-            "catalog",
-            "schema",
-            "_user_agent_entry",
-        ):
-            if key in connection_parameters:
-                raise DbtValidationError(f"The connection parameter `{key}` is reserved.")
-        if "http_headers" in connection_parameters:
-            http_headers = connection_parameters["http_headers"]
-            if not isinstance(http_headers, dict) or any(
-                not isinstance(key, str) or not isinstance(value, str)
-                for key, value in http_headers.items()
-            ):
-                raise DbtValidationError(
-                    "The connection parameter `http_headers` should be dict of strings: "
-                    f"{http_headers}."
-                )
-        if "_socket_timeout" not in connection_parameters:
-            connection_parameters["_socket_timeout"] = 600
-        self.connection_parameters = connection_parameters
-
-    def validate_creds(self) -> None:
-        for key in ["host", "http_path"]:
-            if not getattr(self, key):
-                raise DbtConfigError(
-                    "The config '{}' is required to connect to Databricks".format(key)
-                )
-        if not self.token and self.auth_type != "oauth":
-            raise DbtConfigError(
-                ("The config `auth_type: oauth` is required when not using access token")
-            )
-
-        if not self.client_id and self.client_secret:
-            raise DbtConfigError(
-                (
-                    "The config 'client_id' is required to connect "
-                    "to Databricks when 'client_secret' is present"
-                )
-            )
-
-    @classmethod
-    def get_invocation_env(cls) -> Optional[str]:
-        invocation_env = os.environ.get(DBT_DATABRICKS_INVOCATION_ENV)
-        if invocation_env:
-            # Thrift doesn't allow nested () so we need to ensure
-            # that the passed user agent is valid.
-            if not DBT_DATABRICKS_INVOCATION_ENV_REGEX.search(invocation_env):
-                raise DbtValidationError(f"Invalid invocation environment: {invocation_env}")
-        return invocation_env
-
-    @classmethod
-    def get_all_http_headers(cls, user_http_session_headers: Dict[str, str]) -> Dict[str, str]:
-        http_session_headers_str: Optional[str] = os.environ.get(
-            DBT_DATABRICKS_HTTP_SESSION_HEADERS
-        )
-
-        http_session_headers_dict: Dict[str, str] = (
-            {
-                k: v if isinstance(v, str) else json.dumps(v)
-                for k, v in json.loads(http_session_headers_str).items()
-            }
-            if http_session_headers_str is not None
-            else {}
-        )
-
-        intersect_http_header_keys = (
-            user_http_session_headers.keys() & http_session_headers_dict.keys()
-        )
-
-        if len(intersect_http_header_keys) > 0:
-            raise DbtValidationError(
-                f"Intersection with reserved http_headers in keys: {intersect_http_header_keys}"
-            )
-
-        http_session_headers_dict.update(user_http_session_headers)
-
-        return http_session_headers_dict
-
-    @property
-    def type(self) -> str:
-        return "databricks"
-
-    @property
-    def unique_field(self) -> str:
-        return cast(str, self.host)
-
-    def connection_info(self, *, with_aliases: bool = False) -> Iterable[Tuple[str, Any]]:
-        as_dict = self.to_dict(omit_none=False)
-        connection_keys = set(self._connection_keys(with_aliases=with_aliases))
-        aliases: List[str] = []
-        if with_aliases:
-            aliases = [k for k, v in self._ALIASES.items() if v in connection_keys]
-        for key in itertools.chain(self._connection_keys(with_aliases=with_aliases), aliases):
-            if key in as_dict:
-                yield key, as_dict[key]
-
-    def _connection_keys(self, *, with_aliases: bool = False) -> Tuple[str, ...]:
-        # Assuming `DatabricksCredentials.connection_info(self, *, with_aliases: bool = False)`
-        # is called from only:
-        #
-        # - `Profile` with `with_aliases=True`
-        # - `DebugTask` without `with_aliases` (`False` by default)
-        #
-        # Thus, if `with_aliases` is `True`, `DatabricksCredentials._connection_keys` should return
-        # the internal key names; otherwise it can use aliases to show in `dbt debug`.
-        connection_keys = ["host", "http_path", "schema"]
-        if with_aliases:
-            connection_keys.insert(2, "database")
-        elif self.database:
-            connection_keys.insert(2, "catalog")
-        if self.session_properties:
-            connection_keys.append("session_properties")
-        return tuple(connection_keys)
-
-    @classmethod
-    def extract_cluster_id(cls, http_path: str) -> Optional[str]:
-        m = EXTRACT_CLUSTER_ID_FROM_HTTP_PATH_REGEX.match(http_path)
-        if m:
-            return m.group(1).strip()
-        else:
-            return None
-
-    @property
-    def cluster_id(self) -> Optional[str]:
-        return self.extract_cluster_id(self.http_path)  # type: ignore[arg-type]
-
-    def authenticate(self, in_provider: Optional[TCredentialProvider]) -> TCredentialProvider:
-        self.validate_creds()
-        host: str = self.host or ""
-        if self._credentials_provider:
-            return self._provider_from_dict()  # type: ignore
-        if in_provider:
-            if isinstance(in_provider, m2m_auth) or isinstance(in_provider, token_auth):
-                self._credentials_provider = in_provider.as_dict()
-            return in_provider
-
-        provider: TCredentialProvider
-        # dbt will spin up multiple threads. This has to be sync. So lock here
-        self._lock.acquire()
-        try:
-            if self.token:
-                provider = token_auth(self.token)
-                self._credentials_provider = provider.as_dict()
-                return provider
-
-            if self.client_id and self.client_secret:
-                provider = m2m_auth(
-                    host=host,
-                    client_id=self.client_id or "",
-                    client_secret=self.client_secret or "",
-                )
-                self._credentials_provider = provider.as_dict()
-                return provider
-
-            oauth_client = OAuthClient(
-                host=host,
-                client_id=self.client_id if self.client_id else CLIENT_ID,
-                client_secret="",
-                redirect_url=REDIRECT_URL,
-                scopes=SCOPES,
-            )
-            # optional branch. Try and keep going if it does not work
-            try:
-                # try to get cached credentials
-                credsdict = self.get_sharded_password("dbt-databricks", host)
-
-                if credsdict:
-                    provider = SessionCredentials.from_dict(oauth_client, json.loads(credsdict))
-                    # if refresh token is expired, this will throw
-                    try:
-                        if provider.token().valid:
-                            return provider
-                    except Exception as e:
-                        logger.debug(e)
-                        # whatever it is, get rid of the cache
-                        self.delete_sharded_password("dbt-databricks", host)
-
-            # error with keyring. Maybe machine has no password persistency
-            except Exception as e:
-                logger.debug(e)
-                logger.info("could not retrieved saved token")
-
-            # no token, go fetch one
-            consent = oauth_client.initiate_consent()
-
-            provider = consent.launch_external_browser()
-            # save for later
-            self._credentials_provider = provider.as_dict()
-            try:
-                self.set_sharded_password(
-                    "dbt-databricks", host, json.dumps(self._credentials_provider)
-                )
-            # error with keyring. Maybe machine has no password persistency
-            except Exception as e:
-                logger.debug(e)
-                logger.info("could not save token")
-
-            return provider
-
-        finally:
-            self._lock.release()
-
-    def set_sharded_password(self, service_name: str, username: str, password: str) -> None:
-        max_size = MAX_NT_PASSWORD_SIZE
-
-        # if not Windows or "small" password, stick to the default
-        if os.name != "nt" or len(password) < max_size:
-            keyring.set_password(service_name, username, password)
-        else:
-            logger.debug(f"password is {len(password)} characters, sharding it.")
-
-            password_shards = [
-                password[i : i + max_size] for i in range(0, len(password), max_size)
-            ]
-            shard_info = {
-                "sharded_password": True,
-                "shard_count": len(password_shards),
-            }
-
-            # store the "shard info" as the "base" password
-            keyring.set_password(service_name, username, json.dumps(shard_info))
-            # then store all shards with the shard number as postfix
-            for i, s in enumerate(password_shards):
-                keyring.set_password(service_name, f"{username}__{i}", s)
-
-    def get_sharded_password(self, service_name: str, username: str) -> Optional[str]:
-        password = keyring.get_password(service_name, username)
-
-        # check for "shard info" stored as json
-        try:
-            password_as_dict = json.loads(str(password))
-            if password_as_dict.get("sharded_password"):
-                # if password was stored shared, reconstruct it
-                shard_count = int(password_as_dict.get("shard_count"))
-
-                password = ""
-                for i in range(shard_count):
-                    password += str(keyring.get_password(service_name, f"{username}__{i}"))
-        except ValueError:
-            pass
-
-        return password
-
-    def delete_sharded_password(self, service_name: str, username: str) -> None:
-        password = keyring.get_password(service_name, username)
-
-        # check for "shard info" stored as json. If so delete all shards
-        try:
-            password_as_dict = json.loads(str(password))
-            if password_as_dict.get("sharded_password"):
-                shard_count = int(password_as_dict.get("shard_count"))
-                for i in range(shard_count):
-                    keyring.delete_password(service_name, f"{username}__{i}")
-        except ValueError:
-            pass
-
-        # delete "base" password
-        keyring.delete_password(service_name, username)
-
-    def _provider_from_dict(self) -> Optional[TCredentialProvider]:
-        if self.token:
-            return token_auth.from_dict(self._credentials_provider)
-
-        if self.client_id and self.client_secret:
-            return m2m_auth.from_dict(
-                host=self.host or "",
-                client_id=self.client_id or "",
-                client_secret=self.client_secret or "",
-                raw=self._credentials_provider or {"token": {}},
-            )
-
-        oauth_client = OAuthClient(
-            host=self.host or "",
-            client_id=CLIENT_ID,
-            client_secret="",
-            redirect_url=REDIRECT_URL,
-            scopes=SCOPES,
-        )
-
-        return SessionCredentials.from_dict(
-            client=oauth_client, raw=self._credentials_provider or {"token": {}}
-        )
-
-
 class DatabricksSQLConnectionWrapper:
     """Wrap a Databricks SQL connector in a way that no-ops transactions"""
 
diff --git a/dbt/adapters/databricks/credentials.py b/dbt/adapters/databricks/credentials.py
new file mode 100644
index 00000000..567bf6df
--- /dev/null
+++ b/dbt/adapters/databricks/credentials.py
@@ -0,0 +1,393 @@
+import itertools
+import json
+import os
+import re
+import threading
+from dataclasses import dataclass
+from typing import Any
+from typing import cast
+from typing import Dict
+from typing import Iterable
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import keyring
+from databricks.sdk.core import CredentialsProvider
+from databricks.sdk.oauth import OAuthClient
+from databricks.sdk.oauth import SessionCredentials
+from dbt.adapters.contracts.connection import Credentials
+from dbt.adapters.databricks.auth import m2m_auth
+from dbt.adapters.databricks.auth import token_auth
+from dbt.adapters.databricks.logging import logger
+from dbt_common.exceptions import DbtConfigError
+from dbt_common.exceptions import DbtValidationError
+
+CATALOG_KEY_IN_SESSION_PROPERTIES = "databricks.catalog"
+DBT_DATABRICKS_INVOCATION_ENV = "DBT_DATABRICKS_INVOCATION_ENV"
+DBT_DATABRICKS_INVOCATION_ENV_REGEX = re.compile("^[A-z0-9\\-]+$")
+EXTRACT_CLUSTER_ID_FROM_HTTP_PATH_REGEX = re.compile(r"/?sql/protocolv1/o/\d+/(.*)")
+DBT_DATABRICKS_HTTP_SESSION_HEADERS = "DBT_DATABRICKS_HTTP_SESSION_HEADERS"
+
+REDIRECT_URL = "http://localhost:8020"
+CLIENT_ID = "dbt-databricks"
+SCOPES = ["all-apis", "offline_access"]
+MAX_NT_PASSWORD_SIZE = 1280
+
+TCredentialProvider = Union[CredentialsProvider, SessionCredentials]
+
+
+@dataclass
+class DatabricksCredentials(Credentials):
+    database: Optional[str] = None  # type: ignore[assignment]
+    schema: Optional[str] = None  # type: ignore[assignment]
+    host: Optional[str] = None
+    http_path: Optional[str] = None
+    token: Optional[str] = None
+    client_id: Optional[str] = None
+    client_secret: Optional[str] = None
+    session_properties: Optional[Dict[str, Any]] = None
+    connection_parameters: Optional[Dict[str, Any]] = None
+    auth_type: Optional[str] = None
+
+    # Named compute resources specified in the profile. Used for
+    # creating a connection when a model specifies a compute resource.
+    compute: Optional[Dict[str, Any]] = None
+
+    connect_retries: int = 1
+    connect_timeout: Optional[int] = None
+    retry_all: bool = False
+    connect_max_idle: Optional[int] = None
+
+    _credentials_provider: Optional[Dict[str, Any]] = None
+    _lock = threading.Lock()  # to avoid concurrent auth
+
+    _ALIASES = {
+        "catalog": "database",
+        "target_catalog": "target_database",
+    }
+
+    @classmethod
+    def __pre_deserialize__(cls, data: Dict[Any, Any]) -> Dict[Any, Any]:
+        data = super().__pre_deserialize__(data)
+        if "database" not in data:
+            data["database"] = None
+        return data
+
+    def __post_init__(self) -> None:
+        if "." in (self.schema or ""):
+            raise DbtValidationError(
+                f"The schema should not contain '.': {self.schema}\n"
+                "If you are trying to set a catalog, please use `catalog` instead.\n"
+            )
+
+        session_properties = self.session_properties or {}
+        if CATALOG_KEY_IN_SESSION_PROPERTIES in session_properties:
+            if self.database is None:
+                self.database = session_properties[CATALOG_KEY_IN_SESSION_PROPERTIES]
+                del session_properties[CATALOG_KEY_IN_SESSION_PROPERTIES]
+            else:
+                raise DbtValidationError(
+                    f"Got duplicate keys: (`{CATALOG_KEY_IN_SESSION_PROPERTIES}` "
+                    'in session_properties) all map to "database"'
+                )
+        self.session_properties = session_properties
+
+        if self.database is not None:
+            database = self.database.strip()
+            if not database:
+                raise DbtValidationError(f"Invalid catalog name : `{self.database}`.")
+            self.database = database
+        else:
+            self.database = "hive_metastore"
+
+        connection_parameters = self.connection_parameters or {}
+        for key in (
+            "server_hostname",
+            "http_path",
+            "access_token",
+            "client_id",
+            "client_secret",
+            "session_configuration",
+            "catalog",
+            "schema",
+            "_user_agent_entry",
+        ):
+            if key in connection_parameters:
+                raise DbtValidationError(f"The connection parameter `{key}` is reserved.")
+        if "http_headers" in connection_parameters:
+            http_headers = connection_parameters["http_headers"]
+            if not isinstance(http_headers, dict) or any(
+                not isinstance(key, str) or not isinstance(value, str)
+                for key, value in http_headers.items()
+            ):
+                raise DbtValidationError(
+                    "The connection parameter `http_headers` should be dict of strings: "
+                    f"{http_headers}."
+                )
+        if "_socket_timeout" not in connection_parameters:
+            connection_parameters["_socket_timeout"] = 600
+        self.connection_parameters = connection_parameters
+
+    def validate_creds(self) -> None:
+        for key in ["host", "http_path"]:
+            if not getattr(self, key):
+                raise DbtConfigError(
+                    "The config '{}' is required to connect to Databricks".format(key)
+                )
+        if not self.token and self.auth_type != "oauth":
+            raise DbtConfigError(
+                ("The config `auth_type: oauth` is required when not using access token")
+            )
+
+        if not self.client_id and self.client_secret:
+            raise DbtConfigError(
+                (
+                    "The config 'client_id' is required to connect "
+                    "to Databricks when 'client_secret' is present"
+                )
+            )
+
+    @classmethod
+    def get_invocation_env(cls) -> Optional[str]:
+        invocation_env = os.environ.get(DBT_DATABRICKS_INVOCATION_ENV)
+        if invocation_env:
+            # Thrift doesn't allow nested () so we need to ensure
+            # that the passed user agent is valid.
+            if not DBT_DATABRICKS_INVOCATION_ENV_REGEX.search(invocation_env):
+                raise DbtValidationError(f"Invalid invocation environment: {invocation_env}")
+        return invocation_env
+
+    @classmethod
+    def get_all_http_headers(cls, user_http_session_headers: Dict[str, str]) -> Dict[str, str]:
+        http_session_headers_str: Optional[str] = os.environ.get(
+            DBT_DATABRICKS_HTTP_SESSION_HEADERS
+        )
+
+        http_session_headers_dict: Dict[str, str] = (
+            {
+                k: v if isinstance(v, str) else json.dumps(v)
+                for k, v in json.loads(http_session_headers_str).items()
+            }
+            if http_session_headers_str is not None
+            else {}
+        )
+
+        intersect_http_header_keys = (
+            user_http_session_headers.keys() & http_session_headers_dict.keys()
+        )
+
+        if len(intersect_http_header_keys) > 0:
+            raise DbtValidationError(
+                f"Intersection with reserved http_headers in keys: {intersect_http_header_keys}"
+            )
+
+        http_session_headers_dict.update(user_http_session_headers)
+
+        return http_session_headers_dict
+
+    @property
+    def type(self) -> str:
+        return "databricks"
+
+    @property
+    def unique_field(self) -> str:
+        return cast(str, self.host)
+
+    def connection_info(self, *, with_aliases: bool = False) -> Iterable[Tuple[str, Any]]:
+        as_dict = self.to_dict(omit_none=False)
+        connection_keys = set(self._connection_keys(with_aliases=with_aliases))
+        aliases: List[str] = []
+        if with_aliases:
+            aliases = [k for k, v in self._ALIASES.items() if v in connection_keys]
+        for key in itertools.chain(self._connection_keys(with_aliases=with_aliases), aliases):
+            if key in as_dict:
+                yield key, as_dict[key]
+
+    def _connection_keys(self, *, with_aliases: bool = False) -> Tuple[str, ...]:
+        # Assuming `DatabricksCredentials.connection_info(self, *, with_aliases: bool = False)`
+        # is called from only:
+        #
+        # - `Profile` with `with_aliases=True`
+        # - `DebugTask` without `with_aliases` (`False` by default)
+        #
+        # Thus, if `with_aliases` is `True`, `DatabricksCredentials._connection_keys` should return
+        # the internal key names; otherwise it can use aliases to show in `dbt debug`.
+        connection_keys = ["host", "http_path", "schema"]
+        if with_aliases:
+            connection_keys.insert(2, "database")
+        elif self.database:
+            connection_keys.insert(2, "catalog")
+        if self.session_properties:
+            connection_keys.append("session_properties")
+        return tuple(connection_keys)
+
+    @classmethod
+    def extract_cluster_id(cls, http_path: str) -> Optional[str]:
+        m = EXTRACT_CLUSTER_ID_FROM_HTTP_PATH_REGEX.match(http_path)
+        if m:
+            return m.group(1).strip()
+        else:
+            return None
+
+    @property
+    def cluster_id(self) -> Optional[str]:
+        return self.extract_cluster_id(self.http_path)  # type: ignore[arg-type]
+
+    def authenticate(self, in_provider: Optional[TCredentialProvider]) -> TCredentialProvider:
+        self.validate_creds()
+        host: str = self.host or ""
+        if self._credentials_provider:
+            return self._provider_from_dict()  # type: ignore
+        if in_provider:
+            if isinstance(in_provider, m2m_auth) or isinstance(in_provider, token_auth):
+                self._credentials_provider = in_provider.as_dict()
+            return in_provider
+
+        provider: TCredentialProvider
+        # dbt will spin up multiple threads. This has to be sync. So lock here
+        self._lock.acquire()
+        try:
+            if self.token:
+                provider = token_auth(self.token)
+                self._credentials_provider = provider.as_dict()
+                return provider
+
+            if self.client_id and self.client_secret:
+                provider = m2m_auth(
+                    host=host,
+                    client_id=self.client_id or "",
+                    client_secret=self.client_secret or "",
+                )
+                self._credentials_provider = provider.as_dict()
+                return provider
+
+            oauth_client = OAuthClient(
+                host=host,
+                client_id=self.client_id if self.client_id else CLIENT_ID,
+                client_secret="",
+                redirect_url=REDIRECT_URL,
+                scopes=SCOPES,
+            )
+            # optional branch. Try and keep going if it does not work
+            try:
+                # try to get cached credentials
+                credsdict = self.get_sharded_password("dbt-databricks", host)
+
+                if credsdict:
+                    provider = SessionCredentials.from_dict(oauth_client, json.loads(credsdict))
+                    # if refresh token is expired, this will throw
+                    try:
+                        if provider.token().valid:
+                            return provider
+                    except Exception as e:
+                        logger.debug(e)
+                        # whatever it is, get rid of the cache
+                        self.delete_sharded_password("dbt-databricks", host)
+
+            # error with keyring. Maybe machine has no password persistency
+            except Exception as e:
+                logger.debug(e)
+                logger.info("could not retrieved saved token")
+
+            # no token, go fetch one
+            consent = oauth_client.initiate_consent()
+
+            provider = consent.launch_external_browser()
+            # save for later
+            self._credentials_provider = provider.as_dict()
+            try:
+                self.set_sharded_password(
+                    "dbt-databricks", host, json.dumps(self._credentials_provider)
+                )
+            # error with keyring. Maybe machine has no password persistency
+            except Exception as e:
+                logger.debug(e)
+                logger.info("could not save token")
+
+            return provider
+
+        finally:
+            self._lock.release()
+
+    def set_sharded_password(self, service_name: str, username: str, password: str) -> None:
+        max_size = MAX_NT_PASSWORD_SIZE
+
+        # if not Windows or "small" password, stick to the default
+        if os.name != "nt" or len(password) < max_size:
+            keyring.set_password(service_name, username, password)
+        else:
+            logger.debug(f"password is {len(password)} characters, sharding it.")
+
+            password_shards = [
+                password[i : i + max_size] for i in range(0, len(password), max_size)
+            ]
+            shard_info = {
+                "sharded_password": True,
+                "shard_count": len(password_shards),
+            }
+
+            # store the "shard info" as the "base" password
+            keyring.set_password(service_name, username, json.dumps(shard_info))
+            # then store all shards with the shard number as postfix
+            for i, s in enumerate(password_shards):
+                keyring.set_password(service_name, f"{username}__{i}", s)
+
+    def get_sharded_password(self, service_name: str, username: str) -> Optional[str]:
+        password = keyring.get_password(service_name, username)
+
+        # check for "shard info" stored as json
+        try:
+            password_as_dict = json.loads(str(password))
+            if password_as_dict.get("sharded_password"):
+                # if password was stored shared, reconstruct it
+                shard_count = int(password_as_dict.get("shard_count"))
+
+                password = ""
+                for i in range(shard_count):
+                    password += str(keyring.get_password(service_name, f"{username}__{i}"))
+        except ValueError:
+            pass
+
+        return password
+
+    def delete_sharded_password(self, service_name: str, username: str) -> None:
+        password = keyring.get_password(service_name, username)
+
+        # check for "shard info" stored as json. If so delete all shards
+        try:
+            password_as_dict = json.loads(str(password))
+            if password_as_dict.get("sharded_password"):
+                shard_count = int(password_as_dict.get("shard_count"))
+                for i in range(shard_count):
+                    keyring.delete_password(service_name, f"{username}__{i}")
+        except ValueError:
+            pass
+
+        # delete "base" password
+        keyring.delete_password(service_name, username)
+
+    def _provider_from_dict(self) -> Optional[TCredentialProvider]:
+        if self.token:
+            return token_auth.from_dict(self._credentials_provider)
+
+        if self.client_id and self.client_secret:
+            return m2m_auth.from_dict(
+                host=self.host or "",
+                client_id=self.client_id or "",
+                client_secret=self.client_secret or "",
+                raw=self._credentials_provider or {"token": {}},
+            )
+
+        oauth_client = OAuthClient(
+            host=self.host or "",
+            client_id=CLIENT_ID,
+            client_secret="",
+            redirect_url=REDIRECT_URL,
+            scopes=SCOPES,
+        )
+
+        return SessionCredentials.from_dict(
+            client=oauth_client, raw=self._credentials_provider or {"token": {}}
+        )
diff --git a/dbt/adapters/databricks/python_submissions.py b/dbt/adapters/databricks/python_submissions.py
index 3f33a7e1..2d874fa9 100644
--- a/dbt/adapters/databricks/python_submissions.py
+++ b/dbt/adapters/databricks/python_submissions.py
@@ -15,9 +15,9 @@ from urllib3.util.retry import Retry
 from dbt.adapters.base import PythonJobHelper
 from dbt.adapters.databricks import utils
 from dbt.adapters.databricks.__version__ import version
-from dbt.adapters.databricks.connections import BearerAuth
-from dbt.adapters.databricks.connections import DatabricksCredentials
-from dbt.adapters.databricks.connections import TCredentialProvider
+from dbt.adapters.databricks.auth import BearerAuth
+from dbt.adapters.databricks.credentials import DatabricksCredentials
+from dbt.adapters.databricks.credentials import TCredentialProvider
 from dbt.adapters.databricks.logging import logger
 
 


Test patch:
diff --git a/tests/unit/python/test_python_submissions.py b/tests/unit/python/test_python_submissions.py
index 4fafb84f..27f1eb5f 100644
--- a/tests/unit/python/test_python_submissions.py
+++ b/tests/unit/python/test_python_submissions.py
@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import Mock
 
-from dbt.adapters.databricks.connections import DatabricksCredentials
+from dbt.adapters.databricks.credentials import DatabricksCredentials
 from dbt.adapters.databricks.python_submissions import BaseDatabricksHelper
 from dbt.adapters.databricks.python_submissions import DBContext
 
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index bb5d4c53..2c7ca4ee 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -5,18 +5,17 @@ from typing import Dict
 from typing import Optional
 from unittest import mock
 
-import pytest
-from agate import Row
-
 import dbt.adapters.exceptions
 import dbt.flags as flags
+import pytest
+from agate import Row
 from dbt.adapters.databricks import __version__
 from dbt.adapters.databricks import DatabricksAdapter
 from dbt.adapters.databricks import DatabricksRelation
 from dbt.adapters.databricks.column import DatabricksColumn
-from dbt.adapters.databricks.connections import CATALOG_KEY_IN_SESSION_PROPERTIES
-from dbt.adapters.databricks.connections import DBT_DATABRICKS_HTTP_SESSION_HEADERS
-from dbt.adapters.databricks.connections import DBT_DATABRICKS_INVOCATION_ENV
+from dbt.adapters.databricks.credentials import CATALOG_KEY_IN_SESSION_PROPERTIES
+from dbt.adapters.databricks.credentials import DBT_DATABRICKS_HTTP_SESSION_HEADERS
+from dbt.adapters.databricks.credentials import DBT_DATABRICKS_INVOCATION_ENV
 from dbt.adapters.databricks.impl import check_not_found_error
 from dbt.adapters.databricks.impl import get_identifier_list_string
 from dbt.config import RuntimeConfig
diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py
index d619c384..cd2f171c 100644
--- a/tests/unit/test_auth.py
+++ b/tests/unit/test_auth.py
@@ -3,7 +3,7 @@ import unittest
 import keyring.backend
 import pytest
 
-from dbt.adapters.databricks.connections import DatabricksCredentials
+from dbt.adapters.databricks.credentials import DatabricksCredentials
 
 
 @pytest.mark.skip(reason="Need to mock requests to OIDC")
diff --git a/tests/unit/test_compute_config.py b/tests/unit/test_compute_config.py
index 12920f0e..91d70b0f 100644
--- a/tests/unit/test_compute_config.py
+++ b/tests/unit/test_compute_config.py
@@ -2,6 +2,7 @@ import unittest
 
 import dbt.exceptions
 from dbt.adapters.databricks import connections
+from dbt.adapters.databricks.credentials import DatabricksCredentials
 from dbt.contracts.graph import model_config
 from dbt.contracts.graph import nodes
 
@@ -15,7 +16,7 @@ class TestDatabricksConnectionHTTPPath(unittest.TestCase):
 
     def test_get_http_path_model(self):
         default_path = "my_http_path"
-        creds = connections.DatabricksCredentials(http_path=default_path)
+        creds = DatabricksCredentials(http_path=default_path)
 
         path = connections._get_http_path(None, creds)
         self.assertEqual(default_path, path)
@@ -72,7 +73,7 @@ class TestDatabricksConnectionHTTPPath(unittest.TestCase):
 
     def test_get_http_path_seed(self):
         default_path = "my_http_path"
-        creds = connections.DatabricksCredentials(http_path=default_path)
+        creds = DatabricksCredentials(http_path=default_path)
 
         path = connections._get_http_path(None, creds)
         self.assertEqual(default_path, path)
@@ -129,7 +130,7 @@ class TestDatabricksConnectionHTTPPath(unittest.TestCase):
 
     def test_get_http_path_snapshot(self):
         default_path = "my_http_path"
-        creds = connections.DatabricksCredentials(http_path=default_path)
+        creds = DatabricksCredentials(http_path=default_path)
 
         path = connections._get_http_path(None, creds)
         self.assertEqual(default_path, path)
diff --git a/tests/unit/test_idle_config.py b/tests/unit/test_idle_config.py
index eac7d9d2..1e317e2c 100644
--- a/tests/unit/test_idle_config.py
+++ b/tests/unit/test_idle_config.py
@@ -1,12 +1,12 @@
-import unittest
-
-import dbt.exceptions
+import pytest
 from dbt.adapters.databricks import connections
+from dbt.adapters.databricks.credentials import DatabricksCredentials
 from dbt.contracts.graph import model_config
 from dbt.contracts.graph import nodes
+from dbt_common.exceptions import DbtRuntimeError
 
 
-class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
+class TestDatabricksConnectionMaxIdleTime:
     """Test the various cases for determining a specified warehouse."""
 
     errMsg = (
@@ -14,11 +14,11 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
     )
 
     def test_get_max_idle_default(self):
-        creds = connections.DatabricksCredentials()
+        creds = DatabricksCredentials()
 
         # No node and nothing specified in creds
         time = connections._get_max_idle_time(None, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
 
         node = nodes.ModelNode(
             relation_name="a_relation",
@@ -37,31 +37,31 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
 
         # node has no configuration so should get back default
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
 
         # empty configuration should return default
         node.config = model_config.ModelConfig()
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
 
         # node with no extras in configuration should return default
         node.config._extra = {}
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
 
         # node that specifies a compute with no corresponding definition should return default
         node.config._extra["databricks_compute"] = "foo"
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
 
         creds.compute = {}
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
 
         # if alternate compute doesn't specify a max time should return default
         creds.compute = {"foo": {}}
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(connections.DEFAULT_MAX_IDLE_TIME, time)
+        assert connections.DEFAULT_MAX_IDLE_TIME == time
         # with self.assertRaisesRegex(
         #     dbt.exceptions.DbtRuntimeError,
         #     self.errMsg,
@@ -74,11 +74,11 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
 
     def test_get_max_idle_creds(self):
         creds_idle_time = 77
-        creds = connections.DatabricksCredentials(connect_max_idle=creds_idle_time)
+        creds = DatabricksCredentials(connect_max_idle=creds_idle_time)
 
         # No node so value should come from creds
         time = connections._get_max_idle_time(None, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
         node = nodes.ModelNode(
             relation_name="a_relation",
@@ -97,36 +97,36 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
 
         # node has no configuration so should get value from creds
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
         # empty configuration should get value from creds
         node.config = model_config.ModelConfig()
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
         # node with no extras in configuration should get value from creds
         node.config._extra = {}
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
         # node that specifies a compute with no corresponding definition should get value from creds
         node.config._extra["databricks_compute"] = "foo"
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
         creds.compute = {}
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
         # if alternate compute doesn't specify a max time should get value from creds
         creds.compute = {"foo": {}}
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(creds_idle_time, time)
+        assert creds_idle_time == time
 
     def test_get_max_idle_compute(self):
         creds_idle_time = 88
         compute_idle_time = 77
-        creds = connections.DatabricksCredentials(connect_max_idle=creds_idle_time)
+        creds = DatabricksCredentials(connect_max_idle=creds_idle_time)
         creds.compute = {"foo": {"connect_max_idle": compute_idle_time}}
 
         node = nodes.SnapshotNode(
@@ -149,12 +149,12 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
         node.config._extra = {"databricks_compute": "foo"}
 
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(compute_idle_time, time)
+        assert compute_idle_time == time
 
     def test_get_max_idle_invalid(self):
         creds_idle_time = "foo"
         compute_idle_time = "bar"
-        creds = connections.DatabricksCredentials(connect_max_idle=creds_idle_time)
+        creds = DatabricksCredentials(connect_max_idle=creds_idle_time)
         creds.compute = {"alternate_compute": {"connect_max_idle": compute_idle_time}}
 
         node = nodes.SnapshotNode(
@@ -175,39 +175,39 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
 
         node.config = model_config.SnapshotConfig()
 
-        with self.assertRaisesRegex(
-            dbt.exceptions.DbtRuntimeError,
-            f"{creds_idle_time} is not a valid value for connect_max_idle. "
-            "Must be a number of seconds.",
-        ):
+        with pytest.raises(DbtRuntimeError) as info:
             connections._get_max_idle_time(node, creds)
+        assert (
+            f"{creds_idle_time} is not a valid value for connect_max_idle. "
+            "Must be a number of seconds."
+        ) in str(info.value)
 
         node.config._extra["databricks_compute"] = "alternate_compute"
-        with self.assertRaisesRegex(
-            dbt.exceptions.DbtRuntimeError,
-            f"{compute_idle_time} is not a valid value for connect_max_idle. "
-            "Must be a number of seconds.",
-        ):
+        with pytest.raises(DbtRuntimeError) as info:
             connections._get_max_idle_time(node, creds)
+        assert (
+            f"{compute_idle_time} is not a valid value for connect_max_idle. "
+            "Must be a number of seconds."
+        ) in str(info.value)
 
         creds.compute["alternate_compute"]["connect_max_idle"] = "1.2.3"
-        with self.assertRaisesRegex(
-            dbt.exceptions.DbtRuntimeError,
-            "1.2.3 is not a valid value for connect_max_idle. " "Must be a number of seconds.",
-        ):
+        with pytest.raises(DbtRuntimeError) as info:
             connections._get_max_idle_time(node, creds)
+        assert (
+            "1.2.3 is not a valid value for connect_max_idle. " "Must be a number of seconds."
+        ) in str(info.value)
 
         creds.compute["alternate_compute"]["connect_max_idle"] = "1,002.3"
-        with self.assertRaisesRegex(
-            dbt.exceptions.DbtRuntimeError,
-            "1,002.3 is not a valid value for connect_max_idle. " "Must be a number of seconds.",
-        ):
+        with pytest.raises(DbtRuntimeError) as info:
             connections._get_max_idle_time(node, creds)
+        assert (
+            "1,002.3 is not a valid value for connect_max_idle. " "Must be a number of seconds."
+        ) in str(info.value)
 
     def test_get_max_idle_simple_string_conversion(self):
         creds_idle_time = "12"
         compute_idle_time = "34"
-        creds = connections.DatabricksCredentials(connect_max_idle=creds_idle_time)
+        creds = DatabricksCredentials(connect_max_idle=creds_idle_time)
         creds.compute = {"alternate_compute": {"connect_max_idle": compute_idle_time}}
 
         node = nodes.SnapshotNode(
@@ -229,12 +229,12 @@ class TestDatabricksConnectionMaxIdleTime(unittest.TestCase):
         node.config = model_config.SnapshotConfig()
 
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(float(creds_idle_time), time)
+        assert float(creds_idle_time) == time
 
         node.config._extra["databricks_compute"] = "alternate_compute"
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(float(compute_idle_time), time)
+        assert float(compute_idle_time) == time
 
         creds.compute["alternate_compute"]["connect_max_idle"] = "  56 "
         time = connections._get_max_idle_time(node, creds)
-        self.assertEqual(56, time)
+        assert 56 == time