Cherrypicks to aio connector part9

.github/workflows/build_test.yml

@@ -10,10 +10,7 @@ on:
             - v*
     pull_request:
         branches:
-            - master
-            - main
-            - prep-**
-            - dev/aio-connector
+            - '**'
     workflow_dispatch:
         inputs:
           logLevel:
@@ -24,7 +21,7 @@ on:
             description: "Test scenario tags"
 
 concurrency:
-  # older builds for the same pull request numer or branch should be cancelled
+  # older builds for the same pull request number or branch should be cancelled
   cancel-in-progress: true
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
 

DESCRIPTION.md

@@ -15,9 +15,11 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added handling of PAT provided in `password` field.
   - Improved error message for client-side query cancellations due to timeouts.
   - Added support of GCS regional endpoints.
-  - Added `gcs_use_virtual_endpoints` connection property that forces the usage of the virtual GCS usage. See more: https://cloud.google.com/storage/docs/request-endpoints#xml-api
+  - Added `gcs_use_virtual_endpoints` connection property that forces the usage of the virtual GCS usage. Thanks to this it should be possible to set up private DNS entry for the GCS endpoint. See more: https://cloud.google.com/storage/docs/request-endpoints#xml-api
   - Fixed a bug that caused driver to fail silently on `TO_DATE` arrow to python conversion when invalid date was followed by the correct one.
   - Added `check_arrow_conversion_error_on_every_column` connection property that can be set to `False` to restore previous behaviour in which driver will ignore errors until it occurs in the last column. This flag's purpose is to unblock workflows that may be impacted by the bugfix and will be removed in later releases.
+  - Lower log levels from info to debug for some of the messages to make the output easier to follow.
+  - Allow the connector to inherit a UUID4 generated upstream, provided in statement parameters (field: `requestId`), rather than automatically generate a UUID4 to use for the HTTP Request ID.
 
 - v3.14.0(March 03, 2025)
   - Bumped pyOpenSSL dependency upper boundary from <25.0.0 to <26.0.0.

Jenkinsfile

