[COST-5093] Update mapped tags in EC2 summary table (#5147)

* [COST-5093] Update mapped tags in EC2 compute table * update unit tests * code cleanup * address feedback --------- Co-authored-by: David <[email protected]>
project-koku · Jul 12, 2024 · c08f7c6 · c08f7c6
1 parent 31f7de6
commit c08f7c6
Show file tree

Hide file tree

Showing 7 changed files with 99 additions and 34 deletions.
diff --git a/koku/api/report/test/util/baker_recipes.py b/koku/api/report/test/util/baker_recipes.py
@@ -42,6 +42,20 @@ def decimal_yielder():
     _quantity=10,
 )
 
+aws_ec2_compute_summary = Recipe(
+    "AWSCostEntryLineItemSummaryByEC2Compute",
+    resource_id=cycle(f"i-000000{i}" for i in range(AWS_CONSTANTS.length - 1)),
+    instance_type=cycle(AWS_CONSTANTS["instance_types"]),
+    operating_system=cycle(AWS_CONSTANTS["operating_systems"]),
+    unit=cycle(AWS_CONSTANTS["units"]),
+    region=cycle(AWS_GEOG["regions"]),
+    memory=cycle(AWS_CONSTANTS["memory"]),
+    vcpu=cycle(AWS_CONSTANTS["vcpus"]),
+    cost_category=cycle(AWS_CONSTANTS["cost_category"]),
+    _fill_optional=True,
+    _quantity=AWS_CONSTANTS.length,
+)
+
 azure_daily_summary = Recipe(
     "AzureCostEntryLineItemDailySummary",
     service_name=cycle(AZURE_CONSTANTS["service_names"]),

diff --git a/koku/api/report/test/util/constants.py b/koku/api/report/test/util/constants.py
@@ -68,6 +68,17 @@ def __setitem__(self, item, value):
     None,
     None,
 )
+AWS_VCPUS = (2, 4, 8, 32, 48, 96, 128)
+AWS_MEMORY = ("16 GiB", "8 GiB", "32 GiB", "128 GiB", "512 GiB", "64 GiB", "1024 GiB")
+AWS_OPERATING_SYSTEMS = (
+    "Ubuntu",
+    "Red Hat Enterprise Linux",
+    "Fedora",
+    "Debian",
+    "CentOS",
+    "Oracle Linux",
+    "FreeBSD",
+)
 
 AWS_CONSTANTS = SameLengthDict(
     {
@@ -78,6 +89,9 @@ def __setitem__(self, item, value):
         "resource_ids": AWS_RESOURCE_IDS,
         "units": AWS_UNITS,
         "cost_category": AWS_COST_CATEGORIES,
+        "operating_systems": AWS_OPERATING_SYSTEMS,
+        "vcpus": AWS_VCPUS,
+        "memory": AWS_MEMORY,
     }
 )
 

diff --git a/koku/api/report/test/util/model_bakery_loader.py b/koku/api/report/test/util/model_bakery_loader.py
@@ -253,6 +253,18 @@ def load_aws_data(self, linked_openshift_provider=None, day_list=None):
                         tags=cycle(self.tags),
                         source_uuid=provider.uuid,
                     )
+
+                baker.make_recipe(
+                    "api.report.test.util.aws_ec2_compute_summary",
+                    cost_entry_bill=bill,
+                    usage_account_id=cycle(usage_account_ids),
+                    account_alias=cycle(aliases),
+                    currency_code=self.currency,
+                    usage_start=start_date,
+                    usage_end=end_date,
+                    tags=cycle(self.tags),
+                    source_uuid=provider.uuid,
+                )
         bill_ids = [bill.id for bill in bills]
         with AWSReportDBAccessor(self.schema) as accessor:
             accessor.populate_category_summary_table(bill_ids, self.first_start_date, self.last_end_date)

diff --git a/koku/masu/database/aws_report_db_accessor.py b/koku/masu/database/aws_report_db_accessor.py
@@ -349,7 +349,7 @@ def populate_markup_cost(self, provider_uuid, markup, start_date, end_date, bill
                         source_uuid=provider_uuid, source_type=Provider.PROVIDER_AWS, **date_filters
                     ).update(markup_cost=(F("unblended_cost") * markup))
 
