From 09abe5f513db6b98e2679cc9b7b0d81ca25f802e Mon Sep 17 00:00:00 2001 From: buua436 Date: Mon, 13 Jul 2026 17:34:28 +0800 Subject: [PATCH] fix: ensure database model indexes (#16860) --- api/db/db_models.py | 107 ++++++++++++++++---------------------------- 1 file changed, 38 insertions(+), 69 deletions(-) diff --git a/api/db/db_models.py b/api/db/db_models.py index b8bff59dd4..dcfb4b1e61 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -43,7 +43,6 @@ from peewee import ( Metadata, Model, TextField, - PrimaryKeyField, ) from playhouse.migrate import MySQLMigrator, PostgresqlMigrator, migrate from playhouse.pool import PooledMySQLDatabase, PooledPostgresqlDatabase @@ -1332,74 +1331,6 @@ class SyncLogs(DataBaseModel): db_table = "sync_logs" -class EvaluationDataset(DataBaseModel): - """Ground truth dataset for RAG evaluation""" - - id = CharField(max_length=32, primary_key=True) - tenant_id = CharField(max_length=32, null=False, index=True, help_text="tenant ID") - name = CharField(max_length=255, null=False, index=True, help_text="dataset name") - description = TextField(null=True, help_text="dataset description") - kb_ids = JSONField(null=False, help_text="knowledge base IDs to evaluate against") - created_by = CharField(max_length=32, null=False, index=True, help_text="creator user ID") - create_time = BigIntegerField(null=False, index=True, help_text="creation timestamp") - update_time = BigIntegerField(null=False, help_text="last update timestamp") - status = IntegerField(null=False, default=1, help_text="1=valid, 0=invalid") - - class Meta: - db_table = "evaluation_datasets" - - -class EvaluationCase(DataBaseModel): - """Individual test case in an evaluation dataset""" - - id = CharField(max_length=32, primary_key=True) - dataset_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_datasets") - question = TextField(null=False, help_text="test question") - reference_answer = TextField(null=True, help_text="optional ground truth answer") - relevant_doc_ids = JSONField(null=True, help_text="expected relevant document IDs") - relevant_chunk_ids = JSONField(null=True, help_text="expected relevant chunk IDs") - metadata = JSONField(null=True, help_text="additional context/tags") - create_time = BigIntegerField(null=False, help_text="creation timestamp") - - class Meta: - db_table = "evaluation_cases" - - -class EvaluationRun(DataBaseModel): - """A single evaluation run""" - - id = CharField(max_length=32, primary_key=True) - dataset_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_datasets") - dialog_id = CharField(max_length=32, null=False, index=True, help_text="dialog configuration being evaluated") - name = CharField(max_length=255, null=False, help_text="run name") - config_snapshot = JSONField(null=False, help_text="dialog config at time of evaluation") - metrics_summary = JSONField(null=True, help_text="aggregated metrics") - status = CharField(max_length=32, null=False, default="PENDING", help_text="PENDING/RUNNING/COMPLETED/FAILED") - created_by = CharField(max_length=32, null=False, index=True, help_text="user who started the run") - create_time = BigIntegerField(null=False, index=True, help_text="creation timestamp") - complete_time = BigIntegerField(null=True, help_text="completion timestamp") - - class Meta: - db_table = "evaluation_runs" - - -class EvaluationResult(DataBaseModel): - """Result for a single test case in an evaluation run""" - - id = CharField(max_length=32, primary_key=True) - run_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_runs") - case_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_cases") - generated_answer = TextField(null=False, help_text="generated answer") - retrieved_chunks = JSONField(null=False, help_text="chunks that were retrieved") - metrics = JSONField(null=False, help_text="all computed metrics") - execution_time = FloatField(null=False, help_text="response time in seconds") - token_usage = JSONField(null=True, help_text="prompt/completion tokens") - create_time = BigIntegerField(null=False, help_text="creation timestamp") - - class Meta: - db_table = "evaluation_results" - - class Memory(DataBaseModel): id = CharField(max_length=32, primary_key=True) name = CharField(max_length=128, null=False, index=False, help_text="Memory name") @@ -1525,6 +1456,43 @@ def alter_db_rename_column(migrator, table_name, old_column_name, new_column_nam pass +def ensure_model_indexes(migrator): + """Create indexes declared by the Peewee models when they are missing.""" + members = inspect.getmembers(sys.modules[__name__], inspect.isclass) + for name, model in members: + if model == DataBaseModel or not issubclass(model, DataBaseModel): + continue + + table_name = model._meta.table_name + expected = {} + for field in model._meta.fields.values(): + if field.primary_key: + continue + if field.index or field.unique: + expected[(field.name,)] = bool(field.unique) + + for columns, unique in model._meta.indexes: + expected[tuple(columns)] = bool(unique) + + if not expected: + continue + + try: + existing = {tuple(index.columns): bool(index.unique) for index in DB.get_indexes(table_name)} + except Exception as ex: + logging.error(f"Failed to inspect indexes on {table_name}: {ex}") + continue + + for columns, unique in expected.items(): + if columns in existing and (not unique or existing[columns]): + continue + try: + migrate(migrator.add_index(table_name, columns, unique=unique)) + logging.info(f"Created {'unique ' if unique else ''}index on {table_name} ({', '.join(columns)})") + except Exception as ex: + logging.error(f"Failed to create {'unique ' if unique else ''}index on {table_name} ({', '.join(columns)}): {ex}") + + def migrate_add_unique_email(migrator): """Deduplicates user emails and add UNIQUE constraint to email column (idempotent)""" # step 0: check existing index state on user.email and prepare for unique constraint @@ -1834,6 +1802,7 @@ def migrate_db(): # this is after re-enabling logging to allow logging changed user emails migrate_add_unique_email(migrator) migrate_model_type_names() + ensure_model_indexes(migrator) def migrate_model_type_names():