Apply changes to async tests and workflows

Author: sfc-gh-pczajka

.github/workflows/build_test.yml

@@ -424,6 +424,13 @@ jobs:
         run: |
           gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" \
           .github/workflows/parameters/public/parameters_${{ matrix.cloud-provider }}.py.gpg > test/parameters.py
+      - name: Setup private key file
+        shell: bash
+        env:
+          PYTHON_PRIVATE_KEY_SECRET: ${{ secrets.PYTHON_PRIVATE_KEY_SECRET }}
+        run: |
+          gpg --quiet --batch --yes --decrypt --passphrase="$PYTHON_PRIVATE_KEY_SECRET" \
+          .github/workflows/parameters/public/rsa_keys/rsa_key_python_${{ matrix.cloud-provider }}.p8.gpg > test/rsa_key_python_${{ matrix.cloud-provider }}.p8
       - name: Download wheel(s)
         uses: actions/download-artifact@v4
         with:

test/integ/aio/conftest.py

@@ -2,9 +2,14 @@
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
 
+import os
 from contextlib import asynccontextmanager
-from test.integ.conftest import get_db_parameters, is_public_testaccount
-from typing import AsyncContextManager, Callable, Generator
+from test.integ.conftest import (
+    _get_private_key_bytes_for_olddriver,
+    get_db_parameters,
+    is_public_testaccount,
+)
+from typing import AsyncContextManager, AsyncGenerator, Callable
 
 import pytest
 
