Fpawlowski/snow 2011595 snowflake connector for python logging presigned ur ls (#2257)

Author: sfc-gh-fpawlowski

.github/workflows/build_test.yml

@@ -178,19 +178,6 @@ jobs:
         run: python -m pip install -U setuptools pip wheel
       - name: Install tox
         run: python -m pip install tox>=4
-
-      - name: Run only the selected test
-        run: python -m tox run -e py39-integ
-        env:
-          PYTEST_ADDOPTS: "test/integ/test_a_put_get_with_aws_token.py::test_put_get_with_aws"
-          PYTHON_VERSION: ${{ matrix.python-version }}
-          TOX_PARALLEL_NO_SPINNER: 1
-          cloud_provider: [aws]
-        shell: bash
-#
-#      - name: Read logs
-#        run: cat ./log_aws.log
-
       - name: Run tests
         run: python -m tox run -e `echo py${PYTHON_VERSION/\./}-{extras,unit,integ,pandas,sso}-ci | sed 's/ /,/g'`
         env:

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
         - id: check-hooks-apply
         - id: check-useless-excludes
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v4.4.0
     hooks:
     -   id: trailing-whitespace
         exclude: >

test/integ/conftest.py

@@ -164,16 +164,22 @@ def init_test_schema(db_parameters) -> Generator[None]:
 
     This is automatically called per test session.
     """
-    ret = db_parameters
-    with snowflake.connector.connect(
-        user=ret["user"],
-        password=ret["password"],
-        host=ret["host"],
-        port=ret["port"],
-        database=ret["database"],
-        account=ret["account"],
-        protocol=ret["protocol"],
-    ) as con:
+    connection_params = {
+        "user": db_parameters["user"],
+        "password": db_parameters["password"],
+        "host": db_parameters["host"],
+        "port": db_parameters["port"],
+        "database": db_parameters["database"],
+        "account": db_parameters["account"],
+        "protocol": db_parameters["protocol"],
+    }
+
+    # Role may be needed when running on preprod, but is not present on Jenkins jobs
+    optional_role = db_parameters.get("role")
+    if optional_role is not None:
+        connection_params.update(role=optional_role)
+
+    with snowflake.connector.connect(**connection_params) as con:
         con.cursor().execute(f"CREATE SCHEMA IF NOT EXISTS {TEST_SCHEMA}")
         yield
         con.cursor().execute(f"DROP SCHEMA IF EXISTS {TEST_SCHEMA}")

test/integ/test_large_result_set.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+import logging
 from unittest.mock import Mock
 
 import pytest
 
+from snowflake.connector.secret_detector import SecretDetector
 from snowflake.connector.telemetry import TelemetryField
 
 NUMBER_OF_ROWS = 50000
@@ -111,8 +113,9 @@ def test_query_large_result_set_n_threads(
 
 @pytest.mark.aws
 @pytest.mark.skipolddriver
-def test_query_large_result_set(conn_cnx, db_parameters, ingest_data):
+def test_query_large_result_set(conn_cnx, db_parameters, ingest_data, caplog):
     """[s3] Gets Large Result set."""
+    caplog.set_level(logging.DEBUG)
     sql = "select * from {name} order by 1".format(name=db_parameters["name"])
     with conn_cnx() as cnx:
         telemetry_data = []
@@ -161,3 +164,17 @@ 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
+                assert (
+                    expected_token_prefix + SecretDetector.SECRET_STARRED_MASK_STR
+                    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"

test/integ/test_a_put_get_with_aws_token.py renamed to test/integ/test_put_get_with_aws_token.py

@@ -4,6 +4,7 @@
 import glob
 import gzip
 import os
+from logging import DEBUG
 
 import pytest
 
@@ -43,6 +44,7 @@
 def test_put_get_with_aws(tmpdir, conn_cnx, from_path, caplog):
     """[s3] Puts and Gets a small text using AWS S3."""
     # create a data file
+    caplog.set_level(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:
@@ -91,22 +93,17 @@ def test_put_get_with_aws(tmpdir, conn_cnx, from_path, caplog):
 
     aws_request_present = False
     expected_token_prefix = "X-Amz-Signature="
-
-    print("\n\n")
-    print(caplog.text)
-    print("\n\n")
-
-    with open("./log_aws.log", "a") as f:
-        f.write(caplog.text)
-
     for line in caplog.text.splitlines():
-        if expected_token_prefix in line:
+        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 + SecretDetector.SECRET_STARRED_MASK_STR in line
+                expected_token_prefix
+                + getattr(SecretDetector, "SECRET_STARRED_MASK_STR", "****")
+                in line
+                or expected_token_prefix not in line
             ), "connectionpool logger is leaking sensitive information"
 
-    # Connection pool is used on GitHub actions, but not always locally
     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/test_put_get_with_azure_token.py

@@ -88,11 +88,13 @@ def test_put_get_with_azure(tmpdir, conn_cnx, from_path, caplog):
     for line in caplog.text.splitlines():
         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 (
-                expected_token_prefix + SecretDetector.SECRET_STARRED_MASK_STR in line
+                expected_token_prefix
+                + getattr(SecretDetector, "SECRET_STARRED_MASK_STR", "****")
+                in line
             ), "connectionpool logger is leaking sensitive information"
 
-    # Connection pool is used on GitHub actions, but not always locally
     assert (
         azure_request_present
     ), "Azure URL was not found in logs, so it can't be assumed that no leaks happened in it"