Alembic genesis (#106)
* chore: create initial migration for alembic, refactor provision/upgrade db scripts to fully leverage alembic * add schema everywhere, remove duplicates * fix: Add default schema public --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
doria
GitHub
parent
d1292aaae4
commit
606aa14411
|
|
@ -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)
|
||||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
41
src/db.py
41
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()
|
||||
|
|
|
|||
Loading…
Reference in New Issue