101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
import asyncio
|
|
from logging.config import fileConfig
|
|
|
|
from sqlalchemy import pool
|
|
from sqlalchemy.engine import Connection
|
|
from sqlalchemy.ext.asyncio import async_engine_from_config
|
|
|
|
from alembic import context
|
|
|
|
# Import your models
|
|
from app.core.database import Base
|
|
from app.models import User, Team, Tree, Session, Attachment, InviteCode
|
|
from app.models.email_verification_token import EmailVerificationToken
|
|
from app.models.tree_embedding import TreeEmbedding
|
|
from app.models.copilot_conversation import CopilotConversation
|
|
from app.models.assistant_chat import AssistantChat
|
|
from app.models.survey_response import SurveyResponse
|
|
from app.models.survey_invite import SurveyInvite
|
|
from app.models.ai_suggestion import AISuggestion # noqa: F401
|
|
from app.models.kb_import import KBImport, KBImportNode # noqa: F401
|
|
from app.models.script_template import ScriptCategory, ScriptTemplate, ScriptGeneration # noqa: F401
|
|
from app.models.psa_connection import PsaConnection # noqa: F401
|
|
from app.models.ai_session import AISession # noqa: F401
|
|
from app.models.ai_session_step import AISessionStep # noqa: F401
|
|
from app.models.psa_post_log import PsaPostLog # noqa: F401
|
|
from app.models.psa_member_mapping import PsaMemberMapping # noqa: F401
|
|
from app.models.script_builder_session import ScriptBuilderSession, ScriptBuilderMessage # noqa: F401
|
|
from app.models.session_branch import SessionBranch # noqa: F401
|
|
from app.models.fork_point import ForkPoint # noqa: F401
|
|
from app.models.session_handoff import SessionHandoff # noqa: F401
|
|
from app.models.session_resolution_output import SessionResolutionOutput # noqa: F401
|
|
from app.core.config import settings
|
|
|
|
# this is the Alembic Config object
|
|
config = context.config
|
|
|
|
# Override sqlalchemy.url with the sync version for migrations
|
|
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL_SYNC)
|
|
|
|
# Interpret the config file for Python logging.
|
|
if config.config_file_name is not None:
|
|
fileConfig(config.config_file_name)
|
|
|
|
# add your model's MetaData object here
|
|
target_metadata = Base.metadata
|
|
|
|
|
|
def run_migrations_offline() -> None:
|
|
"""Run migrations in 'offline' mode."""
|
|
url = config.get_main_option("sqlalchemy.url")
|
|
context.configure(
|
|
url=url,
|
|
target_metadata=target_metadata,
|
|
literal_binds=True,
|
|
dialect_opts={"paramstyle": "named"},
|
|
)
|
|
|
|
with context.begin_transaction():
|
|
context.run_migrations()
|
|
|
|
|
|
def do_run_migrations(connection: Connection) -> None:
|
|
context.configure(connection=connection, target_metadata=target_metadata)
|
|
|
|
with context.begin_transaction():
|
|
context.run_migrations()
|
|
|
|
|
|
async def run_async_migrations() -> None:
|
|
"""Run migrations in 'online' mode with async engine."""
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
connectable = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
poolclass=pool.NullPool,
|
|
)
|
|
|
|
async with connectable.connect() as connection:
|
|
await connection.run_sync(do_run_migrations)
|
|
|
|
await connectable.dispose()
|
|
|
|
|
|
def run_migrations_online() -> None:
|
|
"""Run migrations in 'online' mode."""
|
|
from sqlalchemy import create_engine
|
|
|
|
connectable = create_engine(
|
|
settings.DATABASE_URL_SYNC,
|
|
poolclass=pool.NullPool,
|
|
)
|
|
|
|
with connectable.connect() as connection:
|
|
do_run_migrations(connection)
|
|
|
|
|
|
if context.is_offline_mode():
|
|
run_migrations_offline()
|
|
else:
|
|
run_migrations_online()
|