add execution tags by env vars (#239)

* add execution tags by env vars

Author: saartochner-lumigo

src/lumigo_tracer/auto_tag/auto_tag_event.py

@@ -4,7 +4,7 @@
 
 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
+from lumigo_tracer.lumigo_utils import get_logger, is_api_gw_event, Configuration
 
 AUTO_TAG_API_GW_HEADERS: Optional[List[str]] = (
     str_to_list(os.environ.get("LUMIGO_AUTO_TAG_API_GW_HEADERS", "")) or []
@@ -48,13 +48,25 @@ def auto_tag(event: dict):
                     add_execution_tag(key, headers[key])
 
 
+class ConfigurationHandler(EventAutoTagHandler):
+    @staticmethod
+    def is_supported(event: dict) -> bool:
+        return bool(event) and any(key in event for key in Configuration.auto_tag)
+
+    @staticmethod
+    def auto_tag(event: dict):
+        for key in Configuration.auto_tag:
+            if key in event:
+                add_execution_tag(key, event[key])
+
+
 class AutoTagEvent:
     @staticmethod
     def auto_tag_event(
         event: Optional[Dict] = None, handlers: Optional[List[EventAutoTagHandler]] = None
     ) -> None:
         if event:
-            handlers = handlers or [ApiGWHandler()]
+            handlers = handlers or [ApiGWHandler(), ConfigurationHandler()]
             for handler in handlers:
                 try:
                     if handler.is_supported(event):

src/lumigo_tracer/lumigo_utils.py

@@ -82,6 +82,7 @@
 EDGE_KINESIS_STREAM_NAME = "prod_trc-inges-edge_edge-kinesis-stream"
 STACKTRACE_LINE_TO_DROP = "lumigo_tracer/tracer.py"
 Container = TypeVar("Container", dict, list)
+DEFAULT_AUTO_TAG_KEY = "LUMIGO_AUTO_TAG"
 
 _logger: Dict[str, logging.Logger] = {}
 
@@ -139,6 +140,7 @@ class Configuration:
     edge_kinesis_aws_secret_access_key: Optional[str] = None
     should_scrub_known_services: bool = False
     is_sync_tracer: bool = False
+    auto_tag: List[str] = []
 
     @staticmethod
     def get_max_entry_size(has_error: bool = False) -> int:
@@ -162,6 +164,7 @@ def config(
     edge_kinesis_stream_name: Optional[str] = None,
     edge_kinesis_aws_access_key_id: Optional[str] = None,
     edge_kinesis_aws_secret_access_key: Optional[str] = None,
+    auto_tag: Optional[List[str]] = None,
 ) -> None:
     """
     This function configure the lumigo wrapper.
@@ -181,6 +184,7 @@ def config(
     :param edge_kinesis_stream_name: The name of the Kinesis to push the spans in China region
     :param edge_kinesis_aws_access_key_id: The credentials to push to the Kinesis in China region
     :param edge_kinesis_aws_secret_access_key: The credentials to push to the Kinesis in China region
+    :param auto_tag: The keys from the event that should be used as execution tags.
     """
 
     Configuration.token = token or os.environ.get(LUMIGO_TOKEN_KEY, "")
@@ -242,6 +246,9 @@ def config(
         os.environ.get("LUMIGO_SCRUB_KNOWN_SERVICES") == "true"
     )
     Configuration.is_sync_tracer = os.environ.get(LUMIGO_SYNC_TRACING, "FALSE").lower() == "true"
+    Configuration.auto_tag = auto_tag or os.environ.get(
+        "LUMIGO_AUTO_TAG", DEFAULT_AUTO_TAG_KEY
+    ).split(",")
 
 
 def _is_span_has_error(span: dict) -> bool:

src/test/unit/auto_tag/test_auto_tag_event.py

@@ -1,7 +1,13 @@
+import pytest
+
 from lumigo_tracer.auto_tag import auto_tag_event
-from lumigo_tracer.auto_tag.auto_tag_event import EventAutoTagHandler, AutoTagEvent
+from lumigo_tracer.auto_tag.auto_tag_event import (
+    EventAutoTagHandler,
+    AutoTagEvent,
+    ConfigurationHandler,
+)
 from lumigo_tracer.spans_container import SpansContainer
-from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY
+from lumigo_tracer.lumigo_utils import EXECUTION_TAGS_KEY, Configuration
 
 
 class ExceptionHandler(EventAutoTagHandler):
@@ -278,3 +284,41 @@ def test_auto_tag_key_in_header(monkeypatch):
 
 def set_header_key(monkeypatch, header: str):
     monkeypatch.setattr(auto_tag_event, "AUTO_TAG_API_GW_HEADERS", [header])
+
+
+@pytest.mark.parametrize(
+    "config, event, expected",
+    [
+        (["key1"], {"key1": "value1"}, True),
+        (["key1"], {"key2": "value2"}, False),
+        ([], {}, False),
+    ],
+)
+def test_configuration_handler_is_supported(config, event, expected):
+    Configuration.auto_tag = config
+    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"})
+    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
+
+
+@pytest.mark.parametrize(
+    "value, expected",
+    [
+        (1, "1"),  # int
+        ({"a": "b"}, "{'a': 'b'}"),  # dict
+    ],
+)
+def test_configuration_handler_auto_tag_non_string(value, expected):
+    Configuration.auto_tag = ["key1"]
+
+    ConfigurationHandler.auto_tag({"key1": value})
+
+    tags = SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY]
+    assert {"key": "key1", "value": expected} in tags

src/test/unit/test_lumigo_utils.py

@@ -49,6 +49,7 @@
     InternalState,
     concat_old_body_to_new,
     TRUNCATE_SUFFIX,
+    DEFAULT_AUTO_TAG_KEY,
 )
 import json
 
