feat: introduce alembic migration verification (#238)

* feat: introducer migration verification checks

* fix: move verification to tests/alembic

* feat: add verification steps for all alembic migrations

* fix: isolate test runs and implement all migration tests

* test: parametrize

* fix: CR comments

* fix: Add README

* feat: add precommit hook for validating alembic

* fix: rm pytest-alembic package

* test: create bulk resources to test migration batching

* fix: add latest test

* fix: tests to handle non-standard schema

* chore: Code Rabbit nits

* fix: CR comments 1

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja 2025-10-23 16:24:35 -04:00 committed by GitHub
parent a7d01d17df
commit 5db7b4948c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
36 changed files with 2876 additions and 364 deletions

View File

@ -65,12 +65,31 @@ repos:
# Run main application tests
- id: pytest-main
name: pytest (main app)
entry: uv run pytest tests/
entry: uv run pytest tests/ --ignore=tests/alembic/
language: system
files: ^(src/|tests/).*\.py$
stages: [pre-push]
pass_filenames: false
# Run Alembic tests only when migrations change
- id: pytest-alembic
name: pytest (alembic migrations)
entry: uv run pytest tests/alembic/
language: system
files: ^(migrations/.*\.py|tests/alembic/.*\.py)$
stages: [pre-push]
pass_filenames: false
require_serial: true
# Ensure each alembic migration revision has a corresponding test file
- id: ensure-alembic-coverage
name: ensure alembic migration test coverage
entry: uv run python scripts/ensure_alembic_tests.py
language: system
files: ^(migrations/versions/.*\.py|tests/alembic/revisions/.*\.py)$
stages: [pre-push]
pass_filenames: false
# Run Python SDK tests (if they exist)
- id: pytest-python-sdk
name: pytest (Python SDK)

View File

@ -476,12 +476,10 @@ def downgrade() -> None:
# Make peer_name NOT NULL
op.alter_column("collections", "peer_name", nullable=False, schema=schema)
# Recreate the foreign key constraint for peer_name
if not fk_exists(
"collections", "collections_peer_name_workspace_name_fkey", inspector
):
# Recreate the legacy foreign key constraint for peer_name
if not fk_exists("collections", "fk_collections_peer_name_peers", inspector):
op.create_foreign_key(
"collections_peer_name_workspace_name_fkey",
"fk_collections_peer_name_peers",
"collections",
"peers",
["peer_name", "workspace_name"],
@ -530,7 +528,7 @@ def downgrade() -> None:
# Step 6: Make documents collection_name NOT NULL
op.alter_column("documents", "collection_name", nullable=False, schema=schema)
# Step 6a: Restore peer_name column to documents (set to observed value)
# Step 6a: Restore peer_name column to documents (set to observer value)
if not column_exists("documents", "peer_name", inspector):
op.add_column(
"documents",
@ -637,14 +635,14 @@ def downgrade() -> None:
schema=schema,
)
# Step 11: Recreate the old foreign key constraint from documents to collections
# Step 11: Recreate the legacy foreign key constraint from documents to collections
if not fk_exists(
"documents",
"documents_collection_name_peer_name_workspace_name_fkey",
"fk_documents_collection_name_collections",
inspector,
):
op.create_foreign_key(
"documents_collection_name_peer_name_workspace_name_fkey",
"fk_documents_collection_name_collections",
"documents",
"collections",
["collection_name", "peer_name", "workspace_name"],

View File

@ -435,18 +435,18 @@ def downgrade():
op.alter_column("documents", "user_id", nullable=True, schema=schema)
print("Made app_id and user_id nullable again for documents table")
# Drop foreign key constraints
try:
# Drop foreign key constraints (guarded)
if fk_exists("documents", "documents_user_id_fkey", inspector):
op.drop_constraint("documents_user_id_fkey", "documents", schema=schema)
print("Dropped user_id foreign key for documents table")
except Exception as e:
print(f"Error dropping user_id foreign key for documents table: {e}")
else:
print("user_id foreign key for documents table does not exist; skipping drop")
try:
if fk_exists("documents", "documents_app_id_fkey", inspector):
op.drop_constraint("documents_app_id_fkey", "documents", schema=schema)
print("Dropped app_id foreign key for documents table")
except Exception as e:
print(f"Error dropping app_id foreign key for documents table: {e}")
else:
print("app_id foreign key for documents table does not exist; skipping drop")
# Drop the columns
op.drop_column("documents", "user_id", schema=schema)
@ -464,12 +464,12 @@ def downgrade():
op.alter_column("collections", "app_id", nullable=True, schema=schema)
print("Made app_id nullable again for collections table")
# Drop foreign key constraint
try:
# Drop foreign key constraint (guarded)
if fk_exists("collections", "collections_app_id_fkey", inspector):
op.drop_constraint("collections_app_id_fkey", "collections", schema=schema)
print("Dropped app_id foreign key for collections table")
except Exception as e:
print(f"Error dropping app_id foreign key for collections table: {e}")
else:
print("app_id foreign key for collections table does not exist; skipping drop")
# Drop the column
op.drop_column("collections", "app_id", schema=schema)
@ -484,20 +484,26 @@ def downgrade():
)
print("Dropped index ix_metamessages_app_id")
# Make app_id nullable again
op.alter_column("metamessages", "app_id", nullable=True, schema=schema)
print("Made app_id nullable again for metamessages table")
# Make app_id nullable again (only if column exists)
if column_exists("metamessages", "app_id", inspector):
op.alter_column("metamessages", "app_id", nullable=True, schema=schema)
print("Made app_id nullable again for metamessages table")
else:
print("metamessages.app_id column does not exist; skipping alter nullable")
# Drop foreign key constraint
try:
# Drop foreign key constraint (guarded)
if fk_exists("metamessages", "metamessages_app_id_fkey", inspector):
op.drop_constraint("metamessages_app_id_fkey", "metamessages", schema=schema)
print("Dropped app_id foreign key for metamessages table")
except Exception as e:
print(f"Error dropping app_id foreign key for metamessages table: {e}")
else:
print("app_id foreign key for metamessages table does not exist; skipping drop")
# Drop the column
op.drop_column("metamessages", "app_id", schema=schema)
print("Dropped app_id column from metamessages table")
# Drop the column (only if present)
if column_exists("metamessages", "app_id", inspector):
op.drop_column("metamessages", "app_id", schema=schema)
print("Dropped app_id column from metamessages table")
else:
print("metamessages.app_id column does not exist; skipping drop column")
# 2. Messages table
@ -514,18 +520,18 @@ def downgrade():
op.alter_column("messages", "user_id", nullable=True, schema=schema)
print("Made app_id and user_id nullable again for messages table")
# Drop foreign key constraints
try:
# Drop foreign key constraints (guarded)
if fk_exists("messages", "messages_user_id_fkey", inspector):
op.drop_constraint("messages_user_id_fkey", "messages", schema=schema)
print("Dropped user_id foreign key for messages table")
except Exception as e:
print(f"Error dropping user_id foreign key for messages table: {e}")
else:
print("user_id foreign key for messages table does not exist; skipping drop")
try:
if fk_exists("messages", "messages_app_id_fkey", inspector):
op.drop_constraint("messages_app_id_fkey", "messages", schema=schema)
print("Dropped app_id foreign key for messages table")
except Exception as e:
print(f"Error dropping app_id foreign key for messages table: {e}")
else:
print("app_id foreign key for messages table does not exist; skipping drop")
# Drop the columns
op.drop_column("messages", "user_id", schema=schema)
@ -543,12 +549,12 @@ def downgrade():
op.alter_column("sessions", "app_id", nullable=True, schema=schema)
print("Made app_id nullable again for sessions table")
# Drop foreign key constraint
try:
# Drop foreign key constraint (guarded)
if fk_exists("sessions", "sessions_app_id_fkey", inspector):
op.drop_constraint("sessions_app_id_fkey", "sessions", schema=schema)
print("Dropped app_id foreign key for sessions table")
except Exception as e:
print(f"Error dropping app_id foreign key for sessions table: {e}")
else:
print("app_id foreign key for sessions table does not exist; skipping drop")
# Drop the column
op.drop_column("sessions", "app_id", schema=schema)

View File

@ -183,3 +183,46 @@ def downgrade() -> None:
if column_exists("active_queue_sessions", "work_unit_data", inspector):
op.drop_column("active_queue_sessions", "work_unit_data", schema=schema)
# Re-add columns to active_queue_sessions
if not column_exists("active_queue_sessions", "session_id", inspector):
op.add_column(
"active_queue_sessions",
sa.Column("session_id", sa.TEXT(), nullable=True),
schema=schema,
)
if not column_exists("active_queue_sessions", "sender_name", inspector):
op.add_column(
"active_queue_sessions",
sa.Column("sender_name", sa.TEXT(), nullable=True),
schema=schema,
)
if not column_exists("active_queue_sessions", "target_name", inspector):
op.add_column(
"active_queue_sessions",
sa.Column("target_name", sa.TEXT(), nullable=True),
schema=schema,
)
if not column_exists("active_queue_sessions", "task_type", inspector):
op.add_column(
"active_queue_sessions",
sa.Column("task_type", sa.TEXT(), nullable=False),
schema=schema,
)
# Re-add unique constraint
if not constraint_exists(
"active_queue_sessions",
"unique_active_queue_session",
"unique",
inspector,
):
op.create_unique_constraint(
"unique_active_queue_session",
"active_queue_sessions",
["session_id", "sender_name", "target_name", "task_type"],
schema=schema,
)

View File

@ -40,14 +40,15 @@ def upgrade() -> None:
server_default=sa.func.now(),
),
# Foreign key constraints
sa.ForeignKeyConstraint(["message_id"], ["messages.public_id"]),
sa.ForeignKeyConstraint(["workspace_name"], ["workspaces.name"]),
sa.ForeignKeyConstraint(["message_id"], [f"{schema}.messages.public_id"]),
sa.ForeignKeyConstraint(["workspace_name"], [f"{schema}.workspaces.name"]),
sa.ForeignKeyConstraint(
["session_name", "workspace_name"],
["sessions.name", "sessions.workspace_name"],
[f"{schema}.sessions.name", f"{schema}.sessions.workspace_name"],
),
sa.ForeignKeyConstraint(
["peer_name", "workspace_name"], ["peers.name", "peers.workspace_name"]
["peer_name", "workspace_name"],
[f"{schema}.peers.name", f"{schema}.peers.workspace_name"],
),
schema=schema,
)

View File

