Review fixes - add secret filter to external libraries + tests

Author: sfc-gh-pczajka

src/snowflake/connector/externals_utils/externals_setup.py

@@ -9,6 +9,9 @@
     "snowflake.connector.vendored.urllib3",
     "botocore",
     "boto3",
+    "aiohttp",  # this should not break even if [aio] extra is not installed - in such case logger will remain unused
+    "aiobotocore",
+    "aioboto3",
 ]
 # TODO: after migration to the external urllib3 from the vendored one (SNOW-2041970),
 #  we should change filters here immediately to the below module's logger:

test/integ/aio/test_large_result_set_async.py

@@ -5,10 +5,11 @@
 
 from __future__ import annotations
 
-from unittest.mock import Mock
+import logging
 
 import pytest
 
+from snowflake.connector.secret_detector import SecretDetector
 from snowflake.connector.telemetry import TelemetryField
 
 NUMBER_OF_ROWS = 50000
@@ -18,7 +19,9 @@
 
 @pytest.fixture()
 async def ingest_data(request, conn_cnx, db_parameters):
-    async with conn_cnx() as cnx:
+    async with conn_cnx(
+        session_parameters={"python_connector_query_result_format": "json"},
+    ) as cnx:
         await cnx.cursor().execute(
             """
     create or replace table {name} (
@@ -89,7 +92,12 @@ async def test_query_large_result_set_n_threads(
     conn_cnx, db_parameters, ingest_data, num_threads
 ):
     sql = "select * from {name} order by 1".format(name=db_parameters["name"])
-    async with conn_cnx(client_prefetch_threads=num_threads) as cnx:
+    async with conn_cnx(
+        client_prefetch_threads=num_threads,
+        session_parameters={
+            "python_connector_query_result_format": "json",
+        },
+    ) as cnx:
         assert cnx.client_prefetch_threads == num_threads
         results = []
         async for rec in await cnx.cursor().execute(sql):
@@ -102,13 +110,26 @@ async def test_query_large_result_set_n_threads(
 
 @pytest.mark.aws
 @pytest.mark.skipolddriver
-async def test_query_large_result_set(conn_cnx, db_parameters, ingest_data):
+async def test_query_large_result_set(conn_cnx, db_parameters, ingest_data, caplog):
     """[s3] Gets Large Result set."""
+    caplog.set_level(logging.DEBUG)
+    caplog.set_level(logging.DEBUG, logger="snowflake.connector.vendored.urllib3")
+    caplog.set_level(
+        logging.DEBUG, logger="snowflake.connector.vendored.urllib3.connectionpool"
+    )
+    caplog.set_level(logging.DEBUG, logger="aiohttp")
+    caplog.set_level(logging.DEBUG, logger="aiohttp.client")
     sql = "select * from {name} order by 1".format(name=db_parameters["name"])
-    async with conn_cnx() as cnx:
+    async with conn_cnx(
+        session_parameters={
+            "python_connector_query_result_format": "json",
+        }
+    ) as cnx:
         telemetry_data = []
-        add_log_mock = Mock()
-        add_log_mock.side_effect = lambda datum: telemetry_data.append(datum)
+
+        async def add_log_mock(datum):
+            telemetry_data.append(datum)
+
         cnx._telemetry.add_log_to_batch = add_log_mock
 
         result2 = []
@@ -152,3 +173,20 @@ async def test_query_large_result_set(conn_cnx, db_parameters, ingest_data):
                 "Expected three telemetry logs (one per query) "
                 "for log type {}".format(field.value)
             )
+
+        aws_request_present = False
+        expected_token_prefix = "X-Amz-Signature="
+        for line in caplog.text.splitlines():
+            if expected_token_prefix in line:
+                aws_request_present = True
+                # getattr is used to stay compatible with old driver - before SECRET_STARRED_MASK_STR was added
+                assert (
+                    expected_token_prefix
+                    + getattr(SecretDetector, "SECRET_STARRED_MASK_STR", "****")
+                    in line
+                ), "connectionpool logger is leaking sensitive information"
+
+        # If no AWS request appeared in logs, we cannot assert masking here.
+        assert (
+            aws_request_present
+        ), "AWS URL was not found in logs, so it can't be assumed that no leaks happened in it"

test/integ/aio/test_put_get_with_aws_token_async.py

@@ -7,12 +7,15 @@
 
 import glob
 import gzip
+import logging
 import os
 
 import pytest
 from aiohttp import ClientResponseError
 
 from snowflake.connector.constants import UTF8
+from snowflake.connector.file_transfer_agent import SnowflakeS3ProgressPercentage
+from snowflake.connector.secret_detector import SecretDetector
 
 try:  # pragma: no cover
     from snowflake.connector.aio._file_transfer_agent import SnowflakeFileMeta
@@ -38,9 +41,10 @@
 @pytest.mark.parametrize(
     "from_path", [True, pytest.param(False, marks=pytest.mark.skipolddriver)]
 )
-async def test_put_get_with_aws(tmpdir, aio_connection, from_path):
+async def test_put_get_with_aws(tmpdir, aio_connection, from_path, caplog):
     """[s3] Puts and Gets a small text using AWS S3."""
     # create a data file
+    caplog.set_level(logging.DEBUG)
     fname = str(tmpdir.join("test_put_get_with_aws_token.txt.gz"))
     original_contents = "123,test1\n456,test2\n"
     with gzip.open(fname, "wb") as f:
@@ -60,6 +64,8 @@ async def test_put_get_with_aws(tmpdir, aio_connection, from_path):
             f"%{table_name}",
             from_path,
             sql_options=" auto_compress=true parallel=30",
+            _put_callback=SnowflakeS3ProgressPercentage,
+            _get_callback=SnowflakeS3ProgressPercentage,
             file_stream=file_stream,
         )
         rec = await csr.fetchone()
@@ -71,22 +77,44 @@ async def test_put_get_with_aws(tmpdir, aio_connection, from_path):
             f"copy into @%{table_name} from {table_name} "
             "file_format=(type=csv compression='gzip')"
         )
-        await csr.execute(f"get @%{table_name} file://{tmp_dir}")
+        await csr.execute(
+            f"get @%{table_name} file://{tmp_dir}",
+            _put_callback=SnowflakeS3ProgressPercentage,
+            _get_callback=SnowflakeS3ProgressPercentage,
+        )
         rec = await csr.fetchone()
         assert rec[0].startswith("data_"), "A file downloaded by GET"
         assert rec[1] == 36, "Return right file size"
         assert rec[2] == "DOWNLOADED", "Return DOWNLOADED status"
         assert rec[3] == "", "Return no error message"
     finally:
-        await csr.execute(f"drop table {table_name}")
+        await csr.execute(f"drop table if exists {table_name}")
         if file_stream:
             file_stream.close()
+        await aio_connection.close()
 
     files = glob.glob(os.path.join(tmp_dir, "data_*"))
     with gzip.open(files[0], "rb") as fd:
         contents = fd.read().decode(UTF8)
     assert original_contents == contents, "Output is different from the original file"
 
+    aws_request_present = False
+    expected_token_prefix = "X-Amz-Signature="
+    for line in caplog.text.splitlines():
+        if ".amazonaws." in line:
+            aws_request_present = True
+            # getattr is used to stay compatible with old driver - before SECRET_STARRED_MASK_STR was added
+            assert (
+                expected_token_prefix
+                + getattr(SecretDetector, "SECRET_STARRED_MASK_STR", "****")
+                in line
+                or expected_token_prefix not in line
+            ), "connectionpool logger is leaking sensitive information"
+
+    assert (
+        aws_request_present
+    ), "AWS URL was not found in logs, so it can't be assumed that no leaks happened in it"
+
 
 @pytest.mark.skipolddriver
 async def test_put_with_invalid_token(tmpdir, aio_connection):
@@ -141,3 +169,4 @@ async def test_put_with_invalid_token(tmpdir, aio_connection):
             await client.upload_chunk(0)
     finally:
         await csr.execute(f"drop table if exists {table_name}")
+        await aio_connection.close()

test/integ/aio/test_put_get_with_azure_token_async.py

@@ -20,6 +20,7 @@
     SnowflakeAzureProgressPercentage,
     SnowflakeProgressPercentage,
 )
+from snowflake.connector.secret_detector import SecretDetector
 
 try:
     from snowflake.connector.util_text import random_string
@@ -86,13 +87,24 @@ async def test_put_get_with_azure(tmpdir, aio_connection, from_path, caplog):
     finally:
         if file_stream:
             file_stream.close()
-        await csr.execute(f"drop table {table_name}")
+        await csr.execute(f"drop table if exists {table_name}")
+        await aio_connection.close()
 
+    azure_request_present = False
+    expected_token_prefix = "sig="
     for line in caplog.text.splitlines():
-        if "blob.core.windows.net" in line:
+        if "blob.core.windows.net" in line and expected_token_prefix in line:
+            azure_request_present = True
+            # getattr is used to stay compatible with old driver - before SECRET_STARRED_MASK_STR was added
             assert (
-                "sig=" not in line
+                expected_token_prefix
+                + getattr(SecretDetector, "SECRET_STARRED_MASK_STR", "****")
+                in line
             ), "connectionpool logger is leaking sensitive information"
+
+    assert (
+        azure_request_present
+    ), "Azure URL was not found in logs, so it can't be assumed that no leaks happened in it"
     files = glob.glob(os.path.join(tmp_dir, "data_*"))
     with gzip.open(files[0], "rb") as fd:
         contents = fd.read().decode(UTF8)
@@ -141,6 +153,7 @@ async def run(csr, sql):
         assert rows == number_of_files * number_of_lines, "Number of rows"
     finally:
         await run(csr, "drop table if exists {name}")
+        await aio_connection.close()
 
 
 async def test_put_copy_duplicated_files_azure(tmpdir, aio_connection):
@@ -216,6 +229,7 @@ async def run(csr, sql):
         assert rows == number_of_files * number_of_lines, "Number of rows"
     finally:
         await run(csr, "drop table if exists {name}")
+        await aio_connection.close()
 
 
 async def test_put_get_large_files_azure(tmpdir, aio_connection):
@@ -280,3 +294,4 @@ async def run(cnx, sql):
         assert all([rec[2] == "DOWNLOADED" for rec in all_recs])
     finally:
         await run(aio_connection, "RM @~/{dir}")
+        await aio_connection.close()