Use database sequences for readable IDs

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

backend/.env.example

@@ -10,7 +10,7 @@ POSTGRES_USER=postgres
 POSTGRES_PASSWORD=change-me
 POSTGRES_HOST=aspsirp.com
 POSTGRES_PORT=5432
-POSTGRES_CONN_MAX_AGE=60
+POSTGRES_CONN_MAX_AGE=0
 POSTGRES_CONN_HEALTH_CHECKS=true
 
 # Redis cache settings.

backend/apps/common/management/commands/generate_perf_data.py

@@ -39,6 +39,7 @@ from apps.cases.models import (
     CaseStatus,
     CaseVerdict,
 )
+from apps.common.readable_ids import sync_readable_id_sequence
 from apps.enrichments.models import Enrichment, EnrichmentProvider, EnrichmentType
 from apps.knowledge.models import Knowledge, KnowledgeSource
 from apps.playbooks.models import Playbook, PlaybookJobStatus
@@ -372,9 +373,25 @@ class Command(BaseCommand):
             user_ids=user_ids,
             batch_size=batch_size,
         )
+        self.sync_readable_id_sequences(readable_offsets, scale)
 
         self.stdout.write(self.style.SUCCESS(f"Generated performance data for run {run_id}."))
 
+    def sync_readable_id_sequences(self, readable_offsets, scale):
+        sequence_targets = {
+            "case": scale.cases,
+            "alert": scale.alerts,
+            "artifact": scale.artifacts,
+            "enrichment": scale.enrichments,
+            "knowledge": min(scale.knowledge, scale.cases),
+            "playbook": scale.playbooks,
+        }
+        for prefix, count in sequence_targets.items():
+            if count <= 0:
+                continue
+            sync_readable_id_sequence(prefix, readable_offsets[prefix] + count - 1)
+        self.stdout.write("Synchronized readable ID sequences.")
+
     def ensure_perf_users(self, now):
         User = get_user_model()
         users = []

backend/apps/common/migrations/0001_readable_id_sequences.py

@@ -0,0 +1,52 @@
+from django.db import migrations
+
+
+READABLE_ID_TARGETS = (
+    ("case", "cases", "case_id", "readable_id_case_seq"),
+    ("alert", "alerts", "alert_id", "readable_id_alert_seq"),
+    ("artifact", "artifacts", "artifact_id", "readable_id_artifact_seq"),
+    ("enrichment", "enrichments", "enrichment_id", "readable_id_enrichment_seq"),
+    ("playbook", "playbooks", "playbook_id", "readable_id_playbook_seq"),
+    ("knowledge", "knowledge", "knowledge_id", "readable_id_knowledge_seq"),
+)
+
+
+def create_and_sync_sequences(apps, schema_editor):
+    with schema_editor.connection.cursor() as cursor:
+        for prefix, table_name, field_name, sequence_name in READABLE_ID_TARGETS:
+            cursor.execute(f"CREATE SEQUENCE IF NOT EXISTS {sequence_name}")
+            cursor.execute(
+                f"""
+                SELECT COALESCE(MAX(substring({field_name} FROM %s)::bigint), 0)
+                FROM {table_name}
+                WHERE {field_name} ~ %s
+                """,
+                [f"^{prefix}_([0-9]+)$", f"^{prefix}_[0-9]+$"],
+            )
+            max_number = cursor.fetchone()[0]
+            if max_number:
+                cursor.execute("SELECT setval(%s::regclass, %s, true)", [sequence_name, max_number])
+            else:
+                cursor.execute("SELECT setval(%s::regclass, 1, false)", [sequence_name])
+
+
+def drop_sequences(apps, schema_editor):
+    with schema_editor.connection.cursor() as cursor:
+        for _prefix, _table_name, _field_name, sequence_name in reversed(READABLE_ID_TARGETS):
+            cursor.execute(f"DROP SEQUENCE IF EXISTS {sequence_name}")
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("cases", "0002_case_case_created_id_idx"),
+        ("alerts", "0002_alert_alert_created_id_idx_and_more"),
+        ("artifacts", "0002_artifact_artifact_created_id_idx"),
+        ("enrichments", "0002_remove_mcp_provider_choice"),
+        ("playbooks", "0001_initial"),
+        ("knowledge", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.RunPython(create_and_sync_sequences, reverse_code=drop_sequences),
+    ]