@ -83,7 +83,7 @@ def upgrade() -> None:
sa.CheckConstraint("length(name) <= 512", name="name_length"),
sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
sa.ForeignKeyConstraint(
["app_id"], ["apps.public_id"], name=op.f("fk_users_app_id_apps")
["app_id"], [f"{schema}.apps.public_id"], name=op.f("fk_users_app_id_apps")
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_users")),
sa.UniqueConstraint("name", "app_id", name="unique_name_app_user"),
@ -130,7 +130,9 @@ def upgrade() -> None:
sa.CheckConstraint("length(public_id) = 21", name="public_id_length"),
sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
sa.ForeignKeyConstraint(
["user_id"], ["users.public_id"], name=op.f("fk_sessions_user_id_users")
["user_id"],
[f"{schema}.users.public_id"],
name=op.f("fk_sessions_user_id_users"),
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_sessions")),
sa.UniqueConstraint("public_id", name=op.f("uq_sessions_public_id")),
@ -186,7 +188,7 @@ def upgrade() -> None:
sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
sa.ForeignKeyConstraint(
["session_id"],
["sessions.public_id"],
[f"{schema}.sessions.public_id"],
name=op.f("fk_messages_session_id_sessions"),
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_messages")),
@ -246,7 +248,7 @@ def upgrade() -> None:
sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
sa.ForeignKeyConstraint(
["message_id"],
["messages.public_id"],
[f"{schema}.messages.public_id"],
name=op.f("fk_metamessages_message_id_messages"),
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_metamessages")),
@ -308,7 +310,9 @@ def upgrade() -> None:
sa.CheckConstraint("length(name) <= 512", name="name_length"),
sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
sa.ForeignKeyConstraint(
["user_id"], ["users.public_id"], name=op.f("fk_collections_user_id_users")
["user_id"],
[f"{schema}.users.public_id"],
name=op.f("fk_collections_user_id_users"),
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_collections")),
sa.UniqueConstraint("name", "user_id", name="unique_name_collection_user"),
@ -372,7 +376,7 @@ def upgrade() -> None:
sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
sa.ForeignKeyConstraint(
["collection_id"],
["collections.public_id"],
[f"{schema}.collections.public_id"],
name=op.f("fk_documents_collection_id_collections"),
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_documents")),
@ -412,7 +416,9 @@ def upgrade() -> None:
sa.Column("payload", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
sa.Column("processed", sa.Boolean(), nullable=False, server_default="false"),
sa.ForeignKeyConstraint(
["session_id"], ["sessions.id"], name=op.f("fk_queue_session_id_sessions")
["session_id"],
[f"{schema}.sessions.id"],
name=op.f("fk_queue_session_id_sessions"),
),
sa.PrimaryKeyConstraint("id", name=op.f("pk_queue")),
schema=schema,
@ -437,7 +443,7 @@ def upgrade() -> None:
),
sa.ForeignKeyConstraint(
["session_id"],
["sessions.id"],
[f"{schema}.sessions.id"],
name=op.f("fk_active_queue_sessions_session_id_sessions"),
),
sa.PrimaryKeyConstraint("session_id", name=op.f("pk_active_queue_sessions")),

View File

@ -13,7 +13,7 @@ import sqlalchemy as sa
import tiktoken
from alembic import op
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from sqlalchemy import Inspector, text
from sqlalchemy.dialects import postgresql
from migrations.utils import (
@ -114,8 +114,117 @@ def downgrade() -> None:
# Step 11: Restore foreign keys
restore_foreign_keys(schema)
# Step 12: Readd metamessages table
if not table_exists("metamessages", inspector):
op.create_table(
"metamessages",
sa.Column(
"id",
sa.BigInteger(),
sa.Identity(always=False),
primary_key=True,
autoincrement=True,
index=True,
nullable=False,
),
sa.Column("public_id", sa.TEXT(), nullable=False, index=True, unique=True),
sa.Column("content", sa.TEXT(), nullable=True),
sa.Column(
"message_id",
sa.TEXT(),
sa.ForeignKey("messages.public_id"),
nullable=True,
index=True,
),
sa.Column(
"created_at",
sa.DateTime(timezone=True),
nullable=False,
index=True,
server_default=sa.text("now()"),
),
sa.Column("label", sa.TEXT(), nullable=False, index=True),
sa.Column(
"session_id",
sa.TEXT(),
sa.ForeignKey("sessions.public_id"),
nullable=True,
index=True,
),
sa.Column(
"user_id",
sa.TEXT(),
sa.ForeignKey("users.public_id"),
nullable=True,
index=True,
),
sa.Column(
"app_id",
sa.TEXT(),
sa.ForeignKey("apps.public_id"),
nullable=True,
index=True,
),
sa.Column(
"metadata",
postgresql.JSONB(astext_type=sa.Text()),
nullable=False,
server_default=sa.text("'{}'::jsonb"),
),
sa.Index(
"idx_metamessages_lookup",
"label",
sa.text("id DESC"),
postgresql_include=["public_id", "message_id", "created_at"],
),
sa.Index(
"idx_metamessages_user_lookup",
"user_id",
"label",
sa.text("id DESC"),
),
sa.Index(
"idx_metamessages_session_lookup",
"session_id",
"label",
sa.text("id DESC"),
),
sa.Index(
"idx_metamessages_message_lookup",
"message_id",
"label",
sa.text("id DESC"),
),
sa.CheckConstraint("length(public_id) = 21", name="public_id_length"),
sa.CheckConstraint(
"public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"
),
sa.CheckConstraint("length(content) <= 65535", name="content_length"),
sa.CheckConstraint("length(label) <= 512", name="label_length"),
# Added constraints to ensure consistency
sa.CheckConstraint(
"(message_id IS NULL) OR (session_id IS NOT NULL)",
name="message_requires_session",
),
sa.ForeignKeyConstraint(
["session_id"],
[f"{schema}.sessions.public_id"],
"fk_metamessages_session_id_sessions",
),
sa.ForeignKeyConstraint(
["user_id"],
[f"{schema}.users.public_id"],
"fk_metamessages_user_id_users",
),
sa.ForeignKeyConstraint(
["app_id"],
[f"{schema}.apps.public_id"],
"fk_metamessages_app_id_apps",
),
)
def rename_tables(schema: str, inspector) -> None:
def rename_tables(schema: str, inspector: Inspector) -> None:
"""Rename apps->workspaces and users->peers tables."""
if inspector.has_table("apps", schema=schema):
op.rename_table("apps", "workspaces", schema=schema)
@ -123,7 +232,7 @@ def rename_tables(schema: str, inspector) -> None:
op.rename_table("users", "peers", schema=schema)
def update_workspaces_table(schema: str, inspector) -> None:
def update_workspaces_table(schema: str, inspector: Inspector) -> None:
"""Update workspaces table (formerly apps)."""
# Add configuration column
@ -176,7 +285,7 @@ def update_workspaces_table(schema: str, inspector) -> None:
)
def update_peers_table(schema: str, inspector) -> None:
def update_peers_table(schema: str, inspector: Inspector) -> None:
"""Update peers table (formerly users)."""
# Add configuration column
@ -266,9 +375,10 @@ def update_peers_table(schema: str, inspector) -> None:
)
op.drop_column("peers", "app_id", schema=schema)
op.alter_column("peers", "workspace_name", nullable=False, schema=schema)
def update_sessions_table(schema: str, inspector) -> None:
def update_sessions_table(schema: str, inspector: Inspector) -> None:
"""Update sessions table."""
# Add configuration column
@ -304,6 +414,7 @@ def update_sessions_table(schema: str, inspector) -> None:
f"UPDATE {schema}.sessions SET workspace_name = workspaces.name FROM {schema}.workspaces WHERE sessions.app_id = workspaces.id"
)
)
op.alter_column("sessions", "workspace_name", nullable=False, schema=schema)
op.add_column(
"sessions",
@ -365,7 +476,7 @@ def update_sessions_table(schema: str, inspector) -> None:
)
def create_and_populate_session_peers_table(schema: str, inspector) -> None:
def create_and_populate_session_peers_table(schema: str, inspector: Inspector) -> None:
"""Create and populate session_peers table."""
# Create session_peers table
@ -400,11 +511,11 @@ def create_and_populate_session_peers_table(schema: str, inspector) -> None:
),
sa.ForeignKeyConstraint(
["peer_name", "workspace_name"],
["peers.name", "peers.workspace_name"],
[f"{schema}.peers.name", f"{schema}.peers.workspace_name"],
),
sa.ForeignKeyConstraint(
["session_name", "workspace_name"],
["sessions.name", "sessions.workspace_name"],
[f"{schema}.sessions.name", f"{schema}.sessions.workspace_name"],
),
sa.PrimaryKeyConstraint("workspace_name", "session_name", "peer_name"),
)
@ -456,7 +567,7 @@ def create_and_populate_session_peers_table(schema: str, inspector) -> None:
)
def update_messages_table(schema: str, inspector) -> None:
def update_messages_table(schema: str, inspector: Inspector) -> None:
"""Update messages table."""
# Add new columns
@ -618,7 +729,7 @@ def update_messages_table(schema: str, inspector) -> None:
backfill_token_counts(schema)
def update_collections_table(schema: str, inspector) -> None:
def update_collections_table(schema: str, inspector: Inspector) -> None:
"""Update collections table."""
# Add new columns
@ -741,7 +852,7 @@ def update_collections_table(schema: str, inspector) -> None:
)
def update_documents_table(schema: str, inspector) -> None:
def update_documents_table(schema: str, inspector: Inspector) -> None:
"""Update documents table."""
# Add new columns
@ -873,7 +984,9 @@ def update_documents_table(schema: str, inspector) -> None:
)
def update_queue_and_active_queue_sessions_tables(schema: str, inspector) -> None:
def update_queue_and_active_queue_sessions_tables(
schema: str, inspector: Inspector
) -> None:
"""Update queue and active_queue_sessions tables."""
# Drop foreign key constraints first, before changing column types
@ -903,7 +1016,7 @@ def update_queue_and_active_queue_sessions_tables(schema: str, inspector) -> Non
# Get the mapping of old session.id (integer) to new session.id (text, which is public_id)
# At this point, sessions table still has both id (integer) and public_id (text) columns
session_id_mapping = {}
session_id_mapping: dict[int, str] = {}
if table_exists("sessions", inspector):
sessions_mapping = connection.execute(
sa.text(f"SELECT id, public_id FROM {schema}.sessions")
@ -1081,7 +1194,7 @@ def backfill_token_counts(schema: str) -> None:
break
# Calculate token counts and update messages in batches
batch_updates = []
batch_updates: list[tuple[int, int]] = []
for message_id, content in messages:
token_count = _count_tokens(content)
batch_updates.append((message_id, token_count))
@ -1110,7 +1223,7 @@ def backfill_token_counts(schema: str) -> None:
offset += batch_size
def restore_app_user_columns(schema: str, inspector) -> None:
def restore_app_user_columns(schema: str, inspector: Inspector) -> None:
"""Restore app_id and user_id columns to peers and sessions."""
# Add app_id back to peers
if not column_exists("peers", "app_id", inspector):
@ -1154,8 +1267,7 @@ def restore_app_user_columns(schema: str, inspector) -> None:
schema=schema,
)
# Populate user_id from session_peers
# The user peer is the one that existed before agent peers were created
# Agent peers have names that match their IDs (both are nanoids), user peers have different names and IDs
# Prefer user peers (where id != name), but fall back to any peer if needed
op.execute(
sa.text(f"""
UPDATE {schema}.sessions SET user_id = (
@ -1172,7 +1284,7 @@ def restore_app_user_columns(schema: str, inspector) -> None:
op.alter_column("sessions", "user_id", nullable=False, schema=schema)
def restore_documents_table(schema: str, inspector) -> None:
def restore_documents_table(schema: str, inspector: Inspector) -> None:
"""Restore documents table to pre-peer paradigm state."""
# Add back id column as primary key
if not column_exists("documents", "temp_id", inspector):
@ -1306,7 +1418,7 @@ def restore_documents_table(schema: str, inspector) -> None:
)
def restore_collections_table(schema: str, inspector) -> None:
def restore_collections_table(schema: str, inspector: Inspector) -> None:
"""Restore collections table to pre-peer paradigm state."""
# Add back id column as primary key
if not column_exists("collections", "temp_id", inspector):
@ -1433,7 +1545,7 @@ def restore_collections_table(schema: str, inspector) -> None:
)
def restore_messages_table(schema: str, inspector) -> None:
def restore_messages_table(schema: str, inspector: Inspector) -> None:
"""Restore messages table to pre-peer paradigm state."""
# Add back old columns
if not column_exists("messages", "session_id", inspector):
@ -1553,8 +1665,20 @@ def restore_messages_table(schema: str, inspector) -> None:
referent_schema=schema,
)
op.create_index(
"idx_messages_session_lookup",
"messages",
["session_id", "id"],
postgresql_include=[
"public_id",
"is_user",
"created_at",
],
schema=schema,
)
def restore_sessions_table(schema: str, inspector) -> None:
def restore_sessions_table(schema: str, inspector: Inspector) -> None:
"""Restore sessions table to pre-peer paradigm state."""
# Add back id column as BigInteger primary key
if not column_exists("sessions", "temp_id", inspector):
@ -1624,8 +1748,12 @@ def restore_sessions_table(schema: str, inspector) -> None:
"public_id_format", "sessions", "public_id ~ '^[A-Za-z0-9_-]+$'", schema=schema
)
op.create_index(
"idx_sessions_user_lookup", "sessions", ["user_id", "public_id"], schema=schema
)
def restore_peers_table(schema: str, inspector) -> None:
def restore_peers_table(schema: str, inspector: Inspector) -> None:
"""Restore peers table to pre-user paradigm state."""
# Add back id column as BigInteger primary key
if not column_exists("peers", "temp_id", inspector):
@ -1703,7 +1831,7 @@ def restore_peers_table(schema: str, inspector) -> None:
)
def restore_workspaces_table(schema: str, inspector) -> None:
def restore_workspaces_table(schema: str, inspector: Inspector) -> None:
"""Restore workspaces table to pre-peer paradigm state (apps)."""
# Add back id column as BigInteger primary key
if not column_exists("workspaces", "temp_id", inspector):
@ -1764,14 +1892,16 @@ def restore_workspaces_table(schema: str, inspector) -> None:
)
def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> None:
def restore_queue_and_active_queue_sessions_tables(
schema: str, inspector: Inspector
) -> None:
"""Restore queue and active_queue_sessions tables to pre-peer paradigm state."""
connection = op.get_bind()
# Create reverse mapping from session.public_id (text) back to session.id (BigInteger)
# At this point in downgrade, sessions table still has both id (BigInteger) and public_id (text)
session_id_reverse_mapping = {}
session_id_reverse_mapping: dict[str, int] = {}
if table_exists("sessions", inspector):
sessions_mapping = connection.execute(
sa.text(f"SELECT id, public_id FROM {schema}.sessions")
@ -1781,26 +1911,27 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No
session_id_reverse_mapping[text_id] = big_int_id
# Update queue table
if table_exists("queue", inspector) and session_id_reverse_mapping:
if table_exists("queue", inspector):
# Get current session_id values in queue table (they are text now)
queue_session_ids = connection.execute(
sa.text(
f"SELECT DISTINCT session_id FROM {schema}.queue WHERE session_id IS NOT NULL"
)
).fetchall()
# Convert session_id values back to BigInteger
for (session_id,) in queue_session_ids:
if session_id in session_id_reverse_mapping:
old_id = session_id_reverse_mapping[session_id]
connection.execute(
sa.text(
f"UPDATE {schema}.queue SET session_id = :old_id WHERE session_id = :new_id"
),
{"old_id": str(old_id), "new_id": str(session_id)},
if session_id_reverse_mapping:
queue_session_ids = connection.execute(
sa.text(
f"SELECT DISTINCT session_id FROM {schema}.queue WHERE session_id IS NOT NULL"
)
).fetchall()
# Change column type back to BigInteger
# Convert session_id values back to BigInteger
for (session_id,) in queue_session_ids:
if session_id in session_id_reverse_mapping:
old_id = session_id_reverse_mapping[session_id]
connection.execute(
sa.text(
f"UPDATE {schema}.queue SET session_id = :old_id WHERE session_id = :new_id"
),
{"old_id": str(old_id), "new_id": str(session_id)},
)
# Change column type back to BigInteger (always)
op.alter_column(
"queue",
"session_id",
@ -1810,36 +1941,40 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No
)
# Update active_queue_sessions table
if table_exists("active_queue_sessions", inspector) and session_id_reverse_mapping:
# Get current session_id values in active_queue_sessions table (they are text now)
active_queue_session_ids = connection.execute(
sa.text(
f"SELECT DISTINCT session_id FROM {schema}.active_queue_sessions WHERE session_id IS NOT NULL"
)
).fetchall()
if table_exists("active_queue_sessions", inspector):
if session_id_reverse_mapping:
# Get current session_id values in active_queue_sessions table (they are text now)
active_queue_session_ids = connection.execute(
sa.text(
f"SELECT DISTINCT session_id FROM {schema}.active_queue_sessions WHERE session_id IS NOT NULL"
)
).fetchall()
if constraint_exists(
"active_queue_sessions", "pk_active_queue_sessions", "primary", inspector
):
op.drop_constraint(
"pk_active_queue_sessions",
if constraint_exists(
"active_queue_sessions",
type_="primary",
schema=schema,
)
# Convert session_id values back to BigInteger
for (session_id,) in active_queue_session_ids:
if session_id in session_id_reverse_mapping:
old_id = session_id_reverse_mapping[session_id]
connection.execute(
sa.text(
f"UPDATE {schema}.active_queue_sessions SET session_id = :old_id WHERE session_id = :new_id"
),
{"old_id": str(old_id), "new_id": str(session_id)},
"pk_active_queue_sessions",
"primary",
inspector,
):
op.drop_constraint(
"pk_active_queue_sessions",
"active_queue_sessions",
type_="primary",
schema=schema,
)
# Change column type back to BigInteger
# Convert session_id values back to BigInteger
for (session_id,) in active_queue_session_ids:
if session_id in session_id_reverse_mapping:
old_id = session_id_reverse_mapping[session_id]
connection.execute(
sa.text(
f"UPDATE {schema}.active_queue_sessions SET session_id = :old_id WHERE session_id = :new_id"
),
{"old_id": str(old_id), "new_id": str(session_id)},
)
# Change column type back to BigInteger (always)
op.alter_column(
"active_queue_sessions",
"session_id",
@ -1848,12 +1983,16 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No
postgresql_using="session_id::bigint",
)
op.create_primary_key(
"pk_active_queue_sessions",
"active_queue_sessions",
["session_id"],
schema=schema,
)
# Ensure primary key exists on session_id for pre-peer shape
if not constraint_exists(
"active_queue_sessions", "pk_active_queue_sessions", "primary", inspector
):
op.create_primary_key(
"pk_active_queue_sessions",
"active_queue_sessions",
["session_id"],
schema=schema,
)
if constraint_exists(
"active_queue_sessions",
@ -1877,29 +2016,10 @@ def restore_queue_and_active_queue_sessions_tables(schema: str, inspector) -> No
if column_exists("active_queue_sessions", "task_type", inspector):
op.drop_column("active_queue_sessions", "task_type", schema=schema)
# Restore foreign key constraints
if table_exists("queue", inspector):
op.create_foreign_key(
"fk_queue_session_id_sessions",
"queue",
"sessions",
["session_id"],
["id"],
referent_schema=schema,
)
if table_exists("active_queue_sessions", inspector):
op.create_foreign_key(
"fk_active_queue_sessions_session_id_sessions",
"active_queue_sessions",
"sessions",
["session_id"],
["id"],
referent_schema=schema,
)
# Defer restoring foreign key constraints to restore_foreign_keys()
def restore_table_names(schema: str, inspector) -> None:
def restore_table_names(schema: str, inspector: Inspector) -> None:
"""Restore table names: workspaces->apps and peers->users."""
if inspector.has_table("workspaces", schema=schema):
op.rename_table("workspaces", "apps", schema=schema)
@ -1957,6 +2077,23 @@ def restore_foreign_keys(schema: str) -> None:
["public_id"],
referent_schema=schema,
)
# Restore queue/session FKs last, after all type changes are complete
op.create_foreign_key(
"fk_queue_session_id_sessions",
"queue",
"sessions",
["session_id"],
["id"],
referent_schema=schema,
)
op.create_foreign_key(
"fk_active_queue_sessions_session_id_sessions",
"active_queue_sessions",
"sessions",
["session_id"],
["id"],
referent_schema=schema,
)
op.create_foreign_key(
"fk_sessions_app_id_apps",
"sessions",

View File

@ -81,7 +81,7 @@ extend-immutable-calls = ["fastapi.Depends"]
[tool.pytest.ini_options]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "session"
addopts = "--strict-markers --cov=src/ --cov=sdks/python/src/honcho --cov-report=term-missing"
addopts = "--strict-markers --cov=src/ --cov=sdks/python/src/honcho --cov-report=term-missing --ignore=tests/alembic"
testpaths = ["tests"]
pythonpath = ["src"]

View File

@ -0,0 +1,68 @@
#!/usr/bin/env uv run python
"""
Script that validates that all alembic migration revisions have a corresponding test file.
Note that this script is actively used within our precommit hooks and should not be removed.
If this script is moved, the corresponding precommit hook will need to be updated.
"""
import sys
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
MIGRATIONS_DIR = PROJECT_ROOT / "migrations" / "versions"
TESTS_DIR = PROJECT_ROOT / "tests" / "alembic" / "revisions"
def main() -> int:
missing: list[str] = []
# Gather migration base names (without extension)
migration_files = [
p for p in MIGRATIONS_DIR.glob("*.py") if p.name != "__init__.py"
]
migration_basenames = {p.stem for p in migration_files}
print(f"Migration basenames: {migration_basenames}")
# Gather test file base names mapped back to migration names by stripping leading 'test_'
test_files = [p for p in TESTS_DIR.glob("test_*.py") if p.name != "__init__.py"]
test_targets = {p.stem.removeprefix("test_") for p in test_files}
for migration_basename in sorted(migration_basenames):
if migration_basename not in test_targets:
missing.append(migration_basename)
if missing:
print(
"Missing Alembic tests for the following migration revisions:",
file=sys.stderr,
)
for name in missing:
print(f" - {name}", file=sys.stderr)
print(
"\nExpected test files under tests/alembic/revisions named as:",
file=sys.stderr,
)
for name in missing:
print(f" - tests/alembic/revisions/test_{name}.py", file=sys.stderr)
print("\nScaffold helper commands:", file=sys.stderr)
for name in missing:
print(
f" - python -m tests.alembic.scaffold {name.split('_', 1)[0]}",
file=sys.stderr,
)
return 1
# Optional: warn if there are tests without corresponding migrations (stale tests)
stale_tests = sorted(test_targets - migration_basenames)
if stale_tests:
print(
"Warning: Found tests without corresponding migration files (stale?):",
file=sys.stderr,
)
for name in stale_tests:
print(f" - tests/alembic/revisions/test_{name}.py", file=sys.stderr)
return 0
if __name__ == "__main__":
raise SystemExit(main())

View File

@ -1,4 +1,4 @@
#!/usr/bin/env python
#!/usr/bin/env uv run python
"""
Utility script to generate a JWT secret for use in the .env file.
This uses the same logic as the automatically generated version in security.py.

View File

@ -1,3 +1,4 @@
#!/usr/bin/env uv run python
"""
Script to generate embeddings for existing messages that don't already have embeddings.

View File

@ -113,7 +113,7 @@ class Peer(Base):
DateTime(timezone=True), index=True, default=func.now()
)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), index=True
ForeignKey("workspaces.name"), index=True, nullable=False
)
configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)
@ -149,7 +149,7 @@ class Session(Base):
)
messages = relationship("Message", back_populates="session")
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), index=True
ForeignKey("workspaces.name"), index=True, nullable=False
)
configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)

29
tests/alembic/README.md Normal file
View File

@ -0,0 +1,29 @@
These tests validate Alembic migrations end-to-end for structure, order, and correctness.
They ensure reversibility, expected schema and data, and integration with the registry and pipeline.
The key components are the verifier, the test pipeline, the registry, and the revisions under test.
### Verifier
- The verifier runs checks when specific revisions are applied and reverted.
- It validates the schema after upgrade, verifies data migrations such as defaults, backfills, and transforms, and confirms reversibility after downgrade.
- Assertions are grouped per-revision or feature, and helpers use the SQLAlchemy inspector to introspect the database.
### Test Pipeline
- The pipeline orchestrates the database lifecycle by creating a database, applying upgrades and downgrades, running verifications, and tearing down resources.
- It typically starts from base, upgrades to the revision immediately before a target revision, seeds the DB, runs the target migration, and then validates the schema + data
- It relies on shared fixtures such as `engine`, `connection`, and `alembic_config`, and it ensures isolation per test
### Registry
- The registry declares revisions and test metadata used to drive scenarios.
- It defines ordering and selection, attaches verifier callbacks to specific revisions or ranges, and centralizes per revision expectations.
### Revisions
- Revisions are the migration scripts under `alembic/revisions`.
- Each revision should provide functions decorated with `register_before_upgrade()` and `register_after_upgrade()`. These are used to validate schemas and data before and after a migration is run
### Running the Tests
- Tests can be run all together `pytest tests/alembic` or individually `pytest tests/alembic -k "revision_number"`. For example, to run the test against a1b2c3d4e5f6_initial_schema.py, you would run the command `pytest tests/alembic -k "a1b2c3d4e5f6"`

View File

@ -0,0 +1,5 @@
"""Test package for Alembic migration scenarios."""
from . import revisions
__all__ = ["revisions"]

107
tests/alembic/conftest.py Normal file
View File

@ -0,0 +1,107 @@
"""Pytest configuration for Alembic-focused tests."""
from __future__ import annotations
import os
from collections.abc import Generator
from pathlib import Path
import pytest
from alembic.config import Config
from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from sqlalchemy.engine.url import URL
from sqlalchemy_utils import (
create_database, # pyright: ignore
database_exists, # pyright: ignore
drop_database, # pyright: ignore
)
from src.config import settings
from tests.conftest import CONNECTION_URI
ALEMBIC_CONFIG_PATH = Path(__file__).resolve().parents[2] / "alembic.ini"
ALEMBIC_TEST_DB_URL: URL = CONNECTION_URI.set(database="alembic_migration_tests")
@pytest.fixture(scope="session", autouse=True)
def configure_alembic_settings(alembic_database: str) -> Generator[None, None, None]:
"""Point application settings at the Alembic test database."""
previous_uri = settings.DB.CONNECTION_URI
os.environ["DB_CONNECTION_URI"] = alembic_database
settings.DB.CONNECTION_URI = alembic_database
try:
yield
finally:
settings.DB.CONNECTION_URI = previous_uri
if previous_uri:
os.environ["DB_CONNECTION_URI"] = previous_uri
else:
os.environ.pop("DB_CONNECTION_URI", None)
@pytest.fixture
def alembic_cfg(alembic_database: str) -> Config:
"""Provide an Alembic Config bound to the alembic test database."""
cfg = Config(str(ALEMBIC_CONFIG_PATH))
cfg.set_main_option(
"script_location", str(ALEMBIC_CONFIG_PATH.parent / "migrations")
)
cfg.set_main_option("sqlalchemy.url", alembic_database)
return cfg
@pytest.fixture(scope="session")
def alembic_database() -> Generator[str, None, None]:
"""Provision a dedicated DB for Alembic verification tests."""
assert ALEMBIC_TEST_DB_URL.database == "alembic_migration_tests", (
"Can't set up Alembic test database fixture. "
+ "ALEMBIC_TEST_DB_URL.database is {ALEMBIC_TEST_DB_URL.database}, "
+ "expected 'alembic_migration_tests'. "
)
if database_exists(ALEMBIC_TEST_DB_URL):
drop_database(ALEMBIC_TEST_DB_URL) # start fresh
create_database(ALEMBIC_TEST_DB_URL)
engine = create_engine(str(ALEMBIC_TEST_DB_URL))
try:
with engine.begin() as conn:
conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
yield str(ALEMBIC_TEST_DB_URL)
finally:
engine.dispose()
if database_exists(ALEMBIC_TEST_DB_URL):
drop_database(ALEMBIC_TEST_DB_URL)
@pytest.fixture
def alembic_engine(alembic_database: str) -> Generator[Engine, None, None]:
"""Yield an engine bound to the Alembic test database."""
engine = create_engine(alembic_database, pool_pre_ping=True)
try:
yield engine
finally:
engine.dispose()
@pytest.fixture(autouse=True)
def reset_schema_between_tests(
alembic_engine: Engine,
) -> Generator[None, None, None]:
"""Drop and recreate the schema for a clean slate each test (fast reset)."""
schema = settings.DB.SCHEMA
def _reset_schema() -> None:
with alembic_engine.begin() as conn:
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
conn.execute(text(f'CREATE SCHEMA "{schema}"'))
_reset_schema()
yield
_reset_schema()

73
tests/alembic/registry.py Normal file
View File

@ -0,0 +1,73 @@
"""Registry for migration-specific prepare and verification hooks."""
from __future__ import annotations
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from tests.alembic.verifier import MigrationVerifier
BeforeUpgradeHook = Callable[["MigrationVerifier"], None]
AfterUpgradeHook = Callable[["MigrationVerifier"], None]
@dataclass(slots=True)
class RevisionHooks:
"""Container for lifecycle hooks tied to a revision."""
before_upgrade: BeforeUpgradeHook | None = None
after_upgrade: AfterUpgradeHook | None = None
_REGISTRY: dict[str, RevisionHooks] = {}
def register_before_upgrade(
revision: str,
) -> Callable[[BeforeUpgradeHook], BeforeUpgradeHook]:
"""Register a callable to execute before upgrading to ``revision``."""
def decorator(func: BeforeUpgradeHook) -> BeforeUpgradeHook:
hooks = _REGISTRY.setdefault(revision, RevisionHooks())
if hooks.before_upgrade is not None:
raise ValueError(
f"before_upgrade hook already registered for revision {revision}"
)
hooks.before_upgrade = func
return func
return decorator
def register_after_upgrade(
revision: str,
) -> Callable[[AfterUpgradeHook], AfterUpgradeHook]:
"""Register a callable to execute after upgrading to ``revision``."""
def decorator(func: AfterUpgradeHook) -> AfterUpgradeHook:
hooks = _REGISTRY.setdefault(revision, RevisionHooks())
if hooks.after_upgrade is not None:
raise ValueError(
f"after_upgrade hook already registered for revision {revision}"
)
hooks.after_upgrade = func
return func
return decorator
def get_registered_hooks() -> Mapping[str, RevisionHooks]:
"""Return a snapshot of the registered revision hooks."""
return dict(_REGISTRY)
__all__ = [
"RevisionHooks",
"register_before_upgrade",
"register_after_upgrade",
"get_registered_hooks",
]

View File

@ -0,0 +1,35 @@
"""Register revision-specific hooks for migration verification."""
from . import (
test_05486ce795d5_make_session_name_required_on_messages,
test_08894082221a_replace_collection_name_with_observer_,
test_20f89a421aff_rename_metamessage_type_to_label,
test_66e63cf2cf77_add_indexes_to_documents_table,
test_76ffba56fe8c_add_error_field_to_queueitem,
test_88b0fb10906f_add_webhooks_table,
test_556a16564f50_add_user_id_and_app_id_to_tables,
test_564ba40505c5_add_session_name_column_to_documents,
test_917195d9b5e9_add_messageembedding_table,
test_a1b2c3d4e5f6_initial_schema,
test_b765d82110bd_change_metamessages_to_user_level_with_,
test_bb6fb3a7a643_add_message_seq_in_session_column,
test_c3828084f472_add_indexes_for_messages_and_,
test_d429de0e5338_adopt_peer_paradigm,
)
__all__ = [
"test_05486ce795d5_make_session_name_required_on_messages",
"test_08894082221a_replace_collection_name_with_observer_",
"test_20f89a421aff_rename_metamessage_type_to_label",
"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",
"test_76ffba56fe8c_add_error_field_to_queueitem",
"test_88b0fb10906f_add_webhooks_table",
"test_917195d9b5e9_add_messageembedding_table",
"test_a1b2c3d4e5f6_initial_schema",
"test_b765d82110bd_change_metamessages_to_user_level_with_",
"test_bb6fb3a7a643_add_message_seq_in_session_column",
"test_c3828084f472_add_indexes_for_messages_and_",
"test_d429de0e5338_adopt_peer_paradigm",
]

View File

@ -0,0 +1,104 @@
"""Hooks for revision 05486ce795d5 (session_name required)."""
from __future__ import annotations
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
WORKSPACE_ID = generate_nanoid()
WORKSPACE_NAME = "workspace-name"
PEER_ID = generate_nanoid()
PEER_NAME = "peer-name"
SESSION_ID = generate_nanoid()
SESSION_NAME = "session-name"
MESSAGE_ID = generate_nanoid()
MESSAGE_CONTENT = "message-content"
@register_before_upgrade("05486ce795d5")
def seed_orphaned_message(verifier: MigrationVerifier) -> None:
"""Insert a message without a session_name to exercise the migration."""
schema = verifier.schema
connection = verifier.conn
connection.execute(
text(
f"""
INSERT INTO "{schema}"."workspaces" ("id", "name")
VALUES (:workspace_id, :workspace_name)
"""
),
{"workspace_id": WORKSPACE_ID, "workspace_name": WORKSPACE_NAME},
)
connection.execute(
text(
f"""
INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name")
VALUES (:peer_id, :peer_name, :workspace_name)
"""
),
{
"peer_id": PEER_ID,
"peer_name": PEER_NAME,
"workspace_name": WORKSPACE_NAME,
},
)
connection.execute(
text(
f"""
INSERT INTO "{schema}"."messages"
("public_id", "workspace_name", "peer_name", "session_name", "content")
VALUES (:public_id, :workspace_name, :peer_name, NULL, :content)
"""
),
{
"public_id": MESSAGE_ID,
"workspace_name": WORKSPACE_NAME,
"peer_name": PEER_NAME,
"content": MESSAGE_CONTENT,
},
)
@register_after_upgrade("05486ce795d5")
def verify_session_name_enforced(verifier: MigrationVerifier) -> None:
"""Ensure session_name is populated and constrained after the migration."""
verifier.assert_column_exists("messages", "session_name", nullable=False)
schema = verifier.schema
conn = verifier.conn
session = conn.execute(
text(
f'SELECT "name", "workspace_name" FROM "{schema}"."sessions" '
+ 'WHERE "name" = :session_name AND "workspace_name" = :workspace_name'
),
{
"session_name": PEER_NAME,
"workspace_name": WORKSPACE_NAME,
},
).one_or_none()
assert (
session is not None
), "Session with expected name and workspace does not exist"
message = conn.execute(
text(
f'SELECT "public_id", "content" FROM "{schema}"."messages" '
+ 'WHERE "session_name" = :session_name AND "workspace_name" = :workspace_name AND "peer_name" = :peer_name'
),
{
"session_name": PEER_NAME,
"workspace_name": WORKSPACE_NAME,
"peer_name": PEER_NAME,
},
).one_or_none()
assert message is not None, "Message with expected session_name does not exist"
assert message.content == MESSAGE_CONTENT

View File

@ -0,0 +1,361 @@
"""Hooks for revision 08894082221a (observer/observed refactor)."""
from __future__ import annotations
from nanoid import generate as generate_nanoid
from sqlalchemy import inspect, text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
COLLECTION_ID: str = generate_nanoid()
GLOBAL_REP_COLLECTION_NAME = "global_representation"
DOCUMENT_ID = generate_nanoid()
DOCUMENT_NAME = "document-name"
LEGACY_PEER_ID = generate_nanoid()
LEGACY_PEER_NAME = "legacy-peer"
WORKSPACE_ID = generate_nanoid()
WORKSPACE_NAME = "workspace-name"
# Based on migration logic: when name='global_representation', observer=peer_name and observed=peer_name
EXPECTED_OBSERVED = LEGACY_PEER_NAME
@register_before_upgrade("08894082221a")
def prepare_observer_observed(verifier: MigrationVerifier) -> None:
schema = verifier.schema
conn = verifier.conn
# Create workspace
conn.execute(
text(
f'INSERT INTO "{schema}"."workspaces" ("id", "name") '
+ "VALUES (:workspace_id, :workspace_name)"
),
{"workspace_id": WORKSPACE_ID, "workspace_name": WORKSPACE_NAME},
)
# Create peer
conn.execute(
text(
f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") '
+ "VALUES (:peer_id, :peer_name, :workspace_name)"
),
{
"peer_id": LEGACY_PEER_ID,
"peer_name": LEGACY_PEER_NAME,
"workspace_name": WORKSPACE_NAME,
},
)
# Create collection with the old schema (has 'name' field)
conn.execute(
text(
f'INSERT INTO "{schema}"."collections" '
+ '("id", "name", "peer_name", "workspace_name") '
+ "VALUES (:collection_id, :collection_name, :peer_name, :workspace_name)"
),
{
"collection_id": COLLECTION_ID,
"collection_name": GLOBAL_REP_COLLECTION_NAME,
"peer_name": LEGACY_PEER_NAME,
"workspace_name": WORKSPACE_NAME,
},
)
# Create document with the old schema (has 'collection_name' field) without session_name
conn.execute(
text(
f'INSERT INTO "{schema}"."documents" '
+ '("id", "collection_name", "peer_name", "workspace_name", "content") '
+ "VALUES (:document_id, :collection_name, :peer_name, :workspace_name, :content)"
),
{
"document_id": DOCUMENT_ID,
"collection_name": GLOBAL_REP_COLLECTION_NAME,
"peer_name": LEGACY_PEER_NAME,
"workspace_name": WORKSPACE_NAME,
"content": "test content",
},
)
# Bulk create many documents with the old schema (no session_name)
bulk_stmt = text(
f'INSERT INTO "{schema}"."documents" '
+ '("id", "collection_name", "peer_name", "workspace_name", "content") '
+ "VALUES (:document_id, :collection_name, :peer_name, :workspace_name, :content)"
)
params_template = {
"collection_name": GLOBAL_REP_COLLECTION_NAME,
"peer_name": LEGACY_PEER_NAME,
"workspace_name": WORKSPACE_NAME,
"content": "test content",
}
total_docs = 120_000
batch_size = 10_000
for start in range(0, total_docs, batch_size):
end = min(start + batch_size, total_docs)
batch_params = [
{"document_id": generate_nanoid(), **params_template}
for _ in range(start, end)
]
conn.execute(bulk_stmt, batch_params)
# Speed up bulk seeding
conn.execute(text("SET LOCAL synchronous_commit = OFF"))
# -----------------------------
# Group 1: self-observation collections (global_representation)
# observer = peer_name, observed = peer_name
# -----------------------------
peer_insert_stmt = text(
f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") '
+ "VALUES (:peer_id, :peer_name, :workspace_name)"
)
collection_insert_stmt = text(
f'INSERT INTO "{schema}"."collections" ("id", "name", "peer_name", "workspace_name") '
+ "VALUES (:collection_id, :name, :peer_name, :workspace_name)"
)
total = 100_000
batch_size = 10_000
for start in range(0, total, batch_size):
end = min(start + batch_size, total)
# Insert peers for group 1
peers_params = [
{
"peer_id": generate_nanoid(),
"peer_name": f"selfpeer_{i}",
"workspace_name": WORKSPACE_NAME,
}
for i in range(start, end)
]
conn.execute(peer_insert_stmt, peers_params)
# Insert collections for group 1
collections_params = [
{
"collection_id": generate_nanoid(),
"name": GLOBAL_REP_COLLECTION_NAME,
"peer_name": f"selfpeer_{i}",
"workspace_name": WORKSPACE_NAME,
}
for i in range(start, end)
]
conn.execute(collection_insert_stmt, collections_params)
# -----------------------------
# Group 2: prefix pattern (observer_observed)
# observer = 'prefix_observer', observed extracted from name suffix
# -----------------------------
# Create the fixed observer peer
conn.execute(
peer_insert_stmt,
{
"peer_id": generate_nanoid(),
"peer_name": "prefix_observer",
"workspace_name": WORKSPACE_NAME,
},
)
for start in range(0, total, batch_size):
end = min(start + batch_size, total)
# Observed peers for group 2
peers_params = [
{
"peer_id": generate_nanoid(),
"peer_name": f"obs_prefix_{i}",
"workspace_name": WORKSPACE_NAME,
}
for i in range(start, end)
]
conn.execute(peer_insert_stmt, peers_params)
# Collections for group 2 (name = 'prefix_observer_obs_prefix_<i>')
collections_params = [
{
"collection_id": generate_nanoid(),
"name": f"prefix_observer_obs_prefix_{i}",
"peer_name": "prefix_observer",
"workspace_name": WORKSPACE_NAME,
}
for i in range(start, end)
]
conn.execute(collection_insert_stmt, collections_params)
# -----------------------------
# Group 3: suffix pattern (observed_observer)
# observer = 'suffix_observer', observed extracted from name prefix
# -----------------------------
# Create the fixed observer peer
conn.execute(
peer_insert_stmt,
{
"peer_id": generate_nanoid(),
"peer_name": "suffix_observer",
"workspace_name": WORKSPACE_NAME,
},
)
for start in range(0, total, batch_size):
end = min(start + batch_size, total)
# Observed peers for group 3
peers_params = [
{
"peer_id": generate_nanoid(),
"peer_name": f"obs_suffix_{i}",
"workspace_name": WORKSPACE_NAME,
}
for i in range(start, end)
]
conn.execute(peer_insert_stmt, peers_params)
# Collections for group 3 (name = 'obs_suffix_<i>_suffix_observer')
collections_params = [
{
"collection_id": generate_nanoid(),
"name": f"obs_suffix_{i}_suffix_observer",
"peer_name": "suffix_observer",
"workspace_name": WORKSPACE_NAME,
}
for i in range(start, end)
]
conn.execute(collection_insert_stmt, collections_params)
@register_after_upgrade("08894082221a")
def verify_observer_observed_migration(verifier: MigrationVerifier) -> None:
"""Assert that collections/documents now rely on observer/observed fields."""
inspector = inspect(verifier.conn)
collection_columns = {
col["name"]
for col in inspector.get_columns("collections", schema=verifier.schema)
}
assert "observer" in collection_columns
assert "observed" in collection_columns
assert "name" not in collection_columns
document_columns = {
col["name"]
for col in inspector.get_columns("documents", schema=verifier.schema)
}
assert "observer" in document_columns
assert "observed" in document_columns
assert "collection_name" not in document_columns
verifier.assert_indexes_exist(
[
("collections", "idx_collections_observer"),
("collections", "idx_collections_observed"),
("documents", "idx_documents_observer"),
("documents", "idx_documents_observed"),
]
)
verifier.assert_constraint_exists(
"collections", "unique_observer_observed_collection", "unique"
)
verifier.assert_constraint_exists(
"documents", "documents_observer_observed_workspace_name_fkey", "foreign_key"
)
collection = verifier.conn.execute(
text(
'SELECT "observer", "observed", "workspace_name" '
+ f'FROM "{verifier.schema}"."collections" '
+ 'WHERE "id" = :collection_id'
),
{"collection_id": COLLECTION_ID},
).one()
assert collection.observer == LEGACY_PEER_NAME
assert collection.observed == EXPECTED_OBSERVED
assert collection.workspace_name == WORKSPACE_NAME
document = verifier.conn.execute(
text(
'SELECT "observer", "observed", "workspace_name", "session_name" '
+ f'FROM "{verifier.schema}"."documents" '
+ 'WHERE "id" = :document_id'
),
{"document_id": DOCUMENT_ID},
).one()
assert document.observer == LEGACY_PEER_NAME
assert document.observed == EXPECTED_OBSERVED
assert document.workspace_name == WORKSPACE_NAME
assert document.session_name == "__global_observations__"
# Verify no NULLs remain in documents.session_name
verifier.assert_no_nulls("documents", "session_name")
# Verify collections mapping for Group 1 (self-observation)
count_self = verifier.conn.execute(
text(
"SELECT COUNT(1) FROM "
+ f'"{verifier.schema}"."collections" '
+ 'WHERE "workspace_name" = :ws '
+ 'AND "observer" LIKE :prefix '
+ 'AND "observed" = "observer"'
),
{"ws": WORKSPACE_NAME, "prefix": "selfpeer_%"},
).scalar()
assert count_self == 100_000
# Verify collections mapping for Group 2 (prefix pattern observer_observed)
count_prefix = verifier.conn.execute(
text(
"SELECT COUNT(1) FROM "
+ f'"{verifier.schema}"."collections" '
+ 'WHERE "workspace_name" = :ws '
+ 'AND "observer" = :observer '
+ 'AND "observed" LIKE :obs_prefix'
),
{
"ws": WORKSPACE_NAME,
"observer": "prefix_observer",
"obs_prefix": "obs_prefix_%",
},
).scalar()
assert count_prefix == 100_000
distinct_prefix_observed = verifier.conn.execute(
text(
'SELECT COUNT(DISTINCT "observed") FROM '
+ f'"{verifier.schema}"."collections" '
+ 'WHERE "workspace_name" = :ws AND "observer" = :observer'
),
{"ws": WORKSPACE_NAME, "observer": "prefix_observer"},
).scalar()
assert distinct_prefix_observed == 100_000
# Verify collections mapping for Group 3 (suffix pattern observed_observer)
count_suffix = verifier.conn.execute(
text(
"SELECT COUNT(1) FROM "
+ f'"{verifier.schema}"."collections" '
+ 'WHERE "workspace_name" = :ws '
+ 'AND "observer" = :observer '
+ 'AND "observed" LIKE :obs_prefix'
),
{
"ws": WORKSPACE_NAME,
"observer": "suffix_observer",
"obs_prefix": "obs_suffix_%",
},
).scalar()
assert count_suffix == 100_000
distinct_suffix_observed = verifier.conn.execute(
text(
'SELECT COUNT(DISTINCT "observed") FROM '
+ f'"{verifier.schema}"."collections" '
+ 'WHERE "workspace_name" = :ws AND "observer" = :observer'
),
{"ws": WORKSPACE_NAME, "observer": "suffix_observer"},
).scalar()
assert distinct_suffix_observed == 100_000

View File

@ -0,0 +1,123 @@
"""Hooks for revision 20f89a421aff (metamessage label rename)."""
from __future__ import annotations
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
APP_ID = generate_nanoid()
USER_ID = generate_nanoid()
SESSION_ID = generate_nanoid()
MESSAGE_ID = generate_nanoid()
METAMESSAGE_ID = generate_nanoid()
@register_before_upgrade("20f89a421aff")
def prepare_metamessage_label(verifier: MigrationVerifier) -> None:
OLD_INDEXES = (
("metamessages", "idx_metamessages_lookup"),
("metamessages", "idx_metamessages_user_lookup"),
("metamessages", "idx_metamessages_session_lookup"),
("metamessages", "idx_metamessages_message_lookup"),
)
verifier.assert_indexes_exist(OLD_INDEXES)
verifier.assert_constraint_exists(
"metamessages", "metamessage_type_length", "check"
)
schema = verifier.schema
connection = verifier.conn
inspector = verifier.get_inspector()
connection.execute(
text(
f'INSERT INTO "{schema}"."apps" ("public_id", "name") '
+ "VALUES (:id, :name)"
),
{"id": APP_ID, "name": "verify-app"},
)
connection.execute(
text(
f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") '
+ "VALUES (:id, :name, :app_id)"
),
{"id": USER_ID, "name": "verify-user", "app_id": APP_ID},
)
connection.execute(
text(
f'INSERT INTO "{schema}"."sessions" '
+ '("public_id", "user_id", "app_id", "is_active") '
+ "VALUES (:id, :user_id, :app_id, true)"
),
{"id": SESSION_ID, "user_id": USER_ID, "app_id": APP_ID},
)
connection.execute(
text(
f'INSERT INTO "{schema}"."messages" '
+ '("public_id", "session_id", "app_id", "user_id", "is_user", "content") '
+ "VALUES (:id, :session_id, :app_id, :user_id, true, :content)"
),
{
"id": MESSAGE_ID,
"session_id": SESSION_ID,
"app_id": APP_ID,
"user_id": USER_ID,
"content": "legacy metamessage seed",
},
)
connection.execute(
text(
f'INSERT INTO "{schema}"."metamessages" '
+ '("public_id", "metamessage_type", "content", "message_id", "user_id", "session_id", "app_id") '
+ "VALUES (:id, :type, :content, :message_id, :user_id, :session_id, :app_id)"
),
{
"id": METAMESSAGE_ID,
"type": "seed",
"content": "legacy metamessage",
"message_id": MESSAGE_ID,
"user_id": USER_ID,
"session_id": SESSION_ID,
"app_id": APP_ID,
},
)
columns = {
col["name"] for col in inspector.get_columns("metamessages", schema=schema)
}
assert "metamessage_type" in columns
assert "label" not in columns
@register_after_upgrade("20f89a421aff")
def verify_metamessage_label(verifier: MigrationVerifier) -> None:
NEW_INDEXES = (
("metamessages", "idx_metamessages_lookup"),
("metamessages", "idx_metamessages_user_lookup"),
("metamessages", "idx_metamessages_session_lookup"),
("metamessages", "idx_metamessages_message_lookup"),
)
verifier.assert_indexes_exist(NEW_INDEXES)
verifier.assert_column_exists("metamessages", "label", nullable=False)
verifier.assert_constraint_exists("metamessages", "label_length", "check")
verifier.assert_constraint_exists(
"metamessages", "metamessage_type_length", "check", exists=False
)
verifier.assert_column_exists("metamessages", "metamessage_type", exists=False)
row = verifier.conn.execute(
text(
f'SELECT "label" FROM "{verifier.schema}"."metamessages" '
+ 'WHERE "public_id" = :public_id'
),
{"public_id": METAMESSAGE_ID},
).one()
assert row.label == "seed"

View File

@ -0,0 +1,185 @@
"""Hooks for revision 556a16564f50 (propagate app/user identifiers)."""
from __future__ import annotations
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
APP_ID = generate_nanoid()
USER_ID = generate_nanoid()
SESSION_ID = generate_nanoid()
MESSAGE_ID = generate_nanoid()
METAMESSAGE_ID = generate_nanoid()
COLLECTION_ID = generate_nanoid()
DOCUMENT_ID = generate_nanoid()
@register_before_upgrade("556a16564f50")
def prepare_user_app_ids(verifier: MigrationVerifier) -> None:
schema = verifier.schema
connection = verifier.conn
# Create app
connection.execute(
text(
f'INSERT INTO "{schema}"."apps" ("public_id", "name") '
+ "VALUES (:app_id, :name)"
),
{"app_id": APP_ID, "name": "test-app"},
)
# Create user
connection.execute(
text(
f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") '
+ "VALUES (:user_id, :name, :app_id)"
),
{"user_id": USER_ID, "name": "test-user", "app_id": APP_ID},
)
# Create session
connection.execute(
text(
f'INSERT INTO "{schema}"."sessions" ("public_id", "user_id", "is_active") '
+ "VALUES (:session_id, :user_id, true)"
),
{"session_id": SESSION_ID, "user_id": USER_ID},
)
# Create message
connection.execute(
text(
f'INSERT INTO "{schema}"."messages" ("public_id", "session_id", "is_user", "content") '
+ "VALUES (:message_id, :session_id, true, :content)"
),
{
"message_id": MESSAGE_ID,
"session_id": SESSION_ID,
"content": "test-content",
},
)
# Create metamessage
connection.execute(
text(
f'INSERT INTO "{schema}"."metamessages" ("public_id", "user_id", "content", "metamessage_type") '
+ "VALUES (:metamessage_id, :user_id, :content, :metamessage_type)"
),
{
"metamessage_id": METAMESSAGE_ID,
"user_id": USER_ID,
"content": "test-content",
"metamessage_type": "test-type",
},
)
# Create collection
connection.execute(
text(
f'INSERT INTO "{schema}"."collections" ("public_id", "name", "user_id") '
+ "VALUES (:collection_id, :name, :user_id)"
),
{
"collection_id": COLLECTION_ID,
"name": "test-collection",
"user_id": USER_ID,
},
)
# Create document
connection.execute(
text(
f'INSERT INTO "{schema}"."documents" ("public_id", "content", "collection_id") '
+ "VALUES (:document_id, :content, :collection_id)"
),
{
"document_id": DOCUMENT_ID,
"content": "test-document-content",
"collection_id": COLLECTION_ID,
},
)
@register_after_upgrade("556a16564f50")
def verify_app_and_user_ids(verifier: MigrationVerifier) -> None:
"""Ensure app/user identifiers are present after the upgrade."""
verifier.assert_column_exists("sessions", "app_id", nullable=False)
verifier.assert_indexes_exist([("sessions", "ix_sessions_app_id")])
verifier.assert_column_exists("messages", "app_id", nullable=False)
verifier.assert_column_exists("messages", "user_id", nullable=False)
verifier.assert_indexes_exist(
[
("messages", "ix_messages_app_id"),
("messages", "ix_messages_user_id"),
]
)
verifier.assert_column_exists("metamessages", "app_id", nullable=False)
verifier.assert_indexes_exist([("metamessages", "ix_metamessages_app_id")])
verifier.assert_column_exists("collections", "app_id", nullable=False)
verifier.assert_indexes_exist([("collections", "ix_collections_app_id")])
verifier.assert_column_exists("documents", "app_id", nullable=False)
verifier.assert_column_exists("documents", "user_id", nullable=False)
verifier.assert_indexes_exist(
[
("documents", "ix_documents_app_id"),
("documents", "ix_documents_user_id"),
]
)
conn = verifier.conn
schema = verifier.schema
session_row = conn.execute(
text(
f'SELECT "app_id" FROM "{schema}"."sessions" '
+ 'WHERE "public_id" = :session_public_id'
),
{"session_public_id": SESSION_ID},
).one()
assert session_row.app_id == APP_ID
message_row = conn.execute(
text(
f'SELECT "app_id", "user_id" FROM "{schema}"."messages" '
+ 'WHERE "public_id" = :message_public_id'
),
{"message_public_id": MESSAGE_ID},
).one()
assert message_row.app_id == APP_ID
assert message_row.user_id == USER_ID
metamessage_row = conn.execute(
text(
f'SELECT "app_id" FROM "{schema}"."metamessages" '
+ 'WHERE "public_id" = :public_id'
),
{"public_id": METAMESSAGE_ID},
).one()
assert metamessage_row.app_id == APP_ID
collection_row = conn.execute(
text(
f'SELECT "app_id" FROM "{schema}"."collections" '
+ 'WHERE "public_id" = :public_id'
),
{"public_id": COLLECTION_ID},
).one()
assert collection_row.app_id == APP_ID
document_row = conn.execute(
text(
f'SELECT "app_id", "user_id" FROM "{schema}"."documents" '
+ 'WHERE "public_id" = :public_id'
),
{"public_id": DOCUMENT_ID},
).one()
assert document_row.app_id == APP_ID
assert document_row.user_id == USER_ID

View File

@ -0,0 +1,211 @@
"""Hooks for revision 564ba40505c5 (documents session_name column)."""
from __future__ import annotations
import json
import time
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
WORKSPACE_ID = generate_nanoid()
WORKSPACE_NAME = generate_nanoid()
PEER_ID = generate_nanoid()
PEER_NAME = generate_nanoid()
SESSION_ID = generate_nanoid()
SESSION_NAME = generate_nanoid()
COLLECTION_ID = generate_nanoid()
COLLECTION_NAME = generate_nanoid()
DOCUMENT_ID = generate_nanoid()
@register_before_upgrade("564ba40505c5")
def seed_document_session_metadata(verifier: MigrationVerifier) -> None:
schema = verifier.schema
connection = verifier.conn
verifier.assert_column_exists("documents", "session_name", exists=False)
verifier.assert_indexes_not_exist([("documents", "idx_documents_session_name")])
verifier.assert_constraint_exists(
"documents", "fk_documents_session_workspace", "foreign_key", exists=False
)
# Seed workspaces, peers, and collections
connection.execute(
text(
f"""
INSERT INTO "{schema}"."workspaces" ("id", "name")
VALUES (:workspace_id, :workspace_name)
"""
),
{"workspace_id": WORKSPACE_ID, "workspace_name": WORKSPACE_NAME},
)
connection.execute(
text(
f"""
INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name")
VALUES (:peer_id, :peer_name, :workspace_name)
"""
),
{
"peer_id": PEER_ID,
"peer_name": PEER_NAME,
"workspace_name": WORKSPACE_NAME,
},
)
connection.execute(
text(
f"""
INSERT INTO "{schema}"."sessions"
("id", "name", "workspace_name", "is_active")
VALUES (:session_id, :session_name, :workspace_name, true)
"""
),
{
"session_id": SESSION_ID,
"session_name": SESSION_NAME,
"workspace_name": WORKSPACE_NAME,
},
)
# Create session_peers entry so downgrade can find user_id
connection.execute(
text(
f"""
INSERT INTO "{schema}"."session_peers"
("workspace_name", "session_name", "peer_name")
VALUES (:workspace_name, :session_name, :peer_name)
"""
),
{
"workspace_name": WORKSPACE_NAME,
"session_name": SESSION_NAME,
"peer_name": PEER_NAME,
},
)
connection.execute(
text(
f"""
INSERT INTO "{schema}"."collections"
("id", "name", "peer_name", "workspace_name")
VALUES (:collection_id, :collection_name, :peer_name, :workspace_name)
"""
),
{
"collection_id": COLLECTION_ID,
"collection_name": COLLECTION_NAME,
"peer_name": PEER_NAME,
"workspace_name": WORKSPACE_NAME,
},
)
# Bulk insert 500k documents with session_name in internal_metadata to test batched backfill
print("[564ba40505c5] Starting bulk insert of 500k documents...", flush=True)
t0 = time.perf_counter()
connection.execute(text("SET LOCAL synchronous_commit = OFF"))
connection.execute(
text(
f'INSERT INTO "{schema}"."documents" '
+ '("id", "collection_name", "peer_name", "workspace_name", "content", "internal_metadata") '
+ "SELECT "
+ " 'bulk-' || substr(md5(random()::text), 1, 16), "
+ f" '{COLLECTION_NAME}', "
+ f" '{PEER_NAME}', "
+ f" '{WORKSPACE_NAME}', "
+ " 'seed content', "
+ f" jsonb_build_object('session_name', '{SESSION_NAME}') "
+ "FROM generate_series(1, :n)"
),
{"n": 500_000},
)
t1 = time.perf_counter()
print(f"[564ba40505c5] Bulk insert completed in {t1 - t0:.2f}s", flush=True)
# Add document with session_name in internal_metadata
internal_metadata = json.dumps(
{
"session_name": SESSION_NAME,
}
)
connection.execute(
text(
f"""
INSERT INTO "{schema}"."documents"
(
"id",
"collection_name",
"peer_name",
"workspace_name",
"content",
"internal_metadata"
)
VALUES (
:document_id,
:collection_name,
:peer_name,
:workspace_name,
:content,
:internal_metadata
)
"""
),
{
"document_id": DOCUMENT_ID,
"collection_name": COLLECTION_NAME,
"peer_name": PEER_NAME,
"workspace_name": WORKSPACE_NAME,
"content": "seed content",
"internal_metadata": internal_metadata,
},
)
@register_after_upgrade("564ba40505c5")
def verify_document_session_column(verifier: MigrationVerifier) -> None:
"""Ensure the session_name column reflects migrated metadata."""
verifier.assert_column_exists("documents", "session_name", nullable=True)
verifier.assert_indexes_exist([("documents", "idx_documents_session_name")])
verifier.assert_constraint_exists(
"documents", "fk_documents_session_workspace", "foreign_key"
)
# Quick diagnostics: counts before assertions
total_docs = (
verifier.conn.execute(
text(f'SELECT COUNT(*) FROM "{verifier.schema}"."documents"')
).scalar()
or 0
)
null_sessions = (
verifier.conn.execute(
text(
f'SELECT COUNT(*) FROM "{verifier.schema}"."documents" '
+ 'WHERE "session_name" IS NULL'
)
).scalar()
or 0
)
print(
f"[564ba40505c5] Documents total={total_docs}, session_name NULLs={null_sessions}",
flush=True,
)
# All documents that contained session_name in internal_metadata should now have session_name populated
verifier.assert_no_nulls("documents", "session_name")
# Sanity-check seeded document preserved expected session_name
row = verifier.conn.execute(
text(
f'SELECT "session_name" FROM "{verifier.schema}"."documents" '
+ 'WHERE "id" = :document_id'
),
{"document_id": DOCUMENT_ID},
).one()
assert row.session_name == SESSION_NAME

View File

@ -0,0 +1,16 @@
"""Hooks for revision 66e63cf2cf77 (documents HNSW index)."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("66e63cf2cf77")
def prepare_documents_hnsw(verifier: MigrationVerifier) -> None:
verifier.assert_indexes_not_exist([("documents", "idx_documents_embedding_hnsw")])
@register_after_upgrade("66e63cf2cf77")
def verify_documents_hnsw(verifier: MigrationVerifier) -> None:
verifier.assert_indexes_exist([("documents", "idx_documents_embedding_hnsw")])

View File

@ -0,0 +1,44 @@
"""Hooks for revision 76ffba56fe8c (queue created_at/error columns)."""
from __future__ import annotations
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("76ffba56fe8c")
def prepare_queue_created_at(verifier: MigrationVerifier) -> None:
verifier.assert_column_exists("queue", "created_at", nullable=False, exists=False)
verifier.assert_column_exists("queue", "error", exists=False)
verifier.assert_indexes_not_exist([("queue", "ix_queue_created_at")])
conn = verifier.conn
schema = verifier.schema
# Bulk insert 500k additional queue items efficiently to exercise batched backfill
conn.execute(text("SET LOCAL synchronous_commit = OFF"))
conn.execute(
text(
f'INSERT INTO "{schema}"."queue" '
+ '("session_id", "work_unit_key", "task_type", "payload", "processed") '
+ "SELECT NULL, "
+ " 'seed-batch-queue-item-' || gs::text, "
+ " 'representation', "
+ " jsonb_build_object('marker','bulk'), "
+ " false "
+ "FROM generate_series(1, :n) AS gs"
),
{"n": 500_000},
)
@register_after_upgrade("76ffba56fe8c")
def verify_queue_created_at(verifier: MigrationVerifier) -> None:
verifier.assert_column_exists("queue", "created_at", nullable=False)
verifier.assert_column_exists("queue", "error")
verifier.assert_indexes_exist([("queue", "ix_queue_created_at")])
# Ensure all rows have non-null created_at after backfill
verifier.assert_no_nulls("queue", "created_at")

View File

@ -0,0 +1,115 @@
"""Hooks for revision 88b0fb10906f (webhooks and queue updates)."""
from __future__ import annotations
import json
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
PAYLOAD_MARKER = generate_nanoid()
WORKSPACE_NAME = generate_nanoid()
SESSION_NAME = generate_nanoid()
SENDER_NAME = generate_nanoid()
TARGET_NAME = generate_nanoid()
@register_before_upgrade("88b0fb10906f")
def prepare_webhooks(verifier: MigrationVerifier) -> None:
# Validate objects introduced by the migration are absent
verifier.assert_table_exists("webhook_endpoints", exists=False)
verifier.assert_column_exists("queue", "task_type", exists=False)
verifier.assert_column_exists("queue", "work_unit_key", exists=False)
verifier.assert_column_exists(
"active_queue_sessions", "work_unit_key", exists=False
)
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_work_unit_key", "unique", exists=False
)
# Validate objects removed by the migration still exist
verifier.assert_column_exists("active_queue_sessions", "session_id")
verifier.assert_column_exists("active_queue_sessions", "sender_name")
verifier.assert_column_exists("active_queue_sessions", "target_name")
verifier.assert_column_exists("active_queue_sessions", "task_type")
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_active_queue_session", "unique"
)
# Bulk insert 500k queue items to test batched backfill
verifier.conn.execute(text("SET LOCAL synchronous_commit = OFF"))
verifier.conn.execute(
text(
f'INSERT INTO "{verifier.schema}"."queue" ("payload") '
+ "SELECT jsonb_build_object("
+ " 'marker', 'bulk-' || substr(md5(random()::text), 1, 21),"
+ " 'task_type', 'representation',"
+ f" 'workspace_name', '{WORKSPACE_NAME}',"
+ f" 'session_name', '{SESSION_NAME}',"
+ f" 'sender_name', '{SENDER_NAME}',"
+ f" 'target_name', '{TARGET_NAME}'"
+ ") FROM generate_series(1, :n)"
),
{"n": 500_000},
)
payload = json.dumps(
{
"marker": PAYLOAD_MARKER,
"task_type": "representation",
"workspace_name": WORKSPACE_NAME,
"session_name": SESSION_NAME,
"sender_name": SENDER_NAME,
"target_name": TARGET_NAME,
}
)
# Add single queue item to test expected data transformation
verifier.conn.execute(
text(f'INSERT INTO "{verifier.schema}"."queue" ("payload") VALUES (:payload)'),
{"payload": payload},
)
@register_after_upgrade("88b0fb10906f")
def verify_webhooks_and_queue(verifier: MigrationVerifier) -> None:
"""Validate queue enrichment and webhook table creation."""
verifier.assert_table_exists("webhook_endpoints")
verifier.assert_column_exists("queue", "task_type", nullable=False)
verifier.assert_column_exists("queue", "work_unit_key", nullable=False)
verifier.assert_column_exists("active_queue_sessions", "work_unit_key")
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_work_unit_key", "unique"
)
verifier.assert_column_exists("active_queue_sessions", "session_id", exists=False)
verifier.assert_column_exists("active_queue_sessions", "sender_name", exists=False)
verifier.assert_column_exists("active_queue_sessions", "target_name", exists=False)
verifier.assert_column_exists("active_queue_sessions", "task_type", exists=False)
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_active_queue_session", "unique", exists=False
)
# Ensure no NULLs remain after backfill across all rows
verifier.assert_no_nulls("queue", "task_type")
verifier.assert_no_nulls("queue", "work_unit_key")
row = verifier.conn.execute(
text(
f'SELECT "task_type", "work_unit_key" FROM "{verifier.schema}"."queue" '
+ "WHERE payload->>'marker' = :marker"
),
{"marker": PAYLOAD_MARKER},
).one()
assert row.task_type == "representation"
expected_key = (
f"representation:{WORKSPACE_NAME}:{SESSION_NAME}:{SENDER_NAME}:{TARGET_NAME}"
)
assert row.work_unit_key == expected_key

