Snow 2043816 kerberos proxy auth python

DESCRIPTION.md

@@ -7,6 +7,13 @@ https://docs.snowflake.com/
 Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python
 
 # Release Notes
+- v3.17(TBD)
+  - Bumped numpy dependency from <2.1.0 to <=2.2.4
+  - Added Windows support for Python 3.13.
+  - Add `bulk_upload_chunks` parameter to `write_pandas` function. Setting this parameter to True changes the behaviour of write_pandas function to first write all the data chunks to the local disk and then perform the wildcard upload of the chunks folder to the stage. In default behaviour the chunks are being saved, uploaded and deleted one by one.
+  - Add `headers_customizers` parameter to the `connect` function. Setting this parameter allows enriching outgoing request headers using a list of customizers. Only header enrichment is supported — modifying query parameters or overwriting the existing headers is not allowed.
+  - Added support for new authentication mechanism PAT with external session ID
+
 - v3.16.1(TBD)
   - Added in-band OCSP exception telemetry.
   - Added `APPLICATION_PATH` within `CLIENT_ENVIRONMENT` to distinguish between multiple scripts using the PythonConnector in the same environment.

src/snowflake/connector/connection.py

@@ -20,7 +20,16 @@
 from logging import getLogger
 from threading import Lock
 from types import TracebackType
-from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence
+from typing import (
+    Any,
+    Callable,
+    Generator,
+    Iterable,
+    Iterator,
+    MutableSequence,
+    NamedTuple,
+    Sequence,
+)
 from uuid import UUID
 
 from cryptography.hazmat.backends import default_backend
@@ -102,6 +111,11 @@
     ER_NOT_IMPLICITY_SNOWFLAKE_DATATYPE,
 )
 from .errors import DatabaseError, Error, OperationalError, ProgrammingError
+from .http_interceptor import (
+    HeadersCustomizer,
+    HeadersCustomizerInterceptor,
+    HttpInterceptor,
+)
 from .log_configuration import EasyLoggingConfigPython
 from .network import (
     DEFAULT_AUTHENTICATOR,
@@ -116,6 +130,9 @@
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
     WORKLOAD_IDENTITY_AUTHENTICATOR,
+    HttpAdapterFactory,
+    InterceptingAdapter,
+    ProxySupportAdapter,
     ReauthenticationRequest,
     SnowflakeRestful,
 )
@@ -368,6 +385,10 @@ def _get_private_bytes_from_file(
         True,
         bool,
     ),  # SNOW-XXXXX: remove the check_arrow_conversion_error_on_every_column flag