@@ -38,10 +38,10 @@ timestamps {
       stage('Test') {
           try {
           def commit_hash = "main" // default which we want to override
-          def bptp_tag = "bptp-built"
+          def bptp_tag = "bptp-stable"
           def response = authenticatedGithubCall("https://api.github.com/repos/snowflakedb/snowflake/git/ref/tags/${bptp_tag}")
           commit_hash = response.object.sha
-          // Append the bptp-built commit sha to params
+          // Append the bptp-stable commit sha to params
           params += [string(name: 'svn_revision', value: commit_hash)]
           } catch(Exception e) {
           println("Exception computing commit hash from: ${response}")

benchmark/benchmark_unit_converter.py

@@ -1,7 +1,5 @@
 #!/usr/bin/env python
-#
-# Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
-#
+
 from __future__ import annotations
 
 from logging import getLogger

samples/auth_by_key_pair_from_file.py

@@ -1,7 +1,4 @@
 #!/usr/bin/env python
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
 """
 This sample shows how to implement a key pair authentication plugin
 which reads private key from a file

setup.cfg

@@ -44,6 +44,8 @@ python_requires = >=3.9
 packages = find_namespace:
 install_requires =
     asn1crypto>0.24.0,<2.0.0
+    boto3>=1.0
+    botocore>=1.0
     cffi>=1.9,<2.0.0
     cryptography>=3.1.0
     pyOpenSSL>=22.0.0,<25.0.0
@@ -98,3 +100,4 @@ secure-local-storage =
     keyring>=23.1.0,<26.0.0
 aio =
     aiohttp
+    aioboto3

setup.py

@@ -1,7 +1,4 @@
 #!/usr/bin/env python
-#
-# Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
-#
 
 import os
 import sys

src/snowflake/connector/__init__.py

@@ -1,8 +1,4 @@
 #!/usr/bin/env python
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
-
 # Python Db API v2
 #
 from __future__ import annotations
@@ -16,6 +12,8 @@
 import logging
 from logging import NullHandler
 
+from snowflake.connector.externals_utils.externals_setup import setup_external_libraries
+
 from .connection import SnowflakeConnection
 from .cursor import DictCursor
 from .dbapi import (
@@ -48,6 +46,7 @@
 from .version import VERSION
 
 logging.getLogger(__name__).addHandler(NullHandler())
+setup_external_libraries()
 
 
 @wraps(SnowflakeConnection.__init__)

src/snowflake/connector/_query_context_cache.py

@@ -1,6 +1,3 @@
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
 from __future__ import annotations
 
 from functools import total_ordering

src/snowflake/connector/_sql_util.py

@@ -1,7 +1,3 @@
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
-
 from __future__ import annotations
 
 import re

src/snowflake/connector/_utils.py

@@ -1,13 +1,10 @@
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
-
 from __future__ import annotations
 
 import string
 from enum import Enum
 from random import choice
 from threading import Timer
+from uuid import UUID
 
 
 class TempObjectType(Enum):
@@ -33,6 +30,8 @@ class TempObjectType(Enum):
     "PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS"
 )
 
+REQUEST_ID_STATEMENT_PARAM_NAME = "requestId"
+
 
 def generate_random_alphanumeric(length: int = 10) -> str:
     return "".join(choice(ALPHANUMERIC) for _ in range(length))
@@ -46,6 +45,21 @@ def get_temp_type_for_object(use_scoped_temp_objects: bool) -> str:
     return SCOPED_TEMPORARY_STRING if use_scoped_temp_objects else TEMPORARY_STRING
 
 
+def is_uuid4(str_or_uuid: str | UUID) -> bool:
+    """Check whether provided string str is a valid UUID version4."""
+    if isinstance(str_or_uuid, UUID):
+        return str_or_uuid.version == 4
+
+    if not isinstance(str_or_uuid, str):
+        return False
+
+    try:
+        uuid_str = str(UUID(str_or_uuid, version=4))
+    except ValueError:
+        return False
+    return uuid_str == str_or_uuid
+
+
 class _TrackedQueryCancellationTimer(Timer):
     def __init__(self, interval, function, args=None, kwargs=None):
         super().__init__(interval, function, args, kwargs)

src/snowflake/connector/aio/__init__.py

@@ -1,7 +1,3 @@
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
-
 from __future__ import annotations
 
 from ._connection import SnowflakeConnection

src/snowflake/connector/aio/_azure_storage_client.py

@@ -1,7 +1,3 @@
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
-
 from __future__ import annotations
 
 import base64
@@ -15,7 +11,6 @@
 
 import aiohttp
 
-from ..azure_storage_client import AzureCredentialFilter
 from ..azure_storage_client import (
     SnowflakeAzureRestClient as SnowflakeAzureRestClientSync,
 )
@@ -37,8 +32,6 @@
 
 logger = getLogger(__name__)
 
-getLogger("aiohttp").addFilter(AzureCredentialFilter())
-
 
 class SnowflakeAzureRestClient(
     SnowflakeStorageClientAsync, SnowflakeAzureRestClientSync
@@ -49,7 +42,6 @@ def __init__(
         credentials: StorageCredential | None,
         chunk_size: int,
         stage_info: dict[str, Any],
-        use_s3_regional_url: bool = False,
         unsafe_file_write: bool = False,
     ) -> None:
         SnowflakeAzureRestClientSync.__init__(

src/snowflake/connector/aio/_build_upload_agent.py

@@ -1,7 +1,5 @@
 #!/usr/bin/env python
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
+
 
 from __future__ import annotations
 

src/snowflake/connector/aio/_connection.py

@@ -1,6 +1,3 @@
-#
-# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-#
 from __future__ import annotations
 
 import asyncio
@@ -35,6 +32,7 @@
 from ..connection import _get_private_bytes_from_file
 from ..constants import (
     _CONNECTIVITY_ERR_MSG,
+    ENV_VAR_EXPERIMENTAL_AUTHENTICATION,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
@@ -55,6 +53,7 @@
     ER_CONNECTION_IS_CLOSED,
     ER_FAILED_TO_CONNECT_TO_DB,
     ER_INVALID_VALUE,
+    ER_INVALID_WIF_SETTINGS,
 )
 from ..network import (
     DEFAULT_AUTHENTICATOR,
@@ -64,14 +63,17 @@
     PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
+    WORKLOAD_IDENTITY_AUTHENTICATOR,
     ReauthenticationRequest,
 )
 from ..sqlstate import SQLSTATE_CONNECTION_NOT_EXISTS, SQLSTATE_FEATURE_NOT_SUPPORTED
 from ..telemetry import TelemetryData, TelemetryField
 from ..time_util import get_time_millis
 from ..util_text import split_statements
+from ..wif_util import AttestationProvider
 from ._cursor import SnowflakeCursor
 from ._description import CLIENT_NAME
+from ._direct_file_operation_utils import FileOperationParser, StreamDownloader
 from ._network import SnowflakeRestful
 from ._telemetry import TelemetryClient
 from ._time_util import HeartBeatTimer
@@ -87,6 +89,7 @@
     AuthByPlugin,
     AuthByUsrPwdMfa,
     AuthByWebBrowser,
+    AuthByWorkloadIdentity,
 )
 
 logger = getLogger(__name__)
@@ -116,6 +119,10 @@ def __init__(
         # check SNOW-1218851 for long term improvement plan to refactor ocsp code
         atexit.register(self._close_at_exit)
 
+        # Set up the file operation parser and stream downloader.
+        self._file_operation_parser = FileOperationParser(self)
+        self._stream_downloader = StreamDownloader(self)
+
     def __enter__(self):
         # async connection does not support sync context manager
         raise TypeError(
@@ -301,8 +308,6 @@ async def __open_connection(self):
                     backoff_generator=self._backoff_generator,
                 )
             elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
-                if not self._token and self._password:
-                    self._token = self._password
                 self.auth_class = AuthByPAT(self._token)
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
@@ -320,6 +325,29 @@ async def __open_connection(self):
                     timeout=self.login_timeout,
                     backoff_generator=self._backoff_generator,
                 )
+            elif self._authenticator == WORKLOAD_IDENTITY_AUTHENTICATOR:
+                if ENV_VAR_EXPERIMENTAL_AUTHENTICATION not in os.environ:
+                    Error.errorhandler_wrapper(
+                        self,
+                        None,
+                        ProgrammingError,
+                        {
+                            "msg": f"Please set the '{ENV_VAR_EXPERIMENTAL_AUTHENTICATION}' environment variable to use the '{WORKLOAD_IDENTITY_AUTHENTICATOR}' authenticator.",
+                            "errno": ER_INVALID_WIF_SETTINGS,
+                        },
+                    )
+                # Standardize the provider enum.
+                if self._workload_identity_provider and isinstance(
+                    self._workload_identity_provider, str
+                ):
+                    self._workload_identity_provider = AttestationProvider.from_string(
+                        self._workload_identity_provider
+                    )
+                self.auth_class = AuthByWorkloadIdentity(
+                    provider=self._workload_identity_provider,
+                    token=self._token,
+                    entra_resource=self._workload_identity_entra_resource,
+                )
             else:
                 # okta URL, e.g., https://<account>.okta.com/
                 self.auth_class = AuthByOkta(
@@ -767,7 +795,7 @@ async def close(self, retry: bool = True) -> None:
             await self._cancel_heartbeat()
 
             # close telemetry first, since it needs rest to send remaining data
-            logger.info("closed")
+            logger.debug("closed")
 
             await self._telemetry.close(
                 send_on_close=bool(retry and self.telemetry_enabled)
@@ -776,15 +804,15 @@ async def close(self, retry: bool = True) -> None:
                 await self._all_async_queries_finished()
                 and not self._server_session_keep_alive
             ):
-                logger.info("No async queries seem to be running, deleting session")
+                logger.debug("No async queries seem to be running, deleting session")
                 try:
                     await self.rest.delete_session(retry=retry)
                 except Exception as e:
                     logger.debug(
                         "Exception encountered in deleting session. ignoring...: %s", e
                     )
             else:
-                logger.info(
+                logger.debug(
                     "There are {} async queries still running, not deleting session".format(
                         len(self._async_sfqids)
                     )