View File

@ -0,0 +1,26 @@
"""Hooks for revision 917195d9b5e9 (message embeddings table)."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
INDEXES = (
("message_embeddings", "idx_message_embeddings_message_id"),
("message_embeddings", "idx_message_embeddings_workspace_name"),
("message_embeddings", "idx_message_embeddings_session_name"),
("message_embeddings", "idx_message_embeddings_peer_name"),
("message_embeddings", "idx_message_embeddings_created_at"),
("message_embeddings", "idx_message_embeddings_embedding_hnsw"),
)
@register_before_upgrade("917195d9b5e9")
def prepare_message_embeddings(verifier: MigrationVerifier) -> None:
verifier.assert_table_exists("message_embeddings", exists=False)
@register_after_upgrade("917195d9b5e9")
def verify_message_embeddings_table(verifier: MigrationVerifier) -> None:
verifier.assert_table_exists("message_embeddings")
verifier.assert_indexes_exist(INDEXES)

View File

@ -0,0 +1,39 @@
"""Hooks for initial schema revision a1b2c3d4e5f6."""
from __future__ import annotations
from nanoid import generate as generate_nanoid
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
APP_ID = generate_nanoid()
APP_NAME = "seed-app"
USER_ID = generate_nanoid()
USER_NAME = "seed-user"
SESSION_ID = generate_nanoid()
SESSION_NAME = "seed-session"
MESSAGE_PUBLIC_ID = generate_nanoid()
MESSAGE_CONTENT = "seed message"
METAMESSAGE_PUBLIC_ID = generate_nanoid()
METAMESSAGE_CONTENT = "seed metamessage"
COLLECTION_ID = generate_nanoid()
COLLECTION_NAME = "seed-collection"
DOCUMENT_ID = generate_nanoid()
DOCUMENT_CONTENT = "Seed document"
@register_before_upgrade("a1b2c3d4e5f6")
def prepare_initial(_verifier: MigrationVerifier) -> None:
pass
@register_after_upgrade("a1b2c3d4e5f6")
def verify_initial_schema(verifier: MigrationVerifier) -> None:
verifier.assert_table_exists("apps", exists=True)
verifier.assert_table_exists("users", exists=True)
verifier.assert_table_exists("sessions", exists=True)
verifier.assert_table_exists("messages", exists=True)
verifier.assert_table_exists("collections", exists=True)
verifier.assert_table_exists("documents", exists=True)
verifier.assert_table_exists("metamessages", exists=True)

View File

@ -0,0 +1,114 @@
"""Hooks for revision b765d82110bd (metamessages user-level migration)."""
from __future__ import annotations
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
METAMESSAGE_ID = generate_nanoid()
USER_ID = generate_nanoid()
SESSION_ID = generate_nanoid()
MESSAGE_ID = generate_nanoid()
_INDEXES = (
("metamessages", "idx_metamessages_user_lookup"),
("metamessages", "idx_metamessages_session_lookup"),
("metamessages", "idx_metamessages_message_lookup"),
)
@register_before_upgrade("b765d82110bd")
def prepare_metamessages(verifier: MigrationVerifier) -> None:
schema = verifier.schema
connection = verifier.conn
# Create app
APP_ID = generate_nanoid()
connection.execute(
text(
f'INSERT INTO "{schema}"."apps" ("public_id", "name") '
+ "VALUES (:app_id, :name)"
),
{"app_id": APP_ID, "name": "test-app"},
)
# Create user
connection.execute(
text(
f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") '
+ "VALUES (:user_id, :name, :app_id)"
),
{"user_id": USER_ID, "name": "test-user", "app_id": APP_ID},
)
# Create session
connection.execute(
text(
f'INSERT INTO "{schema}"."sessions" ("public_id", "user_id", "is_active") '
+ "VALUES (:session_id, :user_id, true)"
),
{"session_id": SESSION_ID, "user_id": USER_ID},
)
# Create message
connection.execute(
text(
f'INSERT INTO "{schema}"."messages" ("public_id", "session_id", "is_user", "content") '
+ "VALUES (:message_id, :session_id, true, :content)"
),
{
"message_id": MESSAGE_ID,
"session_id": SESSION_ID,
"content": "test-content",
},
)
# Create metamessage
connection.execute(
text(
f'INSERT INTO "{schema}"."metamessages" ("public_id", "message_id", "content", "metamessage_type") '
+ "VALUES (:metamessage_id, :message_id, :content, :metamessage_type)"
),
{
"metamessage_id": METAMESSAGE_ID,
"message_id": MESSAGE_ID,
"content": "test-content",
"metamessage_type": "test-type",
},
)
@register_after_upgrade("b765d82110bd")
def verify_metamessages_promoted(verifier: MigrationVerifier) -> None:
"""Confirm metamessages rows migrated to user-level shape."""
verifier.assert_column_exists("metamessages", "user_id", nullable=False)
verifier.assert_column_exists("metamessages", "session_id")
verifier.assert_column_exists("metamessages", "message_id", nullable=True)
verifier.assert_indexes_exist(_INDEXES)
verifier.assert_constraint_exists(
"metamessages", "fk_metamessages_user_id_users", "foreign_key"
)
verifier.assert_constraint_exists(
"metamessages", "fk_metamessages_session_id_sessions", "foreign_key"
)
verifier.assert_constraint_exists(
"metamessages", "message_requires_session", "check"
)
verifier.assert_no_nulls("metamessages", "user_id")
row = verifier.conn.execute(
text(
'SELECT "user_id", "session_id", "message_id", "metamessage_type", "content" '
+ f'FROM "{verifier.schema}"."metamessages" '
+ 'WHERE "public_id" = :public_id'
),
{"public_id": METAMESSAGE_ID},
).one()
assert row.user_id == USER_ID
assert row.session_id == SESSION_ID
assert row.message_id == MESSAGE_ID
assert row.metamessage_type == "test-type"
assert row.content == "test-content"

View File

@ -0,0 +1,16 @@
"""Hooks for revision bb6fb3a7a643 (add_message_seq_in_session_column)."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("bb6fb3a7a643")
def prepare_add_message_seq_in_session_column(_verifier: MigrationVerifier) -> None:
"""Seed state and assertions before upgrading to bb6fb3a7a643."""
@register_after_upgrade("bb6fb3a7a643")
def verify_add_message_seq_in_session_column(_verifier: MigrationVerifier) -> None:
"""Add assertions validating the effects of bb6fb3a7a643."""

