break Azure records out into a kafka message per hour

project-koku · Dec 12, 2023 · 3545f3d · 3545f3d
1 parent 930268e
commit 3545f3d
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 18 deletions.
diff --git a/koku/subs/subs_data_messenger.py b/koku/subs/subs_data_messenger.py
@@ -8,8 +8,10 @@
 import os
 import uuid
 from collections import defaultdict
+from datetime import timedelta
 from tempfile import mkdtemp
 
+from dateutil import parser
 from django.conf import settings
 
 from api.common import log_json
@@ -106,21 +108,41 @@ def process_and_send_subs_message(self, upload_keys):
                         instance_id = self.determine_azure_instance_id(row)
                         if not instance_id:
                             continue
-                        row["subs_resource_id"] = instance_id
-                    # row["subs_product_ids"] is a string of numbers separated by '-' to be sent as a list
-                    msg = self.build_subs_msg(
-                        row["subs_resource_id"],
-                        row["subs_account"],
-                        row["subs_start_time"],
-                        row["subs_end_time"],
-                        row["subs_vcpu"],
-                        row["subs_sla"],
-                        row["subs_usage"],
-                        row["subs_role"],
-                        row["subs_product_ids"].split("-"),
-                    )
-                    self.send_kafka_message(msg)
-                    msg_count += 1
+                        # Azure is daily records but subs need hourly records
+                        start = parser.parse(row["subs_start_time"])
+                        LOG.info(f"start\n{start}\n{type(start)}")
+                        for i in range(int(row["subs_usage_quantity"])):
+                            end = start + timedelta(hours=1)
+                            msg = self.build_subs_msg(
+                                instance_id,
+                                row["subs_account"],
+                                str(start),
+                                str(end),
+                                row["subs_vcpu"],
+                                row["subs_sla"],
+                                row["subs_usage"],
+                                row["subs_role"],
+                                row["subs_product_ids"].split("-"),
+                            )
+                            # move to the next hour in the range
+                            start = end
+                            self.send_kafka_message(msg)
+                            msg_count += 1
+                    else:
+                        # row["subs_product_ids"] is a string of numbers separated by '-' to be sent as a list
+                        msg = self.build_subs_msg(
+                            row["subs_resource_id"],
+                            row["subs_account"],
+                            row["subs_start_time"],
+                            row["subs_end_time"],
+                            row["subs_vcpu"],
+                            row["subs_sla"],
+                            row["subs_usage"],
+                            row["subs_role"],
+                            row["subs_product_ids"].split("-"),
+                        )
+                        self.send_kafka_message(msg)
+                        msg_count += 1
             LOG.info(
                 log_json(
                     self.tracing_id,
@@ -158,7 +180,7 @@ def build_subs_msg(
             "role": role,
             "sla": sla,
             "usage": usage,
-            "billing_provider": "aws",
+            "billing_provider": self.provider_type.lower(),
             "billing_account_id": billing_account_id,
         }
         return bytes(json.dumps(subs_json), "utf-8")
diff --git a/koku/subs/test/test_subs_data_messenger.py b/koku/subs/test/test_subs_data_messenger.py
@@ -226,6 +226,7 @@ def test_process_and_send_subs_message_azure_with_id(
                 "subs_usage": "Production",
                 "subs_sla": "Premium",
                 "subs_role": "Red Hat Enterprise Linux Server",
+                "subs_usage_quantity": "4",
                 "subs_product_ids": "479-70",
                 "subs_instance": "",
                 "source": self.azure_provider.uuid,
@@ -236,8 +237,8 @@ def test_process_and_send_subs_message_azure_with_id(
         with patch("builtins.open", mock_op):
             self.azure_messenger.process_and_send_subs_message(upload_keys)
         mock_azure_id.assert_called_once()
-        mock_msg_builder.assert_called_once()
-        mock_producer.assert_called_once()
+        self.assertEqual(mock_msg_builder.call_count, 4)
+        self.assertEqual(mock_producer.call_count, 4)
 
     @patch("subs.subs_data_messenger.SUBSDataMessenger.determine_azure_instance_id")
     @patch("subs.subs_data_messenger.os.remove")

diff --git a/koku/subs/trino_sql/azure_subs_pre_or_clause.sql b/koku/subs/trino_sql/azure_subs_pre_or_clause.sql
@@ -5,6 +5,7 @@ SELECT
   json_extract_scalar(lower(additionalinfo), '$.vcpus') as subs_vcpu,
   COALESCE(NULLIF(subscriptionid, ''), subscriptionguid) as subs_account,
   regexp_extract(COALESCE(NULLIF(resourceid, ''), instancename), '([^/]+$)') as subs_resource_id,
+  CAST(ceil(coalesce(nullif(quantity, 0), usagequantity)) AS INTEGER) as subs_usage_quantity,
   CASE lower(json_extract_scalar(lower(tags), '$.com_redhat_rhel_variant'))
     WHEN 'workstation' THEN 'Red Hat Enterprise Linux Workstation'
     ELSE 'Red Hat Enterprise Linux Server'
@@ -34,3 +35,5 @@ WHERE
     AND metercategory = 'Virtual Machines'
     AND json_extract_scalar(lower(additionalinfo), '$.vcpus') IS NOT NULL
     AND json_extract_scalar(lower(lower(tags)), '$.com_redhat_rhel') IS NOT NULL
+    -- ensure there is usage
+    AND ceil(coalesce(nullif(quantity, 0), usagequantity)) > 0