-    def update_line_item_daily_summary_with_tag_mapping(self, start_date, end_date, bill_ids=None):
+    def update_line_item_daily_summary_with_tag_mapping(self, start_date, end_date, bill_ids=None, table_name=None):
         """
         Updates the line item daily summary table with tag mapping pieces.
 
@@ -368,14 +368,15 @@ def update_line_item_daily_summary_with_tag_mapping(self, start_date, end_date,
                 LOG.debug("No tag mappings for AWS.")
                 return
 
-        table_name = self._table_map["line_item_daily_summary"]
-        sql = pkgutil.get_data("masu.database", "sql/aws/aws_tag_mapping_update_daily_summary.sql")
+        table_name = table_name if table_name else self._table_map["line_item_daily_summary"]
+        sql = pkgutil.get_data("masu.database", "sql/aws/aws_tag_mapping_update_summary_tables.sql")
         sql = sql.decode("utf-8")
         sql_params = {
             "start_date": start_date,
             "end_date": end_date,
             "bill_ids": bill_ids,
             "schema": self.schema,
+            "table": table_name,
         }
         self._prepare_and_execute_raw_sql_query(table_name, sql, sql_params)
 

diff --git a/.../aws_tag_mapping_update_daily_summary.sql → ...aws_tag_mapping_update_summary_tables.sql b/.../aws_tag_mapping_update_daily_summary.sql → ...aws_tag_mapping_update_summary_tables.sql
@@ -37,7 +37,7 @@ cte_update_tag_keys as (
                 )
         END as update_tags
     FROM
-        {{schema | sqlsafe}}.reporting_awscostentrylineitem_daily_summary AS lids
+        {{schema | sqlsafe}}.{{table | sqlsafe}} AS lids
     CROSS JOIN (
         SELECT
             jsonb_object_agg(child_key, parent_key) AS tag_map
@@ -51,7 +51,7 @@ cte_update_tag_keys as (
         AND lids.cost_entry_bill_id in {{ bill_ids | inclause }}
         {% endif %}
 )
-UPDATE {{schema | sqlsafe}}.reporting_awscostentrylineitem_daily_summary AS lids
+UPDATE {{schema | sqlsafe}}.{{table | sqlsafe}} AS lids
 SET tags = update_data.update_tags
 FROM cte_update_tag_keys as update_data
 WHERE lids.uuid = update_data.uuid

diff --git a/koku/masu/processor/aws/aws_report_parquet_summary_updater.py b/koku/masu/processor/aws/aws_report_parquet_summary_updater.py
@@ -57,9 +57,10 @@ def update_summary_tables(self, start_date, end_date, **kwargs):
 
         """
         start_date, end_date = self._get_sql_inputs(start_date, end_date)
+        ec2_compute_summary_table = AWS_CUR_TABLE_MAP["ec2_compute_summary"]
 
         with schema_context(self._schema):
-            partition_summary_tables = UI_SUMMARY_TABLES + (AWS_CUR_TABLE_MAP["ec2_compute_summary"],)
+            partition_summary_tables = (*UI_SUMMARY_TABLES, ec2_compute_summary_table)
             self._handle_partitions(self._schema, partition_summary_tables, start_date, end_date)
 
         with CostModelDBAccessor(self._schema, self._provider.uuid) as cost_model_accessor:
@@ -118,7 +119,7 @@ def update_summary_tables(self, start_date, end_date, **kwargs):
                     self._provider.uuid,
                     month_start_date,
                     end_date,
-                    table=AWS_CUR_TABLE_MAP["ec2_compute_summary"],
+                    table=ec2_compute_summary_table,
                     filters={"source_uuid": self._provider.uuid},
                 )
 
@@ -127,6 +128,11 @@ def update_summary_tables(self, start_date, end_date, **kwargs):
                     self._provider.uuid, start_date, current_bill_id, markup_value
                 )
 
+                # Update mapped tags in EC2 compute summary table
+                accessor.update_line_item_daily_summary_with_tag_mapping(
+                    month_start_date, end_date, bill_ids, table_name=ec2_compute_summary_table
+                )
+
             for bill in bills:
                 if bill.summary_data_creation_datetime is None:
                     bill.summary_data_creation_datetime = timezone.now()

diff --git a/koku/masu/test/database/test_aws_report_db_accessor.py b/koku/masu/test/database/test_aws_report_db_accessor.py
@@ -34,6 +34,7 @@
 from reporting.provider.all.models import TagMapping
 from reporting.provider.aws.models import AWSCostEntryBill
 from reporting.provider.aws.models import AWSCostEntryLineItemDailySummary
+from reporting.provider.aws.models import AWSCostEntryLineItemSummaryByEC2Compute
 
 
 class AWSReportDBAccessorTest(MasuTestCase):
