From e678206ded8c62d9d247adb1cef8415262d02c3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Fern=C3=A1ndez=20Poyatos?=
 <victor@prowler.com>
Date: Tue, 17 Jun 2025 14:58:38 +0200
Subject: [PATCH] feat(migrations): add tools for partitioned indexes

---
 api/src/backend/api/db_utils.py   |   3 +
 api/src/backend/api/operations.py | 140 ++++++++++++++++++++++++++++++
 2 files changed, 143 insertions(+)
 create mode 100644 api/src/backend/api/operations.py

diff --git a/api/src/backend/api/db_utils.py b/api/src/backend/api/db_utils.py
index d163a02fd3..754c6953d3 100644
--- a/api/src/backend/api/db_utils.py
+++ b/api/src/backend/api/db_utils.py
@@ -250,6 +250,7 @@ def register_enum(apps, schema_editor, enum_class):  # noqa: F841
         register_adapter(enum_class, enum_adapter)
 
 
+# DEPRECATED
 def _should_create_index_on_partition(
     partition_name: str, all_partitions: bool = False
 ) -> bool:
@@ -316,6 +317,7 @@ def _should_create_index_on_partition(
         return True
 
 
+# DEPRECATED
 def create_index_on_partitions(
     apps,  # noqa: F841
     schema_editor,
@@ -381,6 +383,7 @@ def create_index_on_partitions(
             schema_editor.execute(sql)
 
 
+# DEPRECATED
 def drop_index_on_partitions(
     apps,  # noqa: F841
     schema_editor,
diff --git a/api/src/backend/api/operations.py b/api/src/backend/api/operations.py
new file mode 100644
index 0000000000..ba3e0e4e7d
--- /dev/null
+++ b/api/src/backend/api/operations.py
@@ -0,0 +1,140 @@
+import re
+from datetime import datetime, timezone
+
+from django.db import connection
+from django.db.migrations.operations.base import Operation
+
+
+class CreatePartitionedIndex(Operation):
+    reversible = True
+
+    def __init__(
+        self,
+        parent_table: str,
+        index_name: str,
+        columns: str,
+        method: str = "BTREE",
+        where: str = "",
+        all_partitions: bool = False,
+        create_parent_index: bool = True,
+    ):
+        self.parent_table = parent_table
+        self.index_name = index_name
+        self.columns = columns
+        self.method = method
+        self.where = where
+        self.all_partitions = all_partitions
+        self.create_parent_index = create_parent_index
+
+    def state_forwards(self, app_label, state):
+        pass  # No state change
+
+    def database_forwards(self, app_label, schema_editor, from_state, to_state):
+        parent_index_name = f"{self.index_name}"
+        where_sql = f" WHERE {self.where}" if self.where else ""
+
+        with connection.cursor() as cursor:
+            cursor.execute(
+                """
+                SELECT inhrelid::regclass::text
+                FROM pg_inherits
+                WHERE inhparent = %s::regclass
+                """,
+                [self.parent_table],
+            )
+            partitions = [row[0] for row in cursor.fetchall()]
+
+        if self.create_parent_index:
+            sql = (
+                f"CREATE INDEX IF NOT EXISTS {parent_index_name} "
+                f"ON ONLY {self.parent_table} USING {self.method} ({self.columns})"
+                f"{where_sql};"
+            )
+            schema_editor.execute(sql)
+
+        for partition in partitions:
+            if self._should_create_index_on_partition(partition, self.all_partitions):
+                child_index_name = f"{partition.replace('.', '_')}_{self.index_name}"
+                create_sql = (
+                    f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {child_index_name} "
+                    f"ON {partition} USING {self.method} ({self.columns})"
+                    f"{where_sql};"
+                )
+                schema_editor.execute(create_sql)
+
+                if self.create_parent_index:
+                    attach_sql = (
+                        f"ALTER INDEX {parent_index_name} "
+                        f"ATTACH PARTITION {child_index_name};"
+                    )
+                    try:
+                        schema_editor.execute(attach_sql)
+                    except Exception as e:
+                        print(
+                            f"Warning: Could not attach index {child_index_name}: {e}"
+                        )
+
+    def database_backwards(self, app_label, schema_editor, from_state, to_state):
+        if self.create_parent_index:
+            parent_index_name = self.index_name
+            drop_parent_sql = f"DROP INDEX IF EXISTS {parent_index_name};"
+            schema_editor.execute(drop_parent_sql)
+
+        with connection.cursor() as cursor:
+            cursor.execute(
+                """
+                SELECT inhrelid::regclass::text
+                FROM pg_inherits
+                WHERE inhparent = %s::regclass
+                """,
+                [self.parent_table],
+            )
+            partitions = [row[0] for row in cursor.fetchall()]
+
+        for partition in partitions:
+            idx_name = f"{partition.replace('.', '_')}_{self.index_name}"
+            drop_sql = f"DROP INDEX CONCURRENTLY IF EXISTS {idx_name};"
+            schema_editor.execute(drop_sql)
+
+    def describe(self):
+        return f"Create partitioned index {self.index_name} on {self.parent_table}"
+
+    def _should_create_index_on_partition(
+        self, partition_name: str, all_partitions: bool
+    ) -> bool:
+        if all_partitions:
+            return True
+
+        date_pattern = r"(\d{4})_([a-z]{3})$"
+        match = re.search(date_pattern, partition_name)
+        if not match:
+            return True
+
+        try:
+            year_str, month_abbr = match.groups()
+            year = int(year_str)
+            month_map = {
+                "jan": 1,
+                "feb": 2,
+                "mar": 3,
+                "apr": 4,
+                "may": 5,
+                "jun": 6,
+                "jul": 7,
+                "aug": 8,
+                "sep": 9,
+                "oct": 10,
+                "nov": 11,
+                "dec": 12,
+            }
+            month = month_map.get(month_abbr.lower())
+            if month is None:
+                return True
+
+            partition_date = datetime(year, month, 1, tzinfo=timezone.utc)
+            current_month_start = datetime.now(timezone.utc).replace(
+                day=1, hour=0, minute=0, second=0, microsecond=0
+            )
+            return partition_date >= current_month_start
+        except Exception:
+            return True