@@ -398,6 +399,24 @@ def test_config_step_function_without_envs(monkeypatch, configuration_value):
     assert Configuration.is_step_function == configuration_value
 
 
+def test_config_lumigo_auto_tag(monkeypatch):
+    monkeypatch.setenv("LUMIGO_AUTO_TAG", "key1,key2")
+    config()
+    assert Configuration.auto_tag == ["key1", "key2"]
+
+
+def test_config_lumigo_no_auto_tag_env(monkeypatch):
+    monkeypatch.delenv("LUMIGO_AUTO_TAG", raising=False)
+    config()
+    assert Configuration.auto_tag == [DEFAULT_AUTO_TAG_KEY]
+
+
+def test_config_lumigo_auto_tag_kwarg(monkeypatch):
+    monkeypatch.delenv("LUMIGO_AUTO_TAG", raising=False)
+    config(auto_tag=["key1", "key2"])
+    assert Configuration.auto_tag == ["key1", "key2"]
+
+
 def test_config_lumigo_domains_scrubber_with_envs(monkeypatch):
     monkeypatch.setenv("LUMIGO_DOMAINS_SCRUBBER", '["lambda.us-west-2.amazonaws.com"]')
     config()

src/test/unit/test_tracer.py

@@ -341,6 +341,21 @@ def lambda_test_function(event, context):
     ]
 
 
+def test_wrapping_with_auto_tags(context):
+    key = "my_key"
+    value = "my_value"
+
+    @lumigo_tracer(auto_tag=[key])
+    def lambda_test_function(event, context):
+        return "ret_value"
+
+    result = lambda_test_function({key: value}, context)
+    assert result == "ret_value"
+    assert SpansContainer.get_span().function_span[EXECUTION_TAGS_KEY] == [
+        {"key": key, "value": value}
+    ]
+
+
 def test_not_jsonable_return(monkeypatch, context):
     @lumigo_tracer()
     def lambda_test_function(event, context):