diff --git a/migrations/versions/29ade7350c19_remove_document_level_valid_constraint.py b/migrations/versions/29ade7350c19_remove_document_level_valid_constraint.py new file mode 100644 index 00000000..5511b938 --- /dev/null +++ b/migrations/versions/29ade7350c19_remove_document_level_valid_constraint.py @@ -0,0 +1,76 @@ +"""remove Document level_valid constraint + +Revision ID: 29ade7350c19 +Revises: b8183c5ffb48 +Create Date: 2025-11-11 12:54:16.586701 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +from migrations.utils import constraint_exists, get_schema + +# revision identifiers, used by Alembic. +revision: str = "29ade7350c19" +down_revision: str | None = "b8183c5ffb48" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = get_schema() + + +def upgrade() -> None: + """Remove the level_valid CHECK constraint from documents table.""" + connection = op.get_bind() + inspector = sa.inspect(connection) + + # Drop CHECK constraint for level + if constraint_exists("documents", "level_valid", "check", inspector): + op.drop_constraint( + "level_valid", + "documents", + type_="check", + schema=schema, + ) + + +def downgrade() -> None: + """Restore the level_valid CHECK constraint to documents table.""" + connection = op.get_bind() + inspector = sa.inspect(connection) + + # Normalize any rows with values outside the legacy domain before recreating constraint + # This ensures downgrade succeeds even if new values like 'implicit' were added after upgrade + batch_size = 5000 + while True: + result = connection.execute( + sa.text( + f""" + WITH batch AS ( + SELECT id + FROM {schema}.documents + WHERE level NOT IN ('explicit', 'deductive') + LIMIT :batch_size + ) + UPDATE {schema}.documents d + SET level = 'explicit' + FROM batch + WHERE d.id = batch.id + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Recreate CHECK constraint for level + if not constraint_exists("documents", "level_valid", "check", inspector): + op.create_check_constraint( + "level_valid", + "documents", + "level IN ('explicit', 'deductive')", + schema=schema, + ) diff --git a/src/models.py b/src/models.py index f0042cd0..eb0cf7ae 100644 --- a/src/models.py +++ b/src/models.py @@ -405,7 +405,6 @@ class Document(Base): CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("length(content) <= 65535", name="content_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), - CheckConstraint("level IN ('explicit', 'deductive')", name="level_valid"), # Composite foreign key constraint for collections ForeignKeyConstraint( ["observer", "observed", "workspace_name"], diff --git a/src/schemas.py b/src/schemas.py index 4d99dd42..8b6068da 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -1,6 +1,6 @@ import datetime import ipaddress -from typing import Annotated, Any, Literal, Self +from typing import Annotated, Any, Self from urllib.parse import urlparse import tiktoken @@ -15,6 +15,7 @@ from pydantic import ( from src.config import settings from src.utils.representation import Representation +from src.utils.types import DocumentLevel RESOURCE_NAME_PATTERN = r"^[a-zA-Z0-9_-]+$" @@ -305,7 +306,7 @@ class DocumentCreate(DocumentBase): session_name: str = Field( description="The session from which the document was derived" ) - level: Literal["explicit", "deductive"] = Field( + level: DocumentLevel = Field( default="explicit", description="The level of the document (explicit or deductive)", ) diff --git a/tests/alembic/revisions/__init__.py b/tests/alembic/revisions/__init__.py index fedc98c5..9703972b 100644 --- a/tests/alembic/revisions/__init__.py +++ b/tests/alembic/revisions/__init__.py @@ -5,6 +5,7 @@ from . import ( test_066e87ca5b07_align_schema_with_declarative_models, test_08894082221a_replace_collection_name_with_observer_, test_20f89a421aff_rename_metamessage_type_to_label, + test_29ade7350c19_remove_document_level_valid_constraint, test_66e63cf2cf77_add_indexes_to_documents_table, test_76ffba56fe8c_add_error_field_to_queueitem, test_88b0fb10906f_add_webhooks_table, @@ -17,6 +18,7 @@ from . import ( test_bb6fb3a7a643_add_message_seq_in_session_column, test_c3828084f472_add_indexes_for_messages_and_, test_d429de0e5338_adopt_peer_paradigm, + test_e9b705f9adf9_add_server_defaults_to_timestamp_, test_ec8f94139b02_codify_workspace_name_and_message_id_in_, ) @@ -25,6 +27,7 @@ __all__ = [ "test_066e87ca5b07_align_schema_with_declarative_models", "test_08894082221a_replace_collection_name_with_observer_", "test_20f89a421aff_rename_metamessage_type_to_label", + "test_29ade7350c19_remove_document_level_valid_constraint", "test_556a16564f50_add_user_id_and_app_id_to_tables", "test_564ba40505c5_add_session_name_column_to_documents", "test_66e63cf2cf77_add_indexes_to_documents_table", @@ -37,5 +40,6 @@ __all__ = [ "test_bb6fb3a7a643_add_message_seq_in_session_column", "test_c3828084f472_add_indexes_for_messages_and_", "test_d429de0e5338_adopt_peer_paradigm", + "test_e9b705f9adf9_add_server_defaults_to_timestamp_", "test_ec8f94139b02_codify_workspace_name_and_message_id_in_", ] diff --git a/tests/alembic/revisions/test_29ade7350c19_remove_document_level_valid_constraint.py b/tests/alembic/revisions/test_29ade7350c19_remove_document_level_valid_constraint.py new file mode 100644 index 00000000..509ae2ae --- /dev/null +++ b/tests/alembic/revisions/test_29ade7350c19_remove_document_level_valid_constraint.py @@ -0,0 +1,20 @@ +"""Hooks for revision 29ade7350c19 (remove_document_level_valid_constraint).""" + +from __future__ import annotations + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + + +@register_before_upgrade("29ade7350c19") +def prepare_remove_document_level_valid_constraint(verifier: MigrationVerifier) -> None: + """Seed state and assertions before upgrading to 29ade7350c19.""" + # Verify CHECK constraint exists before migration (from b8183c5ffb48) + verifier.assert_constraint_exists("documents", "level_valid", "check") + + +@register_after_upgrade("29ade7350c19") +def verify_remove_document_level_valid_constraint(verifier: MigrationVerifier) -> None: + """Add assertions validating the effects of 29ade7350c19.""" + # Verify CHECK constraint no longer exists after migration + verifier.assert_constraint_exists("documents", "level_valid", "check", exists=False)