address comment

snowflakedb · Aug 31, 2022 · b7719d2 · b7719d2
1 parent acca9f8
commit b7719d2
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 75 deletions.
diff --git a/src/snowflake/connector/__init__.py b/src/snowflake/connector/__init__.py
@@ -9,7 +9,6 @@
 apilevel = "2.0"
 threadsafety = 2
 paramstyle = "pyformat"
-log_imported_packages_in_telemetry = True
 
 import logging
 from logging import NullHandler
@@ -91,5 +90,4 @@ def Connect(**kwargs) -> SnowflakeConnection:
     "ROWID",
     # Extended data type (experimental)
     "Json",
-    "log_imported_packages_in_telemetry",
 ]
diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py
@@ -198,8 +198,8 @@ def DefaultConverterClass():
         (type(None), str),
     ),  # Path to connection diag whitelist json
     "log_imported_packages_in_telemetry": (
-        None,
-        (type(None), bool),
+        True,
+        bool,
     ),  # Whether to log imported packages in telemetry
 }
 
@@ -965,12 +965,10 @@ def __config(self, **kwargs):
             )
             self._use_openssl_only = os.environ["SF_USE_OPENSSL_ONLY"] == "True"
 
-        if self._log_imported_packages_in_telemetry is None:
-            import snowflake.connector
-
-            self._log_imported_packages_in_telemetry = (
-                snowflake.connector.log_imported_packages_in_telemetry
-            )
+        if "log_imported_packages_in_telemetry" in kwargs:
+            self._log_imported_packages_in_telemetry = kwargs[
+                "log_imported_packages_in_telemetry"
+            ]
 
     def cmd_query(
         self,

diff --git a/test/integ/test_connection.py b/test/integ/test_connection.py
@@ -1112,69 +1112,76 @@ def test_ocsp_cache_working(conn_cnx):
 
 @pytest.mark.skipolddriver
 def test_imported_packages_telemetry(conn_cnx, capture_sf_telemetry, db_parameters):
-    try:
-        # these imports are not used but for testing
-        import html.parser  # noqa: F401
-        import json  # noqa: F401
-        import multiprocessing as mp  # noqa: F401
-        from datetime import date  # noqa: F401
-        from math import sqrt  # noqa: F401
-
-        with conn_cnx() as conn, capture_sf_telemetry.patch_connection(
-            conn, False
-        ) as telemetry_test:
-            conn._log_telemetry_imported_packages()
-            assert len(telemetry_test.records) > 0
-            for t in telemetry_test.records:
-                if (
-                    t.message[TelemetryField.KEY_TYPE.value]
-                    == TelemetryField.IMPORTED_PACKAGES.value
-                ):
-                    assert "pytest" in t.message["value"]
-                    assert "unittest" in t.message["value"]
-                    assert "json" in t.message["value"]
-                    assert "multiprocessing" in t.message["value"]
-                    assert "html" in t.message["value"]
-                    assert "datetime" in t.message["value"]
-                    assert "math" in t.message["value"]
-                    assert "__main__" not in t.message["value"]
-                    assert CLIENT_NAME == t.message[TelemetryField.KEY_SOURCE.value]
-
-        # test different application
-        new_application_name = "PythonSnowpark"
-        config = {
-            "user": db_parameters["user"],
-            "password": db_parameters["password"],
-            "host": db_parameters["host"],
-            "port": db_parameters["port"],
-            "account": db_parameters["account"],
-            "schema": db_parameters["schema"],
-            "database": db_parameters["database"],
-            "protocol": db_parameters["protocol"],
-            "timezone": "UTC",
-            "application": new_application_name,
-        }
-        with snowflake.connector.connect(
-            **config
-        ) as conn, capture_sf_telemetry.patch_connection(conn, False) as telemetry_test:
-            conn._log_telemetry_imported_packages()
-            assert len(telemetry_test.records) > 0
-            for t in telemetry_test.records:
-                if (
-                    t.message[TelemetryField.KEY_TYPE.value]
-                    == TelemetryField.IMPORTED_PACKAGES.value
-                ):
-                    assert (
-                        new_application_name
-                        == t.message[TelemetryField.KEY_SOURCE.value]
-                    )
+    # these imports are not used but for testing
+    import html.parser  # noqa: F401
+    import json  # noqa: F401
+    import multiprocessing as mp  # noqa: F401
+    from datetime import date  # noqa: F401
+    from math import sqrt  # noqa: F401
+
+    def check_packages(message: str, expected_packages: list[str]) -> bool:
+        return (
+            all([package in message for package in expected_packages])
+            and "__main__" not in message
+        )
 
-        # test opt out
-        snowflake.connector.log_imported_packages_in_telemetry = False
-        with conn_cnx() as conn, capture_sf_telemetry.patch_connection(
-            conn, False
-        ) as telemetry_test:
-            conn._log_telemetry_imported_packages()
-            assert len(telemetry_test.records) == 0
-    finally:
-        snowflake.connector.log_imported_packages_in_telemetry = True
+    packages = [
+        "pytest",
+        "unittest",
+        "json",
+        "multiprocessing",
+        "html",
+        "datetime",
+        "math",
+    ]
+
+    with conn_cnx() as conn, capture_sf_telemetry.patch_connection(
+        conn, False
+    ) as telemetry_test:
+        conn._log_telemetry_imported_packages()
+        assert len(telemetry_test.records) > 0
+        assert any(
+            [
+                t.message[TelemetryField.KEY_TYPE.value]
+                == TelemetryField.IMPORTED_PACKAGES.value
+                and CLIENT_NAME == t.message[TelemetryField.KEY_SOURCE.value]
+                and check_packages(t.message["value"], packages)
+                for t in telemetry_test.records
+            ]
+        )
+
+    # test different application
+    new_application_name = "PythonSnowpark"
+    config = {
+        "user": db_parameters["user"],
+        "password": db_parameters["password"],
+        "host": db_parameters["host"],
+        "port": db_parameters["port"],
+        "account": db_parameters["account"],
+        "schema": db_parameters["schema"],
+        "database": db_parameters["database"],
+        "protocol": db_parameters["protocol"],
+        "timezone": "UTC",
+        "application": new_application_name,
+    }
+    with snowflake.connector.connect(
+        **config
+    ) as conn, capture_sf_telemetry.patch_connection(conn, False) as telemetry_test:
+        conn._log_telemetry_imported_packages()
+        assert len(telemetry_test.records) > 0
+        assert any(
+            [
+                t.message[TelemetryField.KEY_TYPE.value]
+                == TelemetryField.IMPORTED_PACKAGES.value
+                and new_application_name == t.message[TelemetryField.KEY_SOURCE.value]
+                for t in telemetry_test.records
+            ]
+        )
+
+    # test opt out
+    config["log_imported_packages_in_telemetry"] = False
+    with snowflake.connector.connect(
+        **config
+    ) as conn, capture_sf_telemetry.patch_connection(conn, False) as telemetry_test:
+        conn._log_telemetry_imported_packages()
+        assert len(telemetry_test.records) == 0