[COST-4571] Mechanism to delay summary for settings (#4933)

Author: myersCody

koku/api/settings/tags/mapping/utils.py

@@ -0,0 +1,61 @@
+#
+# Copyright 2024 Red Hat Inc.
+# SPDX-License-Identifier: Apache-2.0
+#
+from collections import defaultdict
+
+from django.db.models import F
+from django.db.models import Func
+
+from api.models import Provider
+from api.utils import DateHelper
+from masu.processor.tasks import delayed_summarize_current_month
+from reporting.models import AWSTagsSummary
+from reporting.models import AzureTagsSummary
+from reporting.models import GCPTagsSummary
+from reporting.models import OCITagsSummary
+from reporting.provider.ocp.models import OCPTagsValues
+from reporting.provider.ocp.models import OCPUsageReportPeriod
+
+
+def resummarize_current_month_by_tag_keys(enabled_rows, schema_name):
+    """Creates a mapping to use for resummarizing sources given tag keys
+
+    enabled_rows: List of enabled tag keys rows
+    schema_name: Str of the schema name to be used
+    start_date: datetime of when to start looking for providers that match the tag key
+    """
+    start_date = DateHelper().this_month_start
+    cloud_model_mapping = {
+        Provider.PROVIDER_AWS: AWSTagsSummary,
+        Provider.PROVIDER_AZURE: AzureTagsSummary,
+        Provider.PROVIDER_GCP: GCPTagsSummary,
+        Provider.PROVIDER_OCI: OCITagsSummary,
+    }
+    key_sorting = defaultdict(list)
+    for row in enabled_rows:
+        key_sorting[row.provider_type].append(row.key)
+
+    for provider_type, key_list in key_sorting.items():
+        if model := cloud_model_mapping.get(provider_type):
+            provider_uuids = (
+                model.objects.filter(key__in=key_list, cost_entry_bill__billing_period_start=start_date)
+                .values_list("cost_entry_bill__provider__uuid", flat=True)
+                .distinct()
+            )
+            delayed_summarize_current_month(schema_name, list(provider_uuids), provider_type)
+        elif provider_type == Provider.PROVIDER_OCP:
+            clusters = (
+                OCPTagsValues.objects.filter(key__in=key_list)
+                .annotate(clusters=Func(F("cluster_ids"), function="unnest"))
+                .values_list("clusters", flat=True)
+                .distinct()
+            )
+            provider_uuids = (
+                OCPUsageReportPeriod.objects.filter(cluster_id__in=clusters, report_period_start=start_date)
+                .values_list("provider__uuid", flat=True)
+                .distinct()
+            )
+            delayed_summarize_current_month(schema_name, list(provider_uuids), provider_type)
+
+    return provider_type

koku/api/settings/tags/mapping/view.py

@@ -19,6 +19,7 @@
 from api.settings.tags.mapping.serializers import AddChildSerializer
 from api.settings.tags.mapping.serializers import EnabledTagKeysSerializer
 from api.settings.tags.mapping.serializers import TagMappingSerializer
+from api.settings.tags.mapping.utils import resummarize_current_month_by_tag_keys
 from api.settings.utils import NonValidatedMultipleChoiceFilter
 from api.settings.utils import SettingsFilter
 from reporting.provider.all.models import EnabledTagKeys
@@ -112,9 +113,10 @@ def put(self, request):
         serializer = AddChildSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
         parent_row = EnabledTagKeys.objects.get(uuid=serializer.data.get("parent"))
-        children = EnabledTagKeys.objects.filter(uuid__in=serializer.data.get("children"))
-        tag_mappings = [TagMapping(parent=parent_row, child=child_row) for child_row in children]
+        children_rows = list(EnabledTagKeys.objects.filter(uuid__in=serializer.data.get("children")))
+        tag_mappings = [TagMapping(parent=parent_row, child=child_row) for child_row in children_rows]
         TagMapping.objects.bulk_create(tag_mappings)
+        resummarize_current_month_by_tag_keys(children_rows, request.user.customer.schema_name)
         return Response(status=status.HTTP_204_NO_CONTENT)
 
 
@@ -123,11 +125,11 @@ class SettingsTagMappingChildRemoveView(APIView):
 
     def put(self, request: Request):
         children_uuids = request.data.get("ids", [])
-        if not EnabledTagKeys.objects.filter(uuid__in=children_uuids).exists():
+        enabled_tags = EnabledTagKeys.objects.filter(uuid__in=children_uuids)
+        if not enabled_tags.exists():
             return Response({"detail": "Invalid children UUIDs."}, status=status.HTTP_400_BAD_REQUEST)
-
         TagMapping.objects.filter(child__in=children_uuids).delete()
-
+        resummarize_current_month_by_tag_keys(list(enabled_tags), request.user.customer.schema_name)
         return Response(status=status.HTTP_204_NO_CONTENT)
 
 
@@ -136,9 +138,9 @@ class SettingsTagMappingParentRemoveView(APIView):
 
     def put(self, request: Request):
         parents_uuid = request.data.get("ids", [])
-        if not EnabledTagKeys.objects.filter(uuid__in=parents_uuid).exists():
+        parent_rows = EnabledTagKeys.objects.filter(uuid__in=parents_uuid)
+        if not parent_rows.exists():
             return Response({"detail": "Invalid parents UUIDs."}, status=status.HTTP_400_BAD_REQUEST)
-
         TagMapping.objects.filter(parent__in=parents_uuid).delete()
-
+        resummarize_current_month_by_tag_keys(list(parent_rows), request.user.customer.schema_name)
         return Response({"detail": "Parents deleted successfully."}, status=status.HTTP_204_NO_CONTENT)

koku/api/settings/test/tags/mappings/test_utils.py

@@ -0,0 +1,36 @@
+from django_tenants.utils import tenant_context
+
+from api.models import Provider
+from api.settings.tags.mapping.utils import resummarize_current_month_by_tag_keys
+from masu.test import MasuTestCase
+from reporting.provider.all.models import EnabledTagKeys
+from reporting_common.models import DelayedCeleryTasks
+
+
+class TestTagMappingUtils(MasuTestCase):
+    """Test the utils for Tag mapping"""
+
+    def setUp(self):
+        super().setUp()
+        self.test_matrix = {
+            Provider.PROVIDER_AWS: self.aws_provider.uuid,
+            Provider.PROVIDER_AZURE: self.azure_provider.uuid,
+            Provider.PROVIDER_GCP: self.gcp_provider.uuid,
+            Provider.PROVIDER_OCI: self.oci_provider.uuid,
+            Provider.PROVIDER_OCP: self.ocp_provider.uuid,
+        }
+
+    def test_find_tag_key_providers(self):
+        with tenant_context(self.tenant):
+            for ptype, uuid in self.test_matrix.items():
+                with self.subTest(ptype=ptype, uuid=uuid):
+                    keys = list(EnabledTagKeys.objects.filter(provider_type=ptype))
+                    resummarize_current_month_by_tag_keys(keys, self.schema_name)
+                    self.assertTrue(DelayedCeleryTasks.objects.filter(provider_uuid=uuid).exists())
+
+    def test_multiple_returns(self):
+        with tenant_context(self.tenant):
+            keys = list(EnabledTagKeys.objects.all())
+            resummarize_current_month_by_tag_keys(keys, self.schema_name)
+            for uuid in self.test_matrix.values():
+                self.assertTrue(DelayedCeleryTasks.objects.filter(provider_uuid=uuid).exists())

koku/koku/celery.py

@@ -206,6 +206,15 @@ def validate_cron_expression(expresssion):
     "schedule": crontab(0, 0, day_of_month="15"),
 }
 
