From f01d725b3b762ec1cc4c5a7cb23bc81cecefb456 Mon Sep 17 00:00:00 2001 From: funnywolf Date: Sat, 11 Jul 2026 18:51:10 +0800 Subject: [PATCH] Use database sequences for readable IDs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- backend/.env.example | 2 +- .../management/commands/generate_perf_data.py | 17 ++++++ .../migrations/0001_readable_id_sequences.py | 52 +++++++++++++++++++ backend/apps/common/migrations/__init__.py | 0 backend/apps/common/readable_ids.py | 51 +++++++++++++++--- backend/asp/settings.py | 2 +- 6 files changed, 114 insertions(+), 10 deletions(-) create mode 100644 backend/apps/common/migrations/0001_readable_id_sequences.py create mode 100644 backend/apps/common/migrations/__init__.py diff --git a/backend/.env.example b/backend/.env.example index 6f9d145..de2c13c 100644 --- a/backend/.env.example +++ b/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. diff --git a/backend/apps/common/management/commands/generate_perf_data.py b/backend/apps/common/management/commands/generate_perf_data.py index a7e1ef1..4e79011 100644 --- a/backend/apps/common/management/commands/generate_perf_data.py +++ b/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 = [] diff --git a/backend/apps/common/migrations/0001_readable_id_sequences.py b/backend/apps/common/migrations/0001_readable_id_sequences.py new file mode 100644 index 0000000..de0ad15 --- /dev/null +++ b/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), + ] diff --git a/backend/apps/common/migrations/__init__.py b/backend/apps/common/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/apps/common/readable_ids.py b/backend/apps/common/readable_ids.py index 2675502..bafb32e 100644 --- a/backend/apps/common/readable_ids.py +++ b/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): diff --git a/backend/asp/settings.py b/backend/asp/settings.py index 86a5dd3..aed2839 100644 --- a/backend/asp/settings.py +++ b/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), } }