diff --git a/backend/alembic/versions/c7bf5721733e_add_has_been_indexed_to_.py b/backend/alembic/versions/c7bf5721733e_add_has_been_indexed_to_.py new file mode 100644 index 00000000000..d66970430f3 --- /dev/null +++ b/backend/alembic/versions/c7bf5721733e_add_has_been_indexed_to_.py @@ -0,0 +1,48 @@ +"""Add has_been_indexed to DocumentByConnectorCredentialPair + +Revision ID: c7bf5721733e +Revises: fec3db967bf7 +Create Date: 2025-01-13 12:39:05.831693 + +""" +from alembic import op +import sqlalchemy as sa + +# revision identifiers, used by Alembic. +revision = "c7bf5721733e" +down_revision = "fec3db967bf7" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # assume all existing rows have been indexed, no better approach + op.add_column( + "document_by_connector_credential_pair", + sa.Column("has_been_indexed", sa.Boolean(), nullable=True), + ) + op.execute( + "UPDATE document_by_connector_credential_pair SET has_been_indexed = TRUE" + ) + op.alter_column( + "document_by_connector_credential_pair", + "has_been_indexed", + nullable=False, + ) + + # Add index to optimize get_document_counts_for_cc_pairs query pattern + op.create_index( + "idx_document_cc_pair_counts", + "document_by_connector_credential_pair", + ["connector_id", "credential_id", "has_been_indexed"], + unique=False, + ) + + +def downgrade() -> None: + # Remove the index first before removing the column + op.drop_index( + "idx_document_cc_pair_counts", + table_name="document_by_connector_credential_pair", + ) + op.drop_column("document_by_connector_credential_pair", "has_been_indexed") diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 8759d84e21f..7b47cb13b9d 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -955,6 +955,14 @@ class DocumentByConnectorCredentialPair(Base): "credential_id", unique=False, ), + # Index to optimize get_document_counts_for_cc_pairs query pattern + Index( + "idx_document_cc_pair_counts", + "connector_id", + "credential_id", + "has_been_indexed", + unique=False, + ), )