+# Specify the frequency for checking delayed summary tasks
+DELAYED_TASK_POLLING_MINUTES = ENVIRONMENT.get_value("DELAYED_TASK_POLLING_MINUTES", default="30")
+trigger_delayed_tasks_schedule = crontab(minute=f"*/{DELAYED_TASK_POLLING_MINUTES}")
+app.conf.beat_schedule["delayed_tasks_trigger"] = {
+    "task": "masu.celery.tasks.trigger_delayed_tasks",
+    "schedule": trigger_delayed_tasks_schedule,
+}
+
+
 # Celery timeout if broker is unavailable to avoid blocking indefinitely
 app.conf.broker_transport_options = {"max_retries": 4, "interval_start": 0, "interval_step": 0.5, "interval_max": 3}
 

koku/masu/celery/tasks.py

@@ -48,6 +48,7 @@
 from masu.util.oci.common import OCI_REPORT_TYPES
 from masu.util.ocp.common import OCP_REPORT_TYPES
 from reporting.models import TRINO_MANAGED_TABLES
+from reporting_common.models import DelayedCeleryTasks
 from sources.tasks import delete_source
 
 LOG = logging.getLogger(__name__)
@@ -561,3 +562,9 @@ def get_celery_queue_items(self, queue_name=None, task_name=None):
         decoded_tasks[queue] = task_list
 
     return decoded_tasks
+
+
+@celery_app.task(name="masu.celery.tasks.trigger_delayed_tasks", queue=GET_REPORT_FILES_QUEUE)
+def trigger_delayed_tasks(*args, **kwargs):
+    """Removes the expired records starting the delayed celery tasks."""
+    DelayedCeleryTasks.trigger_delayed_tasks()

koku/masu/processor/tasks.py

@@ -61,6 +61,7 @@
 from masu.util.oci.common import deduplicate_reports_for_oci
 from reporting.ingress.models import IngressReports
 from reporting_common.models import CostUsageReportStatus
+from reporting_common.models import DelayedCeleryTasks
 from reporting_common.states import ManifestState
 from reporting_common.states import ManifestStep
 
@@ -108,6 +109,31 @@
     UPDATE_SUMMARY_TABLES_QUEUE_XL,
 ]
 
