databricks-sql-connector 3.0.3b1__tar.gz → 3.1.0__tar.gz

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/CHANGELOG.md

@@ -1,8 +1,15 @@
 # Release History
 
-# 3.0.3b1 (2024-01-29)
+# 3.1.0 (2024-02-16)
 
-- Fix: Thrift connections would fail if installed `urllib3` was below version `1.26.0`
+- Revert retry-after behavior to be exponential backoff (#349)
+- Support Databricks OAuth on Azure (#351)
+- Support Databricks OAuth on GCP (#338)
+
+# 3.0.3 (2024-02-02)
+
+- Revised docstrings and examples for OAuth (#339)
+- Redact the URL query parameters from the urllib3.connectionpool logs (#341)
 
 # 3.0.2 (2024-01-25)
 
@@ -48,7 +55,7 @@
 
 ## 2.9.2 (2023-08-17)
 
-__Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed. The log changes are incorporated into version 2.9.3 and greater.__
+**Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed. The log changes are incorporated into version 2.9.3 and greater.**
 
 - Other: Add `examples/v3_retries_query_execute.py` (#199)
 - Other: suppress log message when `_enable_v3_retries` is not `True` (#199)
@@ -56,7 +63,7 @@ __Note: this release was yanked from Pypi on 13 September 2023 due to compatibil
 
 ## 2.9.1 (2023-08-11)
 
-__Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed.__
+**Note: this release was yanked from Pypi on 13 September 2023 due to compatibility issues with environments where `urllib3<=2.0.0` were installed.**
 
 - Other: Explicitly pin urllib3 to ^2.0.0 (#191)
 
@@ -109,6 +116,7 @@ __Note: this release was yanked from Pypi on 13 September 2023 due to compatibil
 - Other: Relax sqlalchemy required version as it was unecessarily strict.
 
 ## 2.5.0 (2023-04-14)
+
 - Add support for External Auth providers
 - Fix: Python HTTP proxies were broken
 - Other: All Thrift requests that timeout during connection will be automatically retried
@@ -130,8 +138,8 @@ __Note: this release was yanked from Pypi on 13 September 2023 due to compatibil
 
 ## 2.2.2 (2023-01-03)
 
-- Support custom oauth client id and redirect port 
-- Fix: Add none check on _oauth_persistence in DatabricksOAuthProvider
+- Support custom oauth client id and redirect port
+- Fix: Add none check on \_oauth_persistence in DatabricksOAuthProvider
 
 ## 2.2.1 (2022-11-29)
 
@@ -163,57 +171,71 @@ Huge thanks to @dbaxa for contributing this change!
 
 - Add retry logic for `GetOperationStatus` requests that fail with an `OSError`
 - Reorganised code to use Poetry for dependency management.
+
 ## 2.0.2 (2022-05-04)
+
 - Better exception handling in automatic connection close
 
 ## 2.0.1 (2022-04-21)
+
 - Fixed Pandas dependency in setup.cfg to be >= 1.2.0
 
 ## 2.0.0 (2022-04-19)
+
 - Initial stable release of V2
-- Added better support for complex types, so that in Databricks runtime 10.3+, Arrays, Maps and Structs will get 
+- Added better support for complex types, so that in Databricks runtime 10.3+, Arrays, Maps and Structs will get
   deserialized as lists, lists of tuples and dicts, respectively.
 - Changed the name of the metadata arg to http_headers
 
 ## 2.0.b2 (2022-04-04)
+
 - Change import of collections.Iterable to collections.abc.Iterable to make the library compatible with Python 3.10
 - Fixed bug with .tables method so that .tables works as expected with Unity-Catalog enabled endpoints
 
 ## 2.0.0b1 (2022-03-04)
+
 - Fix packaging issue (dependencies were not being installed properly)
 - Fetching timestamp results will now return aware instead of naive timestamps
 - The client will now default to using simplified error messages
 
 ## 2.0.0b (2022-02-08)
+
 - Initial beta release of V2. V2 is an internal re-write of large parts of the connector to use Databricks edge features. All public APIs from V1 remain.
-- Added Unity Catalog support (pass catalog and / or  schema key word args to the .connect method to select initial schema and catalog)
+- Added Unity Catalog support (pass catalog and / or schema key word args to the .connect method to select initial schema and catalog)
 
 ---
 
 **Note**: The code for versions prior to `v2.0.0b` is not contained in this repository. The below entries are included for reference only.
 
 ---
+
 ## 1.0.0 (2022-01-20)
+
 - Add operations for retrieving metadata
 - Add the ability to access columns by name on result rows
 - Add the ability to provide configuration settings on connect
 
 ## 0.9.4 (2022-01-10)
+
 - Improved logging and error messages.
 
 ## 0.9.3 (2021-12-08)
+
 - Add retries for 429 and 503 HTTP responses.
 
 ## 0.9.2 (2021-12-02)
+
 - (Bug fix) Increased Thrift requirement from 0.10.0 to 0.13.0 as 0.10.0 was in fact incompatible
 - (Bug fix) Fixed error message after query execution failed -SQLSTATE and Error message were misplaced
 
 ## 0.9.1 (2021-09-01)
+
 - Public Preview release, Experimental tag removed
 - minor updates in internal build/packaging
 - no functional changes
 
 ## 0.9.0 (2021-08-04)
+
 - initial (Experimental) release of pyhive-forked connector
 - Python DBAPI 2.0 (PEP-0249), thrift based
 - see docs for more info: https://docs.databricks.com/dev-tools/python-sql-connector.html

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: databricks-sql-connector
-Version: 3.0.3b1
+Version: 3.1.0
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
 Author: Databricks
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: alembic
 Provides-Extra: sqlalchemy
 Requires-Dist: alembic (>=1.0.11,<2.0.0) ; extra == "alembic"

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "3.0.3b1"
+version = "3.1.0"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"
@@ -58,6 +58,7 @@ exclude = ['ttypes\.py$', 'TCLIService\.py$']
 exclude = '/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist|thrift_api)/'
 
 [tool.pytest.ini_options]
+markers = {"reviewed" = "Test case has been reviewed by Databricks"}
 minversion = "6.0"
 log_cli = "false"
 log_cli_level = "INFO"

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sql/__init__.py

@@ -8,6 +8,39 @@ threadsafety = 1  # Threads may share the module, but not connections.
 
 paramstyle = "named"
 
+import re
+
+
+class RedactUrlQueryParamsFilter(logging.Filter):
+    pattern = re.compile(r"(\?|&)([\w-]+)=([^&]+)")
+    mask = r"\1\2=<REDACTED>"
+
+    def __init__(self):
+        super().__init__()
+
+    def redact(self, string):
+        return re.sub(self.pattern, self.mask, str(string))
+
+    def filter(self, record):
+        record.msg = self.redact(str(record.msg))
+        if isinstance(record.args, dict):
+            for k in record.args.keys():
+                record.args[k] = (
+                    self.redact(record.args[k])
+                    if isinstance(record.arg[k], str)
+                    else record.args[k]
+                )
+        else:
+            record.args = tuple(
+                (self.redact(arg) if isinstance(arg, str) else arg)
+                for arg in record.args
+            )
+
+        return True
+
+
+logging.getLogger("urllib3.connectionpool").addFilter(RedactUrlQueryParamsFilter())
+
 
 class DBAPITypeObject(object):
     def __init__(self, *values):
@@ -29,7 +62,7 @@ DATETIME = DBAPITypeObject("timestamp")
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "3.0.3b1"
+__version__ = "3.1.0"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sql/auth/auth.py

@@ -8,12 +8,11 @@ from databricks.sql.auth.authenticators import (
     ExternalAuthProvider,
     DatabricksOAuthProvider,
 )
-from databricks.sql.auth.endpoint import infer_cloud_from_host, CloudType
-from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 
 
 class AuthType(Enum):
     DATABRICKS_OAUTH = "databricks-oauth"
+    AZURE_OAUTH = "azure-oauth"
     # other supported types (access_token, user/pass) can be inferred
     # we can add more types as needed later
 
@@ -51,7 +50,7 @@ class ClientContext:
 def get_auth_provider(cfg: ClientContext):
     if cfg.credentials_provider:
         return ExternalAuthProvider(cfg.credentials_provider)
-    if cfg.auth_type == AuthType.DATABRICKS_OAUTH.value:
+    if cfg.auth_type in [AuthType.DATABRICKS_OAUTH.value, AuthType.AZURE_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
@@ -62,6 +61,7 @@ def get_auth_provider(cfg: ClientContext):
             cfg.oauth_redirect_port_range,
             cfg.oauth_client_id,
             cfg.oauth_scopes,
+            cfg.auth_type,
         )
     elif cfg.access_token is not None:
         return AccessTokenAuthProvider(cfg.access_token)
@@ -87,19 +87,22 @@ def normalize_host_name(hostname: str):
     return f"{maybe_scheme}{hostname}{maybe_trailing_slash}"
 
 
-def get_client_id_and_redirect_port(hostname: str):
+def get_client_id_and_redirect_port(use_azure_auth: bool):
     return (
         (PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_REDIRECT_PORT_RANGE)
-        if infer_cloud_from_host(hostname) == CloudType.AWS
+        if not use_azure_auth
         else (PYSQL_OAUTH_AZURE_CLIENT_ID, PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE)
     )
 
 
 def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
-    (client_id, redirect_port_range) = get_client_id_and_redirect_port(hostname)
+    auth_type = kwargs.get("auth_type")
+    (client_id, redirect_port_range) = get_client_id_and_redirect_port(
+        auth_type == AuthType.AZURE_OAUTH.value
+    )
     cfg = ClientContext(
         hostname=normalize_host_name(hostname),
-        auth_type=kwargs.get("auth_type"),
+        auth_type=auth_type,
         access_token=kwargs.get("access_token"),
         username=kwargs.get("_username"),
         password=kwargs.get("_password"),

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sql/auth/authenticators.py

@@ -18,6 +18,7 @@ class AuthProvider:
 
 HeaderFactory = Callable[[], Dict[str, str]]
 
+
 # In order to keep compatibility with SDK
 class CredentialsProvider(abc.ABC):
     """CredentialsProvider is the protocol (call-side interface)
@@ -69,16 +70,13 @@ class DatabricksOAuthProvider(AuthProvider):
         redirect_port_range: List[int],
         client_id: str,
         scopes: List[str],
+        auth_type: str = "databricks-oauth",
     ):
         try:
-            cloud_type = infer_cloud_from_host(hostname)
-            if not cloud_type:
-                raise NotImplementedError("Cannot infer the cloud type from hostname")
-
-            idp_endpoint = get_oauth_endpoints(cloud_type)
+            idp_endpoint = get_oauth_endpoints(hostname, auth_type == "azure-oauth")
             if not idp_endpoint:
                 raise NotImplementedError(
-                    f"OAuth is not supported for cloud ${cloud_type.value}"
+                    f"OAuth is not supported for host ${hostname}"
                 )
 
             # Convert to the corresponding scopes in the corresponding IdP

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sql/auth/endpoint.py

@@ -1,9 +1,9 @@
 #
 # It implements all the cloud specific OAuth configuration/metadata
 #
-#   Azure: It uses AAD
+#   Azure:  It uses Databricks internal IdP or Azure AD
 #   AWS: It uses Databricks internal IdP
-#   GCP: Not support yet
+#   GCP: It uses Databricks internal IdP
 #
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -21,6 +21,7 @@ class OAuthScope:
 class CloudType(Enum):
     AWS = "aws"
     AZURE = "azure"
+    GCP = "gcp"
 
 
 DATABRICKS_AWS_DOMAINS = [
@@ -34,6 +35,10 @@ DATABRICKS_AZURE_DOMAINS = [
     ".databricks.azure.cn",
     ".databricks.azure.us",
 ]
+DATABRICKS_GCP_DOMAINS = [".gcp.databricks.com"]
+
+# Domain supported by Databricks InHouse OAuth
+DATABRICKS_OAUTH_AZURE_DOMAINS = [".azuredatabricks.net"]
 
 
 # Infer cloud type from Databricks SQL instance hostname
@@ -45,10 +50,20 @@ def infer_cloud_from_host(hostname: str) -> Optional[CloudType]:
         return CloudType.AZURE
     elif any(e for e in DATABRICKS_AWS_DOMAINS if host.endswith(e)):
         return CloudType.AWS
+    elif any(e for e in DATABRICKS_GCP_DOMAINS if host.endswith(e)):
+        return CloudType.GCP
     else:
         return None
 
 
+def is_supported_databricks_oauth_host(hostname: str) -> bool:
+    host = hostname.lower().replace("https://", "").split("/")[0]
+    domains = (
+        DATABRICKS_AWS_DOMAINS + DATABRICKS_GCP_DOMAINS + DATABRICKS_OAUTH_AZURE_DOMAINS
+    )
+    return any(e for e in domains if host.endswith(e))
+
+
 def get_databricks_oidc_url(hostname: str):
     maybe_scheme = "https://" if not hostname.startswith("https://") else ""
     maybe_trailing_slash = "/" if not hostname.endswith("/") else ""
@@ -94,7 +109,7 @@ class AzureOAuthEndpointCollection(OAuthEndpointCollection):
         return "https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration"
 
 
-class AwsOAuthEndpointCollection(OAuthEndpointCollection):
+class InHouseOAuthEndpointCollection(OAuthEndpointCollection):
     def get_scopes_mapping(self, scopes: List[str]) -> List[str]:
         # No scope mapping in AWS
         return scopes.copy()
@@ -108,10 +123,18 @@ class AwsOAuthEndpointCollection(OAuthEndpointCollection):
         return f"{idp_url}/.well-known/oauth-authorization-server"
 
 
-def get_oauth_endpoints(cloud: CloudType) -> Optional[OAuthEndpointCollection]:
-    if cloud == CloudType.AWS:
-        return AwsOAuthEndpointCollection()
+def get_oauth_endpoints(
+    hostname: str, use_azure_auth: bool
+) -> Optional[OAuthEndpointCollection]:
+    cloud = infer_cloud_from_host(hostname)
+
+    if cloud in [CloudType.AWS, CloudType.GCP]:
+        return InHouseOAuthEndpointCollection()
     elif cloud == CloudType.AZURE:
-        return AzureOAuthEndpointCollection()
+        return (
+            InHouseOAuthEndpointCollection()
+            if is_supported_databricks_oauth_host(hostname) and not use_azure_auth
+            else AzureOAuthEndpointCollection()
+        )
     else:
         return None

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sql/auth/retry.py

@@ -118,13 +118,13 @@ class DatabricksRetryPolicy(Retry):
             _total: int = urllib3_kwargs.pop("total")
             _attempts_remaining = _total
 
-        _urllib_kwargs_we_care_about = {
-            "total": _attempts_remaining,
-            "respect_retry_after_header": True,
-            "backoff_factor": self.delay_min,
-            self._allowed_methods_alias: ["POST"],
-            "status_forcelist": [429, 503, *self.force_dangerous_codes],
-        }
+        _urllib_kwargs_we_care_about = dict(
+            total=_attempts_remaining,
+            respect_retry_after_header=True,
+            backoff_factor=self.delay_min,
+            allowed_methods=["POST"],
+            status_forcelist=[429, 503, *self.force_dangerous_codes],
+        )
 
         urllib3_kwargs.update(**_urllib_kwargs_we_care_about)
 
@@ -132,22 +132,6 @@ class DatabricksRetryPolicy(Retry):
             **urllib3_kwargs,  # type: ignore
         )
 
-    @property
-    def _allowed_methods_alias(self) -> str:
-        """Returns `allowed_methods` if installed urllib3 is >=1.26.0
-        Returns `method_whitelist` otherwise."""
-
-        if not hasattr(self, "__allowed_methods_alias"):
-            import urllib3
-            from distutils.version import LooseVersion
-
-            if LooseVersion(urllib3.__version__) >= LooseVersion("1.26"):
-                self.__allowed_methods_alias = "allowed_methods"
-            else:
-                self.__allowed_methods_alias = "method_whitelist"
-
-        return self.__allowed_methods_alias
-
     @classmethod
     def __private_init__(
         cls, retry_start_time: float, command_type: Optional[CommandType], **init_kwargs
@@ -217,22 +201,22 @@ class DatabricksRetryPolicy(Retry):
         # Gather urllib3's current retry state _before_ increment was called
         # These arguments match the function signature for Retry.__init__
         # Note: if we update urllib3 we may need to add/remove arguments from this dict
-        urllib3_init_params = {
-            "total": self.total,
-            "connect": self.connect,
-            "read": self.read,
-            "redirect": self.redirect,
-            "status": self.status,
-            "other": self.other,
-            self._allowed_methods_alias: self.allowed_methods,
-            "status_forcelist": self.status_forcelist,
-            "backoff_factor": self.backoff_factor,  # type: ignore
-            "raise_on_redirect": self.raise_on_redirect,
-            "raise_on_status": self.raise_on_status,
-            "history": self.history,
-            "remove_headers_on_redirect": self.remove_headers_on_redirect,
-            "respect_retry_after_header": self.respect_retry_after_header,
-        }
+        urllib3_init_params = dict(
+            total=self.total,
+            connect=self.connect,
+            read=self.read,
+            redirect=self.redirect,
+            status=self.status,
+            other=self.other,
+            allowed_methods=self.allowed_methods,
+            status_forcelist=self.status_forcelist,
+            backoff_factor=self.backoff_factor,  # type: ignore
+            raise_on_redirect=self.raise_on_redirect,
+            raise_on_status=self.raise_on_status,
+            history=self.history,
+            remove_headers_on_redirect=self.remove_headers_on_redirect,
+            respect_retry_after_header=self.respect_retry_after_header,
+        )
 
         # Update urllib3's current state to reflect the incremented counters
         urllib3_init_params.update(**urllib3_incremented_counters)
@@ -299,8 +283,10 @@ class DatabricksRetryPolicy(Retry):
         """
         retry_after = self.get_retry_after(response)
         if retry_after:
-            self.check_proposed_wait(retry_after)
-            time.sleep(retry_after)
+            backoff = self.get_backoff_time()
+            proposed_wait = max(backoff, retry_after)
+            self.check_proposed_wait(proposed_wait)
+            time.sleep(proposed_wait)
             return True
 
         return False

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sql/client.py

@@ -96,11 +96,8 @@ class Connection:
                 legacy purposes and will be deprecated in a future release. When this parameter is `True` you will see
                 a warning log message. To suppress this log message, set `use_inline_params="silent"`.
             auth_type: `str`, optional
-                `databricks-oauth` : to use oauth with fine-grained permission scopes, set to `databricks-oauth`.
-                This is currently in private preview for Databricks accounts on AWS.
-                This supports User to Machine OAuth authentication for Databricks on AWS with
-                any IDP configured. This is only for interactive python applications and open a browser window.
-                Note this is beta (private preview)
+                `databricks-oauth` : to use Databricks OAuth with fine-grained permission scopes, set to `databricks-oauth`.
+                `azure-oauth` : to use Microsoft Entra ID OAuth flow, set to `azure-oauth`.
 
             oauth_client_id: `str`, optional
                 custom oauth client_id. If not specified, it will use the built-in client_id of databricks-sql-python.
@@ -111,9 +108,9 @@ class Connection:
 
             experimental_oauth_persistence: configures preferred storage for persisting oauth tokens.
                 This has to be a class implementing `OAuthPersistence`.
-                When `auth_type` is set to `databricks-oauth` without persisting the oauth token in a persistence storage
-                the oauth tokens will only be maintained in memory and if the python process restarts the end user
-                will have to login again.
+                When `auth_type` is set to `databricks-oauth` or `azure-oauth` without persisting the oauth token in a
+                persistence storage the oauth tokens will only be maintained in memory and if the python process
+                restarts the end user will have to login again.
                 Note this is beta (private preview)
 
                 For persisting the oauth token in a prod environment you should subclass and implement OAuthPersistence
@@ -609,12 +606,15 @@ class Cursor:
                     "Local file operations are restricted to paths within the configured staging_allowed_local_path"
                 )
 
-        # TODO: Experiment with DBR sending real headers.
-        # The specification says headers will be in JSON format but the current null value is actually an empty list []
+        # May be real headers, or could be json string
+        headers = (
+            json.loads(row.headers) if isinstance(row.headers, str) else row.headers
+        )
+
         handler_args = {
             "presigned_url": row.presignedUrl,
             "local_file": abs_localFile,
-            "headers": json.loads(row.headers or "{}"),
+            "headers": dict(headers) or {},
         }
 
         logger.debug(

databricks_sql_connector-3.1.0/src/databricks/sqlalchemy/test_local/conftest.py

@@ -0,0 +1,44 @@
+import os
+import pytest
+
+
+@pytest.fixture(scope="session")
+def host():
+    return os.getenv("DATABRICKS_SERVER_HOSTNAME")
+
+
+@pytest.fixture(scope="session")
+def http_path():
+    return os.getenv("DATABRICKS_HTTP_PATH")
+
+
+@pytest.fixture(scope="session")
+def access_token():
+    return os.getenv("DATABRICKS_TOKEN")
+
+
+@pytest.fixture(scope="session")
+def ingestion_user():
+    return os.getenv("DATABRICKS_USER")
+
+
+@pytest.fixture(scope="session")
+def catalog():
+    return os.getenv("DATABRICKS_CATALOG")
+
+
+@pytest.fixture(scope="session")
+def schema():
+    return os.getenv("DATABRICKS_SCHEMA", "default")
+
+
+@pytest.fixture(scope="session", autouse=True)
+def connection_details(host, http_path, access_token, ingestion_user, catalog, schema):
+    return {
+        "host": host,
+        "http_path": http_path,
+        "access_token": access_token,
+        "ingestion_user": ingestion_user,
+        "catalog": catalog,
+        "schema": schema,
+    }

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sqlalchemy/test_local/e2e/test_basic.py

@@ -1,6 +1,5 @@
 import datetime
 import decimal
-import os
 from typing import Tuple, Union, List
 from unittest import skipIf
 
@@ -19,7 +18,7 @@ from sqlalchemy.engine import Engine
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
 from sqlalchemy.schema import DropColumnComment, SetColumnComment
-from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
+from sqlalchemy.types import BOOLEAN, DECIMAL, Date, Integer, String
 
 try:
     from sqlalchemy.orm import declarative_base
@@ -49,12 +48,12 @@ def version_agnostic_select(object_to_select, *args, **kwargs):
         return select(object_to_select, *args, **kwargs)
 
 
-def version_agnostic_connect_arguments(catalog=None, schema=None) -> Tuple[str, dict]:
-    HOST = os.environ.get("host")
-    HTTP_PATH = os.environ.get("http_path")
-    ACCESS_TOKEN = os.environ.get("access_token")
-    CATALOG = catalog or os.environ.get("catalog")
-    SCHEMA = schema or os.environ.get("schema")
+def version_agnostic_connect_arguments(connection_details) -> Tuple[str, dict]:
+    HOST = connection_details["host"]
+    HTTP_PATH = connection_details["http_path"]
+    ACCESS_TOKEN = connection_details["access_token"]
+    CATALOG = connection_details["catalog"]
+    SCHEMA = connection_details["schema"]
 
     ua_connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}
 
@@ -77,8 +76,8 @@ def version_agnostic_connect_arguments(catalog=None, schema=None) -> Tuple[str,
 
 
 @pytest.fixture
-def db_engine() -> Engine:
-    conn_string, connect_args = version_agnostic_connect_arguments()
+def db_engine(connection_details) -> Engine:
+    conn_string, connect_args = version_agnostic_connect_arguments(connection_details)
     return create_engine(conn_string, connect_args=connect_args)
 
 
@@ -92,10 +91,11 @@ def run_query(db_engine: Engine, query: Union[str, Text]):
 
 
 @pytest.fixture
-def samples_engine() -> Engine:
-    conn_string, connect_args = version_agnostic_connect_arguments(
-        catalog="samples", schema="nyctaxi"
-    )
+def samples_engine(connection_details) -> Engine:
+    details = connection_details.copy()
+    details["catalog"] = "samples"
+    details["schema"] = "nyctaxi"
+    conn_string, connect_args = version_agnostic_connect_arguments(details)
     return create_engine(conn_string, connect_args=connect_args)
 
 
@@ -141,7 +141,7 @@ def test_connect_args(db_engine):
 def test_pandas_upload(db_engine, metadata_obj):
     import pandas as pd
 
-    SCHEMA = os.environ.get("schema")
+    SCHEMA = "default"
     try:
         df = pd.read_excel(
             "src/databricks/sqlalchemy/test_local/e2e/demo_data/MOCK_DATA.xlsx"
@@ -409,7 +409,9 @@ def test_get_table_names_smoke_test(samples_engine: Engine):
         _names is not None, "get_table_names did not succeed"
 
 
-def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine):
+def test_has_table_across_schemas(
+    db_engine: Engine, samples_engine: Engine, catalog: str, schema: str
+):
     """For this test to pass these conditions must be met:
     - Table samples.nyctaxi.trips must exist
     - Table samples.tpch.customer must exist
@@ -426,9 +428,6 @@ def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine):
         )
 
         # 3) Check for a table within a different catalog
-        other_catalog = os.environ.get("catalog")
-        other_schema = os.environ.get("schema")
-
         # Create a table in a different catalog
         with db_engine.connect() as conn:
             conn.execute(text("CREATE TABLE test_has_table (numbers_are_cool INT);"))
@@ -442,8 +441,8 @@ def test_has_table_across_schemas(db_engine: Engine, samples_engine: Engine):
                 assert samples_engine.dialect.has_table(
                     connection=conn,
                     table_name="test_has_table",
-                    schema=other_schema,
-                    catalog=other_catalog,
+                    schema=schema,
+                    catalog=catalog,
                 )
             finally:
                 conn.execute(text("DROP TABLE test_has_table;"))
@@ -503,12 +502,12 @@ def test_get_columns(db_engine, sample_table: str):
 
 class TestCommentReflection:
     @pytest.fixture(scope="class")
-    def engine(self):
-        HOST = os.environ.get("host")
-        HTTP_PATH = os.environ.get("http_path")
-        ACCESS_TOKEN = os.environ.get("access_token")
-        CATALOG = os.environ.get("catalog")
-        SCHEMA = os.environ.get("schema")
+    def engine(self, connection_details: dict):
+        HOST = connection_details["host"]
+        HTTP_PATH = connection_details["http_path"]
+        ACCESS_TOKEN = connection_details["access_token"]
+        CATALOG = connection_details["catalog"]
+        SCHEMA = connection_details["schema"]
 
         connection_string = f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}"
         connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}

{databricks_sql_connector-3.0.3b1 → databricks_sql_connector-3.1.0}/src/databricks/sqlalchemy/test_local/test_parsing.py

@@ -64,16 +64,16 @@ def test_extract_3l_namespace_from_bad_constraint_string():
         extract_three_level_identifier_from_constraint_string(input)
 
 
-@pytest.mark.parametrize("schema", [None, "some_schema"])
-def test_build_fk_dict(schema):
+@pytest.mark.parametrize("tschema", [None, "some_schema"])
+def test_build_fk_dict(tschema):
     fk_constraint_string = "FOREIGN KEY (`parent_user_id`) REFERENCES `main`.`some_schema`.`users` (`user_id`)"
 
-    result = build_fk_dict("some_fk_name", fk_constraint_string, schema_name=schema)
+    result = build_fk_dict("some_fk_name", fk_constraint_string, schema_name=tschema)
 
     assert result == {
         "name": "some_fk_name",
         "constrained_columns": ["parent_user_id"],
-        "referred_schema": schema,
+        "referred_schema": tschema,
         "referred_table": "users",
         "referred_columns": ["user_id"],
     }

databricks_sql_connector-3.0.3b1/src/databricks/sqlalchemy/_pytest.ini

@@ -1,3 +0,0 @@
-[pytest]
-markers = 
-    reviewed: Test case has been reviewed by databricks