Add max attr check.

newrelic · Oct 18, 2023 · cbc21d0 · cbc21d0
1 parent c55326d
commit cbc21d0
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/newrelic/core/attribute.py b/newrelic/core/attribute.py
@@ -89,7 +89,8 @@
 
 MAX_NUM_USER_ATTRIBUTES = 128
 MAX_ATTRIBUTE_LENGTH = 255
-MAX_ML_ATTRIBUTE_LENGTH = 4096
+MAX_NUM_ML_USER_ATTRIBUTES = 64
+MAX_ML_ATTRIBUTE_LENGTH = 4095
 MAX_64_BIT_INT = 2**63 - 1
 MAX_LOG_MESSAGE_LENGTH = 32768
 

diff --git a/newrelic/core/custom_event.py b/newrelic/core/custom_event.py
@@ -18,7 +18,7 @@
 
 from newrelic.core.attribute import (check_name_is_string, check_name_length,
         process_user_attribute, NameIsNotStringException, NameTooLongException,
-        MAX_NUM_USER_ATTRIBUTES, MAX_ML_ATTRIBUTE_LENGTH)
+        MAX_NUM_USER_ATTRIBUTES, MAX_ML_ATTRIBUTE_LENGTH, MAX_NUM_ML_USER_ATTRIBUTES)
 
 
 _logger = logging.getLogger(__name__)
@@ -104,13 +104,18 @@ def create_custom_event(event_type, params, is_ml_event=False):
         for k, v in params.items():
             if is_ml_event:
                 key, value = process_user_attribute(k, v, max_length=MAX_ML_ATTRIBUTE_LENGTH)
+                if key:
+                    if len(attributes) >= MAX_NUM_ML_USER_ATTRIBUTES:
+                        _logger.debug('Maximum number of attributes already '
+                                'added to event %r. Dropping attribute: %r=%r',
+                                name, key, value)
             else:
                 key, value = process_user_attribute(k, v)
-            if key:
-                if len(attributes) >= MAX_NUM_USER_ATTRIBUTES:
-                    _logger.debug('Maximum number of attributes already '
-                            'added to event %r. Dropping attribute: %r=%r',
-                            name, key, value)
+                if key:
+                    if len(attributes) >= MAX_NUM_USER_ATTRIBUTES:
+                        _logger.debug('Maximum number of attributes already '
+                                'added to event %r. Dropping attribute: %r=%r',
+                                name, key, value)
                 else:
                     attributes[key] = value
     except Exception:

diff --git a/tests/agent_features/test_ml_events.py b/tests/agent_features/test_ml_events.py
@@ -112,6 +112,23 @@ def _test():
     _test()
 
 
+@reset_core_stats_engine()
+def test_record_ml_event_max_num_attrs():
+    too_many_attrs_event = {}
+    for i in range(65):
+        too_many_attrs_event[str(i)] = str(i)
+
+    max_attrs_event = {}
+    for i in range(64):
+        max_attrs_event[str(i)] = str(i)
+
+    @validate_ml_events([(_intrinsics, max_attrs_event)])
+    @background_task()
+    def _test():
+        record_ml_event("LabelEvent", too_many_attrs_event)
+
+    _test()
+
 @pytest.mark.parametrize(
     "params,expected",
     [