switched to detect platforms by adding it to the base_auth data which will get sent to login-request and eventually persisted in the SessionDPO

Author: sfc-gh-sshetkar

src/snowflake/connector/auth/_auth.py

@@ -3,11 +3,16 @@
 import copy
 import json
 import logging
+import os
+import re
 import uuid
+from concurrent.futures.thread import ThreadPoolExecutor
 from datetime import datetime, timezone
 from threading import Thread
 from typing import TYPE_CHECKING, Any, Callable
 
+import boto3
+from botocore.utils import IMDSFetcher
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives.serialization import (
     Encoding,
@@ -53,7 +58,9 @@
 )
 from ..sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED
 from ..token_cache import TokenCache, TokenKey, TokenType
+from ..vendored import requests
 from ..version import VERSION
+from ..wif_util import DEFAULT_ENTRA_SNOWFLAKE_RESOURCE
 from .no_auth import AuthNoAuth
 
 if TYPE_CHECKING:
@@ -101,6 +108,136 @@ def base_auth_data(
         network_timeout: int | None = None,
         socket_timeout: int | None = None,
     ):
+        def detect_platforms() -> list[str]:
+            def is_ec2_instance(timeout=0.5):
+                try:
+                    fetcher = IMDSFetcher(timeout=timeout, num_attempts=2)
+                    document = fetcher._get_request(
+                        "/latest/dynamic/instance-identity/document",
+                        None,
+                        fetcher._fetch_metadata_token(),
+                    )
+                    return bool(document.content)
+                except Exception:
+                    return False
+
+            def is_aws_lambda():
+                return "LAMBDA_TASK_ROOT" in os.environ
+
+            def is_valid_arn_for_wif(arn: str) -> bool:
+                patterns = [
+                    r"^arn:[^:]+:iam::[^:]+:user/.+$",
+                    r"^arn:[^:]+:sts::[^:]+:assumed-role/.+$",
+                ]
+                return any(re.match(p, arn) for p in patterns)
+
+            def has_aws_identity():
+                try:
+                    caller_identity = boto3.client("sts").get_caller_identity()
+                    if not caller_identity or "Arn" not in caller_identity:
+                        return False
+                    else:
+                        return is_valid_arn_for_wif(caller_identity["Arn"])
+                except Exception:
+                    return False
+
+            def is_azure_vm(timeout=0.5):
+                try:
+                    token_resp = requests.get(
+                        "http://169.254.169.254/metadata/instance?api-version=2021-02-01",
+                        headers={"Metadata": "true"},
+                        timeout=timeout,
+                    )
+                    return token_resp.status_code == 200
+                except requests.RequestException:
+                    return False
+
+            def is_azure_function():
+                service_vars = [
+                    "FUNCTIONS_WORKER_RUNTIME",
+                    "FUNCTIONS_EXTENSION_VERSION",
+                    "AzureWebJobsStorage",
+                ]
+                return all(var in os.environ for var in service_vars)
+
+            def is_managed_identity_available_on_azure_vm(
+                resource=DEFAULT_ENTRA_SNOWFLAKE_RESOURCE, timeout=0.5
+            ):
+                endpoint = f"http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource={resource}"
+                headers = {"Metadata": "true"}
+                try:
+                    response = requests.get(endpoint, headers=headers, timeout=timeout)
+                    return response.status_code == 200
+                except requests.RequestException:
+                    return False
+
+            def has_azure_managed_identity(on_azure_vm, on_azure_function):
+                if on_azure_function:
+                    return bool(os.environ.get("IDENTITY_HEADER"))
+                if on_azure_vm:
+                    return is_managed_identity_available_on_azure_vm()
+                return False
+
+            def is_gce_vm(timeout=0.5):
+                try:
+                    response = requests.get(
+                        "http://metadata.google.internal", timeout=timeout
+                    )
+                    return response.headers.get("Metadata-Flavor") == "Google"
+                except requests.RequestException:
+                    return False
+
+            def is_gce_cloud_run_service():
+                service_vars = ["K_SERVICE", "K_REVISION", "K_CONFIGURATION"]
+                return all(var in os.environ for var in service_vars)
+
+            def is_gce_cloud_run_job():
+                job_vars = ["CLOUD_RUN_JOB", "CLOUD_RUN_EXECUTION"]
+                return all(var in os.environ for var in job_vars)
+
+            def has_gcp_identity(timeout=2):
+                try:
+                    response = requests.get(
+                        "http://metadata/computeMetadata/v1/instance/service-accounts/default/email",
+                        headers={"Metadata-Flavor": "Google"},
+                        timeout=timeout,
+                    )
+                    response.raise_for_status()
+                    return bool(response.text)
+                except requests.RequestException:
+                    return False
+
+            def is_github_action():
+                return "GITHUB_ACTIONS" in os.environ
+
+            with ThreadPoolExecutor(max_workers=10) as executor:
+                futures = {
+                    "is_ec2_instance": executor.submit(is_ec2_instance),
+                    "is_aws_lambda": executor.submit(is_aws_lambda),
+                    "has_aws_identity": executor.submit(has_aws_identity),
+                    "is_azure_vm": executor.submit(is_azure_vm),
+                    "is_azure_function": executor.submit(is_azure_function),
+                    "is_gce_vm": executor.submit(is_gce_vm),
+                    "is_gce_cloud_run_service": executor.submit(
+                        is_gce_cloud_run_service
+                    ),
+                    "is_gce_cloud_run_job": executor.submit(is_gce_cloud_run_job),
+                    "has_gcp_identity": executor.submit(has_gcp_identity),
+                    "is_github_action": executor.submit(is_github_action),
+                }
+
+                platforms = {key: future.result() for key, future in futures.items()}
+
+            platforms["azure_managed_identity"] = has_azure_managed_identity(
+                platforms["is_azure_vm"], platforms["is_azure_function"]
+            )
+
+            detected_platforms = [
+                platform for platform, detected in platforms.items() if detected
+            ]
+
+            return detected_platforms
+
         return {
             "data": {
                 "CLIENT_APP_ID": internal_application_name,
@@ -120,6 +257,7 @@ def base_auth_data(
                     "LOGIN_TIMEOUT": login_timeout,
                     "NETWORK_TIMEOUT": network_timeout,
                     "SOCKET_TIMEOUT": socket_timeout,
+                    "PLATFORM": detect_platforms(),
                 },
             },
         }

