adjust async conftest.py

Author: sfc-gh-pczajka

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,16 @@ 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)
+
     connection = SnowflakeConnection(**ret)
     await connection.connect()
     return connection
@@ -80,7 +98,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 +114,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 +134,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 +147,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()