addressing PR review except arch change

Author: abdihakim92x1

metrics/interfaces/management/commands/seed_random.py

@@ -1,6 +1,6 @@
 import random
 import time
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from datetime import date, timedelta
 from decimal import Decimal
 from typing import override
@@ -22,6 +22,11 @@
     Topic,
 )
 from metrics.data.models.core_models.timeseries import CoreTimeSeries
+from validation import enums as validation_enums
+from validation.geography_code import (
+    NATION_GEOGRAPHY_CODES,
+    UNITED_KINGDOM_GEOGRAPHY_CODE,
+)
 
 SCALE_CONFIGS = {
     "small": {"geographies": 5, "metrics": 10, "days": 30},
@@ -37,22 +42,26 @@ def add_arguments(self, parser: CommandParser) -> None:
             "--dataset",
             choices=["cms", "metrics", "both"],
             default="both",
+            help="Which dataset to seed: CMS, metrics, or both.",
         )
         parser.add_argument(
             "--scale",
             choices=["small", "medium", "large"],
             default="small",
+            help="Size of the random metrics dataset to generate.",
         )
         parser.add_argument(
             "--seed",
             type=int,
             required=False,
             default=None,
+            help="Optional random seed for reproducible metric values.",
         )
         parser.add_argument(
             "--truncate-first",
             action="store_true",
             default=False,
+            help="Clear existing metrics tables before seeding to avoid duplicates.",
         )
 
     def handle(self, *args, **options) -> None:
@@ -82,13 +91,18 @@ def handle(self, *args, **options) -> None:
 
         if should_seed_metrics:
             scale_config = SCALE_CONFIGS[scale]
+            self.stderr.write("Seeding metrics dataset...")
             counts = self._seed_metrics_data(
                 scale_config=scale_config,
                 truncate_first=truncate_first,
+                progress_callback=self.stderr.write,
             )
+            self.stderr.write("Metrics dataset seeding complete.")
 
         if should_seed_cms:
+            self.stderr.write("Building CMS site data...")
             call_command("build_cms_site")
+            self.stderr.write("CMS site build complete.")
 
         runtime_seconds = time.perf_counter() - started_at
         self._print_summary(
@@ -101,35 +115,51 @@ def handle(self, *args, **options) -> None:
 
     @classmethod
     def _seed_metrics_data(
-        cls, *, scale_config: dict[str, int], truncate_first: bool
+        cls,
+        *,
+        scale_config: dict[str, int],
+        truncate_first: bool,
+        progress_callback: Callable[[str], None] | None = None,
     ) -> dict[str, int]:
-        if truncate_first:
-            cls._truncate_metrics_data()
+        """Seed supporting metric models and time series rows for the selected scale."""
+        if progress_callback is not None:
+            progress_callback("Preparing metric taxonomy and geography records...")
 
         with transaction.atomic():
+            if truncate_first:
+                cls._truncate_metrics_data()
+
+            (
+                theme_names,
+                sub_theme_rows,
+                topic_rows,
+            ) = cls._build_theme_hierarchy_records()
             themes = cls._bulk_create(
                 Theme,
-                [Theme(name=f"Theme {index + 1}") for index in range(3)],
+                [Theme(name=name) for name in theme_names],
             )
+            themes_by_name = {theme.name: theme for theme in themes}
 
             sub_themes = cls._bulk_create(
                 SubTheme,
                 [
-                    SubTheme(
-                        name=f"SubTheme {index + 1}", theme=themes[index % len(themes)]
-                    )
-                    for index in range(6)
+                    SubTheme(name=name, theme=themes_by_name[theme_name])
+                    for name, theme_name in sub_theme_rows
                 ],
             )
+            sub_themes_by_key = {
+                (sub_theme.name, sub_theme.theme.name): sub_theme
+                for sub_theme in sub_themes
+            }
 
             topics = cls._bulk_create(
                 Topic,
                 [
                     Topic(
-                        name=f"Topic {index + 1}",
-                        sub_theme=sub_themes[index % len(sub_themes)],
+                        name=topic_name,
+                        sub_theme=sub_themes_by_key[(sub_theme_name, theme_name)],
                     )
-                    for index in range(12)
+                    for topic_name, sub_theme_name, theme_name in topic_rows
                 ],
             )
 
@@ -144,29 +174,47 @@ def _seed_metrics_data(
                 ],
             )
 
-            geography_type = GeographyType.objects.create(name="Nation")
+            geography_seed_values = cls._build_geography_seed_values(
+                count=scale_config["geographies"]
+            )
+            geography_type_names = {
+                record["geography_type"] for record in geography_seed_values
+            }
+            geography_types = cls._bulk_create(
+                GeographyType,
+                [GeographyType(name=name) for name in sorted(geography_type_names)],
+            )
+            geography_types_by_name = {
+                geography_type.name: geography_type
+                for geography_type in geography_types
+            }
 
             geographies = cls._bulk_create(
                 Geography,
                 [
                     Geography(
-                        name=f"Area {index + 1}",
-                        geography_code=f"RND{index + 1:04d}",
-                        geography_type=geography_type,
+                        name=record["name"],
+                        geography_code=record["geography_code"],
+                        geography_type=geography_types_by_name[
+                            record["geography_type"]
+                        ],
                     )
-                    for index in range(scale_config["geographies"])
+                    for record in geography_seed_values
                 ],
             )
 
             stratum = Stratum.objects.create(name="All")
             age = Age.objects.create(name="All ages")
 