@@ -365,36 +366,53 @@ def test_delete_aws_hive_partition_by_month(self, mock_trino, mock_table_exist):
     @patch("masu.database.aws_report_db_accessor.is_feature_cost_3592_tag_mapping_enabled")
     def test_update_line_item_daily_summary_with_tag_mapping(self, mock_unleash):
         """
-        This tests the tag mapping feature.
+        Test that mapped tags are updated in aws line item summary tables.
+        After the update, the child tag's key-value data is cleared and the parent tag's data is updated accordingly.
         """
         mock_unleash.return_value = True
         populated_keys = []
-        with schema_context(self.schema):
-            enabled_tags = EnabledTagKeys.objects.filter(provider_type=Provider.PROVIDER_AWS, enabled=True)
-            for enabled_tag in enabled_tags:
-                tag_count = AWSCostEntryLineItemDailySummary.objects.filter(
-                    tags__has_key=enabled_tag.key,
-                    usage_start__gte=self.dh.this_month_start,
-                    usage_start__lte=self.dh.today,
-                ).count()
-                if tag_count > 0:
-                    key_metadata = [enabled_tag.key, enabled_tag, tag_count]
-                    populated_keys.append(key_metadata)
-                if len(populated_keys) == 2:
-                    break
-            parent_key, parent_obj, parent_count = populated_keys[0]
-            child_key, child_obj, child_count = populated_keys[1]
-            TagMapping.objects.create(parent=parent_obj, child=child_obj)
-            self.accessor.update_line_item_daily_summary_with_tag_mapping(self.dh.this_month_start, self.dh.today)
-            expected_parent_count = parent_count + child_count
-            actual_parent_count = AWSCostEntryLineItemDailySummary.objects.filter(
-                tags__has_key=parent_key, usage_start__gte=self.dh.this_month_start, usage_start__lte=self.dh.today
-            ).count()
-            self.assertEqual(expected_parent_count, actual_parent_count)
-            actual_child_count = AWSCostEntryLineItemDailySummary.objects.filter(
-                tags__has_key=child_key, usage_start__gte=self.dh.this_month_start, usage_start__lte=self.dh.today
-            ).count()
-            self.assertEqual(0, actual_child_count)
+
+        table_classes = [AWSCostEntryLineItemDailySummary, AWSCostEntryLineItemSummaryByEC2Compute]
+
+        for table_class in table_classes:
+            with self.subTest(table_class=table_class):
+                populated_keys = []
+                with schema_context(self.schema):
+                    enabled_tags = EnabledTagKeys.objects.filter(provider_type=Provider.PROVIDER_AWS, enabled=True)
+                    for enabled_tag in enabled_tags:
+                        tag_count = table_class.objects.filter(
+                            tags__has_key=enabled_tag.key,
+                            usage_start__gte=self.dh.this_month_start,
+                            usage_start__lte=self.dh.today,
+                        ).count()
+                        if tag_count > 0:
+                            key_metadata = [enabled_tag.key, enabled_tag, tag_count]
+                            populated_keys.append(key_metadata)
+                        if len(populated_keys) == 2:
+                            break
+
+                    parent_key, parent_obj, parent_count = populated_keys[0]
+                    child_key, child_obj, child_count = populated_keys[1]
+                    TagMapping.objects.create(parent=parent_obj, child=child_obj)
+                    self.accessor.update_line_item_daily_summary_with_tag_mapping(
+                        self.dh.this_month_start, self.dh.today, table_name=table_class._meta.db_table
+                    )
+                    expected_parent_count = parent_count + child_count
+                    actual_parent_count = table_class.objects.filter(
+                        tags__has_key=parent_key,
+                        usage_start__gte=self.dh.this_month_start,
+                        usage_start__lte=self.dh.today,
+                    ).count()
+                    self.assertEqual(expected_parent_count, actual_parent_count)
+                    actual_child_count = table_class.objects.filter(
+                        tags__has_key=child_key,
+                        usage_start__gte=self.dh.this_month_start,
+                        usage_start__lte=self.dh.today,
+                    ).count()
+                    self.assertEqual(0, actual_child_count)
+
+                    # Clear TagMapping objects
+                    TagMapping.objects.filter(parent=parent_obj, child=child_obj).delete()
 
     @patch("masu.database.aws_report_db_accessor.is_feature_cost_3592_tag_mapping_enabled")
     def test_populate_ocp_on_aws_tag_information(self, mock_unleash):