snowflake-connector-python 4.4.0__cp314-cp314-win_amd64.whl

snowflake/connector/platform_detection.py

@@ -0,0 +1,555 @@
+from __future__ import annotations
+
+import logging
+import os
+import re
+from concurrent.futures import CancelledError as FutureCancelledError
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from concurrent.futures.thread import ThreadPoolExecutor
+from contextlib import contextmanager
+from enum import Enum
+from functools import cache
+
+from .constants import (
+    ENV_VAR_BOOL_POSITIVE_VALUES_LOWERCASED,
+    ENV_VAR_DISABLE_PLATFORM_DETECTION,
+)
+from .options import boto3, botocore, installed_boto
+
+if installed_boto:
+    Config = botocore.config.Config
+    IMDSFetcher = botocore.utils.IMDSFetcher
+
+from .session_manager import SessionManager, SessionManagerFactory
+from .vendored.requests import RequestException, Timeout
+
+logger = logging.getLogger(__name__)
+
+# Loggers to suppress during platform detection to avoid noise in customer logs
+_LOGGERS_TO_SUPPRESS = [
+    "snowflake.connector.vendored.urllib3.connectionpool",
+    "botocore.utils",
+    "botocore.httpsession",
+    "urllib3.connectionpool",
+]
+
+
+@contextmanager
+def _suppress_platform_detection_logs():
+    """
+    Context manager to temporarily suppress all logs from underlying HTTP libraries during platform detection.
+
+    This prevents noisy DEBUG logs and stack traces from urllib3 and botocore when detecting
+    cloud platforms, which can confuse customers (SNOW-2204396). Our own debug logs are not affected.
+    """
+    original_levels = {}
+    try:
+        # Completely suppress all logs from noisy libraries
+        for logger_name in _LOGGERS_TO_SUPPRESS:
+            lib_logger = logging.getLogger(logger_name)
+            original_levels[logger_name] = lib_logger.level
+            lib_logger.setLevel(logging.CRITICAL + 1)  # Above CRITICAL = no logs at all
+        yield
+    finally:
+        # Restore original log levels
+        for logger_name, level in original_levels.items():
+            logging.getLogger(logger_name).setLevel(level)
+
+
+class _DetectionState(Enum):
+    """Internal enum to represent the detection state of a platform."""
+
+    DETECTED = "detected"
+    NOT_DETECTED = "not_detected"
+    HTTP_TIMEOUT = "timeout"
+    WORKER_TIMEOUT = "worker_timeout"
+
+
+# Result returned when platform detection is disabled via environment variable
+_PLATFORM_DETECTION_DISABLED_RESULT = ["disabled"]
+
+
+def is_ec2_instance(platform_detection_timeout_seconds: float):
+    """
+    Check if the current environment is running on an AWS EC2 instance.
+
+    If we query the AWS Instance Metadata Service (IMDS) for the instance identity document
+    and receive content back, then we assume we are running on an EC2 instance.
+    This function is compatible with IMDSv1 and IMDSv2 since we send the token in the request.
+    It will ignore the token if on IMDSv1 and use the token if on IMDSv2.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for the metadata service request.
+
+    Returns:
+        _DetectionState: DETECTED if running on EC2, NOT_DETECTED otherwise.
+    """
+    if not installed_boto:
+        logger.debug("boto3 is not installed, skipping EC2 instance detection")
+        return _DetectionState.NOT_DETECTED
+
+    try:
+        fetcher = IMDSFetcher(
+            timeout=platform_detection_timeout_seconds, num_attempts=1
+        )
+        document = fetcher._get_request(
+            "/latest/dynamic/instance-identity/document",
+            None,
+            fetcher._fetch_metadata_token(),
+        )
+        return (
+            _DetectionState.DETECTED
+            if document.content
+            else _DetectionState.NOT_DETECTED
+        )
+    except Exception:
+        return _DetectionState.NOT_DETECTED
+
+
+def is_aws_lambda():
+    """
+    Check if the current environment is running in AWS Lambda.
+
+    If we check for the LAMBDA_TASK_ROOT environment variable and it exists,
+    then we assume we are running in AWS Lambda.
+
+    Returns:
+        _DetectionState: DETECTED if LAMBDA_TASK_ROOT env var exists, NOT_DETECTED otherwise.
+    """
+    return (
+        _DetectionState.DETECTED
+        if "LAMBDA_TASK_ROOT" in os.environ
+        else _DetectionState.NOT_DETECTED
+    )
+
+
+def is_valid_arn_for_wif(arn: str) -> bool:
+    """
+    Validate if an AWS ARN is suitable for use with Snowflake's Workload Identity Federation (WIF).
+
+    Args:
+        arn: The AWS ARN string to validate.
+
+    Returns:
+        bool: True if ARN is valid for WIF, False otherwise.
+    """
+    patterns = [
+        r"^arn:[^:]+:iam::[^:]+:user/.+$",
+        r"^arn:[^:]+:sts::[^:]+:assumed-role/.+$",
+    ]
+    return any(re.match(p, arn) for p in patterns)
+
+
+def has_aws_identity(platform_detection_timeout_seconds: float):
+    """
+    Check if the current environment has a valid AWS identity for authentication.
+
+    If we retrieve an ARN from the caller identity and it is a valid WIF ARN,
+    then we assume we have a valid AWS identity for authentication.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for AWS API calls.
+
+    Returns:
+        _DetectionState: DETECTED if valid AWS identity exists, NOT_DETECTED otherwise.
+    """
+    if not installed_boto:
+        logger.debug("boto3 is not installed, skipping AWS identity detection")
+        return _DetectionState.NOT_DETECTED
+
+    try:
+        config = Config(
+            connect_timeout=platform_detection_timeout_seconds,
+            read_timeout=platform_detection_timeout_seconds,
+            retries={"total_max_attempts": 1},
+        )
+        caller_identity = boto3.client("sts", config=config).get_caller_identity()
+        if not caller_identity or "Arn" not in caller_identity:
+            return _DetectionState.NOT_DETECTED
+        return (
+            _DetectionState.DETECTED
+            if is_valid_arn_for_wif(caller_identity["Arn"])
+            else _DetectionState.NOT_DETECTED
+        )
+    except Exception:
+        return _DetectionState.NOT_DETECTED
+
+
+def is_azure_vm(
+    platform_detection_timeout_seconds: float, session_manager: SessionManager
+):
+    """
+    Check if the current environment is running on an Azure Virtual Machine.
+
+    If we query the Azure Instance Metadata Service and receive an HTTP 200 response,
+    then we assume we are running on an Azure VM.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for the metadata service request.
+        session_manager: SessionManager instance for making HTTP requests.
+
+    Returns:
+        _DetectionState: DETECTED if on Azure VM, HTTP_TIMEOUT if request times out,
+                        NOT_DETECTED otherwise.
+    """
+    try:
+        token_resp = session_manager.get(
+            "http://169.254.169.254/metadata/instance?api-version=2021-02-01",
+            headers={"Metadata": "true"},
+            timeout=platform_detection_timeout_seconds,
+        )
+        return (
+            _DetectionState.DETECTED
+            if token_resp.status_code == 200
+            else _DetectionState.NOT_DETECTED
+        )
+    except Timeout:
+        return _DetectionState.HTTP_TIMEOUT
+    except RequestException:
+        return _DetectionState.NOT_DETECTED
+
+
+def is_azure_function():
+    """
+    Check if the current environment is running in Azure Functions.
+
+    If we check for Azure Functions environment variables (FUNCTIONS_WORKER_RUNTIME,
+    FUNCTIONS_EXTENSION_VERSION, AzureWebJobsStorage) and they all exist,
+    then we assume we are running in Azure Functions.
+
+    Returns:
+        _DetectionState: DETECTED if all Azure Functions env vars are present,
+                        NOT_DETECTED otherwise.
+    """
+    service_vars = [
+        "FUNCTIONS_WORKER_RUNTIME",
+        "FUNCTIONS_EXTENSION_VERSION",
+        "AzureWebJobsStorage",
+    ]
+    return (
+        _DetectionState.DETECTED
+        if all(var in os.environ for var in service_vars)
+        else _DetectionState.NOT_DETECTED
+    )
+
+
+def is_managed_identity_available_on_azure_vm(
+    platform_detection_timeout_seconds,
+    session_manager: SessionManager,
+    resource="https://management.azure.com",
+):
+    """
+    Check if Azure Managed Identity is available and accessible on an Azure VM.
+
+    If we attempt to mint an access token from the Azure Instance Metadata Service
+    managed identity endpoint and receive an HTTP 200 response,
+    then we assume managed identity is available.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for the metadata service request.
+        session_manager: SessionManager instance for making HTTP requests.
+        resource: The Azure resource URI to request a token for.
+
+    Returns:
+        _DetectionState: DETECTED if managed identity is available, HTTP_TIMEOUT if request
+                        times out, NOT_DETECTED otherwise.
+    """
+    endpoint = f"http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource={resource}"
+    headers = {"Metadata": "true"}
+    try:
+        response = session_manager.get(
+            endpoint, headers=headers, timeout=platform_detection_timeout_seconds
+        )
+        return (
+            _DetectionState.DETECTED
+            if response.status_code == 200
+            else _DetectionState.NOT_DETECTED
+        )
+    except Timeout:
+        return _DetectionState.HTTP_TIMEOUT
+    except RequestException:
+        return _DetectionState.NOT_DETECTED
+
+
+def is_managed_identity_available_on_azure_function():
+    return bool(os.environ.get("IDENTITY_HEADER"))
+
+
+def has_azure_managed_identity(
+    platform_detection_timeout_seconds: float, session_manager: SessionManager
+):
+    """
+    Determine if Azure Managed Identity is available in the current environment.
+
+    If we are on Azure Functions and the IDENTITY_HEADER environment variable exists,
+    then we assume managed identity is available.
+    If we are on an Azure VM and can mint an access token from the managed identity endpoint,
+    then we assume managed identity is available.
+    Handles Azure Functions first since the checks are faster
+    Handles Azure VM checks second since they involve network calls.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for managed identity checks.
+        session_manager: SessionManager instance for making HTTP requests.
+
+    Returns:
+        _DetectionState: DETECTED if managed identity is available, HTTP_TIMEOUT if
+                        detection timed out, NOT_DETECTED otherwise.
+    """
+    # short circuit early to save on latency and avoid minting an unnecessary token
+    if is_azure_function() == _DetectionState.DETECTED:
+        return (
+            _DetectionState.DETECTED
+            if is_managed_identity_available_on_azure_function()
+            else _DetectionState.NOT_DETECTED
+        )
+    return is_managed_identity_available_on_azure_vm(
+        platform_detection_timeout_seconds, session_manager
+    )
+
+
+def is_gce_vm(
+    platform_detection_timeout_seconds: float, session_manager: SessionManager
+):
+    """
+    Check if the current environment is running on Google Compute Engine (GCE).
+
+    If we query the Google metadata server and receive a response with the
+    "Metadata-Flavor: Google" header, then we assume we are running on GCE.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for the metadata service request.
+        session_manager: SessionManager instance for making HTTP requests.
+
+    Returns:
+        _DetectionState: DETECTED if on GCE, HTTP_TIMEOUT if request times out,
+                        NOT_DETECTED otherwise.
+    """
+    try:
+        response = session_manager.get(
+            "http://metadata.google.internal",
+            timeout=platform_detection_timeout_seconds,
+        )
+        return (
+            _DetectionState.DETECTED
+            if response.headers and response.headers.get("Metadata-Flavor") == "Google"
+            else _DetectionState.NOT_DETECTED
+        )
+    except Timeout:
+        return _DetectionState.HTTP_TIMEOUT
+    except RequestException:
+        return _DetectionState.NOT_DETECTED
+
+
+def is_gcp_cloud_run_service():
+    """
+    Check if the current environment is running in Google Cloud Run service.
+
+    If we check for Cloud Run service environment variables (K_SERVICE, K_REVISION,
+    K_CONFIGURATION) and they all exist, then we assume we are running in Cloud Run service.
+
+    Returns:
+        _DetectionState: DETECTED if all Cloud Run service env vars are present,
+                        NOT_DETECTED otherwise.
+    """
+    service_vars = ["K_SERVICE", "K_REVISION", "K_CONFIGURATION"]
+    return (
+        _DetectionState.DETECTED
+        if all(var in os.environ for var in service_vars)
+        else _DetectionState.NOT_DETECTED
+    )
+
+
+def is_gcp_cloud_run_job():
+    """
+    Check if the current environment is running in Google Cloud Run job.
+
+    If we check for Cloud Run job environment variables (CLOUD_RUN_JOB, CLOUD_RUN_EXECUTION)
+    and they both exist, then we assume we are running in a Cloud Run job.
+
+    Returns:
+        _DetectionState: DETECTED if all Cloud Run job env vars are present,
+                        NOT_DETECTED otherwise.
+    """
+    job_vars = ["CLOUD_RUN_JOB", "CLOUD_RUN_EXECUTION"]
+    return (
+        _DetectionState.DETECTED
+        if all(var in os.environ for var in job_vars)
+        else _DetectionState.NOT_DETECTED
+    )
+
+
+def has_gcp_identity(
+    platform_detection_timeout_seconds: float, session_manager: SessionManager
+):
+    """
+    Check if the current environment has a valid Google Cloud Platform identity.
+
+    If we query the GCP metadata service for the default service account email
+    and receive a non-empty response, then we assume we have a valid GCP identity.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for the metadata service request.
+        session_manager: SessionManager instance for making HTTP requests.
+    Returns:
+        _DetectionState: DETECTED if valid GCP identity exists, HTTP_TIMEOUT if request
+                        times out, NOT_DETECTED otherwise.
+    """
+    try:
+        response = session_manager.get(
+            "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/email",
+            headers={"Metadata-Flavor": "Google"},
+            timeout=platform_detection_timeout_seconds,
+        )
+        return (
+            _DetectionState.DETECTED
+            if response.status_code == 200
+            else _DetectionState.NOT_DETECTED
+        )
+    except Timeout:
+        return _DetectionState.HTTP_TIMEOUT
+    except RequestException:
+        return _DetectionState.NOT_DETECTED
+
+
+def is_github_action():
+    """
+    Check if the current environment is running in GitHub Actions.
+
+    If we check for the GITHUB_ACTIONS environment variable and it exists,
+    then we assume we are running in GitHub Actions.
+
+    Returns:
+        _DetectionState: DETECTED if GITHUB_ACTIONS env var exists, NOT_DETECTED otherwise.
+    """
+    return (
+        _DetectionState.DETECTED
+        if "GITHUB_ACTIONS" in os.environ
+        else _DetectionState.NOT_DETECTED
+    )
+
+
+@cache
+def detect_platforms(
+    platform_detection_timeout_seconds: float | None,
+    session_manager: SessionManager | None = None,
+) -> list[str]:
+    """
+    Detect all potential platforms that the current environment may be running on.
+    Swallows all exceptions and returns an empty list if any exception occurs to not affect main driver functionality.
+
+    Args:
+        platform_detection_timeout_seconds: Timeout value for platform detection requests. Defaults to 0.2 seconds
+                if None is provided.
+        session_manager: SessionManager instance for making HTTP requests. If None, a new instance will be created.
+
+    Returns:
+        list[str]: List of detected platform names. Platforms that timed out (either HTTP timeout
+                  or thread timeout) will have "_timeout" suffix appended to their name.
+                  Returns _PLATFORM_DETECTION_DISABLED_RESULT if the ENV_VAR_DISABLE_PLATFORM_DETECTION
+                  environment variable is set to a value in ENV_VAR_BOOL_POSITIVE_VALUES_LOWERCASED
+                  (case-insensitive). Returns empty list if any exception occurs during detection.
+    """
+    try:
+        # Check if platform detection is disabled via environment variable
+        if (
+            os.environ.get(ENV_VAR_DISABLE_PLATFORM_DETECTION, "").lower()
+            in ENV_VAR_BOOL_POSITIVE_VALUES_LOWERCASED
+        ):
+            logger.debug(
+                "Platform detection disabled via %s environment variable",
+                ENV_VAR_DISABLE_PLATFORM_DETECTION,
+            )
+            return _PLATFORM_DETECTION_DISABLED_RESULT
+
+        if platform_detection_timeout_seconds is None:
+            platform_detection_timeout_seconds = 0.2
+
+        if session_manager is None:
+            # This should never happen - we expect session manager to be passed from the outer scope
+            logger.debug(
+                "No session manager provided. HTTP settings may not be preserved. Using default."
+            )
+            session_manager = SessionManagerFactory.get_manager(
+                use_pooling=False, max_retries=0
+            )
+
+        # HTTP timeout should be slightly shorter than thread timeout to allow HTTP-level
+        # timeouts to occur before thread executor times out. This helps distinguish between
+        # HTTP_TIMEOUT (network issue) and WORKER_TIMEOUT (thread stuck/hung).
+        http_timeout_epsilon = 0.05  # 5% shorter
+        http_timeout = platform_detection_timeout_seconds * (1 - http_timeout_epsilon)
+        threads_timeout = platform_detection_timeout_seconds
+
+        # Suppress noisy logs from underlying HTTP libraries during platform detection
+        with _suppress_platform_detection_logs():
+            # Run environment-only checks synchronously (no network calls, no threading overhead)
+            platforms = {
+                "is_aws_lambda": is_aws_lambda(),
+                "is_azure_function": is_azure_function(),
+                "is_gce_cloud_run_service": is_gcp_cloud_run_service(),
+                "is_gce_cloud_run_job": is_gcp_cloud_run_job(),
+                "is_github_action": is_github_action(),
+            }
+
+            # Run network-calling functions in parallel
+            if platform_detection_timeout_seconds != 0.0:
+                with ThreadPoolExecutor(max_workers=6) as executor:
+                    futures = {
+                        "is_ec2_instance": executor.submit(
+                            is_ec2_instance, http_timeout
+                        ),
+                        "has_aws_identity": executor.submit(
+                            has_aws_identity, http_timeout
+                        ),
+                        "is_azure_vm": executor.submit(
+                            is_azure_vm,
+                            http_timeout,
+                            session_manager,
+                        ),
+                        "has_azure_managed_identity": executor.submit(
+                            has_azure_managed_identity,
+                            http_timeout,
+                            session_manager,
+                        ),
+                        "is_gce_vm": executor.submit(
+                            is_gce_vm,
+                            http_timeout,
+                            session_manager,
+                        ),
+                        "has_gcp_identity": executor.submit(
+                            has_gcp_identity,
+                            http_timeout,
+                            session_manager,
+                        ),
+                    }
+
+                    # Enforce timeout at executor level - all parallel detections must complete
+                    # within threads_timeout
+                    for key, future in futures.items():
+                        try:
+                            platforms[key] = future.result(timeout=threads_timeout)
+                        except (FutureTimeoutError, FutureCancelledError):
+                            # Thread/future timed out at executor level
+                            platforms[key] = _DetectionState.WORKER_TIMEOUT
+                        except Exception:
+                            # Any other error from the thread
+                            platforms[key] = _DetectionState.NOT_DETECTED
+
+            detected_platforms = []
+            for platform_name, detection_state in platforms.items():
+                if detection_state == _DetectionState.DETECTED:
+                    detected_platforms.append(platform_name)
+                elif detection_state in (
+                    _DetectionState.HTTP_TIMEOUT,
+                    _DetectionState.WORKER_TIMEOUT,
+                ):
+                    detected_platforms.append(f"{platform_name}_timeout")
+
+            logger.debug(
+                "Platform detection completed. Detected platforms: %s",
+                detected_platforms,
+            )
+            return detected_platforms
+    except Exception:
+        return []

snowflake/connector/proxy.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+
+def get_proxy_url(
+    proxy_host: str | None,
+    proxy_port: str | None,
+    proxy_user: str | None = None,
+    proxy_password: str | None = None,
+) -> str | None:
+    http_prefix = "http://"
+    https_prefix = "https://"
+
+    if proxy_host and proxy_port:
+        if proxy_host.startswith(http_prefix):
+            host = proxy_host[len(http_prefix) :]
+        elif proxy_host.startswith(https_prefix):
+            host = proxy_host[len(https_prefix) :]
+        else:
+            host = proxy_host
+        auth = (
+            f"{proxy_user or ''}:{proxy_password or ''}@"
+            if proxy_user or proxy_password
+            else ""
+        )
+        return f"{http_prefix}{auth}{host}:{proxy_port}"
+
+    return None