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

snowflake/connector/telemetry_oob.py

@@ -0,0 +1,544 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import datetime
+import json
+import logging
+import uuid
+from collections import namedtuple
+from queue import Queue
+from threading import Lock
+from typing import Any
+
+from .compat import OK
+from .description import CLIENT_NAME, SNOWFLAKE_CONNECTOR_VERSION
+from .secret_detector import SecretDetector
+from .telemetry import TelemetryField, generate_telemetry_data_dict
+from .test_util import ENABLE_TELEMETRY_LOG, rt_plain_logger
+from .vendored import requests
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_BATCH_SIZE = 10
+DEFAULT_NUM_OF_RETRY_TO_TRIGGER_TELEMETRY = 10
+REQUEST_TIMEOUT = 3
+
+TelemetryAPI = namedtuple("TelemetryAPI", ["url", "api_key"])
+TelemetryServer = namedtuple("TelemetryServer", ["name", "url", "api_key"])
+TelemetryEventBase = namedtuple(
+    "TelemetryEventBase", ["name", "tags", "urgent", "value"]
+)
+
+
+class TelemetryAPIEndpoint:
+    SFCTEST = TelemetryAPI(
+        url="https://sfctest.client-telemetry.snowflakecomputing.com/enqueue",
+        api_key="rRNY3EPNsB4U89XYuqsZKa7TSxb9QVX93yNM4tS6",
+    )
+    SFCDEV = TelemetryAPI(
+        url="https://sfcdev.client-telemetry.snowflakecomputing.com/enqueue",
+        api_key="kyTKLWpEZSaJnrzTZ63I96QXZHKsgfqbaGmAaIWf",
+    )
+    PROD = TelemetryAPI(
+        url="https://client-telemetry.snowflakecomputing.com/enqueue",
+        api_key="wLpEKqnLOW9tGNwTjab5N611YQApOb3t9xOnE1rX",
+    )
+
+
+class TelemetryServerDeployments:
+    DEV = TelemetryServer(
+        "dev", TelemetryAPIEndpoint.SFCTEST.url, TelemetryAPIEndpoint.SFCTEST.api_key
+    )
+    REG = TelemetryServer(
+        "reg", TelemetryAPIEndpoint.SFCTEST.url, TelemetryAPIEndpoint.SFCTEST.api_key
+    )
+    QA1 = TelemetryServer(
+        "qa1", TelemetryAPIEndpoint.SFCDEV.url, TelemetryAPIEndpoint.SFCDEV.api_key
+    )
+    PREPROD3 = TelemetryServer(
+        "preprod3", TelemetryAPIEndpoint.SFCDEV.url, TelemetryAPIEndpoint.SFCDEV.api_key
+    )
+    PROD = TelemetryServer(
+        "prod", TelemetryAPIEndpoint.PROD.url, TelemetryAPIEndpoint.PROD.api_key
+    )
+
+
+ENABLED_DEPLOYMENTS = (
+    TelemetryServerDeployments.DEV.name,
+    TelemetryServerDeployments.REG.name,
+    TelemetryServerDeployments.QA1.name,
+    TelemetryServerDeployments.PREPROD3.name,
+    TelemetryServerDeployments.PROD.name,
+)
+
+
+class TelemetryEvent(TelemetryEventBase):
+    """Base class for log and metric telemetry events.
+
+    This class has all of the logic except for the 'type' of the telemetry event.
+    That must be defined by the child class.
+    """
+
+    def get_type(self):
+        """Gets the telemetry event type."""
+        raise NotImplementedError
+
+    def to_dict(self):
+        """Transform this event into a dictionary."""
+        event = dict()
+        event["Name"] = self.name
+        event["Urgent"] = self.urgent
+        event["Value"] = self.value
+        event["Tags"] = self.generate_tags()
+        event.update(
+            {
+                "UUID": str(uuid.uuid4()),
+                "Created_On": datetime.datetime.now(datetime.timezone.utc)
+                .replace(tzinfo=None)
+                .strftime("%Y-%m-%d %H:%M:%S"),
+                "Type": self.get_type(),
+                "SchemaVersion": 1,
+            }
+        )
+        return event
+
+    def get_deployment(self):
+        """Gets the deployment field specified in tags if it exists."""
+        tags = self.tags
+        if tags:
+            for tag in tags:
+                if tag.get("Name", None) == "deployment":
+                    return tag.get("Value")
+
+        return "Unknown"
+
+    def generate_tags(self):
+        """Generates the tags to send as part of the telemetry event. Parts of the tags are user defined."""
+        tags = dict()
+        # Add in tags that were added to the event
+        if self.tags and len(self.tags) > 0:
+            for k, v in self.tags.items():
+                if v is not None:
+                    tags[str(k).lower()] = str(v)
+
+        telemetry = TelemetryService.get_instance()
+        # Add telemetry service generated tags
+        tags[TelemetryField.KEY_OOB_DRIVER.value] = CLIENT_NAME
+        tags[TelemetryField.KEY_OOB_VERSION.value] = str(SNOWFLAKE_CONNECTOR_VERSION)
+        tags[TelemetryField.KEY_OOB_TELEMETRY_SERVER_DEPLOYMENT.value] = (
+            telemetry.deployment.name
+        )
+        tags[TelemetryField.KEY_OOB_CONNECTION_STRING.value] = (
+            telemetry.get_connection_string()
+        )
+        if telemetry.context and len(telemetry.context) > 0:
+            for k, v in telemetry.context.items():
+                if v is not None:
+                    tags["ctx_" + str(k).lower()] = str(v)
+
+        return tags
+
+
+class TelemetryLogEvent(TelemetryEvent):
+    def get_type(self) -> str:
+        return "Log"
+
+
+class TelemetryMetricEvent(TelemetryEvent):
+    def get_type(self) -> str:
+        return "Metric"
+
+
+class TelemetryService:
+    __instance = None
+    # prevents race condition from multiple threads creating Snowflake connections
+    __lock_init = Lock()
+
+    @classmethod
+    def get_instance(cls) -> TelemetryService:
+        """Static access method."""
+        with cls.__lock_init:
+            if cls.__instance is None:
+                cls()
+        return cls.__instance
+
+    def __init__(self) -> None:
+        """Virtually private constructor."""
+        if TelemetryService.__instance is not None:
+            raise Exception("This class is a singleton!")
+        else:
+            TelemetryService.__instance = self
+        self._enabled = False
+        self._queue = Queue()
+        self.batch_size = DEFAULT_BATCH_SIZE
+        self.num_of_retry_to_trigger_telemetry = (
+            DEFAULT_NUM_OF_RETRY_TO_TRIGGER_TELEMETRY
+        )
+        self.context = dict()
+        self.connection_params = dict()
+        self.deployment = TelemetryServerDeployments.PROD
+
+    def __del__(self) -> None:
+        """Tries to flush all events left in the queue. Ignores all exceptions."""
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    @property
+    def enabled(self) -> bool:
+        """Whether the Telemetry service is enabled or not."""
+        return False
+
+    def enable(self) -> None:
+        """Enable Telemetry Service."""
+        self._enabled = False
+
+    def disable(self) -> None:
+        """Disable Telemetry Service."""
+        self._enabled = False
+
+    @property
+    def queue(self):
+        """Get the queue that holds all of the telemetry events."""
+        return self._queue
+
+    @property
+    def context(self) -> dict[str, Any]:
+        """Returns the context of the current connection."""
+        return self._context
+
+    @context.setter
+    def context(self, value) -> None:
+        """Sets the context of the current connection."""
+        self._context = value
+
+    @property
+    def connection_params(self) -> dict[str, Any]:
+        """Returns the connection parameters from the current connection."""
+        return self._connection_params
+
+    @connection_params.setter
+    def connection_params(self, value) -> None:
+        """Sets the connection parameters from the current connection."""
+        self._connection_params = value
+
+    @property
+    def batch_size(self):
+        """Returns the batch size for uploading results."""
+        return self._batch_size
+
+    @batch_size.setter
+    def batch_size(self, value) -> None:
+        """Sets the batch size for uploading results."""
+        self._batch_size = value
+
+    @property
+    def num_of_retry_to_trigger_telemetry(self) -> int:
+        """Returns the number of HTTP retries before we submit a telemetry event."""
+        return self._num_of_retry_to_trigger_telemetry
+
+    @num_of_retry_to_trigger_telemetry.setter
+    def num_of_retry_to_trigger_telemetry(self, value) -> None:
+        """Sets the number of HTTP retries before we submit a telemetry event."""
+        self._num_of_retry_to_trigger_telemetry = value
+
+    @property
+    def deployment(self: TelemetryServer) -> Any | None:
+        """Returns the deployment that we are sending the telemetry information to."""
+        return self._deployment
+
+    @deployment.setter
+    def deployment(self, value) -> None:
+        """Sets the deployment that we are sending the telemetry information to."""
+        self._deployment = value
+
+    def is_deployment_enabled(self) -> bool:
+        """Returns whether or not this deployment is enabled."""
+        return self.deployment.name in ENABLED_DEPLOYMENTS
+
+    def get_connection_string(self):
+        """Returns the URL used to connect to Snowflake."""
+        return (
+            self.connection_params.get("protocol", "")
+            + "://"
+            + self.connection_params.get("host", "")
+            + ":"
+            + str(self.connection_params.get("port", ""))
+        )
+
+    def add(self, event) -> None:
+        """Adds a telemetry event to the queue. If the event is urgent, upload all telemetry events right away."""
+        if not self.enabled:
+            return
+
+        self.queue.put(event)
+        if self.queue.qsize() > self.batch_size or event.urgent:
+            payload = self.export_queue_to_string()
+            if payload is None:
+                return
+            self._upload_payload(payload)
+
+    def flush(self) -> None:
+        """Flushes all telemetry events in the queue and submit them to the back-end."""
+        if not self.enabled:
+            return
+
+        if not self.queue.empty():
+            payload = self.export_queue_to_string()
+            if payload is None:
+                return
+            self._upload_payload(payload)
+
+    def update_context(self, connection_params) -> None:
+        """Updates the telemetry service context. Remove any passwords or credentials."""
+        self.configure_deployment(connection_params)
+        self.context = dict()
+
+        for key, value in connection_params.items():
+            if (
+                "password" not in key
+                and "passcode" not in key
+                and "privateKey" not in key
+            ):
+                self.context[key] = value
+
+    def configure_deployment(self, connection_params) -> None:
+        """Determines which deployment we are sending Telemetry OOB messages to."""
+        self.connection_params = connection_params
+        account = (
+            self.connection_params.get("account")
+            if self.connection_params.get("account")
+            else ""
+        )
+        host = (
+            self.connection_params.get("host")
+            if self.connection_params.get("host")
+            else ""
+        )
+        port = self.connection_params.get("port", None)
+
+        # Set as PROD by default
+        deployment = TelemetryServerDeployments.PROD
+        if "reg" in host or "local" in host:
+            deployment = TelemetryServerDeployments.REG
+            if port == 8080:
+                deployment = TelemetryServerDeployments.DEV
+        elif "qa1" in host or "qa1" in account:
+            deployment = TelemetryServerDeployments.QA1
+        elif "preprod3" in host:
+            deployment = TelemetryServerDeployments.PREPROD3
+
+        self.deployment = deployment
+
+    def log_ocsp_exception(
+        self,
+        event_type,
+        telemetry_data,
+        exception=None,
+        stack_trace=None,
+        tags=None,
+        urgent: bool = False,
+    ) -> None:
+        """Logs an OCSP Exception and adds it to the queue to be uploaded."""
+        if tags is None:
+            tags = dict()
+        try:
+            if self.enabled:
+                event_name = "OCSPException"
+                if exception is not None:
+                    telemetry_data[TelemetryField.KEY_OOB_EXCEPTION_MESSAGE.value] = (
+                        str(exception)
+                    )
+                if stack_trace is not None:
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_EXCEPTION_STACK_TRACE.value
+                    ] = stack_trace
+
+                if tags is None:
+                    tags = dict()
+
+                tags[TelemetryField.KEY_OOB_EVENT_TYPE.value] = event_type
+
+                log_event = TelemetryLogEvent(
+                    name=event_name, tags=tags, urgent=urgent, value=telemetry_data
+                )
+
+                self.add(log_event)
+        except Exception:
+            # Do nothing on exception, just log
+            logger.debug("Failed to log OCSP exception", exc_info=True)
+
+    def log_http_request_error(
+        self,
+        event_name,
+        url,
+        method,
+        sqlstate,
+        errno,
+        response=None,
+        retry_timeout=None,
+        retry_count=None,
+        exception=None,
+        stack_trace=None,
+        tags=None,
+        urgent: bool = False,
+    ) -> None:
+        """Logs an HTTP Request error and adds it to the queue to be uploaded."""
+        if tags is None:
+            tags = dict()
+        try:
+            if self.enabled:
+                response_status_code = -1
+                # This mimics the output of HttpRequestBase.toString() from JBDC
+                telemetry_data = generate_telemetry_data_dict(
+                    from_dict={
+                        TelemetryField.KEY_OOB_REQUEST.value: f"{method} {url}",
+                        TelemetryField.KEY_OOB_SQL_STATE.value: sqlstate,
+                        TelemetryField.KEY_OOB_ERROR_CODE.value: errno,
+                    },
+                    is_oob_telemetry=True,
+                )
+                if response:
+                    telemetry_data[TelemetryField.KEY_OOB_RESPONSE.value] = (
+                        response.json()
+                    )
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_RESPONSE_STATUS_LINE.value
+                    ] = str(response.reason)
+                    if response.status_code:
+                        response_status_code = str(response.status_code)
+                        telemetry_data[
+                            TelemetryField.KEY_OOB_RESPONSE_STATUS_CODE.value
+                        ] = response_status_code
+                if retry_timeout:
+                    telemetry_data[TelemetryField.KEY_OOB_RETRY_TIMEOUT.value] = str(
+                        retry_timeout
+                    )
+                if retry_count:
+                    telemetry_data[TelemetryField.KEY_OOB_RETRY_COUNT.value] = str(
+                        retry_count
+                    )
+                if exception:
+                    telemetry_data[TelemetryField.KEY_OOB_EXCEPTION_MESSAGE.value] = (
+                        str(exception)
+                    )
+                if stack_trace:
+                    telemetry_data[
+                        TelemetryField.KEY_OOB_EXCEPTION_STACK_TRACE.value
+                    ] = stack_trace
+
+                if tags is None:
+                    tags = dict()
+
+                tags[TelemetryField.KEY_OOB_RESPONSE_STATUS_CODE.value] = (
+                    response_status_code
+                )
+                tags[TelemetryField.KEY_OOB_SQL_STATE.value] = str(sqlstate)
+                tags[TelemetryField.KEY_OOB_ERROR_CODE.value] = errno
+
+                log_event = TelemetryLogEvent(
+                    name=event_name, tags=tags, value=telemetry_data, urgent=urgent
+                )
+
+                self.add(log_event)
+        except Exception:
+            # Do nothing on exception, just log
+            logger.debug("Failed to log HTTP request error", exc_info=True)
+
+    def log_general_exception(
+        self,
+        event_name: str,
+        telemetry_data: dict,
+        tags: dict | None = None,
+        urgent: bool | None = False,
+    ) -> None:
+        """Sends any type of exception through OOB telemetry."""
+        if tags is None:
+            tags = dict()
+        try:
+            if self.enabled:
+                log_event = TelemetryLogEvent(
+                    name=event_name, tags=tags, value=telemetry_data, urgent=urgent
+                )
+                self.add(log_event)
+        except Exception:
+            # Do nothing on exception, just log
+            logger.debug("Failed to log general exception", exc_info=True)
+
+    def _upload_payload(self, payload) -> None:
+        """Uploads the JSON-formatted string payload to the telemetry backend.
+
+        Ignore any exceptions that may arise.
+        """
+        success = True
+        response = None
+        try:
+            if not self.is_deployment_enabled():
+                logger.debug("Skip the disabled deployment: %s", self.deployment.name)
+                return
+            logger.debug(f"Sending OOB telemetry data. Payload: {payload}")
+            if ENABLE_TELEMETRY_LOG:
+                # This logger guarantees the payload won't be masked. Testing purpose.
+                rt_plain_logger.debug(f"OOB telemetry data being sent is {payload}")
+
+            # TODO(SNOW-2259522): Telemetry OOB is currently disabled. If Telemetry OOB is to be re-enabled, this HTTP call must be routed through the connection_argument.session_manager.use_session(use_pooling) (so the SessionManager instance attached to the connection which initialization's fail most likely triggered this telemetry log). It would allow to pick up proxy configuration & custom headers (see tickets SNOW-694457 and SNOW-2203079).
+            with requests.Session() as session:
+                headers = {
+                    "Content-type": "application/json",
+                    "x-api-key": self.deployment.api_key,
+                }
+                response = session.post(
+                    self.deployment.url,
+                    data=payload,
+                    headers=headers,
+                    timeout=REQUEST_TIMEOUT,
+                )
+                if (
+                    response.status_code == OK
+                    and json.loads(response.text).get("statusCode", 0) == OK
+                ):
+                    logger.debug(
+                        "telemetry server request success: %d", response.status_code
+                    )
+                else:
+                    logger.debug(
+                        "telemetry server request error: %d", response.status_code
+                    )
+                    success = False
+        except Exception as e:
+            logger.debug(
+                "Telemetry request failed, Exception response: %s, exception: %s",
+                response,
+                str(e),
+            )
+            success = False
+        finally:
+            logger.debug("Telemetry request success=%s", success)
+
+    def export_queue_to_string(self):
+        """Exports all events in the queue into a JSON formatted string with secrets masked."""
+        logs = list()
+        while not self._queue.empty():
+            logs.append(self._queue.get().to_dict())
+        # We may get an exception trying to serialize a python object to JSON
+        try:
+            payload = json.dumps(logs)
+        except Exception:
+            logger.debug(
+                "Failed to generate a JSON dump from the passed in telemetry OOB events. String representation of logs: %s"
+                % str(logs),
+                exc_info=True,
+            )
+            payload = None
+        _, masked_text, _ = SecretDetector.mask_secrets(payload)
+        return masked_text
+
+    def close(self) -> None:
+        """Closes the telemetry service."""
+        self.flush()
+        self.disable()
+
+    def size(self):
+        """Returns the size of the queue."""
+        return self.queue.qsize()

snowflake/connector/test_util.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import logging
+import os
+import platform
+from typing import cast
+
+IS_LINUX = platform.system() == "Linux"
+RUNNING_ON_JENKINS = os.getenv("JENKINS_HOME") is not None
+REGRESSION_TEST_LOG_DIR = os.getenv("CLIENT_LOG_DIR_PATH_DOCKER")
+ENABLE_TELEMETRY_LOG = RUNNING_ON_JENKINS and REGRESSION_TEST_LOG_DIR and IS_LINUX
+rt_plain_logger = None
+
+
+if ENABLE_TELEMETRY_LOG:
+    rt_plain_logger = logging.getLogger("regression.test.plain.logger")
+    rt_plain_logger.setLevel(logging.DEBUG)
+    ch = logging.FileHandler(
+        os.path.join(
+            cast(str, REGRESSION_TEST_LOG_DIR), "snowflake_ssm_rt_telemetry.log"
+        )
+    )
+    ch.setLevel(logging.DEBUG)
+    ch.setFormatter(
+        logging.Formatter(
+            "%(asctime)s - %(threadName)s %(filename)s:%(lineno)d - %(funcName)s() - %(levelname)s - %(message)s"
+        )
+    )
+    rt_plain_logger.addHandler(ch)