diff --git a/migrations/versions/a1b2c3d4e5f6_initial_schema.py b/migrations/versions/a1b2c3d4e5f6_initial_schema.py new file mode 100644 index 00000000..708ca977 --- /dev/null +++ b/migrations/versions/a1b2c3d4e5f6_initial_schema.py @@ -0,0 +1,214 @@ +"""Initial schema creation + +Revision ID: a1b2c3d4e5f6 +Revises: +Create Date: 2024-01-01 00:00:00.000000 + +""" +from collections.abc import Sequence +from os import getenv +from typing import Union + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = 'a1b2c3d4e5f6' +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + schema = getenv("DATABASE_SCHEMA", "public") + + # Create apps table + op.create_table('apps', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('name', sa.Text(), nullable=False), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.CheckConstraint("length(public_id) = 21", name='public_id_length'), + sa.CheckConstraint("length(name) <= 512", name='name_length'), + sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name='public_id_format'), + sa.PrimaryKeyConstraint('id', name=op.f('pk_apps')), + sa.UniqueConstraint('name', name=op.f('uq_apps_name')), + sa.UniqueConstraint('public_id', name=op.f('uq_apps_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_apps_created_at'), 'apps', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_apps_id'), 'apps', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_apps_name'), 'apps', ['name'], unique=False, schema=schema) + op.create_index(op.f('ix_apps_public_id'), 'apps', ['public_id'], unique=False, schema=schema) + + # Create users table + op.create_table('users', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('name', sa.Text(), nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.Column('app_id', sa.Text(), nullable=False), + sa.CheckConstraint("length(public_id) = 21", name='public_id_length'), + 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')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_users')), + sa.UniqueConstraint('name', 'app_id', name='unique_name_app_user'), + sa.UniqueConstraint('public_id', name=op.f('uq_users_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_users_app_id'), 'users', ['app_id'], unique=False, schema=schema) + op.create_index(op.f('ix_users_created_at'), 'users', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_users_id'), 'users', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_users_name'), 'users', ['name'], unique=False, schema=schema) + op.create_index(op.f('ix_users_public_id'), 'users', ['public_id'], unique=False, schema=schema) + + # Create sessions table + op.create_table('sessions', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('is_active', sa.Boolean(), nullable=False, server_default='true'), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.Column('user_id', sa.Text(), nullable=False), + 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')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_sessions')), + sa.UniqueConstraint('public_id', name=op.f('uq_sessions_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_sessions_created_at'), 'sessions', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_sessions_id'), 'sessions', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_sessions_public_id'), 'sessions', ['public_id'], unique=False, schema=schema) + op.create_index(op.f('ix_sessions_user_id'), 'sessions', ['user_id'], unique=False, schema=schema) + + # Create messages table + op.create_table('messages', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('session_id', sa.Text(), nullable=False), + sa.Column('is_user', sa.Boolean(), nullable=False), + sa.Column('content', sa.Text(), nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.CheckConstraint("length(public_id) = 21", name='public_id_length'), + sa.CheckConstraint("length(content) <= 65535", name='content_length'), + sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name='public_id_format'), + sa.ForeignKeyConstraint(['session_id'], ['sessions.public_id'], name=op.f('fk_messages_session_id_sessions')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_messages')), + sa.UniqueConstraint('public_id', name=op.f('uq_messages_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_messages_created_at'), 'messages', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_messages_id'), 'messages', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_messages_public_id'), 'messages', ['public_id'], unique=False, schema=schema) + op.create_index(op.f('ix_messages_session_id'), 'messages', ['session_id'], unique=False, schema=schema) + + # Create metamessages table + op.create_table('metamessages', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('metamessage_type', sa.Text(), nullable=False), + sa.Column('content', sa.Text(), nullable=False), + sa.Column('message_id', sa.Text(), nullable=True), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.CheckConstraint("length(public_id) = 21", name='public_id_length'), + sa.CheckConstraint("length(content) <= 65535", name='content_length'), + sa.CheckConstraint("length(metamessage_type) <= 512", name='metamessage_type_length'), + sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name='public_id_format'), + sa.ForeignKeyConstraint(['message_id'], ['messages.public_id'], name=op.f('fk_metamessages_message_id_messages')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_metamessages')), + sa.UniqueConstraint('public_id', name=op.f('uq_metamessages_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_metamessages_created_at'), 'metamessages', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_metamessages_id'), 'metamessages', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_metamessages_message_id'), 'metamessages', ['message_id'], unique=False, schema=schema) + op.create_index(op.f('ix_metamessages_metamessage_type'), 'metamessages', ['metamessage_type'], unique=False, schema=schema) + op.create_index(op.f('ix_metamessages_public_id'), 'metamessages', ['public_id'], unique=False, schema=schema) + + # Create collections table + op.create_table('collections', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('name', sa.Text(), nullable=False), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.Column('user_id', sa.Text(), nullable=False), + sa.CheckConstraint("length(public_id) = 21", name='public_id_length'), + 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')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_collections')), + sa.UniqueConstraint('name', 'user_id', name='unique_name_collection_user'), + sa.UniqueConstraint('public_id', name=op.f('uq_collections_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_collections_created_at'), 'collections', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_collections_id'), 'collections', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_collections_name'), 'collections', ['name'], unique=False, schema=schema) + op.create_index(op.f('ix_collections_public_id'), 'collections', ['public_id'], unique=False, schema=schema) + op.create_index(op.f('ix_collections_user_id'), 'collections', ['user_id'], unique=False, schema=schema) + + # Create documents table + op.create_table('documents', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('public_id', sa.Text(), nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), nullable=False, server_default='{}'), + sa.Column('content', sa.Text(), nullable=False), + sa.Column('embedding', Vector(1536), nullable=True), + sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.Column('collection_id', sa.Text(), nullable=False), + sa.CheckConstraint("length(public_id) = 21", name='public_id_length'), + sa.CheckConstraint("length(content) <= 65535", name='content_length'), + sa.CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name='public_id_format'), + sa.ForeignKeyConstraint(['collection_id'], ['collections.public_id'], name=op.f('fk_documents_collection_id_collections')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_documents')), + sa.UniqueConstraint('public_id', name=op.f('uq_documents_public_id')), + schema=schema, + ) + op.create_index(op.f('ix_documents_collection_id'), 'documents', ['collection_id'], unique=False, schema=schema) + op.create_index(op.f('ix_documents_created_at'), 'documents', ['created_at'], unique=False, schema=schema) + op.create_index(op.f('ix_documents_id'), 'documents', ['id'], unique=False, schema=schema) + op.create_index(op.f('ix_documents_public_id'), 'documents', ['public_id'], unique=False, schema=schema) + + # Create queue table + op.create_table('queue', + sa.Column('id', sa.BigInteger(), sa.Identity(), nullable=False), + sa.Column('session_id', sa.BigInteger(), nullable=False), + 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')), + sa.PrimaryKeyConstraint('id', name=op.f('pk_queue')), + schema=schema, + ) + op.create_index(op.f('ix_queue_session_id'), 'queue', ['session_id'], unique=False, schema=schema) + + # Create active_queue_sessions table + op.create_table('active_queue_sessions', + sa.Column('session_id', sa.BigInteger(), nullable=False), + sa.Column('last_updated', sa.DateTime(timezone=True), server_default=sa.text('now()'), nullable=False), + sa.ForeignKeyConstraint(['session_id'], ['sessions.id'], name=op.f('fk_active_queue_sessions_session_id_sessions')), + sa.PrimaryKeyConstraint('session_id', name=op.f('pk_active_queue_sessions')), + schema=schema, + ) + op.create_index(op.f('ix_active_queue_sessions_session_id'), 'active_queue_sessions', ['session_id'], unique=False, schema=schema) + + +def downgrade() -> None: + schema = getenv("DATABASE_SCHEMA", "public") + op.drop_table('active_queue_sessions', schema=schema) + op.drop_table('queue', schema=schema) + op.drop_table('documents', schema=schema) + op.drop_table('collections', schema=schema) + op.drop_table('metamessages', schema=schema) + op.drop_table('messages', schema=schema) + op.drop_table('sessions', schema=schema) + op.drop_table('users', schema=schema) + op.drop_table('apps', schema=schema) \ No newline at end of file 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 e4e11a1e..89225890 100644 --- a/migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py +++ b/migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py @@ -7,6 +7,7 @@ Create Date: 2025-04-03 15:32:16.733312 """ from collections.abc import Sequence +from os import getenv from typing import Union import sqlalchemy as sa @@ -22,6 +23,8 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: + schema = getenv("DATABASE_SCHEMA", "public") + conn = op.get_bind() # Check which columns already exist in the metamessages table @@ -30,13 +33,13 @@ def upgrade() -> None: # 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)) + op.add_column("metamessages", sa.Column("user_id", sa.TEXT(), nullable=True), schema=schema) 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)) + op.add_column("metamessages", sa.Column("session_id", sa.TEXT(), nullable=True), schema=schema) print("Added session_id column") else: print("session_id column already exists") @@ -53,7 +56,7 @@ def upgrade() -> None: "metamessages", "users", ["user_id"], - ["public_id"], + ["public_id"] ) print("Created user_id foreign key") except IntegrityError: @@ -75,7 +78,7 @@ def upgrade() -> None: # 3. Make message_id nullable if it's not already try: op.alter_column( - "metamessages", "message_id", existing_type=sa.TEXT(), nullable=True + "metamessages", "message_id", existing_type=sa.TEXT(), nullable=True, schema=schema ) print("Made message_id nullable") except (ProgrammingError, IntegrityError) as e: @@ -102,7 +105,7 @@ def upgrade() -> None: ) ).scalar() - if orphaned_count > 0: + if orphaned_count and orphaned_count > 0: print( f"Found {orphaned_count} orphaned records. Setting message_id to NULL." ) @@ -130,7 +133,6 @@ def upgrade() -> None: # 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) @@ -145,43 +147,25 @@ def upgrade() -> None: 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"]) - create_index_if_not_exists( - "idx_sessions_user_lookup", "sessions", ["user_id", "public_id"] - ) - - # For complex indices - create_index_if_not_exists( - "idx_messages_session_lookup", - "messages", - ["session_id", "id"], - postgresql_include=["public_id", "is_user", "created_at"], - ) - - create_index_if_not_exists( - "idx_metamessages_lookup", - "metamessages", - ["metamessage_type", text("id DESC")], - postgresql_include=["public_id", "message_id", "created_at"], - ) - create_index_if_not_exists( "idx_metamessages_user_lookup", "metamessages", ["user_id", "metamessage_type", text("id DESC")], + schema=schema, ) create_index_if_not_exists( "idx_metamessages_session_lookup", "metamessages", ["session_id", "metamessage_type", text("id DESC")], + schema=schema, ) create_index_if_not_exists( "idx_metamessages_message_lookup", "metamessages", ["message_id", "metamessage_type", text("id DESC")], + schema=schema, ) # 8. Verify the data is ready for the constraint @@ -191,7 +175,7 @@ def upgrade() -> None: ) ).scalar() - if constraint_violations > 0: + if constraint_violations and 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(""" @@ -230,7 +214,7 @@ def upgrade() -> None: text("SELECT COUNT(*) FROM metamessages WHERE user_id IS NULL") ).scalar() - if null_user_count > 0: + if null_user_count and 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(""" @@ -251,7 +235,7 @@ def upgrade() -> None: if null_user_count == 0: try: op.alter_column( - "metamessages", "user_id", existing_type=sa.TEXT(), nullable=False + "metamessages", "user_id", existing_type=sa.TEXT(), nullable=False, schema=schema ) print("Made user_id not nullable") except Exception as e: @@ -263,7 +247,7 @@ def upgrade() -> None: else: try: op.alter_column( - "metamessages", "user_id", existing_type=sa.TEXT(), nullable=False + "metamessages", "user_id", existing_type=sa.TEXT(), nullable=False, schema=schema ) print("Made user_id not nullable") except Exception as e: @@ -271,11 +255,12 @@ def upgrade() -> None: def downgrade() -> None: + schema = getenv("DATABASE_SCHEMA", "public") # Add try-except blocks to handle case where elements don't exist # 1. Remove the check constraint try: - op.drop_constraint("message_requires_session", "metamessages", type_="check") + op.drop_constraint("message_requires_session", "metamessages", type_="check", schema=schema) except Exception as e: print(f"Error dropping message_requires_session constraint: {e}") @@ -284,27 +269,23 @@ def downgrade() -> None: ("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) + op.drop_index(index_name, table_name=table_name, schema=schema) except Exception as e: print(f"Error dropping index {index_name}: {e}") # 3. Remove foreign key constraints try: op.drop_constraint( - "fk_metamessages_session_id_sessions", "metamessages", type_="foreignkey" + "fk_metamessages_session_id_sessions", "metamessages", type_="foreignkey", schema=schema ) 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" + "fk_metamessages_user_id_users", "metamessages", type_="foreignkey", schema=schema ) except Exception as e: print(f"Error dropping user_id foreign key: {e}") @@ -315,18 +296,18 @@ def downgrade() -> None: DELETE FROM metamessages WHERE message_id IS NULL """) op.alter_column( - "metamessages", "message_id", existing_type=sa.TEXT(), nullable=False + "metamessages", "message_id", existing_type=sa.TEXT(), nullable=False, schema=schema ) except Exception as e: print(f"Error making message_id not nullable: {e}") # 5. Drop the new columns try: - op.drop_column("metamessages", "session_id") + op.drop_column("metamessages", "session_id", schema=schema) except Exception as e: print(f"Error dropping session_id column: {e}") try: - op.drop_column("metamessages", "user_id") + op.drop_column("metamessages", "user_id", schema=schema) except Exception as e: print(f"Error dropping user_id column: {e}") diff --git a/migrations/versions/c3828084f472_add_indexes_for_messages_and_.py b/migrations/versions/c3828084f472_add_indexes_for_messages_and_.py index 6d2d974d..6aa0d58b 100644 --- a/migrations/versions/c3828084f472_add_indexes_for_messages_and_.py +++ b/migrations/versions/c3828084f472_add_indexes_for_messages_and_.py @@ -1,7 +1,7 @@ """Add indexes for messages and metamessages for reads Revision ID: c3828084f472 -Revises: +Revises: a1b2c3d4e5f6 Create Date: 2024-12-12 13:41:40.156095 """ @@ -15,7 +15,7 @@ from alembic import op # revision identifiers, used by Alembic. revision: str = "c3828084f472" -down_revision: Union[str, None] = None +down_revision: Union[str, None] = "a1b2c3d4e5f6" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None diff --git a/scripts/migrate_db.py b/scripts/migrate_db.py index 35cca61a..a4811778 100644 --- a/scripts/migrate_db.py +++ b/scripts/migrate_db.py @@ -2,18 +2,16 @@ import os import sys +from alembic import command +from alembic.config import Config + # Add the project root to the path # This assumes the script is run from the scripts directory project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.insert(0, project_root) -from alembic import command # noqa: E402 -from alembic.config import Config # noqa: E402 - -from src.db import scaffold_db # noqa: E402 if __name__ == "__main__": - scaffold_db() # run alembic migrations alembic_cfg = Config("alembic.ini") command.upgrade(alembic_cfg, "head") diff --git a/scripts/provision_db.py b/scripts/provision_db.py index 1ec8bc0b..93633afb 100644 --- a/scripts/provision_db.py +++ b/scripts/provision_db.py @@ -7,8 +7,9 @@ import sys project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) sys.path.insert(0, project_root) -from src.db import scaffold_db # noqa: E402 +from src.db import init_db # noqa: E402 if __name__ == "__main__": - scaffold_db() - print("Database created") + print("Initializing database using Alembic migrations...") + init_db() + print("Database initialized successfully") diff --git a/src/db.py b/src/db.py index f26e8e12..8bfe8dfb 100644 --- a/src/db.py +++ b/src/db.py @@ -1,7 +1,7 @@ import os from dotenv import load_dotenv -from sqlalchemy import MetaData, create_engine, inspect, text +from sqlalchemy import MetaData, create_engine, text from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.orm import declarative_base @@ -30,44 +30,33 @@ SessionLocal = async_sessionmaker( bind=engine, ) -table_schema = os.getenv("DATABASE_SCHEMA") +table_schema = os.getenv("DATABASE_SCHEMA", "public") meta = MetaData() -if table_schema: - meta.schema = table_schema +meta.schema = table_schema Base = declarative_base(metadata=meta) -def scaffold_db(): - """use a sync engine for scaffolding the database. ddl operations are unavailable - with async engines - """ +def init_db(): + """Initialize the database using Alembic migrations""" + from alembic import command + from alembic.config import Config - # Debug: Print all tables that should be created - print("Tables defined in Base.metadata:") - for table in Base.metadata.sorted_tables: - print(f" - {table.name}") - - # Create engine - engine = create_engine( + # Create a sync engine for schema operations + sync_engine = create_engine( os.environ["CONNECTION_URI"], pool_pre_ping=True, echo=os.getenv("SQL_DEBUG", "false").lower() == "true", ) - # Create inspector to check if database exists - inspector = inspect(engine) - + # Create schema if it doesn't exist if table_schema: - with engine.connect() as connection: + with sync_engine.connect() as connection: connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{table_schema}"')) connection.commit() - print(inspector.get_table_names(Base.metadata.schema)) - - # If no tables exist, create them with SQLAlchemy - if not inspector.get_table_names(Base.metadata.schema): - print("No tables found. Creating database schema...") - Base.metadata.create_all(bind=engine) + # Run Alembic migrations + alembic_cfg = Config("alembic.ini") + command.upgrade(alembic_cfg, "head") # Clean up - engine.dispose() + sync_engine.dispose()