From 9dd214b3a63358fc8290664d2493c4506fc05549 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Mon, 5 May 2025 15:39:00 -0400 Subject: [PATCH] chore: Cleanup Migration (#101) * fix: Add safer database migration * chore: Fix typos * chore: Fix typos --- README.md | 4 +- ...change_metamessages_to_user_level_with_.py | 313 ++++++++++++++---- 2 files changed, 244 insertions(+), 73 deletions(-) diff --git a/README.md b/README.md index 803891b7..0a302c44 100644 --- a/README.md +++ b/README.md @@ -103,12 +103,12 @@ The `Session` object represents a set of interactions a `User` has with an **Messages** The `Message` represents an atomic interaction of a `User` in a `Session`. -`Message`s are labed as either a `User` or AI message. +`Message`s are labeled as either a `User` or AI message. #### Collections At a high level a `Collection` is a named group of `Documents`. Developers -familiar with RAG based applications will be familar with these. `Collection`s +familiar with RAG based applications will be familiar with these. `Collection`s store vector embedded data that developers and agents can retrieve against using functions like cosine similarity. diff --git a/migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py b/migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py index 51772848..e4e11a1e 100644 --- a/migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py +++ b/migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py @@ -6,12 +6,13 @@ Create Date: 2025-04-03 15:32:16.733312 """ -from typing import Sequence, Union +from collections.abc import Sequence +from typing import Union -from alembic import op import sqlalchemy as sa +from alembic import op +from sqlalchemy.exc import IntegrityError, ProgrammingError from sqlalchemy.sql import text -from sqlalchemy.dialects import postgresql # revision identifiers, used by Alembic. revision: str = "b765d82110bd" @@ -21,48 +22,128 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: - # 1. Add new columns to metamessages table - op.add_column("metamessages", sa.Column("user_id", sa.TEXT(), nullable=True)) - op.add_column("metamessages", sa.Column("session_id", sa.TEXT(), nullable=True)) + conn = op.get_bind() + + # Check which columns already exist in the metamessages table + inspector = sa.inspect(conn) + existing_columns = [col["name"] for col in inspector.get_columns("metamessages")] + + # 1. Add new columns to metamessages table if they don't exist + if "user_id" not in existing_columns: + op.add_column("metamessages", sa.Column("user_id", sa.TEXT(), nullable=True)) + print("Added user_id column") + else: + print("user_id column already exists") + + if "session_id" not in existing_columns: + op.add_column("metamessages", sa.Column("session_id", sa.TEXT(), nullable=True)) + print("Added session_id column") + else: + print("session_id column already exists") # 2. Create foreign key constraints for the new columns - op.create_foreign_key( - "fk_metamessages_user_id_users", - "metamessages", - "users", - ["user_id"], - ["public_id"], - ) - op.create_foreign_key( - "fk_metamessages_session_id_sessions", - "metamessages", - "sessions", - ["session_id"], - ["public_id"], - ) + # Check if foreign key constraints already exist + foreign_keys = inspector.get_foreign_keys("metamessages") + fk_names = [fk.get("name") for fk in foreign_keys] - # 3. Make message_id nullable - op.alter_column( - "metamessages", "message_id", existing_type=sa.TEXT(), nullable=True - ) + if "fk_metamessages_user_id_users" not in fk_names: + try: + op.create_foreign_key( + "fk_metamessages_user_id_users", + "metamessages", + "users", + ["user_id"], + ["public_id"], + ) + print("Created user_id foreign key") + except IntegrityError: + print("Cannot create user_id foreign key - integrity error") - # 4. Create indices for users, sessions, and messages - only if they don't exist + if "fk_metamessages_session_id_sessions" not in fk_names: + try: + op.create_foreign_key( + "fk_metamessages_session_id_sessions", + "metamessages", + "sessions", + ["session_id"], + ["public_id"], + ) + print("Created session_id foreign key") + except IntegrityError: + print("Cannot create session_id foreign key - integrity error") - conn = op.get_bind() - inspector = sa.inspect(conn) + # 3. Make message_id nullable if it's not already + try: + op.alter_column( + "metamessages", "message_id", existing_type=sa.TEXT(), nullable=True + ) + print("Made message_id nullable") + except (ProgrammingError, IntegrityError) as e: + print(f"Error making message_id nullable: {e}") - # Check and create each index only if it doesn't exist + # 4. Data migration: Fill session_id from messages + try: + op.execute(""" + UPDATE metamessages m + SET session_id = msg.session_id + FROM messages msg + WHERE m.message_id = msg.public_id + AND m.session_id IS NULL + """) + print("Updated session_id values from messages") + except Exception as e: + print(f"Error updating session_id values: {e}") + + # 5. Handle orphaned records - records with message_id but no session_id + try: + orphaned_count = conn.execute( + text( + "SELECT COUNT(*) FROM metamessages WHERE message_id IS NOT NULL AND session_id IS NULL" + ) + ).scalar() + + if orphaned_count > 0: + print( + f"Found {orphaned_count} orphaned records. Setting message_id to NULL." + ) + op.execute(""" + UPDATE metamessages + SET message_id = NULL + WHERE message_id IS NOT NULL AND session_id IS NULL + """) + except Exception as e: + print(f"Error handling orphaned records: {e}") + + # 6. Data migration: Fill user_id from session's user + try: + op.execute(""" + UPDATE metamessages m + SET user_id = s.user_id + FROM sessions s + WHERE m.session_id = s.public_id + AND m.user_id IS NULL + """) + print("Updated user_id values from sessions") + except Exception as e: + print(f"Error updating user_id values from sessions: {e}") + + # 7. Create indices for the various lookups - only if they don't exist existing_indices = { idx["name"] for schema in [None, "public"] for tbl in inspector.get_table_names(schema=schema) + if tbl in ["metamessages", "users", "sessions", "messages"] for idx in inspector.get_indexes(tbl, schema=schema) } # Helper function to create index if not exists def create_index_if_not_exists(index_name, table_name, columns, **kwargs): if index_name not in existing_indices: - op.create_index(index_name, table_name, columns, **kwargs) + try: + op.create_index(index_name, table_name, columns, **kwargs) + print(f"Created index {index_name}") + except Exception as e: + print(f"Error creating index {index_name}: {e}") # Create all indices with the helper function create_index_if_not_exists("idx_users_app_lookup", "users", ["app_id", "public_id"]) @@ -103,59 +184,149 @@ def upgrade() -> None: ["message_id", "metamessage_type", text("id DESC")], ) - # 7. Add the check constraint for message_id and session_id relationship - op.create_check_constraint( - "message_requires_session", - "metamessages", - "(message_id IS NULL) OR (session_id IS NOT NULL)", - ) + # 8. Verify the data is ready for the constraint + constraint_violations = conn.execute( + text( + "SELECT COUNT(*) FROM metamessages WHERE message_id IS NOT NULL AND session_id IS NULL" + ) + ).scalar() - # 8. Update existing data: fill user_id from message's session's user - # This is a complex data migration that requires SQL - op.execute(""" - UPDATE metamessages m - SET user_id = u.public_id, - session_id = s.public_id - FROM messages msg - JOIN sessions s ON msg.session_id = s.public_id - JOIN users u ON s.user_id = u.public_id - WHERE m.message_id = msg.public_id - AND m.user_id IS NULL - """) + if constraint_violations > 0: + print(f"WARNING: {constraint_violations} records would violate the constraint!") + print("Setting these message_ids to NULL to prevent constraint violation") + op.execute(""" + UPDATE metamessages + SET message_id = NULL + WHERE message_id IS NOT NULL AND session_id IS NULL + """) - # 9. Now that data is migrated, make user_id not nullable - op.alter_column("metamessages", "user_id", existing_type=sa.TEXT(), nullable=False) + # 9. Add the check constraint for message_id and session_id relationship + try: + # Check if constraint already exists + constraint_exists = ( + conn.execute( + text( + "SELECT 1 FROM pg_constraint WHERE conname = 'message_requires_session'" + ) + ).fetchone() + is not None + ) + + if not constraint_exists: + op.create_check_constraint( + "message_requires_session", + "metamessages", + "(message_id IS NULL) OR (session_id IS NOT NULL)", + ) + print("Created message_requires_session constraint") + else: + print("message_requires_session constraint already exists") + except Exception as e: + print(f"Error creating constraint: {e}") + + # 10. Now that data is migrated, make user_id not nullable if it has values + # First ensure all records have a user_id + null_user_count = conn.execute( + text("SELECT COUNT(*) FROM metamessages WHERE user_id IS NULL") + ).scalar() + + if null_user_count > 0: + print(f"WARNING: {null_user_count} records have NULL user_id!") + # Try to derive user_id from message for these orphaned records + op.execute(""" + UPDATE metamessages m + SET user_id = u.public_id + FROM messages msg + JOIN sessions s ON msg.session_id = s.public_id + JOIN users u ON s.user_id = u.public_id + WHERE m.message_id = msg.public_id + AND m.user_id IS NULL + """) + + # Check again after fixes + null_user_count = conn.execute( + text("SELECT COUNT(*) FROM metamessages WHERE user_id IS NULL") + ).scalar() + + if null_user_count == 0: + try: + op.alter_column( + "metamessages", "user_id", existing_type=sa.TEXT(), nullable=False + ) + print("Made user_id not nullable") + except Exception as e: + print(f"Error making user_id not nullable: {e}") + else: + print( + f"Still have {null_user_count} records with NULL user_id, cannot make column not nullable" + ) + else: + try: + op.alter_column( + "metamessages", "user_id", existing_type=sa.TEXT(), nullable=False + ) + print("Made user_id not nullable") + except Exception as e: + print(f"Error making user_id not nullable: {e}") def downgrade() -> None: + # Add try-except blocks to handle case where elements don't exist + # 1. Remove the check constraint - op.drop_constraint("message_requires_session", "metamessages", type_="check") + try: + op.drop_constraint("message_requires_session", "metamessages", type_="check") + except Exception as e: + print(f"Error dropping message_requires_session constraint: {e}") # 2. Drop all the new indices - op.drop_index("idx_metamessages_message_lookup", table_name="metamessages") - op.drop_index("idx_metamessages_session_lookup", table_name="metamessages") - op.drop_index("idx_metamessages_user_lookup", table_name="metamessages") - op.drop_index("idx_metamessages_lookup", table_name="metamessages") - op.drop_index("idx_messages_session_lookup", table_name="messages") - op.drop_index("idx_sessions_user_lookup", table_name="sessions") - op.drop_index("idx_users_app_lookup", table_name="users") + for index_name, table_name in [ + ("idx_metamessages_message_lookup", "metamessages"), + ("idx_metamessages_session_lookup", "metamessages"), + ("idx_metamessages_user_lookup", "metamessages"), + ("idx_metamessages_lookup", "metamessages"), + ("idx_messages_session_lookup", "messages"), + ("idx_sessions_user_lookup", "sessions"), + ("idx_users_app_lookup", "users"), + ]: + try: + op.drop_index(index_name, table_name=table_name) + except Exception as e: + print(f"Error dropping index {index_name}: {e}") # 3. Remove foreign key constraints - op.drop_constraint( - "fk_metamessages_session_id_sessions", "metamessages", type_="foreignkey" - ) - op.drop_constraint( - "fk_metamessages_user_id_users", "metamessages", type_="foreignkey" - ) + try: + op.drop_constraint( + "fk_metamessages_session_id_sessions", "metamessages", type_="foreignkey" + ) + except Exception as e: + print(f"Error dropping session_id foreign key: {e}") + + try: + op.drop_constraint( + "fk_metamessages_user_id_users", "metamessages", type_="foreignkey" + ) + except Exception as e: + print(f"Error dropping user_id foreign key: {e}") # 4. Make message_id required again and clean up data if needed - op.execute(""" - DELETE FROM metamessages WHERE message_id IS NULL - """) - op.alter_column( - "metamessages", "message_id", existing_type=sa.TEXT(), nullable=False - ) + try: + op.execute(""" + DELETE FROM metamessages WHERE message_id IS NULL + """) + op.alter_column( + "metamessages", "message_id", existing_type=sa.TEXT(), nullable=False + ) + except Exception as e: + print(f"Error making message_id not nullable: {e}") # 5. Drop the new columns - op.drop_column("metamessages", "session_id") - op.drop_column("metamessages", "user_id") + try: + op.drop_column("metamessages", "session_id") + except Exception as e: + print(f"Error dropping session_id column: {e}") + + try: + op.drop_column("metamessages", "user_id") + except Exception as e: + print(f"Error dropping user_id column: {e}")