add counter metrics

src/firehose_pyio/boto3_client.py

@@ -132,6 +132,10 @@ def put_record_batch(
         except Exception as e:
             raise FirehoseClientError(str(e), get_http_error_code(e))
 
+    def close(self):
+        """Closes underlying endpoint connections."""
+        self.client.close()
+
 
 class FakeFirehoseClient:
     def __init__(self, fake_config: dict):
@@ -155,3 +159,6 @@ def put_record_batch(
             "Encrypted": False,
             "RequestResponses": request_responses,
         }
+
+    def close(self):
+        return

src/firehose_pyio/io.py

@@ -17,8 +17,10 @@
 
 import typing
 import apache_beam as beam
+from apache_beam import metrics
 from apache_beam.pvalue import PCollection
 
+
 from firehose_pyio.boto3_client import FirehoseClient, FakeFirehoseClient
 from firehose_pyio.options import FirehoseOptions
 
@@ -36,6 +38,16 @@ class _FirehoseWriteFn(beam.DoFn):
         fake_config (dict, optional): Config parameters when using FakeFirehoseClient for testing. Defaults to None.
     """
 
+    total_elements_count = metrics.Metrics.counter(
+        "_FirehoseWriteFn", "total_elements_count"
+    )
+    succeeded_elements_count = metrics.Metrics.counter(
+        "_FirehoseWriteFn", "succeeded_elements_count"
+    )
+    failed_elements_count = metrics.Metrics.counter(
+        "_FirehoseWriteFn", "failed_elements_count"
+    )
+
     def __init__(
         self,
         delivery_stream_name: str,
@@ -72,7 +84,7 @@ def start_bundle(self):
     def process(self, element):
         if isinstance(element, tuple):
             element = element[1]
-        loop, failed = 0, []
+        loop, total, failed = 0, len(element), []
         while loop < self.max_retry:
             responses = self.client.put_record_batch(
                 element, self.delivery_stream_name, self.jsonify
@@ -85,10 +97,13 @@ def process(self, element):
             element = failed
             failed = []
             loop += 1
+        self.total_elements_count.inc(total)
+        self.succeeded_elements_count.inc(total - len(failed))
+        self.failed_elements_count.inc(len(failed))
         return failed
 
     def finish_bundle(self):
-        pass
+        self.client.close()
 
 
 class WriteToFirehose(beam.PTransform):

tests/test_io.py

@@ -20,14 +20,15 @@
 from moto import mock_aws
 from moto.core import DEFAULT_ACCOUNT_ID as ACCOUNT_ID
 import apache_beam as beam
+from apache_beam.metrics.metric import MetricsFilter
 from apache_beam.transforms.util import BatchElements
 from apache_beam import GroupIntoBatches
 from apache_beam.options import pipeline_options
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that, equal_to
 
 from firehose_pyio.boto3_client import FirehoseClient, FirehoseClientError
-from firehose_pyio.io import WriteToFirehose
+from firehose_pyio.io import WriteToFirehose, _FirehoseWriteFn
 
 
 def sample_s3_dest_config(bucket_name):
@@ -208,3 +209,83 @@ def test_write_to_firehose_retry_with_failed_elements(self):
                 )
             )
             assert_that(output, equal_to(["four"]))
+
+
+class TestMetrics(unittest.TestCase):
+    def test_metrics_with_no_failed_element(self):
+        pipeline = TestPipeline()
+        output = (
+            pipeline
+            | beam.Create(["one", "two", "three", "four"])
+            | BatchElements(min_batch_size=4)
+            | WriteToFirehose("non-existing-delivery-stream", False, 3, {"num_keep": 2})
+        )
+        assert_that(output, equal_to([]))
+
+        res = pipeline.run()
+        res.wait_until_finish()
+
+        ## verify total_elements_count
+        metric_results = res.metrics().query(
+            MetricsFilter().with_metric(_FirehoseWriteFn.total_elements_count)
+        )
+        total_elements_count = metric_results["counters"][0]
+        self.assertEqual(total_elements_count.key.metric.name, "total_elements_count")
+        self.assertEqual(total_elements_count.committed, 4)
+
+        ## verify succeeded_elements_count
+        metric_results = res.metrics().query(
+            MetricsFilter().with_metric(_FirehoseWriteFn.succeeded_elements_count)
+        )
+        succeeded_elements_count = metric_results["counters"][0]
+        self.assertEqual(
+            succeeded_elements_count.key.metric.name, "succeeded_elements_count"
+        )
+        self.assertEqual(succeeded_elements_count.committed, 4)
+
+        ## verify failed_elements_count
+        metric_results = res.metrics().query(
+            MetricsFilter().with_metric(_FirehoseWriteFn.failed_elements_count)
+        )
+        failed_elements_count = metric_results["counters"][0]
+        self.assertEqual(failed_elements_count.key.metric.name, "failed_elements_count")
+        self.assertEqual(failed_elements_count.committed, 0)
+
+    def test_metrics_with_failed_element(self):
+        pipeline = TestPipeline()
+        output = (
+            pipeline
+            | beam.Create(["one", "two", "three", "four"])
+            | BatchElements(min_batch_size=4)
+            | WriteToFirehose("non-existing-delivery-stream", False, 3, {"num_keep": 1})
+        )
+        assert_that(output, equal_to(["four"]))
+
+        res = pipeline.run()
+        res.wait_until_finish()
+
+        ## verify total_elements_count
+        metric_results = res.metrics().query(
+            MetricsFilter().with_metric(_FirehoseWriteFn.total_elements_count)
+        )
+        total_elements_count = metric_results["counters"][0]
+        self.assertEqual(total_elements_count.key.metric.name, "total_elements_count")
+        self.assertEqual(total_elements_count.committed, 4)
+
+        ## verify succeeded_elements_count
+        metric_results = res.metrics().query(
+            MetricsFilter().with_metric(_FirehoseWriteFn.succeeded_elements_count)
+        )
+        succeeded_elements_count = metric_results["counters"][0]
+        self.assertEqual(
+            succeeded_elements_count.key.metric.name, "succeeded_elements_count"
+        )
+        self.assertEqual(succeeded_elements_count.committed, 3)
+
+        ## verify failed_elements_count
+        metric_results = res.metrics().query(
+            MetricsFilter().with_metric(_FirehoseWriteFn.failed_elements_count)
+        )
+        failed_elements_count = metric_results["counters"][0]
+        self.assertEqual(failed_elements_count.key.metric.name, "failed_elements_count")
+        self.assertEqual(failed_elements_count.committed, 1)