From 05e10aa15b99f6b4effd6ebbdeb26dcace633032 Mon Sep 17 00:00:00 2001
From: "ramas.grikstas" <ramas.grikstas@kilo.health>
Date: Thu, 28 Jul 2022 15:57:03 +0300
Subject: [PATCH] Added additional tags adding possibility

---
 DatadogStorage.php           | 32 ++++++++++++++++++++++++++++++++
 Tests/DatadogStorageTest.php | 20 ++++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 Tests/DatadogStorageTest.php

diff --git a/DatadogStorage.php b/DatadogStorage.php
index 82e841a..7f29344 100644
--- a/DatadogStorage.php
+++ b/DatadogStorage.php
@@ -21,6 +21,11 @@ class DatadogStorage implements StatsStorage
      */
     private $datadog;
 
+    /**
+     * @var array
+     */
+    private $additionalTags = [];
+
     public function __construct($config = 'datadog:')
     {
         if (false === class_exists(DogStatsd::class)) {
@@ -47,6 +52,8 @@ public function pushConsumerStats(ConsumerStats $stats): void
                 'consumerId' => $stats->getConsumerId(),
             ];
 
+            $tags = $this->addAdditionalTags($tags);
+
             if ($stats->getFinishedAtMs()) {
                 $values['finishedAtMs'] = $stats->getFinishedAtMs();
             }
@@ -77,6 +84,8 @@ public function pushSentMessageStats(SentMessageStats $stats): void
             $tags['command'] = $properties[Config::COMMAND];
         }
 
+        $tags = $this->addAdditionalTags($tags);
+
         $this->datadog->increment($this->config['metric.messages.sent'], 1, $tags);
     }
 
@@ -87,6 +96,8 @@ public function pushConsumedMessageStats(ConsumedMessageStats $stats): void
             'status' => $stats->getStatus(),
         ];
 
+        $tags = $this->addAdditionalTags($tags);
+
         if (ConsumedMessageStats::STATUS_FAILED === $stats->getStatus()) {
             $this->datadog->increment($this->config['metric.messages.failed'], 1, $tags);
         }
@@ -170,4 +181,25 @@ private function prepareConfig($config): array
             'metric.consumers.memoryUsage' => 'enqueue.consumers.memoryUsage',
         ], $config);
     }
+
+    public function getAdditionalTags(): array
+    {
+        return $this->additionalTags;
+    }
+
+    public function setAdditionalTags(array $additionalTags): self
+    {
+        $this->additionalTags = $additionalTags;
+
+        return $this;
+    }
+
+    private function addAdditionalTags(array $tags): array
+    {
+        foreach ($this->additionalTags as $tagKey => $tagValue) {
+            $tags[$tagKey] = $tagValue;
+        }
+
+        return $tags;
+    }
 }
diff --git a/Tests/DatadogStorageTest.php b/Tests/DatadogStorageTest.php
new file mode 100644
index 0000000..48c1074
--- /dev/null
+++ b/Tests/DatadogStorageTest.php
@@ -0,0 +1,20 @@
+<?php
+
+declare(strict_types=1);
+
+use Enqueue\Monitoring\DatadogStorage;
+use PHPUnit\Framework\TestCase;
+
+class DatadogStorageTest extends TestCase
+{
+    public function testCanAddTagsCorrect(): void
+    {
+        $additionalTags = ['environment' => 'staging'];
+
+        $storage = new DatadogStorage();
+        $storage->setAdditionalTags($additionalTags);
+
+        $storageTags = $storage->getAdditionalTags();
+        $this->assertEquals($additionalTags, $storageTags);
+    }
+}