Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions src/lumigo_tracer/auto_tag/auto_tag_event.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
from typing import Dict, List, Optional

from lumigo_tracer.user_utils import add_execution_tag
from lumigo_tracer.parsing_utils import str_to_list
from lumigo_tracer.lumigo_utils import get_logger, is_api_gw_event, Configuration
from lumigo_tracer.parsing_utils import str_to_list, safe_get
from lumigo_tracer.lumigo_utils import get_logger, is_api_gw_event, Configuration, warn_client

AUTO_TAG_API_GW_HEADERS: Optional[List[str]] = (
str_to_list(os.environ.get("LUMIGO_AUTO_TAG_API_GW_HEADERS", "")) or []
Expand Down Expand Up @@ -56,8 +56,12 @@ def is_supported(event: dict) -> bool:
@staticmethod
def auto_tag(event: dict):
for key in Configuration.auto_tag:
if key in event:
add_execution_tag(key, event[key])
try:
value = safe_get(event, key.split(".")) # type: ignore
Copy link
Contributor

@nirLumigo nirLumigo Feb 22, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The user won't be able to auto-tag keys that have "." in them. Are we ok with this?
We can use json format instead which is a bit more complicated but will allow using ".".

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes!
we can support escaping, but this is future work IMO.
WDYT?

if value:
add_execution_tag(key, value)
except Exception as err:
warn_client(f"Failed to auto tag key {key}: {err}")


class AutoTagEvent:
Expand Down
44 changes: 38 additions & 6 deletions src/test/unit/auto_tag/test_auto_tag_event.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import datetime

import pytest

from lumigo_tracer.auto_tag import auto_tag_event
Expand Down Expand Up @@ -299,13 +301,43 @@ def test_configuration_handler_is_supported(config, event, expected):
assert ConfigurationHandler.is_supported(event) == expected


def test_configuration_handler_auto_tag():
Configuration.auto_tag = ["key1", "key2", "key3"]
ConfigurationHandler.auto_tag({"key1": "value1", "key2": "value2", "other": "other"})
@pytest.mark.parametrize(
"auto_tag_keys, event, result_tags",
[
( # happy flow non-nested
["key1", "key2", "key3"],
{"key1": "value1", "key2": "value2", "other": "other"},
[{"key": "key1", "value": "value1"}, {"key": "key2", "value": "value2"}],
),
(["key1.key2"], {"key1": "value1"}, []), # not exists inner key
(["key1.key2"], {"other": "other"}, []), # not exists outer key
( # happy flow nested
["key1.key2"],
{"key1": {"key2": "value"}, "key3": "other"},
[{"key": "key1.key2", "value": "value"}],
),
( # happy flow nested multiple keys
["key1.key2", "key3.key4"],
{"key1": {"key2": "value"}, "key3": {"key4": "value2"}, "key5": "other"},
[{"key": "key1.key2", "value": "value"}, {"key": "key3.key4", "value": "value2"}],
),
],
)
def test_configuration_handler_auto_tag(auto_tag_keys, event, result_tags):
Configuration.auto_tag = auto_tag_keys
ConfigurationHandler.auto_tag(event)
tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
assert len(tags) == len(result_tags)
for tag in result_tags:
assert tag in tags


def test_configuration_handler_auto_tag_failure(capsys):
Configuration.auto_tag = [None, "key2"]
ConfigurationHandler.auto_tag({"key1": datetime, "key2": "value"})
tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
assert len(tags) == 2
assert {"key": "key1", "value": "value1"} in tags
assert {"key": "key2", "value": "value2"} in tags
assert tags == [{"key": "key2", "value": "value"}]
assert "Failed to auto tag" in capsys.readouterr().out


@pytest.mark.parametrize(
Expand Down