SNOW-2268606 zero timeout disables endpoint-based cloud platform detection (#2490)

Author: sfc-gh-mmishchenko

src/snowflake/connector/auth/_auth.py

@@ -195,7 +195,7 @@ def authenticate(
             self._rest._connection.login_timeout,
             self._rest._connection._network_timeout,
             self._rest._connection._socket_timeout,
-            self._rest._connection._platform_detection_timeout_seconds,
+            self._rest._connection.platform_detection_timeout_seconds,
             session_manager=self._rest.session_manager.clone(use_pooling=False),
         )
 

src/snowflake/connector/auth/okta.py

@@ -167,7 +167,9 @@ def _step1(
             conn._internal_application_version,
             conn._ocsp_mode(),
             conn.login_timeout,
-            conn._network_timeout,
+            conn.network_timeout,
+            conn.socket_timeout,
+            conn.platform_detection_timeout_seconds,
             session_manager=conn._session_manager.clone(use_pooling=False),
         )
 

src/snowflake/connector/auth/webbrowser.py

@@ -456,12 +456,14 @@ def _get_sso_url(
         body = Auth.base_auth_data(
             user,
             account,
-            conn._rest._connection.application,
-            conn._rest._connection._internal_application_name,
-            conn._rest._connection._internal_application_version,
-            conn._rest._connection._ocsp_mode(),
-            conn._rest._connection.login_timeout,
-            conn._rest._connection._network_timeout,
+            conn.application,
+            conn._internal_application_name,
+            conn._internal_application_version,
+            conn._ocsp_mode(),
+            conn.login_timeout,
+            conn.network_timeout,
+            conn.socket_timeout,
+            conn.platform_detection_timeout_seconds,
             session_manager=conn.rest.session_manager.clone(use_pooling=False),
         )
 

src/snowflake/connector/platform_detection.py

@@ -417,33 +417,36 @@ def detect_platforms(
         }
 
         # Run network-calling functions in parallel
-        with ThreadPoolExecutor(max_workers=6) as executor:
-            futures = {
-                "is_ec2_instance": executor.submit(
-                    is_ec2_instance, platform_detection_timeout_seconds
-                ),
-                "has_aws_identity": executor.submit(
-                    has_aws_identity, platform_detection_timeout_seconds
-                ),
-                "is_azure_vm": executor.submit(
-                    is_azure_vm, platform_detection_timeout_seconds, session_manager
-                ),
-                "has_azure_managed_identity": executor.submit(
-                    has_azure_managed_identity,
-                    platform_detection_timeout_seconds,
-                    session_manager,
-                ),
-                "is_gce_vm": executor.submit(
-                    is_gce_vm, platform_detection_timeout_seconds, session_manager
-                ),
-                "has_gcp_identity": executor.submit(
-                    has_gcp_identity,
-                    platform_detection_timeout_seconds,
-                    session_manager,
-                ),
-            }
-
-            platforms.update({key: future.result() for key, future in futures.items()})
+        if platform_detection_timeout_seconds != 0.0:
+            with ThreadPoolExecutor(max_workers=6) as executor:
+                futures = {
+                    "is_ec2_instance": executor.submit(
+                        is_ec2_instance, platform_detection_timeout_seconds
+                    ),
+                    "has_aws_identity": executor.submit(
+                        has_aws_identity, platform_detection_timeout_seconds
+                    ),
+                    "is_azure_vm": executor.submit(
+                        is_azure_vm, platform_detection_timeout_seconds, session_manager
+                    ),
+                    "has_azure_managed_identity": executor.submit(
+                        has_azure_managed_identity,
+                        platform_detection_timeout_seconds,
+                        session_manager,
+                    ),
+                    "is_gce_vm": executor.submit(
+                        is_gce_vm, platform_detection_timeout_seconds, session_manager
+                    ),
+                    "has_gcp_identity": executor.submit(
+                        has_gcp_identity,
+                        platform_detection_timeout_seconds,
+                        session_manager,
+                    ),
+                }
+
+                platforms.update(
+                    {key: future.result() for key, future in futures.items()}
+                )
 
         detected_platforms = []
         for platform_name, detection_state in platforms.items():

test/integ/test_connection.py

@@ -204,6 +204,37 @@ def test_platform_detection_timeout(conn_cnx):
         assert cnx.platform_detection_timeout_seconds == 2.5
 
 
+@pytest.mark.skipolddriver
+def test_platform_detection_zero_timeout(conn_cnx):
+    """Tests platform detection with timeout set to zero.
+
+    The expectation is that it mustn't do diagnostic requests at all.
+    """
+    with (
+        mock.patch(
+            "snowflake.connector.platform_detection.is_ec2_instance"
+        ) as is_ec2_instance,
+        mock.patch(
+            "snowflake.connector.platform_detection.has_aws_identity"
+        ) as has_aws_identity,
+        mock.patch("snowflake.connector.platform_detection.is_azure_vm") as is_azure_vm,
+        mock.patch(
+            "snowflake.connector.platform_detection.has_azure_managed_identity"
+        ) as has_azure_managed_identity,
+        mock.patch("snowflake.connector.platform_detection.is_gce_vm") as is_gce_vm,
+        mock.patch(
+            "snowflake.connector.platform_detection.has_gcp_identity"
+        ) as has_gcp_identity,
+    ):
+        with conn_cnx(platform_detection_timeout_seconds=0):
+            assert not is_ec2_instance.called
+            assert not has_aws_identity.called
+            assert not is_azure_vm.called
+            assert not has_azure_managed_identity.called
+            assert not is_gce_vm.called
+            assert not has_gcp_identity.called
+
+
 def test_bad_db(conn_cnx):
     """Attempts to use a bad DB."""
     with conn_cnx(database="baddb") as cnx:
@@ -1119,9 +1150,10 @@ def check_packages(message: str, expected_packages: list[str]) -> bool:
         "math",
     ]
 
-    with conn_cnx() as conn, capture_sf_telemetry.patch_connection(
-        conn, False
-    ) as telemetry_test:
+    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(
@@ -1136,10 +1168,13 @@ def check_packages(message: str, expected_packages: list[str]) -> bool:
 
     # test different application
     new_application_name = "PythonSnowpark"
-    with conn_cnx(
-        timezone="UTC",
-        application=new_application_name,
-    ) as conn, capture_sf_telemetry.patch_connection(conn, False) as telemetry_test:
+    with (
+        conn_cnx(
+            timezone="UTC",
+            application=new_application_name,
+        ) 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(
@@ -1152,11 +1187,14 @@ def check_packages(message: str, expected_packages: list[str]) -> bool:
         )
 
     # test opt out
-    with conn_cnx(
-        timezone="UTC",
-        application=new_application_name,
-        log_imported_packages_in_telemetry=False,
-    ) as conn, capture_sf_telemetry.patch_connection(conn, False) as telemetry_test:
+    with (
+        conn_cnx(
+            timezone="UTC",
+            application=new_application_name,
+            log_imported_packages_in_telemetry=False,
+        ) as conn,
+        capture_sf_telemetry.patch_connection(conn, False) as telemetry_test,
+    ):
         conn._log_telemetry_imported_packages()
         assert len(telemetry_test.records) == 0
 
@@ -1293,9 +1331,10 @@ def test_ocsp_mode_insecure_mode_and_disable_ocsp_checks_match(
     conn_cnx, is_public_test, is_local_dev_setup, caplog
 ):
     caplog.set_level(logging.DEBUG, "snowflake.connector.ocsp_snowflake")
-    with conn_cnx(
-        insecure_mode=True, disable_ocsp_checks=True
-    ) as conn, conn.cursor() as cur:
+    with (
+        conn_cnx(insecure_mode=True, disable_ocsp_checks=True) as conn,
+        conn.cursor() as cur,
+    ):
         assert cur.execute("select 1").fetchall() == [(1,)]
         assert "snowflake.connector.ocsp_snowflake" not in caplog.text
         if is_public_test or is_local_dev_setup:
@@ -1311,9 +1350,10 @@ def test_ocsp_mode_insecure_mode_and_disable_ocsp_checks_mismatch_ocsp_disabled(
     conn_cnx, is_public_test, is_local_dev_setup, caplog
 ):
     caplog.set_level(logging.DEBUG, "snowflake.connector.ocsp_snowflake")
-    with conn_cnx(
-        insecure_mode=False, disable_ocsp_checks=True
-    ) as conn, conn.cursor() as cur:
+    with (
+        conn_cnx(insecure_mode=False, disable_ocsp_checks=True) as conn,
+        conn.cursor() as cur,
+    ):
         assert cur.execute("select 1").fetchall() == [(1,)]
         assert "snowflake.connector.ocsp_snowflake" not in caplog.text
         if is_public_test or is_local_dev_setup:
@@ -1329,9 +1369,10 @@ def test_ocsp_mode_insecure_mode_and_disable_ocsp_checks_mismatch_ocsp_enabled(
     conn_cnx, is_public_test, is_local_dev_setup, caplog
 ):
     caplog.set_level(logging.DEBUG, "snowflake.connector.ocsp_snowflake")
-    with conn_cnx(
-        insecure_mode=True, disable_ocsp_checks=False
-    ) as conn, conn.cursor() as cur:
+    with (
+        conn_cnx(insecure_mode=True, disable_ocsp_checks=False) as conn,
+        conn.cursor() as cur,
+    ):
         assert cur.execute("select 1").fetchall() == [(1,)]
         if is_public_test or is_local_dev_setup:
             assert "snowflake.connector.ocsp_snowflake" in caplog.text
@@ -1430,9 +1471,10 @@ def test_disable_telemetry(conn_cnx, caplog):
 
     # set session parameters to false
     with caplog.at_level(logging.DEBUG):
-        with conn_cnx(
-            session_parameters={"CLIENT_TELEMETRY_ENABLED": False}
-        ) as conn, conn.cursor() as cur:
+        with (
+            conn_cnx(session_parameters={"CLIENT_TELEMETRY_ENABLED": False}) as conn,
+            conn.cursor() as cur,
+        ):
             cur.execute("select 1").fetchall()
             assert not conn.telemetry_enabled and not conn._telemetry._log_batch
             # this enable won't work as the session parameter is set to false

test/unit/test_detect_platforms.py

@@ -26,6 +26,24 @@ def unavailable_metadata_service_with_request_exception(unavailable_metadata_ser
     return unavailable_metadata_service
 
 
+@pytest.fixture
+def labels_detected_by_endpoints():
+    return {
+        "is_ec2_instance",
+        "is_ec2_instance_timeout",
+        "has_aws_identity",
+        "has_aws_identity_timeout",
+        "is_azure_vm",
+        "is_azure_vm_timeout",
+        "has_azure_managed_identity",
+        "has_azure_managed_identity_timeout",
+        "is_gce_vm",
+        "is_gce_vm_timeout",
+        "has_gcp_identity",
+        "has_gcp_identity_timeout",
+    }
+
+
 @pytest.mark.xdist_group(name="serial_tests")
 class TestDetectPlatforms:
     @pytest.fixture(autouse=True)
@@ -288,3 +306,24 @@ def test_gce_cloud_run_job_missing_cloud_run_job(
     ):
         result = detect_platforms(platform_detection_timeout_seconds=None)
         assert "is_gce_cloud_run_job" not in result
+
+    def test_zero_platform_detection_timeout_disables_endpoints_detection_on_cloud(
+        self,
+        fake_azure_vm_metadata_service,
+        fake_azure_function_metadata_service,
+        fake_gce_metadata_service,
+        fake_gce_cloud_run_service_metadata_service,
+        fake_gce_cloud_run_job_metadata_service,
+        fake_github_actions_metadata_service,
+        labels_detected_by_endpoints,
+    ):
+        result = detect_platforms(platform_detection_timeout_seconds=0)
+        assert not labels_detected_by_endpoints.intersection(result)
+
+    def test_zero_platform_detection_timeout_disables_endpoints_detection_out_of_cloud(
+        self,
+        unavailable_metadata_service_with_request_exception,
+        labels_detected_by_endpoints,
+    ):
+        result = detect_platforms(platform_detection_timeout_seconds=0)
+        assert not labels_detected_by_endpoints.intersection(result)