+UPDATE_SUMMARY_TABLES_TASK = "masu.processor.tasks.update_summary_tables"
+
+
+def delayed_summarize_current_month(schema_name: str, provider_uuids: list, provider_type: str):
+    """Delay Resummarize provider data for the current month."""
+    queue = UPDATE_SUMMARY_TABLES_QUEUE
+    if is_customer_large(schema_name):
+        queue = UPDATE_SUMMARY_TABLES_QUEUE_XL
+
+    for provider_uuid in provider_uuids:
+        id = DelayedCeleryTasks.create_or_reset_timeout(
+            task_name=UPDATE_SUMMARY_TABLES_TASK,
+            task_args=[schema_name],
+            task_kwargs={
+                "provider_type": provider_type,
+                "provider_uuid": str(provider_uuid),
+                "start_date": str(DateHelper().this_month_start),
+            },
+            provider_uuid=provider_uuid,
+            queue_name=queue,
+        )
+        if schema_name == settings.QE_SCHEMA:
+            # bypass the wait for QE
+            id.delete()
+
 
 def record_all_manifest_files(manifest_id, report_files, tracing_id):
     """Store all report file names for manifest ID."""
@@ -427,7 +453,7 @@ def summarize_reports(  # noqa: C901
                 ).apply_async(queue=queue_name or fallback_queue)
 
 
-@celery_app.task(name="masu.processor.tasks.update_summary_tables", queue=UPDATE_SUMMARY_TABLES_QUEUE)  # noqa: C901
+@celery_app.task(name=UPDATE_SUMMARY_TABLES_TASK, queue=UPDATE_SUMMARY_TABLES_QUEUE)  # noqa: C901
 def update_summary_tables(  # noqa: C901
     schema,
     provider_type,

koku/reporting_common/test_reporting_common.py

@@ -3,12 +3,22 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 """Test Reporting Common."""
+from unittest.mock import patch
+
 from django.utils import timezone
+from django_tenants.utils import schema_context
 
+from api.models import Provider
+from api.utils import DateHelper
+from masu.processor.tasks import delayed_summarize_current_month
+from masu.processor.tasks import UPDATE_SUMMARY_TABLES_QUEUE
+from masu.processor.tasks import UPDATE_SUMMARY_TABLES_QUEUE_XL
+from masu.processor.tasks import UPDATE_SUMMARY_TABLES_TASK
 from masu.test import MasuTestCase
 from reporting_common.models import CombinedChoices
 from reporting_common.models import CostUsageReportManifest
 from reporting_common.models import CostUsageReportStatus
+from reporting_common.models import DelayedCeleryTasks
 
 
 class TestCostUsageReportStatus(MasuTestCase):
@@ -107,3 +117,42 @@ def test_set_failed_status(self):
         stats.update_status(CombinedChoices.FAILED)
         self.assertIsNotNone(stats.failed_status)
         self.assertEqual(stats.status, CombinedChoices.FAILED)
+
+    @patch("masu.processor.tasks.is_customer_large")
+    def test_delayed_summarize_current_month(self, mock_large_customer):
+        mock_large_customer.return_value = False
+        test_matrix = {
+            Provider.PROVIDER_AWS: self.aws_provider,
+            Provider.PROVIDER_AZURE: self.azure_provider,
+            Provider.PROVIDER_GCP: self.gcp_provider,
+            Provider.PROVIDER_OCI: self.oci_provider,
+            Provider.PROVIDER_OCP: self.ocp_provider,
+        }
+        count = 0
+        for test_provider_type, test_provider in test_matrix.items():
+            with self.subTest(test_provider_type=test_provider_type, test_provider=test_provider):
+                with schema_context(self.schema):
+                    delayed_summarize_current_month(self.schema_name, [test_provider.uuid], test_provider_type)
+                    count += 1
+                    self.assertEqual(DelayedCeleryTasks.objects.all().count(), count)
+                    db_entry = DelayedCeleryTasks.objects.get(provider_uuid=test_provider.uuid)
+                    self.assertEqual(db_entry.task_name, UPDATE_SUMMARY_TABLES_TASK)
+                    self.assertTrue(
+                        db_entry.task_kwargs,
+                        {
+                            "provider_type": test_provider_type,
+                            "provider_uuid": str(test_provider.uuid),
+                            "start_date": str(DateHelper().this_month_start),
+                        },
+                    )
+
+                    self.assertEqual(db_entry.task_args, [self.schema_name])
+                    self.assertEqual(db_entry.queue_name, UPDATE_SUMMARY_TABLES_QUEUE)
+
+    @patch("masu.processor.tasks.is_customer_large")
+    def test_large_customer(self, mock_large_customer):
+        mock_large_customer.return_value = True
+        delayed_summarize_current_month(self.schema_name, [self.aws_provider.uuid], Provider.PROVIDER_AWS)
+        with schema_context(self.schema):
+            db_entry = DelayedCeleryTasks.objects.get(provider_uuid=self.aws_provider.uuid)
+            self.assertEqual(db_entry.queue_name, UPDATE_SUMMARY_TABLES_QUEUE_XL)