From 09cd9e3b069645a87c8520a8d61f19bc1850cad3 Mon Sep 17 00:00:00 2001 From: Dmitrii Cherkasov Date: Wed, 6 Dec 2023 22:59:44 -0800 Subject: [PATCH] ODSC-38697: Implement Telemetry decorator. --- ads/telemetry/telemetry.py | 220 ++++++++++++++++-- .../default_setup/telemetry/test_telemetry.py | 88 +++++++ 2 files changed, 285 insertions(+), 23 deletions(-) create mode 100644 tests/unitary/default_setup/telemetry/test_telemetry.py diff --git a/ads/telemetry/telemetry.py b/ads/telemetry/telemetry.py index 57fa434c9..bcec8ebca 100644 --- a/ads/telemetry/telemetry.py +++ b/ads/telemetry/telemetry.py @@ -5,22 +5,117 @@ # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ import os +import re +from dataclasses import dataclass from enum import Enum, auto -from typing import Any, Dict, Optional +from functools import wraps +from typing import Any, Callable, Dict, Optional import ads.config from ads import __version__ from ads.common import logger +TELEMETRY_ARGUMENT_NAME = "telemetry" + + LIBRARY = "Oracle-ads" EXTRA_USER_AGENT_INFO = "EXTRA_USER_AGENT_INFO" USER_AGENT_KEY = "additional_user_agent" UNKNOWN = "UNKNOWN" +DELIMITER = "&" + + +def update_oci_client_config(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """ + Adds user agent information to the signer config if it is not setup yet. + + Parameters + ---------- + config: Dict + The signer configuration. + + Returns + ------- + Dict + The updated configuration. + """ + + try: + config = config or {} + if not config.get(USER_AGENT_KEY): + config.update( + { + USER_AGENT_KEY: ( + f"{LIBRARY}/version={__version__}#" + f"surface={Surface.surface().name}#" + f"api={os.environ.get(EXTRA_USER_AGENT_INFO,UNKNOWN) or UNKNOWN}" + ) + } + ) + except Exception as ex: + logger.debug(ex) + + return config + + +def telemetry( + entry_point: str = "", + name: str = "", + environ_variable: str = EXTRA_USER_AGENT_INFO, +) -> Callable: + """ + The telemetry decorator. + Injects the Telemetry object into the `kwargs` arguments of the decorated function. + This is essential for adding additional information to the telemetry from within the + decorated function. Eventually this information will be merged into the `additional_user_agent`. + + Important Note: The telemetry decorator exclusively updates the specified environment + variable and does not perform any additional actions. + " + + Parameters + ---------- + entry_point: str + The entry point of the telemetry. + Example: "plugin=project&action=run" + name: str + The name of the telemetry. + environ_variable: (str, optional). Defaults to `EXTRA_USER_AGENT_INFO`. + The name of the environment variable to capture the telemetry sequence. + + Examples + -------- + >>> @telemetry(entry_point="plugin=project&action=run", name="ads") + ... def test_function(**kwargs) + ... telemetry = kwargs.get("telemetry") + ... telemetry.add("param=hello_world") + ... print(telemetry) + + >>> test_function() + ... "ads&plugin=project&action=run¶m=hello_world" + """ + + def decorator(func: Callable) -> Callable: + @wraps(func) + def wrapper(*args, **kwargs) -> Any: + telemetry = Telemetry(name=name, environ_variable=environ_variable).begin( + entry_point + ) + try: + return func(*args, **{**kwargs, **{TELEMETRY_ARGUMENT_NAME: telemetry}}) + except: + raise + finally: + telemetry.restore() + + return wrapper + + return decorator class Surface(Enum): """ - An Enum class for labeling the surface where ADS is being used. + An Enum class used to label the surface where ADS is being utilized. """ WORKSTATION = auto() @@ -53,28 +148,107 @@ def surface(cls): return surface -def update_oci_client_config(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: - """Adds user agent information to the config if it is not setup yet. +@dataclass +class Telemetry: + """ + This class is designed to capture a telemetry sequence and store it in the specified + environment variable. By default the `EXTRA_USER_AGENT_INFO` environment variable is used. - Returns - ------- - Dict - The updated configuration. + Attributes + ---------- + name: (str, optional). Default to empty string. + The name of the telemetry. The very beginning of the telemetry string. + environ_variable: (str, optional). Defaults to `EXTRA_USER_AGENT_INFO`. + The name of the environment variable to capture the telemetry sequence. """ - try: - config = config or {} - if not config.get(USER_AGENT_KEY): - config.update( - { - USER_AGENT_KEY: ( - f"{LIBRARY}/version={__version__}#" - f"surface={Surface.surface().name}#" - f"api={os.environ.get(EXTRA_USER_AGENT_INFO,UNKNOWN) or UNKNOWN}" - ) - } - ) - except Exception as ex: - logger.debug(ex) + name: str = "" + environ_variable: str = EXTRA_USER_AGENT_INFO - return config + def __post_init__(self): + self.name = self._prepare(self.name) + self._original_value = os.environ.get(self.environ_variable) + os.environ[self.environ_variable] = "" + + def restore(self) -> "Telemetry": + """Restores the original value of the environment variable. + + Returns + ------- + self: Telemetry + An instance of the Telemetry. + """ + os.environ[self.environ_variable] = self._original_value + return self + + def clean(self) -> "Telemetry": + """Cleans the associated environment variable. + + Returns + ------- + self: Telemetry + An instance of the Telemetry. + """ + os.environ[self.environ_variable] = "" + return self + + def _begin(self): + self.clean() + os.environ[self.environ_variable] = self.name + + def begin(self, value: str = "") -> "Telemetry": + """ + This method should be invoked at the start of telemetry sequence capture. + It resets the value of the associated environment variable. + + Parameters + ---------- + value: (str, optional). Defaults to empty string. + The value that need to be added to the telemetry. + + Returns + ------- + self: Telemetry + An instance of the Telemetry. + """ + return self.clean().add(self.name).add(value) + + def add(self, value: str) -> "Telemetry": + """Appends the new value to the telemetry data. + + Parameters + ---------- + value: str + The value that need to be added to the telemetry. + + Returns + ------- + self: Telemetry + An instance of the Telemetry. + """ + if not os.environ.get(self.environ_variable): + self._begin() + + if value: + current_value = os.environ.get(self.environ_variable, "") + new_value = self._prepare(value) + + if new_value not in current_value: + os.environ[self.environ_variable] = ( + f"{current_value}{DELIMITER}{new_value}" + if current_value + else new_value + ) + return self + + def print(self) -> None: + """Prints the telemetry sequence from environment variable.""" + print(f"{self.environ_variable} = {os.environ.get(self.environ_variable)}") + + def _prepare(self, value: str): + """Replaces the special characters with the `_` in the input string.""" + return ( + re.sub("[^a-zA-Z0-9\.\-\_\&\=]", "_", re.sub(r"\s+", " ", value)) + if value + else "" + ) diff --git a/tests/unitary/default_setup/telemetry/test_telemetry.py b/tests/unitary/default_setup/telemetry/test_telemetry.py new file mode 100644 index 000000000..4f7e519d1 --- /dev/null +++ b/tests/unitary/default_setup/telemetry/test_telemetry.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python +# -*- coding: utf-8; -*- + +# Copyright (c) 2023 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ + +import os +from unittest.mock import patch + +import pytest + +from ads.telemetry import Telemetry + + +class TestTelemetry: + """Tests the Telemetry. + Class to capture telemetry sequence into the environment variable. + """ + + def setup_method(self): + self.telemetry = Telemetry(name="test.api") + + @patch.dict(os.environ, {}, clear=True) + def test_init(self): + """Ensures initializing Telemetry passes.""" + self.telemetry = Telemetry("test.api") + assert self.telemetry.name == "test.api" + assert self.telemetry.environ_variable in os.environ + assert os.environ[self.telemetry.environ_variable] == "" + + @patch.dict(os.environ, {}, clear=True) + def test_add(self): + """Tests adding the new value to the telemetry.""" + self.telemetry.begin() + self.telemetry.add("key=value").add("new_key=new_value") + assert ( + os.environ[self.telemetry.environ_variable] + == "test.api&key=value&new_key=new_value" + ) + + @patch.dict(os.environ, {}, clear=True) + def test_begin(self): + """Tests cleaning the value of the associated environment variable.""" + self.telemetry.begin("key=value") + assert os.environ[self.telemetry.environ_variable] == "test.api&key=value" + + @patch.dict(os.environ, {}, clear=True) + def test_clean(self): + """Ensures that telemetry associated environment variable can be cleaned.""" + self.telemetry.begin() + self.telemetry.add("key=value").add("new_key=new_value") + assert ( + os.environ[self.telemetry.environ_variable] + == "test.api&key=value&new_key=new_value" + ) + self.telemetry.clean() + assert os.environ[self.telemetry.environ_variable] == "" + + @patch.dict(os.environ, {"EXTRA_USER_AGENT_INFO": "some_existing_value"}, clear=True) + def test_restore(self): + """Ensures that telemetry associated environment variable can be restored to the original value.""" + telemetry = Telemetry(name="test.api") + telemetry.begin() + telemetry.add("key=value").add("new_key=new_value") + assert ( + os.environ[telemetry.environ_variable] + == "test.api&key=value&new_key=new_value" + ) + telemetry.restore() + assert os.environ[telemetry.environ_variable] == "some_existing_value" + + @pytest.mark.parametrize( + "NAME,INPUT_DATA,EXPECTED_RESULT", + [ + ("test.api", "key=va~!@#$%^*()_+lue", "key=va____________lue"), + ("test.api", "key=va lue", "key=va_lue"), + ("", "key=va123***lue", "key=va123___lue"), + ("", "", ""), + ], + ) + @patch.dict(os.environ, {}, clear=True) + def test__prepare(self, NAME, INPUT_DATA, EXPECTED_RESULT): + """Tests replacing special characters in the telemetry input value.""" + telemetry = Telemetry(name=NAME) + telemetry.begin(INPUT_DATA) + expected_result = f"{NAME}&{EXPECTED_RESULT}" if NAME else EXPECTED_RESULT + assert os.environ[telemetry.environ_variable] == expected_result +