+            if progress_callback is not None:
+                progress_callback("Generating Core/API time series rows...")
             core_count, api_count = cls._seed_time_series_rows(
                 metrics=metrics,
                 geographies=geographies,
                 stratum=stratum,
                 age=age,
                 days=scale_config["days"],
+                progress_callback=progress_callback,
             )
 
         return {
@@ -181,6 +229,7 @@ def _seed_metrics_data(
 
     @classmethod
     def _truncate_metrics_data(cls) -> None:
+        """Delete all seeded metrics-related rows in dependency-safe order."""
         APITimeSeries.objects.all().delete()
         CoreTimeSeries.objects.all().delete()
         Metric.objects.all().delete()
@@ -201,6 +250,7 @@ def _seed_time_series_rows(
         stratum: Stratum,
         age: Age,
         days: int,
+        progress_callback: Callable[[str], None] | None = None,
     ) -> tuple[int, int]:
         frequency = TimePeriod.Weekly.value
         today = date.today()
@@ -210,8 +260,11 @@ def _seed_time_series_rows(
         api_rows: list[APITimeSeries] = []
         core_count = 0
         api_count = 0
+        total_metrics = len(metrics)
+        total_row_count = total_metrics * len(geographies) * days
+        log_interval = max(1, total_metrics // 10) if total_metrics else 1
 
-        for metric in metrics:
+        for metric_index, metric in enumerate(metrics, start=1):
             topic = metric.topic
             sub_theme = topic.sub_theme
             theme = sub_theme.theme
@@ -281,6 +334,19 @@ def _seed_time_series_rows(
                         api_count += len(api_rows)
                         api_rows = []
 
+            if (
+                progress_callback is not None
+                and (
+                    metric_index == total_metrics
+                    or metric_index % log_interval == 0
+                )
+            ):
+                processed_row_count = metric_index * len(geographies) * days
+                progress_callback(
+                    f"Processed {metric_index}/{total_metrics} metrics "
+                    f"({processed_row_count:,}/{total_row_count:,} row groups)."
+                )
+
         if core_rows:
             CoreTimeSeries.objects.bulk_create(core_rows, batch_size=batch_size)
             core_count += len(core_rows)
@@ -289,12 +355,105 @@ def _seed_time_series_rows(
             APITimeSeries.objects.bulk_create(api_rows, batch_size=batch_size)
             api_count += len(api_rows)
 
+        if progress_callback is not None:
+            progress_callback(
+                "Inserted "
+                f"{core_count:,} CoreTimeSeries rows and "
+                f"{api_count:,} APITimeSeries rows."
+            )
+
         return core_count, api_count
 
     @staticmethod
     def _bulk_create(model, records: Iterable):
+        """Materialise and bulk insert a sequence of model instances."""
         return model.objects.bulk_create(list(records))
 
+    @classmethod
+    def _build_theme_hierarchy_records(
+        cls,
+    ) -> tuple[list[str], list[tuple[str, str]], list[tuple[str, str, str]]]:
+        child_to_parent: dict[str, str] = {}
+        normalised_to_child: dict[str, str] = {}
+        parent_by_name = validation_enums.ParentTheme.__members__
+
+        for child_theme_group in validation_enums.ChildTheme:
+            resolved_parent = (
+                parent_by_name[child_theme_group.name].value
+                if child_theme_group.name in parent_by_name
+                else validation_enums.ParentTheme.INFECTIOUS_DISEASE.value
+            )
+            for sub_theme_name in child_theme_group.return_list():
+                child_to_parent[sub_theme_name] = resolved_parent
+                normalised_to_child[cls._normalise_key(sub_theme_name)] = (
+                    sub_theme_name
+                )
+
+        topic_rows: list[tuple[str, str, str]] = []
+        sub_theme_pairs: set[tuple[str, str]] = set()
+        for topic_group in validation_enums.Topic:
+            normalised_topic_group = cls._normalise_key(topic_group.name)
+            sub_theme_name = normalised_to_child.get(normalised_topic_group)
+            if sub_theme_name is None:
+                continue
+
+            parent_theme_name = child_to_parent[sub_theme_name]
+            sub_theme_pairs.add((sub_theme_name, parent_theme_name))
+            for topic_value in topic_group.return_list():
+                topic_rows.append((topic_value, sub_theme_name, parent_theme_name))
+
+        theme_names = sorted({parent_name for _, parent_name in sub_theme_pairs})
+        sub_theme_rows = sorted(
+            sub_theme_pairs,
+            key=lambda value: (value[1], value[0]),
+        )
+        return theme_names, sub_theme_rows, topic_rows
+
+    @classmethod
+    def _build_geography_seed_values(cls, *, count: int) -> list[dict[str, str]]:
+        geographies: list[dict[str, str]] = [
+            {
+                "name": "United Kingdom",
+                "geography_code": UNITED_KINGDOM_GEOGRAPHY_CODE,
+                "geography_type": (
+                    validation_enums.GeographyType.UNITED_KINGDOM.value
+                ),
+            }
+        ]
+
+        geographies.extend(
+            {
+                "name": name,
+                "geography_code": code,
+                "geography_type": validation_enums.GeographyType.NATION.value,
+            }
+            for name, code in NATION_GEOGRAPHY_CODES.items()
+        )
+
+        if len(geographies) >= count:
+            return geographies[:count]
+
+        extra_required = count - len(geographies)
+        geographies.extend(
+            {
+                "name": cls._format_enum_name(ltla.name),
+                "geography_code": ltla.value,
+                "geography_type": (
+                    validation_enums.GeographyType.LOWER_TIER_LOCAL_AUTHORITY.value
+                ),
+            }
+            for ltla in list(validation_enums.LTLAs)[:extra_required]
+        )
+        return geographies[:count]
+
+    @staticmethod
+    def _normalise_key(value: str) -> str:
+        return value.lower().replace("-", "_")
+
+    @staticmethod
+    def _format_enum_name(value: str) -> str:
+        return value.replace("_", " ").title()
+
     def _print_summary(
         self,
         *,