Adding disk_size.

project-koku · Sep 9, 2024 · 264072a · 264072a
1 parent b88134a
commit 264072a
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 0 deletions.
diff --git a/nise/generators/gcp/compute_engine_generator.py b/nise/generators/gcp/compute_engine_generator.py
@@ -58,6 +58,8 @@ class ComputeEngineGenerator(GCPGenerator):
 
     LABELS = (([{"key": "vm_key_proj2", "value": "vm_label_proj2"}]), ([]))
 
+    EFFECTIVE_PRICE = ("0.0114", "0.0417", "0.1922", "0.0452")
+
     def __init__(self, start_date, end_date, currency, project, attributes=None):  # noqa: C901
         """Initialize the cloud storage generator."""
         super().__init__(start_date, end_date, currency, project, attributes)
@@ -123,6 +125,9 @@ def _update_data(self, row):  # noqa: C901
         row["currency"] = self._currency
         row["labels"] = self.determine_labels(self.LABELS)
         row["system_labels"] = self.determine_system_labels(sku[3])
+        row["price.effective_price"] = choice(self.EFFECTIVE_PRICE)
+        row["disk_size"] = self._generate_disk_size(row["cost"], row["price.effective_price"])
+
         if self.resource_level:
             resource = self._generate_resource(
                 self._resource_name, self._resource_global_name, self.project.get("region")
@@ -183,6 +188,10 @@ def _update_data(self, row):  # noqa: C901
         month = datetime.strptime(row.get("usage_start_time")[:7], "%Y-%m").month
         invoice["month"] = f"{year}{month:02d}"
         row["invoice"] = invoice
+        effective_price = choice(self.EFFECTIVE_PRICE)
+        row["price"] = {"effective_price": effective_price}
+        row["disk_size"] = self._generate_disk_size(row["cost"], row["price"]["effective_price"])
+
         if self.resource_level:
             resource = self._generate_resource()
             row["resource"] = resource

diff --git a/nise/generators/gcp/gcp_generator.py b/nise/generators/gcp/gcp_generator.py
@@ -57,6 +57,7 @@
     "cost_type",
     "partition_date",
     "price.effective_price",
+    "disk_size",
 )
 
 GCP_RESOURCE_COLUMNS = ("resource.name", "resource.global_name")
@@ -320,3 +321,14 @@ def _generate_hourly_data(self, **kwargs):
             row = self._init_data_row(start, end)
             row = self._update_data(row)
             yield row
+
+    def _generate_disk_size(self, cost, effective_price):
+        """Generate disk size."""
+        hours_in_month = 744
+        disk_size = 0
+
+        if cost and effective_price:
+            rate = float(effective_price) / hours_in_month
+            disk_size = float(cost) / rate
+
+        return disk_size
diff --git a/nise/upload.py b/nise/upload.py
@@ -269,6 +269,13 @@ def gcp_bucket_to_dataset(gcp_bucket_name, file_name, dataset_name, table_name,
                 "mode": "NULLABLE",
             },
             {"name": "cost_type", "type": "STRING", "mode": "NULLABLE"},
+            {
+                "name": "price",
+                "type": "RECORD",
+                "fields": [{"name": "effective_price", "type": "STRING", "mode": "NULLABLE"}],
+                "mode": "NULLABLE",
+            },
+            {"name": "disk_size", "type": "STRING", "mode": "NULLABLE"},
             {
                 "name": "adjustment_info",
                 "type": "RECORD",