src/snowflake/connector/connection.py

@@ -18,13 +18,11 @@
 from functools import partial
 from io import StringIO
 from logging import getLogger
-from threading import Lock, Thread
+from threading import Lock
 from types import TracebackType
 from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence
 from uuid import UUID
 
-import boto3
-from botocore.utils import IMDSFetcher
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
@@ -121,8 +119,7 @@
 from .time_util import HeartBeatTimer, get_time_millis
 from .url_util import extract_top_level_domain_from_hostname
 from .util_text import construct_hostname, parse_account, split_statements
-from .vendored import requests
-from .wif_util import DEFAULT_ENTRA_SNOWFLAKE_RESOURCE, AttestationProvider
+from .wif_util import AttestationProvider
 
 DEFAULT_CLIENT_PREFETCH_THREADS = 4
 MAX_CLIENT_PREFETCH_THREADS = 10
@@ -277,10 +274,6 @@ def _get_private_bytes_from_file(
         True,
         bool,
     ),  # Whether to log imported packages in telemetry
-    "log_platform_in_telemetry": (
-        True,
-        bool,
-    ),  # Whether to log platform in telemetry
     "disable_query_context_cache": (
         False,
         bool,
@@ -387,133 +380,6 @@ class TypeAndBinding(NamedTuple):
     binding: str | None
 
 
-def detect_platforms() -> list[str]:
-    def is_ec2_instance(timeout=0.5):
-        try:
-            fetcher = IMDSFetcher(timeout=timeout, num_attempts=2)
-            document = fetcher._get_request(
-                "/latest/dynamic/instance-identity/document",
-                None,
-                fetcher._fetch_metadata_token(),
-            )
-            return bool(document.content)
-        except Exception:
-            return False
-
-    def is_aws_lambda():
-        return "LAMBDA_TASK_ROOT" in os.environ
-
-    def is_valid_arn_for_wif(arn: str) -> bool:
-        patterns = [
-            r"^arn:[^:]+:iam::[^:]+:user/.+$",
-            r"^arn:[^:]+:sts::[^:]+:assumed-role/.+$",
-        ]
-        return any(re.match(p, arn) for p in patterns)
-
-    def has_aws_identity():
-        try:
-            caller_identity = boto3.client("sts").get_caller_identity()
-            if not caller_identity or "Arn" not in caller_identity:
-                return False
-            else:
-                return is_valid_arn_for_wif(caller_identity["Arn"])
-        except Exception:
-            return False
-
-    def is_azure_vm(timeout=0.5):
-        try:
-            token_resp = requests.get(
-                "http://169.254.169.254/metadata/instance?api-version=2021-02-01",
-                headers={"Metadata": "true"},
-                timeout=timeout,
-            )
-            return token_resp.status_code == 200
-        except requests.RequestException:
-            return False
-
-    def is_azure_function():
-        service_vars = [
-            "FUNCTIONS_WORKER_RUNTIME",
-            "FUNCTIONS_EXTENSION_VERSION",
-            "AzureWebJobsStorage",
-        ]
-        return all(var in os.environ for var in service_vars)
-
-    def is_managed_identity_available_on_azure_vm(
-        resource=DEFAULT_ENTRA_SNOWFLAKE_RESOURCE, timeout=0.5
-    ):
-        endpoint = f"http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource={resource}"
-        headers = {"Metadata": "true"}
-        try:
-            response = requests.get(endpoint, headers=headers, timeout=timeout)
-            return response.status_code == 200
-        except requests.RequestException:
-            return False
-
-    def has_azure_managed_identity(on_azure_vm, on_azure_function):
-        if on_azure_function:
-            return bool(os.environ.get("IDENTITY_HEADER"))
-        if on_azure_vm:
-            return is_managed_identity_available_on_azure_vm()
-        return False
-
-    def is_gce_vm(timeout=0.5):
-        try:
-            response = requests.get("http://metadata.google.internal", timeout=timeout)
-            return response.headers.get("Metadata-Flavor") == "Google"
-        except requests.RequestException:
-            return False
-
-    def is_gce_cloud_run_service():
-        service_vars = ["K_SERVICE", "K_REVISION", "K_CONFIGURATION"]
-        return all(var in os.environ for var in service_vars)
-
-    def is_gce_cloud_run_job():
-        job_vars = ["CLOUD_RUN_JOB", "CLOUD_RUN_EXECUTION"]
-        return all(var in os.environ for var in job_vars)
-
-    def has_gcp_identity(timeout=2):
-        try:
-            response = requests.get(
-                "http://metadata/computeMetadata/v1/instance/service-accounts/default/email",
-                headers={"Metadata-Flavor": "Google"},
-                timeout=timeout,
-            )
-            response.raise_for_status()
-            return bool(response.text)
-        except requests.RequestException:
-            return False
-
-    def is_github_action():
-        return "GITHUB_ACTIONS" in os.environ
-
-    with ThreadPoolExecutor(max_workers=10) as executor:
-        futures = {
-            "is_ec2_instance": executor.submit(is_ec2_instance),
-            "is_aws_lambda": executor.submit(is_aws_lambda),
-            "has_aws_identity": executor.submit(has_aws_identity),
-            "is_azure_vm": executor.submit(is_azure_vm),
-            "is_azure_function": executor.submit(is_azure_function),
-            "is_gce_vm": executor.submit(is_gce_vm),
-            "is_gce_cloud_run_service": executor.submit(is_gce_cloud_run_service),
-            "is_gce_cloud_run_job": executor.submit(is_gce_cloud_run_job),
-            "has_gcp_identity": executor.submit(has_gcp_identity),
-            "is_github_action": executor.submit(is_github_action),
-        }
-
-        platforms = {key: future.result() for key, future in futures.items()}
-
-    platforms["azure_managed_identity"] = has_azure_managed_identity(
-        platforms["is_azure_vm"], platforms["is_azure_function"]
-    )
-
-    detected_platforms = [
-        platform for platform, detected in platforms.items() if detected
-    ]
-
-    return detected_platforms
-
-
 class SnowflakeConnection:
     """Implementation of the connection object for the Snowflake Database.
 
@@ -682,8 +548,6 @@ def __init__(
 
         # get the imported modules from sys.modules
         self._log_telemetry_imported_packages()
-        # log the platform of the client
-        Thread(target=self._log_telemetry_platform_info(), daemon=True).start()
         # check SNOW-1218851 for long term improvement plan to refactor ocsp code
         atexit.register(self._close_at_exit)
 
@@ -2339,20 +2203,6 @@ def _log_telemetry_imported_packages(self) -> None:
                 )
             )
 
-    def _log_telemetry_platform_info(self) -> None:
-        if self._log_platform_in_telemetry:
-            ts = get_time_millis()
-            self._log_telemetry(
-                TelemetryData.from_telemetry_data_dict(
-                    from_dict={
-                        TelemetryField.KEY_TYPE.value: TelemetryField.PLATFORM_INFO.value,
-                        TelemetryField.KEY_VALUE.value: str(detect_platforms()),
-                    },
-                    timestamp=ts,
-                    connection=self,
-                )
-            )
-
     def is_valid(self) -> bool:
         """This function tries to answer the question: Is this connection still good for sending queries?
         Attempts to validate the connections both on the TCP/IP and Session levels."""

src/snowflake/connector/telemetry.py

@@ -37,8 +37,6 @@ class TelemetryField(Enum):
     PANDAS_WRITE = "client_write_pandas"
     # imported packages along with client
     IMPORTED_PACKAGES = "client_imported_packages"
-    # platform information describing where the client is running (AWS EC2, GCP VM, Azure function, etc)
-    PLATFORM_INFO = "client_platform_info"
     # multi-statement usage
     MULTI_STATEMENT = "client_multi_statement_query"
     # Keys for telemetry data sent through either in-band or out-of-band telemetry