136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
"""add external qa followup state fields
|
|
|
|
Revision ID: 20260327_add_external_qa_followup_state
|
|
Revises: 20260327_cleanup_model_routing_legacy_names
|
|
Create Date: 2026-03-27
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
try:
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
except Exception: # pragma: no cover
|
|
JSONB = sa.JSON
|
|
|
|
|
|
revision = "20260327_add_external_qa_followup_state"
|
|
down_revision = "20260327_cleanup_model_routing_legacy_names"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def _json_type(bind):
|
|
return JSONB() if bind.dialect.name == "postgresql" else sa.JSON()
|
|
|
|
|
|
def _table_exists(bind, table_name: str) -> bool:
|
|
return sa.inspect(bind).has_table(table_name)
|
|
|
|
|
|
def _column_exists(bind, table_name: str, column_name: str) -> bool:
|
|
inspector = sa.inspect(bind)
|
|
try:
|
|
columns = inspector.get_columns(table_name)
|
|
except Exception:
|
|
return False
|
|
return any(str(col.get("name")) == column_name for col in columns)
|
|
|
|
|
|
def _index_exists(bind, table_name: str, index_name: str) -> bool:
|
|
inspector = sa.inspect(bind)
|
|
try:
|
|
indexes = inspector.get_indexes(table_name)
|
|
except Exception:
|
|
return False
|
|
return any(idx.get("name") == index_name for idx in indexes)
|
|
|
|
|
|
def _create_index_if_not_exists(
|
|
bind,
|
|
index_name: str,
|
|
table_name: str,
|
|
columns: list[str],
|
|
*,
|
|
unique: bool = False,
|
|
) -> None:
|
|
if not _index_exists(bind, table_name, index_name):
|
|
op.create_index(index_name, table_name, columns, unique=unique)
|
|
|
|
|
|
def _drop_index_if_exists(bind, index_name: str, table_name: str) -> None:
|
|
if _index_exists(bind, table_name, index_name):
|
|
op.drop_index(index_name, table_name=table_name)
|
|
|
|
|
|
def upgrade() -> None:
|
|
bind = op.get_bind()
|
|
|
|
if _table_exists(bind, "qa_sessions"):
|
|
if not _column_exists(bind, "qa_sessions", "external_conversation_key"):
|
|
with op.batch_alter_table("qa_sessions") as batch_op:
|
|
batch_op.add_column(sa.Column("external_conversation_key", sa.VARCHAR(length=160), nullable=True))
|
|
_create_index_if_not_exists(
|
|
bind,
|
|
"ix_qa_sessions_external_conversation_key",
|
|
"qa_sessions",
|
|
["external_conversation_key"],
|
|
)
|
|
_create_index_if_not_exists(
|
|
bind,
|
|
"uq_qa_sessions_tenant_external_conversation_key",
|
|
"qa_sessions",
|
|
["tenant_id", "external_conversation_key"],
|
|
unique=True,
|
|
)
|
|
|
|
if _table_exists(bind, "qa_session_states"):
|
|
if not _column_exists(bind, "qa_session_states", "pending_clarification"):
|
|
with op.batch_alter_table("qa_session_states") as batch_op:
|
|
batch_op.add_column(sa.Column("pending_clarification", _json_type(bind), nullable=True))
|
|
if not _column_exists(bind, "qa_session_states", "pending_clarification_expires_at"):
|
|
with op.batch_alter_table("qa_session_states") as batch_op:
|
|
batch_op.add_column(sa.Column("pending_clarification_expires_at", sa.TIMESTAMP(), nullable=True))
|
|
_create_index_if_not_exists(
|
|
bind,
|
|
"ix_qa_session_states_pending_clarification_expires_at",
|
|
"qa_session_states",
|
|
["pending_clarification_expires_at"],
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
bind = op.get_bind()
|
|
|
|
if _table_exists(bind, "qa_session_states"):
|
|
_drop_index_if_exists(
|
|
bind,
|
|
"ix_qa_session_states_pending_clarification_expires_at",
|
|
"qa_session_states",
|
|
)
|
|
if _column_exists(bind, "qa_session_states", "pending_clarification_expires_at") or _column_exists(
|
|
bind, "qa_session_states", "pending_clarification"
|
|
):
|
|
with op.batch_alter_table("qa_session_states") as batch_op:
|
|
if _column_exists(bind, "qa_session_states", "pending_clarification_expires_at"):
|
|
batch_op.drop_column("pending_clarification_expires_at")
|
|
if _column_exists(bind, "qa_session_states", "pending_clarification"):
|
|
batch_op.drop_column("pending_clarification")
|
|
|
|
if _table_exists(bind, "qa_sessions"):
|
|
_drop_index_if_exists(
|
|
bind,
|
|
"uq_qa_sessions_tenant_external_conversation_key",
|
|
"qa_sessions",
|
|
)
|
|
_drop_index_if_exists(
|
|
bind,
|
|
"ix_qa_sessions_external_conversation_key",
|
|
"qa_sessions",
|
|
)
|
|
if _column_exists(bind, "qa_sessions", "external_conversation_key"):
|
|
with op.batch_alter_table("qa_sessions") as batch_op:
|
|
batch_op.drop_column("external_conversation_key")
|