+    "headers_customizers": (
+        None,
+        (type(None), MutableSequence[HeadersCustomizer]),
+    ),
     "external_session_id": (
         None,
         str,
@@ -468,6 +489,7 @@ class SnowflakeConnection:
         token_file_path: The file path of the token file. If both token and token_file_path are provided, the token in token_file_path will be used.
         unsafe_file_write: When true, files downloaded by GET will be saved with 644 permissions. Otherwise, files will be saved with safe - owner-only permissions: 600.
         check_arrow_conversion_error_on_every_column: When true, the error check after the conversion from arrow to python types will happen for every column in the row. This is a new behaviour which fixes the bug that caused the type errors to trigger silently when occurring at any place other than last column in a row. To revert the previous (faulty) behaviour, please set this flag to false.
+        headers_customizer: List of headers customizers (HeadersCustomizer class). Setting this parameter allows enriching outgoing request headers using a list of customizers. Only header enrichment is supported — modifying query parameters or overwriting the existing headers is not allowed.
     """
 
     OCSP_ENV_LOCK = Lock()
@@ -888,7 +910,61 @@ def check_arrow_conversion_error_on_every_column(self) -> bool:
     def check_arrow_conversion_error_on_every_column(self, value: bool) -> bool:
         self._check_arrow_conversion_error_on_every_column = value
 
+    @property
+    def request_interceptors(self) -> MutableSequence[HttpInterceptor]:
+        return self._request_interceptors
+
+    @property
+    def headers_customizers(self) -> MutableSequence[HeadersCustomizer]:
+        return self._headers_customizers
+
+    @headers_customizers.setter
+    def headers_customizers(self, value: MutableSequence[HeadersCustomizer]) -> None:
+        self._headers_customizers = value
+        request_interceptors = self._create_interceptor_for_headers_customizers(value)
+        self._request_interceptors = (
+            [
+                request_interceptors,
+            ]
+            if request_interceptors
+            else []
+        )
+
+    def add_headers_customizer(
+        self, new_customizer: HeadersCustomizer
+    ) -> SnowflakeConnection:
+        """
+        Builder method to add a single headers customizer to the list of headers customizers.
+        """
+        if new_customizer in self._headers_customizers or not new_customizer:
+            return self
+
+        self._headers_customizers.append(new_customizer)
+        self._request_interceptors.append(
+            HeadersCustomizerInterceptor([new_customizer])
+        )
+        return self
+
+    def clear_headers_customizers(self) -> None:
+        self._headers_customizers.clear()
+        self._request_interceptors[:] = [
+            interceptor
+            for interceptor in self._request_interceptors
+            if not isinstance(interceptor, HeadersCustomizerInterceptor)
+        ]
+
+    @staticmethod
+    def _create_interceptor_for_headers_customizers(
+        headers_customizers: MutableSequence[HeadersCustomizer],
+    ) -> HeadersCustomizerInterceptor | None:
+        return (
+            HeadersCustomizerInterceptor(headers_customizers)
+            if headers_customizers
+            else None
+        )
+
     def connect(self, **kwargs) -> None:
+        ...
         """Establishes connection to Snowflake."""
         logger.debug("connect")
         if len(kwargs) > 0:
@@ -1178,7 +1254,6 @@ def __open_connection(self):
                 ):
                     raise TypeError("auth_class must be a child class of AuthByKeyPair")
                     # TODO: add telemetry for custom auth
-                self.auth_class = self.auth_class
             elif self._authenticator == DEFAULT_AUTHENTICATOR:
                 self.auth_class = AuthByDefault(
                     password=self._password,
@@ -1419,6 +1494,37 @@ def __config(self, **kwargs):
             if "host" not in kwargs:
                 self._host = construct_hostname(kwargs.get("region"), self._account)
 
+        self._headers_customizers = kwargs.get("headers_customizers", [])
+        if self._headers_customizers:
+            header_customizer_interceptor = (
+                self._create_interceptor_for_headers_customizers(
+                    self._headers_customizers
+                )
+            )
+            request_interceptors = (
+                [
+                    header_customizer_interceptor,
+                ]
+                if header_customizer_interceptor
+                else []
+            )
+            HttpAdapterFactory.register_for_connection(
+                connection=self,
+                adapter_cls=InterceptingAdapter,
+                interceptors=request_interceptors,
+            )
+        else:
+            # Default adapter
+            HttpAdapterFactory.register_for_connection(
+                connection=self, adapter_cls=ProxySupportAdapter
+            )
+            self._request_interceptors = []
+
+        if self._headers_customizers:
+            logger.info(
+                f"{len(self._headers_customizers)} custom headers customizers were provided. Requests will be enriched according to the defined conditions."
+            )
+
         logger.info(
             f"Connecting to {_DOMAIN_NAME_MAP.get(extract_top_level_domain_from_hostname(self._host), 'GLOBAL')} Snowflake domain"
         )

src/snowflake/connector/http_interceptor.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from enum import Enum, auto
+from typing import Any, Generator, Iterable, MutableSequence, NamedTuple
+
+METHODS = {
+    "GET",
+    "PUT",
+    "POST",
+    "HEAD",
+    "DELETE",
+}
+
+Headers = dict[str, Any]
+
+logger = logging.getLogger(__name__)
+
+
+class RequestDTO(NamedTuple):
+    url: str | bytes
+    method: str
+    headers: Headers
+
+
+class ConflictDTO(NamedTuple):
+    original_key: str
+    value: Any
+    had_conflict: bool = True
+
+
+class HeadersCustomizer(ABC):
+    @abstractmethod
+    def applies_to(self, request: RequestDTO):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def is_invoked_once(self):
+        """recommended to increase performance"""
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_new_headers(self, request: RequestDTO) -> Headers:
+        raise NotImplementedError()
+
+
+class HttpInterceptor(ABC):
+    class InterceptionHook(Enum):
+        BEFORE_RETRY = auto()
+        BEFORE_REQUEST_ISSUED = auto()
+
+    @abstractmethod
+    def intercept_on(
+        self, hook: HttpInterceptor.InterceptionHook, request: RequestDTO
+    ) -> RequestDTO:
+        raise NotImplementedError()
+
+
+class HeadersCustomizerInterceptor(HttpInterceptor):
+    def __init__(
+        self,
+        headers_customizers: Iterable[HeadersCustomizer] | None = None,
+        static_headers_customizers: Iterable[HeadersCustomizer] | None = None,
+        dynamic_headers_customizers: Iterable[HeadersCustomizer] | None = None,
+    ):
+        if headers_customizers is not None:
+            self._static_headers_customizers, self._dynamic_headers_customizers = (
+                self.split_customizers(headers_customizers)
+            )
+        else:
+            self._static_headers_customizers = static_headers_customizers or []
+            self._dynamic_headers_customizers = dynamic_headers_customizers or []
+
+    @staticmethod
+    def split_customizers(
+        headers_customizers: Iterable[HeadersCustomizer] | None,
+    ) -> tuple[MutableSequence[HeadersCustomizer], MutableSequence[HeadersCustomizer]]:
+        static, dynamic = [], []
+        if headers_customizers:
+            for customizer in headers_customizers:
+                if customizer.is_invoked_once():
+                    static.append(customizer)
+                else:
+                    dynamic.append(customizer)
+        return static, dynamic
+
+    @staticmethod
+    def iter_non_conflicting_headers(
+        original_headers: Headers, other_headers: Headers, case_sensitive: bool = False
+    ) -> Generator[ConflictDTO]:
+        """
+        Yields (key, value, is_conflicting) for each key in other_headers,
+        telling you if it conflicts with original_headers.
+        """
+        original_keys = (
+            set(original_headers)
+            if case_sensitive
+            else {k.lower() for k in original_headers}
+        )
+
+        for key, value in other_headers.items():
+            comp_key = key if case_sensitive else key.lower()
+            if comp_key in original_keys:
+                yield ConflictDTO(key, value, True)
+            else:
+                yield ConflictDTO(key, value, False)
+
+    def intercept_on(
+        self, hook: HttpInterceptor.InterceptionHook, request: RequestDTO
+    ) -> RequestDTO:
+        if hook is HttpInterceptor.InterceptionHook.BEFORE_REQUEST_ISSUED:
+            customizers_to_apply = (
+                self._static_headers_customizers + self._dynamic_headers_customizers
+            )
+            return self._handle_headers_customization(request, customizers_to_apply)
+        elif hook is HttpInterceptor.InterceptionHook.BEFORE_RETRY:
+            return self._handle_headers_customization(
+                request, self._dynamic_headers_customizers
+            )
+        return request
+
+    def _handle_headers_customization(
+        self, request: RequestDTO, headers_customizers: Iterable[HeadersCustomizer]
+    ) -> RequestDTO:
+        # copy preventing mutation in the registered customizer
+        result_headers = dict(request.headers) if request.headers else {}
+
+        for header_customizer in headers_customizers:
+            try:
+                if header_customizer.applies_to(request):
+                    additional_headers = header_customizer.get_new_headers(request)
+
+                    for key, value, is_conflicting in self.iter_non_conflicting_headers(
+                        result_headers, additional_headers
+                    ):
+                        if is_conflicting:
+                            logger.warning(
+                                f"Overwriting header '{key}' detected. Skipping this key."
+                            )
+                        else:
+                            result_headers[key] = value
+            except Exception as ex:
+                # Custom logic failure is treated as non-fatal for the connection
+                logger.warning("Unable to customize headers: %s. Skipping...", ex)
+
+        return RequestDTO(
+            url=request.url,
+            method=request.method,
+            headers=result_headers,
+        )