View File

@ -0,0 +1,23 @@
"""Hooks for revision c3828084f472 (read indexes)."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
_INDEXES = (
("users", "idx_users_app_lookup"),
("sessions", "idx_sessions_user_lookup"),
("messages", "idx_messages_session_lookup"),
("metamessages", "idx_metamessages_lookup"),
)
@register_before_upgrade("c3828084f472")
def prepare_read_indexes(verifier: MigrationVerifier) -> None:
verifier.assert_indexes_not_exist(_INDEXES)
@register_after_upgrade("c3828084f472")
def verify_read_indexes(verifier: MigrationVerifier) -> None:
verifier.assert_indexes_exist(_INDEXES)

View File

@ -0,0 +1,286 @@
"""Hooks for revision d429de0e5338 (adopt peer paradigm)."""
from __future__ import annotations
import json
from nanoid import generate as generate_nanoid
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
APP_ID = generate_nanoid()
APP_NAME = "app-name"
USER_ID = generate_nanoid()
USER_NAME = "user-name"
SESSION_ID = generate_nanoid()
MESSAGE_ID = generate_nanoid()
MESSAGE_CONTENT = "message-content"
COLLECTION_ID = generate_nanoid()
COLLECTION_NAME = "collection-name"
DOCUMENT_ID = generate_nanoid()
QUEUE_MARKER = "seed"
@register_before_upgrade("d429de0e5338")
def prepare_peer_paradigm(verifier: MigrationVerifier) -> None:
"""Seed legacy tables so the migration can transform real data."""
schema = verifier.schema
conn = verifier.conn
conn.execute(
text(
f'INSERT INTO "{schema}"."apps" ("public_id", "name") '
+ "VALUES (:app_id, :app_name)"
),
{"app_id": APP_ID, "app_name": APP_NAME},
)
conn.execute(
text(
f'INSERT INTO "{schema}"."users" ("public_id", "name", "app_id") '
+ "VALUES (:user_id, :user_name, :app_id)"
),
{"user_id": USER_ID, "user_name": USER_NAME, "app_id": APP_ID},
)
session_result = conn.execute(
text(
f'INSERT INTO "{schema}"."sessions" '
+ '("public_id", "user_id", "app_id", "is_active") '
+ "VALUES (:session_id, :user_id, :app_id, true) RETURNING id"
),
{"session_id": SESSION_ID, "user_id": USER_ID, "app_id": APP_ID},
).one()
SESSION_DB_ID = session_result.id
conn.execute(
text(
f'INSERT INTO "{schema}"."messages" '
+ '("public_id", "session_id", "is_user", "content", "user_id", "app_id") '
+ "VALUES (:message_id, :session_id, true, :content, :user_id, :app_id)"
),
{
"message_id": MESSAGE_ID,
"session_id": SESSION_ID,
"content": MESSAGE_CONTENT,
"user_id": USER_ID,
"app_id": APP_ID,
},
)
conn.execute(
text(
f'INSERT INTO "{schema}"."collections" '
+ '("public_id", "name", "user_id", "app_id") '
+ "VALUES (:collection_id, :collection_name, :user_id, :app_id)"
),
{
"collection_id": COLLECTION_ID,
"collection_name": COLLECTION_NAME,
"user_id": USER_ID,
"app_id": APP_ID,
},
)
conn.execute(
text(
f'INSERT INTO "{schema}"."documents" '
+ '("public_id", "collection_id", "user_id", "app_id", "content") '
+ "VALUES (:document_id, :collection_id, :user_id, :app_id, :content)"
),
{
"document_id": DOCUMENT_ID,
"collection_id": COLLECTION_ID,
"user_id": USER_ID,
"app_id": APP_ID,
"content": "document-content",
},
)
queue_payload = json.dumps({"marker": QUEUE_MARKER})
conn.execute(
text(
f'INSERT INTO "{schema}"."queue" '
+ '("session_id", "payload", "processed") '
+ "VALUES (:session_db_id, :payload, false)"
),
{
"session_db_id": SESSION_DB_ID,
"payload": queue_payload,
},
)
@register_after_upgrade("d429de0e5338")
def verify_peer_paradigm(verifier: MigrationVerifier) -> None:
"""Validate the large-scale peer paradigm migration."""
inspector = verifier.get_inspector()
assert inspector.has_table("workspaces", schema=verifier.schema)
assert not inspector.has_table("apps", schema=verifier.schema)
verifier.assert_column_exists("workspaces", "configuration", nullable=False)
verifier.assert_column_exists("workspaces", "internal_metadata", nullable=False)
verifier.assert_column_exists("peers", "workspace_name", nullable=False)
verifier.assert_column_exists("peers", "configuration", nullable=False)
verifier.assert_column_exists("sessions", "name", nullable=False)
verifier.assert_column_exists("sessions", "workspace_name", nullable=False)
verifier.assert_column_exists(
"messages", "session_name", nullable=True
) # later changed back to required
verifier.assert_column_exists("messages", "workspace_name", nullable=False)
verifier.assert_column_exists("messages", "peer_name", nullable=False)
verifier.assert_column_exists("collections", "peer_name", nullable=False)
verifier.assert_column_exists("collections", "workspace_name", nullable=False)
verifier.assert_column_exists("collections", "internal_metadata", nullable=False)
verifier.assert_column_exists("documents", "peer_name", nullable=False)
verifier.assert_column_exists("documents", "workspace_name", nullable=False)
verifier.assert_column_exists("documents", "collection_name", nullable=False)
verifier.assert_column_exists("documents", "internal_metadata", nullable=False)
verifier.assert_column_exists("active_queue_sessions", "id", nullable=False)
verifier.assert_column_exists("active_queue_sessions", "session_id", nullable=True)
verifier.assert_column_exists("active_queue_sessions", "sender_name", nullable=True)
verifier.assert_column_exists("active_queue_sessions", "target_name", nullable=True)
verifier.assert_column_exists("active_queue_sessions", "task_type", nullable=False)
verifier.assert_column_exists("messages", "session_id", exists=False)
verifier.assert_column_exists("messages", "user_id", exists=False)
verifier.assert_column_exists("messages", "app_id", exists=False)
verifier.assert_column_exists("messages", "is_user", exists=False)
verifier.assert_indexes_exist(
[
("messages", "idx_messages_session_lookup"),
("messages", "ix_messages_peer_name"),
("messages", "ix_messages_workspace_name"),
]
)
conn = verifier.conn
schema = verifier.schema
workspace = conn.execute(
text(
'SELECT "id", "name", "configuration", "internal_metadata" '
+ f'FROM "{schema}"."workspaces" WHERE "id" = :workspace_id'
),
{"workspace_id": APP_ID},
).one()
assert workspace.name == APP_NAME
assert workspace.configuration == {}
assert workspace.internal_metadata == {}
peer = conn.execute(
text(
'SELECT "id", "name", "workspace_name", "configuration", "internal_metadata" '
+ f'FROM "{schema}"."peers" WHERE "id" = :peer_id'
),
{"peer_id": USER_ID},
).one()
assert peer.name == USER_NAME
assert peer.workspace_name == APP_NAME
assert peer.configuration == {}
assert peer.internal_metadata == {}
session = conn.execute(
text(
'SELECT "id", "name", "workspace_name", "configuration", "internal_metadata" '
+ f'FROM "{schema}"."sessions" WHERE "id" = :session_id'
),
{"session_id": SESSION_ID},
).one()
assert session.name == SESSION_ID
assert session.workspace_name == APP_NAME
assert session.configuration == {}
assert session.internal_metadata == {}
session_peer = conn.execute(
text(
f'SELECT "peer_name" FROM "{schema}"."session_peers" '
+ 'WHERE "workspace_name" = :workspace_name '
+ 'AND "session_name" = :session_name AND "peer_name" = :peer_name'
),
{
"workspace_name": APP_NAME,
"session_name": SESSION_ID,
"peer_name": USER_NAME,
},
).one()
assert session_peer.peer_name == USER_NAME
message = conn.execute(
text(
'SELECT "session_name", "workspace_name", "peer_name", "token_count", "internal_metadata" '
+ f'FROM "{schema}"."messages" WHERE "public_id" = :message_id'
),
{"message_id": MESSAGE_ID},
).one()
assert message.session_name == SESSION_ID
assert message.workspace_name == APP_NAME
assert message.peer_name == USER_NAME
assert message.token_count >= 0
assert message.internal_metadata == {}
collection = conn.execute(
text(
'SELECT "id", "name", "peer_name", "workspace_name", "internal_metadata" '
+ f'FROM "{schema}"."collections" WHERE "id" = :collection_id'
),
{"collection_id": COLLECTION_ID},
).one()
assert collection.name == COLLECTION_NAME
assert collection.peer_name == USER_NAME
assert collection.workspace_name == APP_NAME
assert collection.internal_metadata == {}
document = conn.execute(
text(
'SELECT "peer_name", "workspace_name", "collection_name", "internal_metadata" '
+ f'FROM "{schema}"."documents" WHERE "id" = :document_id'
),
{"document_id": DOCUMENT_ID},
).one()
assert document.peer_name == USER_NAME
assert document.workspace_name == APP_NAME
assert document.collection_name == COLLECTION_NAME
assert document.internal_metadata == {}
queue_row = conn.execute(
text(
f'SELECT "session_id", "payload" FROM "{schema}"."queue" '
+ "WHERE payload->>'marker' = :marker"
),
{"marker": QUEUE_MARKER},
).one()
assert queue_row.session_id == SESSION_ID
verifier.assert_constraint_exists(
"messages", "fk_messages_session_name_sessions", "foreign_key"
)
verifier.assert_constraint_exists(
"messages", "fk_messages_peer_name_peers", "foreign_key"
)
verifier.assert_constraint_exists(
"messages", "fk_messages_workspace_name_workspaces", "foreign_key"
)
verifier.assert_constraint_exists(
"collections", "fk_collections_peer_name_peers", "foreign_key"
)
verifier.assert_constraint_exists(
"collections", "fk_collections_workspace_name_workspaces", "foreign_key"
)
verifier.assert_constraint_exists(
"peers", "fk_peers_workspace_name_workspaces", "foreign_key"
)
verifier.assert_constraint_exists("peers", "unique_name_workspace_peer", "unique")
verifier.assert_constraint_exists("sessions", "unique_session_name", "unique")
verifier.assert_constraint_exists(
"collections", "unique_name_collection_peer", "unique"
)
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_active_queue_session", "unique"
)
verifier.assert_table_exists("metamessages", exists=False)

148
tests/alembic/scaffold.py Normal file
View File

@ -0,0 +1,148 @@
"""Utility to scaffold migration hook test modules."""
from __future__ import annotations
import argparse
import re
from dataclasses import dataclass
from pathlib import Path
from textwrap import dedent
PROJECT_ROOT = Path(__file__).resolve().parents[2]
MIGRATIONS_DIR = PROJECT_ROOT / "migrations" / "versions"
TESTS_REVISION_DIR = Path(__file__).resolve().parent / "revisions"
REVISION_INIT_PATH = TESTS_REVISION_DIR / "__init__.py"
TEMPLATE = '''"""Hooks for revision {revision}{slug_note}."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("{revision}")
def prepare_{identifier}(_verifier: MigrationVerifier) -> None:
"""Seed state and assertions before upgrading to {revision}."""
@register_after_upgrade("{revision}")
def verify_{identifier}(_verifier: MigrationVerifier) -> None:
"""Add assertions validating the effects of {revision}."""
'''
HEADER = '"""Register revision-specific hooks for migration verification."""'
@dataclass(slots=True)
class MigrationInfo:
revision: str
slug: str
path: Path
@property
def identifier(self) -> str:
"""Return a Python-safe identifier derived from the migration slug."""
base = re.sub(r"[^0-9a-zA-Z]+", "_", self.slug)
base = base.strip("_").lower() or "revision"
if base[0].isdigit():
base = f"revision_{base}"
return base
@property
def test_filename(self) -> str:
return f"test_{self.path.stem}.py"
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=("Generate a revision test stub with before/after upgrade hooks.")
)
parser.add_argument(
"revision",
help="Revision id (e.g. a1b2c3d4e5f6) or path to a migration file.",
)
return parser.parse_args()
def resolve_migration(path_or_revision: str) -> MigrationInfo:
candidate = Path(path_or_revision)
if candidate.suffix == ".py" and candidate.exists():
return MigrationInfo(
revision=candidate.stem.split("_", 1)[0],
slug=candidate.stem.split("_", 1)[1] if "_" in candidate.stem else "",
path=candidate.resolve(),
)
matches = sorted(MIGRATIONS_DIR.glob(f"{path_or_revision}_*.py"))
if not matches:
raise FileNotFoundError(
f"Could not find migration matching {path_or_revision!r} in {MIGRATIONS_DIR}."
)
if len(matches) > 1:
options = ", ".join(match.name for match in matches)
raise ValueError(
f"Revision prefix {path_or_revision!r} matches multiple migrations: {options}."
+ " Provide a more specific revision or path."
)
match = matches[0]
stem = match.stem
if "_" in stem:
revision, slug = stem.split("_", 1)
else:
revision, slug = stem, ""
return MigrationInfo(revision=revision, slug=slug, path=match.resolve())
def build_template(info: MigrationInfo) -> str:
slug_note = f" ({info.slug})" if info.slug else ""
content = TEMPLATE.format(
revision=info.revision,
identifier=info.identifier,
slug_note=slug_note,
)
return dedent(content).rstrip() + "\n"
def write_stub(info: MigrationInfo) -> Path:
TESTS_REVISION_DIR.mkdir(parents=True, exist_ok=True)
target = TESTS_REVISION_DIR / info.test_filename
target.write_text(build_template(info), encoding="utf-8")
return target
def refresh_revision_init() -> None:
modules = sorted(
path.stem for path in TESTS_REVISION_DIR.glob("test_*.py") if path.is_file()
)
import_block = "\n".join(f" {module}," for module in modules)
all_block = "\n".join(f' "{module}",' for module in modules)
new_content = (
dedent(
f"""{HEADER}
from . import (
{import_block}
)
__all__ = [
{all_block}
]
"""
).rstrip()
+ "\n"
)
REVISION_INIT_PATH.write_text(new_content, encoding="utf-8")
def main() -> None:
args = parse_args()
info = resolve_migration(args.revision)
target = write_stub(info)
refresh_revision_init()
print(f"Created {target.relative_to(PROJECT_ROOT)}")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,85 @@
"""Test pipeline for running alembic migrations and corresponding test hooks in order."""
from __future__ import annotations
import pytest
from alembic import command
from alembic.config import Config
from alembic.script import ScriptDirectory
from sqlalchemy import Engine
from tests.alembic.conftest import ALEMBIC_CONFIG_PATH
from tests.alembic.registry import get_registered_hooks
from tests.alembic.verifier import MigrationVerifier
def _load_revision_sequence() -> tuple[str, ...]:
"""Read the Alembic script directory to produce the linear revision order."""
script = ScriptDirectory.from_config(Config(str(ALEMBIC_CONFIG_PATH)))
revisions = list(script.walk_revisions()) # newest -> oldest
revisions.reverse()
return tuple(revision.revision for revision in revisions)
REVISION_SEQUENCE: tuple[str, ...] = _load_revision_sequence()
REVISION_PARAMS = [
pytest.param(revision, id=f"{index:02d}_{revision}")
for index, revision in enumerate(REVISION_SEQUENCE, start=1)
]
def _test_single_revision(
revision: str,
alembic_cfg: Config,
alembic_engine: Engine,
) -> None:
"""
Test a single migration revision upgrade.
For each revision:
- Migrate to the previous revision
- Run the before_upgrade hook to seed and validate the state of the DB before the revision
- Migrate to the current revision
- Run the after_upgrade hook to validate the state of the DB after the revision
"""
hooks_map = get_registered_hooks()
revision_order = list(REVISION_SEQUENCE)
# Find the previous revision in the chain
revision_index = revision_order.index(revision)
previous_revision = (
revision_order[revision_index - 1] if revision_index > 0 else "base"
)
# Migrate up to the previous revision using a shared connection
with alembic_engine.begin() as conn:
alembic_cfg.attributes["connection"] = conn
command.upgrade(alembic_cfg, previous_revision)
# Run before_upgrade hook if it exists
hooks = hooks_map.get(revision)
if hooks and hooks.before_upgrade:
with alembic_engine.begin() as conn:
verifier = MigrationVerifier(conn, revision)
hooks.before_upgrade(verifier)
# Migrate to the current revision using the same pattern
with alembic_engine.begin() as conn:
alembic_cfg.attributes["connection"] = conn
command.upgrade(alembic_cfg, revision)
# Run after_upgrade hook if it exists
if hooks and hooks.after_upgrade:
with alembic_engine.begin() as conn:
verifier = MigrationVerifier(conn, revision)
hooks.after_upgrade(verifier)
@pytest.mark.parametrize("revision", REVISION_PARAMS)
def test_migration_revision(
revision: str,
alembic_cfg: Config,
alembic_engine: Engine,
) -> None:
_test_single_revision(revision, alembic_cfg, alembic_engine)

144
tests/alembic/verifier.py Normal file
View File

@ -0,0 +1,144 @@
"""Utilities for asserting alembic migration behaviour inside tests."""
from __future__ import annotations
from collections.abc import Sequence
from sqlalchemy import inspect, text
from sqlalchemy.engine import Connection
from sqlalchemy.engine.reflection import Inspector
from migrations.utils import get_schema
class MigrationVerifier:
"""Helper to run reusable assertions against the migrated schema."""
def __init__(self, connection: Connection, revision: str):
self.conn: Connection = connection
self.revision: str = revision
self.schema: str = get_schema()
self._inspector: Inspector | None = None
def assert_table_exists(self, table: str, *, exists: bool = True) -> None:
"""Assert that a table exists in the schema"""
tables = self.get_inspector().get_table_names(schema=self.schema)
if exists:
assert table in tables
else:
assert table not in tables
def assert_column_exists(
self,
table: str,
column: str,
*,
exists: bool = True,
nullable: bool | None = None,
) -> None:
"""Assert that a column exists in the schema"""
columns = self.get_inspector().get_columns(table, schema=self.schema)
col_names = [c["name"] for c in columns]
if exists:
assert column in col_names
else:
assert column not in col_names
if nullable is not None:
if column not in col_names:
# Column absence was asserted above; nothing further to verify
return
column_info = next(col for col in columns if col["name"] == column)
actual_nullable = column_info.get("nullable", True)
assert (
actual_nullable == nullable
), f"Column {table}.{column} nullability is {actual_nullable}; expected {nullable}"
def assert_column_type(self, table: str, column: str, expected_type: type) -> None:
"""Assert that a column has the expected type"""
columns = self.get_inspector().get_columns(table, schema=self.schema)
column_info = next((col for col in columns if col["name"] == column), None)
assert (
column_info is not None
), f"Column {table}.{column} not found after migration {self.revision}"
actual_type = column_info["type"]
assert isinstance(
actual_type, expected_type
), f"Column {table}.{column} has type {type(actual_type).__name__}; expected {expected_type.__name__}"
def assert_no_nulls(self, table: str, column: str) -> None:
"""Assert that a column has no null values"""
result = self.conn.execute(
text(
f'SELECT COUNT(*) FROM "{self.schema}"."{table}" '
+ f'WHERE "{column}" IS NULL'
)
)
count = result.scalar() or 0
assert (
count == 0
), f"Found {count} NULL values in {table}.{column} after migration {self.revision}"
def assert_constraint_exists(
self,
table: str,
constraint_name: str,
constraint_type: str,
*,
exists: bool = True,
) -> None:
"""Assert that a constraint exists in the schema"""
names = self.fetch_constraints(table, constraint_type)
if exists:
assert constraint_name in names
else:
assert constraint_name not in names
def assert_indexes_exist(self, checks: Sequence[tuple[str, str]]) -> None:
"""Assert that indexes exist in the schema"""
for table_name, index_name in checks:
indexes = self.get_inspector().get_indexes(table_name, schema=self.schema)
names = [idx["name"] for idx in indexes]
assert (
index_name in names
), f"Index {index_name} not found on {table_name} after migration {self.revision}"
def assert_indexes_not_exist(self, checks: Sequence[tuple[str, str]]) -> None:
"""Assert that indexes do not exist in the schema"""
for table_name, index_name in checks:
indexes = self.get_inspector().get_indexes(table_name, schema=self.schema)
names = [idx["name"] for idx in indexes]
assert (
index_name not in names
), f"Index {index_name} still present on {table_name} after migration {self.revision}"
def get_inspector(self) -> Inspector:
"""Get the inspector for the connection"""
if self._inspector is None:
self._inspector = inspect(self.conn)
return self._inspector
def fetch_constraints(self, table: str, constraint_type: str) -> list[str | None]:
"""Collect the names of constraints in a table"""
inspector = self.get_inspector()
if constraint_type == "unique":
constraints = inspector.get_unique_constraints(table, schema=self.schema)
elif constraint_type == "foreign_key":
constraints = inspector.get_foreign_keys(table, schema=self.schema)
elif constraint_type == "check":
constraints = inspector.get_check_constraints(table, schema=self.schema)
elif constraint_type == "primary_key":
constraint = inspector.get_pk_constraint(table, schema=self.schema)
constraints = [constraint] if constraint else []
else:
raise ValueError(f"Unknown constraint type: {constraint_type}")
return [c.get("name") for c in constraints if c]
__all__ = ["MigrationVerifier"]

View File

@ -493,6 +493,10 @@ def mock_tracked_db(db_session: AsyncSession):
patch("src.routers.sessions.tracked_db", mock_tracked_db_context),
patch("src.crud.representation.tracked_db", mock_tracked_db_context),
patch("src.routers.peers.tracked_db", mock_tracked_db_context),
patch("src.dreamer.dreamer.tracked_db", mock_tracked_db_context),
patch("src.dreamer.dream_scheduler.tracked_db", mock_tracked_db_context),
patch("src.dialectic.chat.tracked_db", mock_tracked_db_context),
patch("src.utils.summarizer.tracked_db", mock_tracked_db_context),
):
yield

View File

@ -9,276 +9,216 @@ from sdks.python.src.honcho.peer import Peer
@pytest.mark.asyncio
async def test_sync_page_get_next_page(
async def test_page_get_next_page(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests that SyncPage.get_next_page() works correctly for sync clients.
Tests that Page.get_next_page() works correctly for both sync and async clients.
"""
honcho_client, client_type = client_fixture
if client_type != "sync":
pytest.skip("Test only applicable to sync client")
if client_type == "async":
assert isinstance(honcho_client, AsyncHoncho)
assert isinstance(honcho_client, Honcho)
# Create multiple peers to test pagination
for i in range(15):
peer = await honcho_client.peer(id=f"pagination-test-peer-async-{i}")
await peer.get_metadata() # Create the peer
# Create multiple peers to test pagination
for i in range(15):
peer = honcho_client.peer(id=f"pagination-test-peer-{i}")
peer.get_metadata() # Create the peer
# Get first page
first_page = await honcho_client.get_peers()
assert isinstance(first_page, AsyncPage)
assert len(first_page.items) > 0
# Get first page - SDK doesn't expose size parameter yet, so we get default pagination
first_page = honcho_client.get_peers()
assert isinstance(first_page, SyncPage)
assert len(first_page.items) > 0
# If there's a next page, test get_next_page()
if first_page.has_next_page():
second_page = await first_page.get_next_page()
assert second_page is not None
assert isinstance(second_page, AsyncPage)
assert second_page.page == 2
else:
assert isinstance(honcho_client, Honcho)
# If there's a next page, test get_next_page()
if first_page.has_next_page():
second_page = first_page.get_next_page()
assert second_page is not None
assert isinstance(second_page, SyncPage)
assert second_page.page == 2
# Create multiple peers to test pagination
for i in range(15):
peer = honcho_client.peer(id=f"pagination-test-peer-{i}")
peer.get_metadata() # Create the peer
# Get first page
first_page = honcho_client.get_peers()
assert isinstance(first_page, SyncPage)
assert len(first_page.items) > 0
# If there's a next page, test get_next_page()
if first_page.has_next_page():
second_page = first_page.get_next_page()
assert second_page is not None
assert isinstance(second_page, SyncPage)
assert second_page.page == 2
@pytest.mark.asyncio
async def test_async_page_get_next_page(
async def test_page_transform_preserved_across_pages(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests that AsyncPage.get_next_page() works correctly for async clients.
Tests that transformation function is preserved when getting next page.
"""
honcho_client, client_type = client_fixture
if client_type != "async":
pytest.skip("Test only applicable to async client")
if client_type == "async":
assert isinstance(honcho_client, AsyncHoncho)
assert isinstance(honcho_client, AsyncHoncho)
# Create multiple peers to ensure pagination
for i in range(25):
peer = await honcho_client.peer(id=f"transform-test-peer-async-{i}")
await peer.get_metadata()
# Create multiple peers to test pagination
for i in range(15):
peer = await honcho_client.peer(id=f"pagination-test-peer-async-{i}")
await peer.get_metadata() # Create the peer
# Get first page
first_page = await honcho_client.get_peers()
assert isinstance(first_page, AsyncPage)
# Get first page - SDK doesn't expose size parameter yet, so we get default pagination
first_page = await honcho_client.get_peers()
assert isinstance(first_page, AsyncPage)
assert len(first_page.items) > 0
# If there's a next page, test get_next_page()
if first_page.has_next_page():
second_page = await first_page.get_next_page()
assert second_page is not None
assert isinstance(second_page, AsyncPage)
assert second_page.page == 2
@pytest.mark.asyncio
async def test_sync_page_transform_preserved_across_pages(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests that transformation function is preserved when getting next page (sync).
"""
honcho_client, client_type = client_fixture
if client_type != "sync":
pytest.skip("Test only applicable to sync client")
assert isinstance(honcho_client, Honcho)
# Create multiple peers to ensure pagination
for i in range(25):
peer = honcho_client.peer(id=f"transform-test-peer-{i}")
peer.get_metadata()
# Get first page
first_page = honcho_client.get_peers()
assert isinstance(first_page, SyncPage)
# Verify items are Peer instances (transformed)
for item in first_page.items:
assert isinstance(item, Peer)
# Get next page and verify transformation is preserved
if first_page.has_next_page():
second_page = first_page.get_next_page()
assert second_page is not None
for item in second_page.items:
assert isinstance(item, Peer)
@pytest.mark.asyncio
async def test_async_page_transform_preserved_across_pages(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests that transformation function is preserved when getting next page (async).
"""
honcho_client, client_type = client_fixture
if client_type != "async":
pytest.skip("Test only applicable to async client")
assert isinstance(honcho_client, AsyncHoncho)
# Create multiple peers to ensure pagination
for i in range(25):
peer = await honcho_client.peer(id=f"transform-test-peer-async-{i}")
await peer.get_metadata()
# Get first page
first_page = await honcho_client.get_peers()
assert isinstance(first_page, AsyncPage)
# Verify items are AsyncPeer instances (transformed)
for item in first_page.items:
assert isinstance(item, AsyncPeer)
# Get next page and verify transformation is preserved
if first_page.has_next_page():
second_page = await first_page.get_next_page()
assert second_page is not None
for item in second_page.items:
# Verify items are AsyncPeer instances (transformed)
for item in first_page.items:
assert isinstance(item, AsyncPeer)
# Get next page and verify transformation is preserved
if first_page.has_next_page():
second_page = await first_page.get_next_page()
assert second_page is not None
for item in second_page.items:
assert isinstance(item, AsyncPeer)
else:
assert isinstance(honcho_client, Honcho)
@pytest.mark.asyncio
async def test_sync_page_get_next_page_throws_exception_on_last_page(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests that get_next_page() returns None when on the last page (sync).
"""
honcho_client, client_type = client_fixture
# Create multiple peers to ensure pagination
for i in range(25):
peer = honcho_client.peer(id=f"transform-test-peer-{i}")
peer.get_metadata()
if client_type != "sync":
pytest.skip("Test only applicable to sync client")
# Get first page
first_page = honcho_client.get_peers()
assert isinstance(first_page, SyncPage)
assert isinstance(honcho_client, Honcho)
# Verify items are Peer instances (transformed)
for item in first_page.items:
assert isinstance(item, Peer)
# Create just a few peers to ensure we're on the last page
for i in range(3):
peer = honcho_client.peer(id=f"last-page-test-peer-{i}")
peer.get_metadata()
# Get first page
first_page = honcho_client.get_peers()
assert isinstance(first_page, SyncPage)
# Should be on last page (or only page)
if not first_page.has_next_page():
# get_next_page should return None
try:
first_page.get_next_page()
except Exception as e:
assert isinstance(e, RuntimeError)
# Get next page and verify transformation is preserved
if first_page.has_next_page():
second_page = first_page.get_next_page()
assert second_page is not None
for item in second_page.items:
assert isinstance(item, Peer)
@pytest.mark.asyncio
async def test_async_page_get_next_page_throws_exception_on_last_page(
async def test_page_get_next_page_throws_exception_on_last_page(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests that get_next_page() returns None when on the last page (async).
Tests that get_next_page() throws RuntimeError when on the last page.
"""
honcho_client, client_type = client_fixture
if client_type != "async":
pytest.skip("Test only applicable to async client")
if client_type == "async":
assert isinstance(honcho_client, AsyncHoncho)
assert isinstance(honcho_client, AsyncHoncho)
# Create just a few peers to ensure we're on the last page
for i in range(3):
peer = await honcho_client.peer(id=f"last-page-test-peer-async-{i}")
await peer.get_metadata()
# Create just a few peers to ensure we're on the last page
for i in range(3):
peer = await honcho_client.peer(id=f"last-page-test-peer-async-{i}")
await peer.get_metadata()
# Get first page
first_page = await honcho_client.get_peers()
assert isinstance(first_page, AsyncPage)
# Get first page
first_page = await honcho_client.get_peers()
assert isinstance(first_page, AsyncPage)
# Should be on last page (or only page)
if not first_page.has_next_page():
# get_next_page should throw runtime error
try:
await first_page.get_next_page()
except Exception as e:
assert isinstance(e, RuntimeError)
else:
assert isinstance(honcho_client, Honcho)
# Should be on last page (or only page)
if not first_page.has_next_page():
# get_next_page should throw runtime error
try:
await first_page.get_next_page()
except Exception as e:
assert isinstance(e, RuntimeError)
# Create just a few peers to ensure we're on the last page
for i in range(3):
peer = honcho_client.peer(id=f"last-page-test-peer-{i}")
peer.get_metadata()
# Get first page
first_page = honcho_client.get_peers()
assert isinstance(first_page, SyncPage)
# Should be on last page (or only page)
if not first_page.has_next_page():
# get_next_page should throw runtime error
try:
first_page.get_next_page()
except Exception as e:
assert isinstance(e, RuntimeError)
@pytest.mark.asyncio
async def test_sync_page_manual_pagination(
async def test_page_manual_pagination(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests manual pagination with get_next_page (sync).
Tests manual pagination with get_next_page.
"""
honcho_client, client_type = client_fixture
if client_type != "sync":
pytest.skip("Test only applicable to sync client")
if client_type == "async":
assert isinstance(honcho_client, AsyncHoncho)
assert isinstance(honcho_client, Honcho)
# Create enough peers to ensure multiple pages
for i in range(25):
peer = await honcho_client.peer(id=f"manual-pagination-test-peer-async-{i}")
await peer.get_metadata()
# Create enough peers to ensure multiple pages
for i in range(25):
peer = honcho_client.peer(id=f"manual-pagination-test-peer-{i}")
peer.get_metadata()
# Collect items via manual pagination
manual_items = []
manual_page = await honcho_client.get_peers()
page_count = 0
# Collect items via manual pagination
manual_items = []
manual_page = honcho_client.get_peers()
page_count = 0
while manual_page is not None:
page_count += 1
manual_items.extend(manual_page.items) # pyright: ignore
while manual_page is not None:
page_count += 1
manual_items.extend(manual_page.items) # pyright: ignore
if not manual_page.has_next_page():
break
if not manual_page.has_next_page():
break
manual_page = await manual_page.get_next_page()
manual_page = manual_page.get_next_page()
# Should have collected all items
assert len(manual_items) >= 25 # pyright: ignore
# All items should be AsyncPeer instances
assert all(isinstance(item, AsyncPeer) for item in manual_items) # pyright: ignore
else:
assert isinstance(honcho_client, Honcho)
# Should have collected all items
assert len(manual_items) >= 25 # pyright: ignore
# All items should be Peer instances
assert all(isinstance(item, Peer) for item in manual_items) # pyright: ignore
# Create enough peers to ensure multiple pages
for i in range(25):
peer = honcho_client.peer(id=f"manual-pagination-test-peer-{i}")
peer.get_metadata()
# Collect items via manual pagination
manual_items = []
manual_page = honcho_client.get_peers()
page_count = 0
@pytest.mark.asyncio
async def test_async_page_manual_pagination(
client_fixture: tuple[Honcho | AsyncHoncho, str],
):
"""
Tests manual pagination with get_next_page (async).
"""
honcho_client, client_type = client_fixture
while manual_page is not None:
page_count += 1
manual_items.extend(manual_page.items) # pyright: ignore
if client_type != "async":
pytest.skip("Test only applicable to async client")
if not manual_page.has_next_page():
break
assert isinstance(honcho_client, AsyncHoncho)
manual_page = manual_page.get_next_page()
# Create enough peers to ensure multiple pages
for i in range(25):
peer = await honcho_client.peer(id=f"manual-pagination-test-peer-async-{i}")
await peer.get_metadata()
# Collect items via manual pagination
manual_items = []
manual_page = await honcho_client.get_peers()
page_count = 0
while manual_page is not None:
page_count += 1
manual_items.extend(manual_page.items) # pyright: ignore
if not manual_page.has_next_page():
break
manual_page = await manual_page.get_next_page()
# Should have collected all items
assert len(manual_items) >= 25 # pyright: ignore
# All items should be AsyncPeer instances
assert all(isinstance(item, AsyncPeer) for item in manual_items) # pyright: ignore
# Should have collected all items
assert len(manual_items) >= 25 # pyright: ignore
# All items should be Peer instances
assert all(isinstance(item, Peer) for item in manual_items) # pyright: ignore