Support clear custom dimensions and related enhancements (#89)

* Support clear custom dimensions and related enhancements

* Adjusted some doc to make it more clear

* Adjusted some comments

* Adjust one more comment

* Update README.md

Co-authored-by: Himtanaya Bhadada <[email protected]>

* Update README.md

Co-authored-by: Stephen-Bao <[email protected]>
Co-authored-by: Mark Kuhn <[email protected]>
Co-authored-by: Himtanaya Bhadada <[email protected]>

Author: 4 people

README.md

@@ -110,9 +110,9 @@ put_dimensions({ "Operation": "Aggregator" })
 put_dimensions({ "Operation": "Aggregator", "DeviceType": "Actuator" })
 ```
 
-- **set_dimensions**(\*dimensions: Dict[str, str]) -> MetricsLogger
+- **set_dimensions**(\*dimensions: Dict[str, str], use_default: bool = False) -> MetricsLogger
 
-Explicitly override all dimensions. This will remove the default dimensions.
+Explicitly override all dimensions. By default, this will disable the default dimensions, but can be configured using the *keyword-only* parameter `use_default`.
 
 **WARNING**: Every distinct value will result in a new CloudWatch Metric.
 If the cardinality of a particular value is expected to be high, you should consider
@@ -132,6 +132,23 @@ set_dimensions(
 )
 ```
 
+```py
+set_dimensions(
+  { "Operation": "Aggregator" },
+  use_default=True  # default dimensions would be enabled
+)
+```
+
+- **reset_dimensions**(use_default: bool) -> MetricsLogger
+
+Explicitly clear all custom dimensions. The behavior of whether default dimensions should be used can be configured with the `use_default` parameter.
+
+Examples:
+
+```py
+reset_dimensions(False)  # this will clear all custom dimensions as well as disable default dimensions
+```
+
 - **set_namespace**(value: str) -> MetricsLogger
 
 Sets the CloudWatch [namespace](https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/cloudwatch_concepts.html#Namespace) that extracted metrics should be published to. If not set, a default value of aws-embedded-metrics will be used.
@@ -149,7 +166,24 @@ set_namespace("MyApplication")
 
 - **flush**()
 
-Flushes the current MetricsContext to the configured sink and resets all properties, dimensions and metric values. The namespace and default dimensions will be preserved across flushes.
+Flushes the current MetricsContext to the configured sink and resets all properties and metric values. The namespace and default dimensions will be preserved across flushes.
+Custom dimensions are **not** preserved by default, but this behavior can be changed by invoking `logger.flush_preserve_dimensions = True`, so that custom dimensions would be preserved after each flushing thereafter.
+
+Example:
+
+```py
+logger.flush()  # only default dimensions will be preserved after each flush()
+```
+
+```py
+logger.flush_preserve_dimensions = True
+logger.flush()  # custom dimensions and default dimensions will be preserved after each flush()
+```
+
+```py
+logger.reset_dimensions(False)
+logger.flush()  # default dimensions are disabled; no dimensions will be preserved after each flush()
+```
 
 ### Configuration
 

aws_embedded_metrics/logger/metrics_context.py

@@ -87,7 +87,7 @@ def put_dimensions(self, dimension_set: Dict[str, str]) -> None:
 
         self.dimensions.append(dimension_set)
 
-    def set_dimensions(self, dimension_sets: List[Dict[str, str]]) -> None:
+    def set_dimensions(self, dimension_sets: List[Dict[str, str]], use_default: bool = False) -> None:
         """
         Overwrite all dimensions.
         ```
@@ -96,7 +96,7 @@ def set_dimensions(self, dimension_sets: List[Dict[str, str]]) -> None:
             { "k1": "v1", "k2": "v2" }])
         ```
         """
-        self.should_use_default_dimensions = False
+        self.should_use_default_dimensions = use_default
 
         for dimension_set in dimension_sets:
             self.validate_dimension_set(dimension_set)
@@ -114,6 +114,16 @@ def set_default_dimensions(self, default_dimensions: Dict) -> None:
         """
         self.default_dimensions = default_dimensions
 
+    def reset_dimensions(self, use_default: bool) -> None:
+        """
+        Clear all custom dimensions on this MetricsLogger instance. Whether default dimensions should
+        be used can be configured by the input parameter.
+        :param use_default: indicates whether default dimensions should be used
+        """
+        new_dimensions: List[Dict] = []
+        self.dimensions = new_dimensions
+        self.should_use_default_dimensions = use_default
+
     def set_property(self, key: str, value: Any) -> None:
         self.properties[key] = value
 
@@ -149,7 +159,7 @@ def create_copy_with_context(self) -> "MetricsContext":
         new_properties: Dict = {}
         new_properties.update(self.properties)
 
-        # dimensions added with put_dimension will not be copied.
+        # custom dimensions will not be copied.
         # the reason for this is so that you can flush the same scope multiple
         # times without stacking new dimensions. Example:
         #
@@ -168,6 +178,16 @@ def create_copy_with_context(self) -> "MetricsContext":
             self.namespace, new_properties, new_dimensions, new_default_dimensions
         )
 
+    def create_copy_with_context_with_dimensions(self) -> "MetricsContext":
+        """
+        Creates a deep copy of the context excluding metrics.
+        Custom dimensions will be copied, this helps with the reuse of dimension sets.
+        """
+        new_context = self.create_copy_with_context()
+        new_context.dimensions.extend(self.dimensions)
+
+        return new_context
+
     @staticmethod
     def empty() -> "MetricsContext":
         return MetricsContext()

aws_embedded_metrics/logger/metrics_logger.py

@@ -29,6 +29,7 @@ def __init__(
     ):
         self.resolve_environment = resolve_environment
         self.context: MetricsContext = context or MetricsContext.empty()
+        self.flush_preserve_dimensions: bool = False
 
     async def flush(self) -> None:
         # resolve the environment and get the sink
@@ -42,7 +43,8 @@ async def flush(self) -> None:
 
         # accept and reset the context
         sink.accept(self.context)
-        self.context = self.context.create_copy_with_context()
+        self.context = self.context.create_copy_with_context() if not self.flush_preserve_dimensions \
+            else self.context.create_copy_with_context_with_dimensions()
 
     def __configure_context_for_environment(self, env: Environment) -> None:
         default_dimensions = {
@@ -63,8 +65,12 @@ def put_dimensions(self, dimensions: Dict[str, str]) -> "MetricsLogger":
         self.context.put_dimensions(dimensions)
         return self
 
-    def set_dimensions(self, *dimensions: Dict[str, str]) -> "MetricsLogger":
-        self.context.set_dimensions(list(dimensions))
+    def set_dimensions(self, *dimensions: Dict[str, str], use_default: bool = False) -> "MetricsLogger":
+        self.context.set_dimensions(list(dimensions), use_default)
+        return self
+
+    def reset_dimensions(self, use_default: bool) -> "MetricsLogger":
+        self.context.reset_dimensions(use_default)
         return self
 
     def set_namespace(self, namespace: str) -> "MetricsLogger":

tests/logger/test_metrics_logger.py

@@ -191,6 +191,50 @@ async def test_put_dimension(mocker):
     assert dimensions[0][expected_key] == expected_value
 
 
+@pytest.mark.asyncio
+async def test_reset_dimension_with_default_dimension(mocker):
+    # arrange
+    pair1 = ["key1", "val1"]
+    pair2 = ["key2", "val2"]
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    # act
+    logger.put_dimensions({pair1[0]: pair1[1]})
+    logger.reset_dimensions(True)
+    logger.put_dimensions({pair2[0]: pair2[1]})
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    dimensions = context.get_dimensions()
+    assert len(dimensions) == 1
+    assert len(dimensions[0]) == 4
+    assert dimensions[0][pair2[0]] == pair2[1]
+
+
+@pytest.mark.asyncio
+async def test_reset_dimension_without_default_dimension(mocker):
+    # arrange
+    pair1 = ["key1", "val1"]
+    pair2 = ["key2", "val2"]
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    # act
+    logger.put_dimensions({pair1[0]: pair1[1]})
+    logger.reset_dimensions(False)
+    logger.put_dimensions({pair2[0]: pair2[1]})
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    dimensions = context.get_dimensions()
+    assert len(dimensions) == 1
+    assert len(dimensions[0]) == 1
+    assert dimensions[0][pair2[0]] == pair2[1]
+
+
 @pytest.mark.asyncio
 async def test_logger_configures_default_dimensions_on_flush(before, mocker):
     # arrange
@@ -267,6 +311,32 @@ async def test_set_dimensions_overrides_all_dimensions(mocker):
     assert dimensions[expected_key] == expected_value
 
 
+@pytest.mark.asyncio
+async def test_configure_set_dimensions_to_preserve_default_dimensions(mocker):
+    # arrange
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    # setup the typical default dimensions
+    env.get_log_group_name.return_value = fake.word()
+    env.get_name.return_value = fake.word()
+    env.get_type.return_value = fake.word()
+
+    expected_key = fake.word()
+    expected_value = fake.word()
+
+    # act
+    logger.set_dimensions({expected_key: expected_value}, use_default=True)
+    await logger.flush()
+
+    # assert
+    context = get_flushed_context(sink)
+    dimension_sets = context.get_dimensions()
+    assert len(dimension_sets) == 1
+    dimensions = dimension_sets[0]
+    assert len(dimensions) == 4
+    assert dimensions[expected_key] == expected_value
+
+
 @pytest.mark.asyncio
 async def test_can_set_namespace(mocker):
     # arrange
@@ -316,6 +386,59 @@ async def test_context_is_preserved_across_flushes(mocker):
     assert context.metrics[metric_key].values == [1]
 
 
+@pytest.mark.asyncio
+async def test_flush_dont_preserve_dimensions_by_default(mocker):
+    # arrange
+    dimension_key = "Dim"
+    dimension_value = "Value"
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    logger.set_dimensions({dimension_key: dimension_value})
+
+    # act
+    await logger.flush()
+
+    context = sink.accept.call_args[0][0]
+    dimensions = context.get_dimensions()
+    assert len(dimensions) == 1
+    assert dimensions[0][dimension_key] == dimension_value
+
+    await logger.flush()
+
+    context = sink.accept.call_args[0][0]
+    dimensions = context.get_dimensions()
+    assert len(dimensions) == 1
+    assert dimension_key not in dimensions[0]
+
+
+@pytest.mark.asyncio
+async def test_configure_flush_to_preserve_dimensions(mocker):
+    # arrange
+    dimension_key = "Dim"
+    dimension_value = "Value"
+
+    logger, sink, env = get_logger_and_sink(mocker)
+
+    logger.set_dimensions({dimension_key: dimension_value})
+    logger.flush_preserve_dimensions = True
+
+    # act
+    await logger.flush()
+
+    context = sink.accept.call_args[0][0]
+    dimensions = context.get_dimensions()
+    assert len(dimensions) == 1
+    assert dimensions[0][dimension_key] == dimension_value
+
+    await logger.flush()
+
+    context = sink.accept.call_args[0][0]
+    dimensions = context.get_dimensions()
+    assert len(dimensions) == 1
+    assert dimensions[0][dimension_key] == dimension_value
+
+
 # Test helper methods