SNOW-1664063: sync main branch changes into async part (#2081)

Author: sfc-gh-aling

src/snowflake/connector/aio/_azure_storage_client.py

@@ -14,6 +14,7 @@
 
 import aiohttp
 
+from ..azure_storage_client import AzureCredentialFilter
 from ..azure_storage_client import (
     SnowflakeAzureRestClient as SnowflakeAzureRestClientSync,
 )
@@ -25,14 +26,16 @@
 if TYPE_CHECKING:  # pragma: no cover
     from ..file_transfer_agent import SnowflakeFileMeta, StorageCredential
 
-logger = getLogger(__name__)
-
 from ..azure_storage_client import (
     ENCRYPTION_DATA,
     MATDESC,
     TOKEN_EXPIRATION_ERR_MESSAGE,
 )
 
+logger = getLogger(__name__)
+
+getLogger("aiohttp").addFilter(AzureCredentialFilter())
+
 
 class SnowflakeAzureRestClient(
     SnowflakeStorageClientAsync, SnowflakeAzureRestClientSync

src/snowflake/connector/aio/_build_upload_agent.py

@@ -10,6 +10,7 @@
 from typing import TYPE_CHECKING, cast
 
 from snowflake.connector import Error
+from snowflake.connector._utils import get_temp_type_for_object
 from snowflake.connector.bind_upload_agent import BindUploadAgent as BindUploadAgentSync
 from snowflake.connector.errors import BindUploadError
 
@@ -30,7 +31,11 @@ def __init__(
         self.cursor = cast("SnowflakeCursor", cursor)
 
     async def _create_stage(self) -> None:
-        await self.cursor.execute(self._CREATE_STAGE_STMT)
+        create_stage_sql = (
+            f"create or replace {get_temp_type_for_object(self._use_scoped_temp_object)} stage {self._STAGE_NAME} "
+            "file_format=(type=csv field_optionally_enclosed_by='\"')"
+        )
+        await self.cursor.execute(create_stage_sql)
 
     async def upload(self) -> None:
         try:

src/snowflake/connector/aio/_connection.py

@@ -35,6 +35,7 @@
 from ..connection import _get_private_bytes_from_file
 from ..connection_diagnostic import ConnectionDiagnostic
 from ..constants import (
+    _CONNECTIVITY_ERR_MSG,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
@@ -443,6 +444,8 @@ async def _authenticate(self, auth_instance: AuthByPlugin):
                     )
                 except OperationalError as auth_op:
                     if auth_op.errno == ER_FAILED_TO_CONNECT_TO_DB:
+                        if _CONNECTIVITY_ERR_MSG in e.msg:
+                            auth_op.msg += f"\n{_CONNECTIVITY_ERR_MSG}"
                         raise auth_op from e
                     logger.debug("Continuing authenticator specific timeout handling")
                     continue

src/snowflake/connector/aio/_cursor.py

@@ -129,8 +129,10 @@ async def _timebomb_task(self, timeout, query):
             logger.debug("started timebomb in %ss", timeout)
             await asyncio.sleep(timeout)
             await self.__cancel_query(query)
+            return True
         except asyncio.CancelledError:
             logger.debug("cancelled timebomb in timebomb task")
+            return False
 
     async def __cancel_query(self, query) -> None:
         if self._sequence_counter >= 0 and not self.is_closed():
@@ -284,7 +286,10 @@ def interrupt_handler(*_):  # pragma: no cover
                 )
             if self._timebomb is not None:
                 self._timebomb.cancel()
-                self._timebomb = None
+                try:
+                    await self._timebomb
+                except asyncio.CancelledError:
+                    pass
                 logger.debug("cancelled timebomb in finally")
 
         if "data" in ret and "parameters" in ret["data"]:
@@ -674,6 +679,11 @@ async def execute(
             logger.debug(ret)
             err = ret["message"]
             code = ret.get("code", -1)
+            if self._timebomb and self._timebomb.result():
+                err = (
+                    f"SQL execution was cancelled by the client due to a timeout. "
+                    f"Error message received from the server: {err}"
+                )
             if "data" in ret:
                 err += ret["data"].get("errorMessage", "")
             errvalue = {
@@ -1067,6 +1077,7 @@ async def wait_until_ready() -> None:
         self._prefetch_hook = wait_until_ready
 
     async def query_result(self, qid: str) -> SnowflakeCursor:
+        """Query the result of a previously executed query."""
         url = f"/queries/{qid}/result"
         ret = await self._connection.rest.request(url=url, method="get")
         self._sfqid = (

src/snowflake/connector/aio/_network.py

@@ -28,6 +28,7 @@
     urlparse,
 )
 from ..constants import (
+    _CONNECTIVITY_ERR_MSG,
     HTTP_HEADER_ACCEPT,
     HTTP_HEADER_CONTENT_TYPE,
     HTTP_HEADER_SERVICE_NAME,
@@ -798,8 +799,19 @@ async def _request_exec(
             finally:
                 raw_ret.close()  # ensure response is closed
         except aiohttp.ClientSSLError as se:
-            logger.debug("Hit non-retryable SSL error, %s", str(se))
-
+            msg = f"Hit non-retryable SSL error, {str(se)}.\n{_CONNECTIVITY_ERR_MSG}"
+            logger.debug(msg)
+            # the following code is for backward compatibility with old versions of python connector which calls
+            # self._handle_unknown_error to process SSLError
+            Error.errorhandler_wrapper(
+                self._connection,
+                None,
+                OperationalError,
+                {
+                    "msg": msg,
+                    "errno": ER_FAILED_TO_REQUEST,
+                },
+            )
         # TODO: sync feature parity, aiohttp network error handling
         except (
             BadStatusLine,

src/snowflake/connector/aio/auth/_auth.py

@@ -13,7 +13,12 @@
 from typing import TYPE_CHECKING, Any, Callable
 
 from ...auth import Auth as AuthSync
-from ...auth._auth import ID_TOKEN, MFA_TOKEN, delete_temporary_credential
+from ...auth._auth import (
+    AUTHENTICATION_REQUEST_KEY_WHITELIST,
+    ID_TOKEN,
+    MFA_TOKEN,
+    delete_temporary_credential,
+)
 from ...compat import urlencode
 from ...constants import (
     HTTP_HEADER_ACCEPT,
@@ -103,7 +108,6 @@ async def authenticate(
 
         body = copy.deepcopy(body_template)
         # updating request body
-        logger.debug("assertion content: %s", auth_instance.assertion_content)
         await auth_instance.update_body(body)
 
         logger.debug(
@@ -141,7 +145,10 @@ async def authenticate(
 
         logger.debug(
             "body['data']: %s",
-            {k: v for (k, v) in body["data"].items() if k != "PASSWORD"},
+            {
+                k: v if k in AUTHENTICATION_REQUEST_KEY_WHITELIST else "******"
+                for (k, v) in body["data"].items()
+            },
         )
 
         try:

test/integ/aio/test_cursor_async.py

@@ -728,7 +728,11 @@ async def test_timeout_query(conn_cnx):
                     "select seq8() as c1 from table(generator(timeLimit => 60))",
                     timeout=5,
                 )
-            assert err.value.errno == 604, "Invalid error code"
+            assert err.value.errno == 604, (
+                "Invalid error code"
+                and "SQL execution was cancelled by the client due to a timeout"
+                in err.value.msg
+            )
 
 
 async def test_executemany(conn, db_parameters):

test/integ/aio/test_dbapi_async.py

@@ -44,9 +44,8 @@ async def conn_local(request, conn_cnx):
     async def fin():
         await drop_dbapi_tables(conn_cnx)
 
-    request.addfinalizer(fin)
-
-    return conn_cnx
+    yield conn_cnx
+    await fin()
 
 
 async def _paraminsert(cur):

test/integ/aio/test_large_result_set_async.py

@@ -87,8 +87,8 @@ async def fin():
                 "drop table if exists {name}".format(name=db_parameters["name"])
             )
 
-    request.addfinalizer(fin)
-    return first_val, last_val
+    yield first_val, last_val
+    await fin()
 
 
 @pytest.mark.aws

test/integ/aio/test_put_get_async.py

@@ -44,7 +44,7 @@ async def test_utf8_filename(tmp_path, aio_connection):
     await aio_connection.connect()
     cursor = aio_connection.cursor()
     await cursor.execute(f"create temporary stage {stage_name}")
-    (
+    await (
         await cursor.execute(
             "PUT 'file://{}' @{}".format(str(test_file).replace("\\", "/"), stage_name)
         )
@@ -128,7 +128,7 @@ async def test_put_special_file_name(tmp_path, aio_connection):
     cursor = aio_connection.cursor()
     await cursor.execute(f"create temporary stage {stage_name}")
     filename_in_put = str(test_file).replace("\\", "/")
-    (
+    await (
         await cursor.execute(
             f"PUT 'file://{filename_in_put}' @{stage_name}",
         )