backend/apps/common/readable_ids.py

@@ -1,8 +1,16 @@
-from django.db import IntegrityError, transaction
+from django.db import IntegrityError, connection, transaction
 
 
 READABLE_ID_WIDTH = 6
 READABLE_ID_RETRIES = 3
+READABLE_ID_SEQUENCES = {
+    "case": "readable_id_case_seq",
+    "alert": "readable_id_alert_seq",
+    "artifact": "readable_id_artifact_seq",
+    "enrichment": "readable_id_enrichment_seq",
+    "playbook": "readable_id_playbook_seq",
+    "knowledge": "readable_id_knowledge_seq",
+}
 
 
 def format_readable_id(prefix: str, number: int) -> str:
@@ -19,18 +27,45 @@ def parse_readable_id_number(value: str | None, prefix: str) -> int:
     return int(suffix) if suffix.isdigit() else 0
 
 
-def next_readable_id(model_class, field_name: str, prefix: str) -> str:
-    values = model_class.objects.exclude(**{field_name: ""}).values_list(field_name, flat=True)
-    max_number = 0
-    for value in values:
-        max_number = max(max_number, parse_readable_id_number(value, prefix))
-    return format_readable_id(prefix, max_number + 1)
+def readable_id_sequence_name(prefix: str) -> str:
+    try:
+        return READABLE_ID_SEQUENCES[prefix]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported readable ID prefix: {prefix}") from exc
+
+
+def next_readable_id(prefix: str) -> str:
+    sequence_name = readable_id_sequence_name(prefix)
+    with connection.cursor() as cursor:
+        cursor.execute("SELECT nextval(%s::regclass)", [sequence_name])
+        number = cursor.fetchone()[0]
+    return format_readable_id(prefix, number)
+
+
+def sync_readable_id_sequence(prefix: str, minimum_value: int) -> None:
+    if minimum_value < 1:
+        return
+    sequence_name = readable_id_sequence_name(prefix)
+    with connection.cursor() as cursor:
+        cursor.execute(
+            """
+            SELECT last_value
+            FROM pg_sequences
+            WHERE schemaname = current_schema()
+              AND sequencename = %s
+            """,
+            [sequence_name],
+        )
+        row = cursor.fetchone()
+        current_value = row[0] if row and row[0] is not None else 0
+        if current_value < minimum_value:
+            cursor.execute("SELECT setval(%s::regclass, %s, true)", [sequence_name, minimum_value])
 
 
 def assign_readable_id(instance, field_name: str, prefix: str) -> None:
     if getattr(instance, field_name):
         return
-    setattr(instance, field_name, next_readable_id(type(instance), field_name, prefix))
+    setattr(instance, field_name, next_readable_id(prefix))
 
 
 def save_with_readable_id(instance, field_name: str, prefix: str, *args, **kwargs):

backend/asp/settings.py

@@ -112,7 +112,7 @@ DATABASES = {
         "PASSWORD": os.environ.get("POSTGRES_PASSWORD", ""),
         "HOST": os.environ.get("POSTGRES_HOST", "localhost"),
         "PORT": os.environ.get("POSTGRES_PORT", "5432"),
-        "CONN_MAX_AGE": _env_int("POSTGRES_CONN_MAX_AGE", 60, minimum=0),
+        "CONN_MAX_AGE": _env_int("POSTGRES_CONN_MAX_AGE", 0, minimum=0),
         "CONN_HEALTH_CHECKS": _env_bool("POSTGRES_CONN_HEALTH_CHECKS", True),
     }
 }