Skip to content

Commit e8189c5

Browse files
mrDzurbmrDzurb
authored
ODSC-38697: Implement Telemetry decorator. (#484)
2 parents 343f5e0 + 09cd9e3 commit e8189c5

File tree

2 files changed

+285
-23
lines changed

2 files changed

+285
-23
lines changed

‎ads/telemetry/telemetry.py

Lines changed: 197 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -5,22 +5,117 @@
55
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
66

77
import os
8+
import re
9+
from dataclasses import dataclass
810
from enum import Enum, auto
9-
from typing import Any, Dict, Optional
11+
from functools import wraps
12+
from typing import Any, Callable, Dict, Optional
1013

1114
import ads.config
1215
from ads import __version__
1316
from ads.common import logger
1417

18+
TELEMETRY_ARGUMENT_NAME = "telemetry"
19+
20+
1521
LIBRARY = "Oracle-ads"
1622
EXTRA_USER_AGENT_INFO = "EXTRA_USER_AGENT_INFO"
1723
USER_AGENT_KEY = "additional_user_agent"
1824
UNKNOWN = "UNKNOWN"
25+
DELIMITER = "&"
26+
27+
28+
def update_oci_client_config(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
29+
"""
30+
Adds user agent information to the signer config if it is not setup yet.
31+
32+
Parameters
33+
----------
34+
config: Dict
35+
The signer configuration.
36+
37+
Returns
38+
-------
39+
Dict
40+
The updated configuration.
41+
"""
42+
43+
try:
44+
config = config or {}
45+
if not config.get(USER_AGENT_KEY):
46+
config.update(
47+
{
48+
USER_AGENT_KEY: (
49+
f"{LIBRARY}/version={__version__}#"
50+
f"surface={Surface.surface().name}#"
51+
f"api={os.environ.get(EXTRA_USER_AGENT_INFO,UNKNOWN) or UNKNOWN}"
52+
)
53+
}
54+
)
55+
except Exception as ex:
56+
logger.debug(ex)
57+
58+
return config
59+
60+
61+
def telemetry(
62+
entry_point: str = "",
63+
name: str = "",
64+
environ_variable: str = EXTRA_USER_AGENT_INFO,
65+
) -> Callable:
66+
"""
67+
The telemetry decorator.
68+
Injects the Telemetry object into the `kwargs` arguments of the decorated function.
69+
This is essential for adding additional information to the telemetry from within the
70+
decorated function. Eventually this information will be merged into the `additional_user_agent`.
71+
72+
Important Note: The telemetry decorator exclusively updates the specified environment
73+
variable and does not perform any additional actions.
74+
"
75+
76+
Parameters
77+
----------
78+
entry_point: str
79+
The entry point of the telemetry.
80+
Example: "plugin=project&action=run"
81+
name: str
82+
The name of the telemetry.
83+
environ_variable: (str, optional). Defaults to `EXTRA_USER_AGENT_INFO`.
84+
The name of the environment variable to capture the telemetry sequence.
85+
86+
Examples
87+
--------
88+
>>> @telemetry(entry_point="plugin=project&action=run", name="ads")
89+
... def test_function(**kwargs)
90+
... telemetry = kwargs.get("telemetry")
91+
... telemetry.add("param=hello_world")
92+
... print(telemetry)
93+
94+
>>> test_function()
95+
... "ads&plugin=project&action=run&param=hello_world"
96+
"""
97+
98+
def decorator(func: Callable) -> Callable:
99+
@wraps(func)
100+
def wrapper(*args, **kwargs) -> Any:
101+
telemetry = Telemetry(name=name, environ_variable=environ_variable).begin(
102+
entry_point
103+
)
104+
try:
105+
return func(*args, **{**kwargs, **{TELEMETRY_ARGUMENT_NAME: telemetry}})
106+
except:
107+
raise
108+
finally:
109+
telemetry.restore()
110+
111+
return wrapper
112+
113+
return decorator
19114

20115

21116
class Surface(Enum):
22117
"""
23-
An Enum class for labeling the surface where ADS is being used.
118+
An Enum class used to label the surface where ADS is being utilized.
24119
"""
25120

26121
WORKSTATION = auto()
@@ -53,28 +148,107 @@ def surface(cls):
53148
return surface
54149

55150

56-
def update_oci_client_config(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
57-
"""Adds user agent information to the config if it is not setup yet.
151+
@dataclass
152+
class Telemetry:
153+
"""
154+
This class is designed to capture a telemetry sequence and store it in the specified
155+
environment variable. By default the `EXTRA_USER_AGENT_INFO` environment variable is used.
58156
59-
Returns
60-
-------
61-
Dict
62-
The updated configuration.
157+
Attributes
158+
----------
159+
name: (str, optional). Default to empty string.
160+
The name of the telemetry. The very beginning of the telemetry string.
161+
environ_variable: (str, optional). Defaults to `EXTRA_USER_AGENT_INFO`.
162+
The name of the environment variable to capture the telemetry sequence.
63163
"""
64164

65-
try:
66-
config = config or {}
67-
if not config.get(USER_AGENT_KEY):
68-
config.update(
69-
{
70-
USER_AGENT_KEY: (
71-
f"{LIBRARY}/version={__version__}#"
72-
f"surface={Surface.surface().name}#"
73-
f"api={os.environ.get(EXTRA_USER_AGENT_INFO,UNKNOWN) or UNKNOWN}"
74-
)
75-
}
76-
)
77-
except Exception as ex:
78-
logger.debug(ex)
165+
name: str = ""
166+
environ_variable: str = EXTRA_USER_AGENT_INFO
79167

80-
return config
168+
def __post_init__(self):
169+
self.name = self._prepare(self.name)
170+
self._original_value = os.environ.get(self.environ_variable)
171+
os.environ[self.environ_variable] = ""
172+
173+
def restore(self) -> "Telemetry":
174+
"""Restores the original value of the environment variable.
175+
176+
Returns
177+
-------
178+
self: Telemetry
179+
An instance of the Telemetry.
180+
"""
181+
os.environ[self.environ_variable] = self._original_value
182+
return self
183+
184+
def clean(self) -> "Telemetry":
185+
"""Cleans the associated environment variable.
186+
187+
Returns
188+
-------
189+
self: Telemetry
190+
An instance of the Telemetry.
191+
"""
192+
os.environ[self.environ_variable] = ""
193+
return self
194+
195+
def _begin(self):
196+
self.clean()
197+
os.environ[self.environ_variable] = self.name
198+
199+
def begin(self, value: str = "") -> "Telemetry":
200+
"""
201+
This method should be invoked at the start of telemetry sequence capture.
202+
It resets the value of the associated environment variable.
203+
204+
Parameters
205+
----------
206+
value: (str, optional). Defaults to empty string.
207+
The value that need to be added to the telemetry.
208+
209+
Returns
210+
-------
211+
self: Telemetry
212+
An instance of the Telemetry.
213+
"""
214+
return self.clean().add(self.name).add(value)
215+
216+
def add(self, value: str) -> "Telemetry":
217+
"""Appends the new value to the telemetry data.
218+
219+
Parameters
220+
----------
221+
value: str
222+
The value that need to be added to the telemetry.
223+
224+
Returns
225+
-------
226+
self: Telemetry
227+
An instance of the Telemetry.
228+
"""
229+
if not os.environ.get(self.environ_variable):
230+
self._begin()
231+
232+
if value:
233+
current_value = os.environ.get(self.environ_variable, "")
234+
new_value = self._prepare(value)
235+
236+
if new_value not in current_value:
237+
os.environ[self.environ_variable] = (
238+
f"{current_value}{DELIMITER}{new_value}"
239+
if current_value
240+
else new_value
241+
)
242+
return self
243+
244+
def print(self) -> None:
245+
"""Prints the telemetry sequence from environment variable."""
246+
print(f"{self.environ_variable} = {os.environ.get(self.environ_variable)}")
247+
248+
def _prepare(self, value: str):
249+
"""Replaces the special characters with the `_` in the input string."""
250+
return (
251+
re.sub("[^a-zA-Z0-9\.\-\_\&\=]", "_", re.sub(r"\s+", " ", value))
252+
if value
253+
else ""
254+
)

‎tests/unitary/default_setup/telemetry/test_telemetry.py

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
#!/usr/bin/env python
2+
# -*- coding: utf-8; -*-
3+
4+
# Copyright (c) 2023 Oracle and/or its affiliates.
5+
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
6+
7+
import os
8+
from unittest.mock import patch
9+
10+
import pytest
11+
12+
from ads.telemetry import Telemetry
13+
14+
15+
class TestTelemetry:
16+
"""Tests the Telemetry.
17+
Class to capture telemetry sequence into the environment variable.
18+
"""
19+
20+
def setup_method(self):
21+
self.telemetry = Telemetry(name="test.api")
22+
23+
@patch.dict(os.environ, {}, clear=True)
24+
def test_init(self):
25+
"""Ensures initializing Telemetry passes."""
26+
self.telemetry = Telemetry("test.api")
27+
assert self.telemetry.name == "test.api"
28+
assert self.telemetry.environ_variable in os.environ
29+
assert os.environ[self.telemetry.environ_variable] == ""
30+
31+
@patch.dict(os.environ, {}, clear=True)
32+
def test_add(self):
33+
"""Tests adding the new value to the telemetry."""
34+
self.telemetry.begin()
35+
self.telemetry.add("key=value").add("new_key=new_value")
36+
assert (
37+
os.environ[self.telemetry.environ_variable]
38+
== "test.api&key=value&new_key=new_value"
39+
)
40+
41+
@patch.dict(os.environ, {}, clear=True)
42+
def test_begin(self):
43+
"""Tests cleaning the value of the associated environment variable."""
44+
self.telemetry.begin("key=value")
45+
assert os.environ[self.telemetry.environ_variable] == "test.api&key=value"
46+
47+
@patch.dict(os.environ, {}, clear=True)
48+
def test_clean(self):
49+
"""Ensures that telemetry associated environment variable can be cleaned."""
50+
self.telemetry.begin()
51+
self.telemetry.add("key=value").add("new_key=new_value")
52+
assert (
53+
os.environ[self.telemetry.environ_variable]
54+
== "test.api&key=value&new_key=new_value"
55+
)
56+
self.telemetry.clean()
57+
assert os.environ[self.telemetry.environ_variable] == ""
58+
59+
@patch.dict(os.environ, {"EXTRA_USER_AGENT_INFO": "some_existing_value"}, clear=True)
60+
def test_restore(self):
61+
"""Ensures that telemetry associated environment variable can be restored to the original value."""
62+
telemetry = Telemetry(name="test.api")
63+
telemetry.begin()
64+
telemetry.add("key=value").add("new_key=new_value")
65+
assert (
66+
os.environ[telemetry.environ_variable]
67+
== "test.api&key=value&new_key=new_value"
68+
)
69+
telemetry.restore()
70+
assert os.environ[telemetry.environ_variable] == "some_existing_value"
71+
72+
@pytest.mark.parametrize(
73+
"NAME,INPUT_DATA,EXPECTED_RESULT",
74+
[
75+
("test.api", "key=va~!@#$%^*()_+lue", "key=va____________lue"),
76+
("test.api", "key=va lue", "key=va_lue"),
77+
("", "key=va123***lue", "key=va123___lue"),
78+
("", "", ""),
79+
],
80+
)
81+
@patch.dict(os.environ, {}, clear=True)
82+
def test__prepare(self, NAME, INPUT_DATA, EXPECTED_RESULT):
83+
"""Tests replacing special characters in the telemetry input value."""
84+
telemetry = Telemetry(name=NAME)
85+
telemetry.begin(INPUT_DATA)
86+
expected_result = f"{NAME}&{EXPECTED_RESULT}" if NAME else EXPECTED_RESULT
87+
assert os.environ[telemetry.environ_variable] == expected_result
88+

0 commit comments

Comments
 (0)