@@ -44,7 +49,7 @@ async def patch_connection(
         self,
         con: SnowflakeConnection,
         propagate: bool = True,
-    ) -> Generator[TelemetryCaptureHandlerAsync, None, None]:
+    ) -> AsyncGenerator[TelemetryCaptureHandlerAsync, None]:
         original_telemetry = con._telemetry
         new_telemetry = TelemetryCaptureHandlerAsync(
             original_telemetry,
@@ -57,6 +62,9 @@ async def patch_connection(
             con._telemetry = original_telemetry
 
 
+RUNNING_OLD_DRIVER = os.getenv("TOX_ENV_NAME") == "olddriver"
+
+
 @pytest.fixture(scope="session")
 def capture_sf_telemetry_async() -> TelemetryCaptureFixtureAsync:
     return TelemetryCaptureFixtureAsync()
@@ -71,6 +79,22 @@ async def create_connection(connection_name: str, **kwargs) -> SnowflakeConnecti
     """
     ret = get_db_parameters(connection_name)
     ret.update(kwargs)
+
+    # Handle private key authentication for old driver if applicable
+    if RUNNING_OLD_DRIVER and "private_key_file" in ret and "private_key" not in ret:
+        private_key_file = ret.get("private_key_file")
+        if private_key_file:
+            private_key_bytes = _get_private_key_bytes_for_olddriver(private_key_file)
+            ret["authenticator"] = "SNOWFLAKE_JWT"
+            ret["private_key"] = private_key_bytes
+            ret.pop("private_key_file", None)
+
+    # If authenticator is explicitly provided and it's not key-pair based, drop key-pair fields
+    authenticator_value = ret.get("authenticator")
+    if authenticator_value.lower() not in {"key_pair_authenticator", "snowflake_jwt"}:
+        ret.pop("private_key", None)
+        ret.pop("private_key_file", None)
+
     connection = SnowflakeConnection(**ret)
     await connection.connect()
     return connection
@@ -80,7 +104,7 @@ async def create_connection(connection_name: str, **kwargs) -> SnowflakeConnecti
 async def db(
     connection_name: str = "default",
     **kwargs,
-) -> Generator[SnowflakeConnection, None, None]:
+) -> AsyncGenerator[SnowflakeConnection, None]:
     if not kwargs.get("timezone"):
         kwargs["timezone"] = "UTC"
     if not kwargs.get("converter_class"):
@@ -96,7 +120,7 @@ async def db(
 async def negative_db(
     connection_name: str = "default",
     **kwargs,
-) -> Generator[SnowflakeConnection, None, None]:
+) -> AsyncGenerator[SnowflakeConnection, None]:
     if not kwargs.get("timezone"):
         kwargs["timezone"] = "UTC"
     if not kwargs.get("converter_class"):
@@ -116,7 +140,7 @@ def conn_cnx():
 
 
 @pytest.fixture()
-async def conn_testaccount() -> SnowflakeConnection:
+async def conn_testaccount() -> AsyncGenerator[SnowflakeConnection, None]:
     connection = await create_connection("default")
     yield connection
     await connection.close()
@@ -129,18 +153,43 @@ def negative_conn_cnx() -> Callable[..., AsyncContextManager[SnowflakeConnection
 
 
 @pytest.fixture()
-async def aio_connection(db_parameters):
-    cnx = SnowflakeConnection(
-        user=db_parameters["user"],
-        password=db_parameters["password"],
-        host=db_parameters["host"],
-        port=db_parameters["port"],
-        account=db_parameters["account"],
-        database=db_parameters["database"],
-        schema=db_parameters["schema"],
-        warehouse=db_parameters["warehouse"],
-        protocol=db_parameters["protocol"],
-        timezone="UTC",
-    )
-    yield cnx
-    await cnx.close()
+async def aio_connection(db_parameters) -> AsyncGenerator[SnowflakeConnection, None]:
+    # Build connection params supporting both password and key-pair auth depending on environment
+    connection_params = {
+        "user": db_parameters["user"],
+        "host": db_parameters["host"],
+        "port": db_parameters["port"],
+        "account": db_parameters["account"],
+        "database": db_parameters["database"],
+        "schema": db_parameters["schema"],
+        "protocol": db_parameters["protocol"],
+        "timezone": "UTC",
+    }
+
+    # Optional fields
+    warehouse = db_parameters.get("warehouse")
+    if warehouse is not None:
+        connection_params["warehouse"] = warehouse
+
+    role = db_parameters.get("role")
+    if role is not None:
+        connection_params["role"] = role
+
+    if "password" in db_parameters and db_parameters["password"]:
+        connection_params["password"] = db_parameters["password"]
+    elif "private_key_file" in db_parameters:
+        # Use key-pair authentication
+        connection_params["authenticator"] = "SNOWFLAKE_JWT"
+        if RUNNING_OLD_DRIVER:
+            private_key_bytes = _get_private_key_bytes_for_olddriver(
+                db_parameters["private_key_file"]
+            )
+            connection_params["private_key"] = private_key_bytes
+        else:
+            connection_params["private_key_file"] = db_parameters["private_key_file"]
+
+    cnx = SnowflakeConnection(**connection_params)
+    try:
+        yield cnx
+    finally:
+        await cnx.close()

test/integ/aio/test_arrow_result_async.py

@@ -136,7 +136,7 @@ async def structured_type_wrapped_conn(conn_cnx, structured_type_support):
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("datatype", ICEBERG_UNSUPPORTED_TYPES)
+@pytest.mark.parametrize("datatype", sorted(ICEBERG_UNSUPPORTED_TYPES))
 async def test_iceberg_negative(
     datatype, conn_cnx, iceberg_support, structured_type_support
 ):
@@ -834,35 +834,46 @@ async def test_select_vector(conn_cnx, is_public_test):
 
 @pytest.mark.asyncio
 async def test_select_time(conn_cnx):
-    for scale in range(10):
-        await select_time_with_scale(conn_cnx, scale)
-
-
-async def select_time_with_scale(conn_cnx, scale):
+    # Test key scales and meaningful cases in a single table operation
+    # Cover: no fractional seconds, milliseconds, microseconds, nanoseconds
+    scales = [0, 3, 6, 9]  # Key precision levels
     cases = [
-        "00:01:23",
-        "00:01:23.1",
-        "00:01:23.12",
-        "00:01:23.123",
-        "00:01:23.1234",
-        "00:01:23.12345",
-        "00:01:23.123456",
-        "00:01:23.1234567",
-        "00:01:23.12345678",
-        "00:01:23.123456789",
+        "00:01:23",  # Basic time
+        "00:01:23.123456789",  # Max precision
+        "23:59:59.999999999",  # Edge case - max time with max precision
+        "00:00:00.000000001",  # Edge case - min time with min precision
     ]
-    table = "test_arrow_time"
-    column = f"(a time({scale}))"
-    values = (
-        "(-1, NULL), ("
-        + "),(".join([f"{i}, '{c}'" for i, c in enumerate(cases)])
-        + f"), ({len(cases)}, NULL)"
-    )
-    await init(conn_cnx, table, column, values)
-    sql_text = f"select a from {table} order by s"
-    row_count = len(cases) + 2
-    col_count = 1
-    await iterate_over_test_chunk("time", conn_cnx, sql_text, row_count, col_count)
+
+    table = "test_arrow_time_scales"
+
+    # Create columns for selected scales only (init function will add 's number' automatically)
+    columns = ", ".join([f"a{i} time({i})" for i in scales])
+    column_def = f"({columns})"
+
+    # Create values for selected scales - each case tests all scales simultaneously
+    value_rows = []
+    for i, case in enumerate(cases):
+        # Each row has the same time value for all scale columns
+        time_values = ", ".join([f"'{case}'" for _ in scales])
+        value_rows.append(f"({i}, {time_values})")
+
+    # Add NULL rows
+    null_values = ", ".join(["NULL" for _ in scales])
+    value_rows.append(f"(-1, {null_values})")
+    value_rows.append(f"({len(cases)}, {null_values})")
+
+    values = ", ".join(value_rows)
+
+    # Single table creation and test
+    await init(conn_cnx, table, column_def, values)
+
+    # Test each scale column
+    for scale in scales:
+        sql_text = f"select a{scale} from {table} order by s"
+        row_count = len(cases) + 2
+        col_count = 1
+        await iterate_over_test_chunk("time", conn_cnx, sql_text, row_count, col_count)
+
     await finish(conn_cnx, table)
 
 

test/integ/aio/test_autocommit_async.py

@@ -5,8 +5,6 @@
 
 from __future__ import annotations
 
-import snowflake.connector.aio
-
 
 async def exe0(cnx, sql):
     return await cnx.cursor().execute(sql)
@@ -164,7 +162,7 @@ async def exe(cnx, sql):
             )
 
 
-async def test_autocommit_parameters(db_parameters):
+async def test_autocommit_parameters(db_parameters, conn_cnx):
     """Tests autocommit parameter.
 
     Args:
@@ -174,17 +172,7 @@ async def test_autocommit_parameters(db_parameters):
     async def exe(cnx, sql):
         return await cnx.cursor().execute(sql.format(name=db_parameters["name"]))
 
-    async with snowflake.connector.aio.SnowflakeConnection(
-        user=db_parameters["user"],
-        password=db_parameters["password"],
-        host=db_parameters["host"],
-        port=db_parameters["port"],
-        account=db_parameters["account"],
-        protocol=db_parameters["protocol"],
-        schema=db_parameters["schema"],
-        database=db_parameters["database"],
-        autocommit=False,
-    ) as cnx:
+    async with conn_cnx(autocommit=False) as cnx:
         await exe(
             cnx,
             """
@@ -193,17 +181,7 @@ async def exe(cnx, sql):
         )
         await _run_autocommit_off(cnx, db_parameters)
 
-    async with snowflake.connector.aio.SnowflakeConnection(
-        user=db_parameters["user"],
-        password=db_parameters["password"],
-        host=db_parameters["host"],
-        port=db_parameters["port"],
-        account=db_parameters["account"],
-        protocol=db_parameters["protocol"],
-        schema=db_parameters["schema"],
-        database=db_parameters["database"],
-        autocommit=True,
-    ) as cnx:
+    async with conn_cnx(autocommit=True) as cnx:
         await _run_autocommit_on(cnx, db_parameters)
         await exe(
             cnx,