345 lines
12 KiB
Python
345 lines
12 KiB
Python
|
"""initial schema part 2: session and agent tables
|
||
|
|
|
||
|
|
Revision ID: 202602260002
|
||
|
|
Revises: 202602260001
|
||
|
|
Create Date: 2026-02-26 20:11:00
|
||
|
|
"""
|
||
|
|
|
||
|
|
from typing import Sequence, Union
|
||
|
|
|
||
|
|
from alembic import op
|
||
|
|
import sqlalchemy as sa
|
||
|
|
from sqlalchemy.dialects import postgresql
|
||
|
|
|
||
|
|
revision: str = "202602260002"
|
||
|
|
down_revision: Union[str, Sequence[str], None] = "202602260001"
|
||
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
||
|
|
depends_on: Union[str, Sequence[str], None] = None
|
||
|
|
|
||
|
|
|
||
|
|
def upgrade() -> None:
|
||
|
|
op.create_table(
|
||
|
|
"sessions",
|
||
|
|
sa.Column("id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("user_id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("session_type", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column("job_id", sa.UUID(), nullable=True),
|
||
|
|
sa.Column("title", sa.String(length=255), nullable=True),
|
||
|
|
sa.Column("status", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column(
|
||
|
|
"last_activity_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"message_count", sa.Integer(), server_default=sa.text("0"), nullable=False
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"total_tokens", sa.Integer(), server_default=sa.text("0"), nullable=False
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"total_cost",
|
||
|
|
sa.Numeric(precision=12, scale=6),
|
||
|
|
server_default=sa.text("0"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"created_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"updated_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
|
||
|
|
sa.PrimaryKeyConstraint("id"),
|
||
|
|
)
|
||
|
|
op.create_index("ix_sessions_job_id", "sessions", ["job_id"], unique=False)
|
||
|
|
op.execute(
|
||
|
|
"CREATE INDEX ix_sessions_user_session_type_last_activity ON sessions (user_id, session_type, last_activity_at DESC)"
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
"ALTER TABLE sessions ADD CONSTRAINT chk_session_type CHECK (session_type IN ('chat', 'automation'))"
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
"ALTER TABLE sessions ADD CONSTRAINT chk_sessions_type_job_consistency CHECK ((session_type = 'chat' AND job_id IS NULL) OR (session_type = 'automation' AND job_id IS NOT NULL))"
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
"ALTER TABLE sessions ADD CONSTRAINT chk_sessions_status CHECK (status IN ('pending', 'running', 'completed', 'failed'))"
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_sessions_user_id",
|
||
|
|
"sessions",
|
||
|
|
"users",
|
||
|
|
["user_id"],
|
||
|
|
["id"],
|
||
|
|
referent_schema="auth",
|
||
|
|
ondelete="CASCADE",
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_sessions_job_id_user_id",
|
||
|
|
"sessions",
|
||
|
|
"automation_jobs",
|
||
|
|
["job_id", "user_id"],
|
||
|
|
["id", "owner_id"],
|
||
|
|
ondelete="RESTRICT",
|
||
|
|
)
|
||
|
|
op.create_index(
|
||
|
|
"ix_sessions_job_user",
|
||
|
|
"sessions",
|
||
|
|
["job_id", "user_id"],
|
||
|
|
unique=False,
|
||
|
|
)
|
||
|
|
_enable_rls("sessions")
|
||
|
|
|
||
|
|
op.create_table(
|
||
|
|
"messages",
|
||
|
|
sa.Column("id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("session_id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("seq", sa.Integer(), nullable=False),
|
||
|
|
sa.Column("role", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column("content", sa.Text(), nullable=False),
|
||
|
|
sa.Column("model_code", sa.String(length=50), nullable=True),
|
||
|
|
sa.Column("tool_name", sa.String(length=100), nullable=True),
|
||
|
|
sa.Column(
|
||
|
|
"input_tokens", sa.Integer(), nullable=False, server_default=sa.text("0")
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"output_tokens", sa.Integer(), nullable=False, server_default=sa.text("0")
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"cost",
|
||
|
|
sa.Numeric(precision=12, scale=6),
|
||
|
|
nullable=False,
|
||
|
|
server_default=sa.text("0"),
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"currency",
|
||
|
|
sa.String(length=3),
|
||
|
|
nullable=False,
|
||
|
|
server_default=sa.text("'USD'"),
|
||
|
|
),
|
||
|
|
sa.Column("latency_ms", sa.Integer(), nullable=True),
|
||
|
|
sa.Column("metadata", postgresql.JSONB(astext_type=sa.Text()), nullable=True),
|
||
|
|
sa.Column(
|
||
|
|
"created_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"updated_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
|
||
|
|
sa.PrimaryKeyConstraint("id"),
|
||
|
|
sa.UniqueConstraint("session_id", "seq", name="uq_messages_session_seq"),
|
||
|
|
)
|
||
|
|
op.create_index("ix_messages_session_id", "messages", ["session_id"], unique=False)
|
||
|
|
op.execute(
|
||
|
|
"ALTER TABLE messages ADD CONSTRAINT chk_message_role CHECK (role IN ('user', 'assistant', 'system', 'tool'))"
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_messages_session_id",
|
||
|
|
"messages",
|
||
|
|
"sessions",
|
||
|
|
["session_id"],
|
||
|
|
["id"],
|
||
|
|
ondelete="CASCADE",
|
||
|
|
)
|
||
|
|
_enable_rls("messages")
|
||
|
|
|
||
|
|
op.create_table(
|
||
|
|
"user_agents",
|
||
|
|
sa.Column("id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("user_id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("llm_id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("agent_type", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column(
|
||
|
|
"config",
|
||
|
|
postgresql.JSONB(astext_type=sa.Text()),
|
||
|
|
server_default="{}",
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column("status", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column("created_by", sa.UUID(), nullable=True),
|
||
|
|
sa.Column("updated_by", sa.UUID(), nullable=True),
|
||
|
|
sa.Column(
|
||
|
|
"created_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"updated_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
|
||
|
|
sa.PrimaryKeyConstraint("id"),
|
||
|
|
sa.UniqueConstraint("user_id", name="uq_user_agents_user_id"),
|
||
|
|
)
|
||
|
|
op.create_index(
|
||
|
|
"ix_user_agents_agent_type", "user_agents", ["agent_type"], unique=False
|
||
|
|
)
|
||
|
|
op.create_index("ix_user_agents_status", "user_agents", ["status"], unique=False)
|
||
|
|
op.execute(
|
||
|
|
"ALTER TABLE user_agents ADD CONSTRAINT chk_agent_type CHECK (agent_type IN ('INTENT_RECOGNITION', 'TASK_EXECUTION', 'RESULT_REPORTING'))"
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_user_agents_user_id",
|
||
|
|
"user_agents",
|
||
|
|
"users",
|
||
|
|
["user_id"],
|
||
|
|
["id"],
|
||
|
|
referent_schema="auth",
|
||
|
|
ondelete="CASCADE",
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_user_agents_llm_id",
|
||
|
|
"user_agents",
|
||
|
|
"llms",
|
||
|
|
["llm_id"],
|
||
|
|
["id"],
|
||
|
|
ondelete="RESTRICT",
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_user_agents_created_by",
|
||
|
|
"user_agents",
|
||
|
|
"users",
|
||
|
|
["created_by"],
|
||
|
|
["id"],
|
||
|
|
referent_schema="auth",
|
||
|
|
ondelete="SET NULL",
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_user_agents_updated_by",
|
||
|
|
"user_agents",
|
||
|
|
"users",
|
||
|
|
["updated_by"],
|
||
|
|
["id"],
|
||
|
|
referent_schema="auth",
|
||
|
|
ondelete="SET NULL",
|
||
|
|
)
|
||
|
|
_enable_rls("user_agents")
|
||
|
|
|
||
|
|
op.create_table(
|
||
|
|
"memories",
|
||
|
|
sa.Column("id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("owner_id", sa.UUID(), nullable=False),
|
||
|
|
sa.Column("agent_id", sa.UUID(), nullable=True),
|
||
|
|
sa.Column("memory_type", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column("title", sa.String(length=255), nullable=True),
|
||
|
|
sa.Column("content", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
|
||
|
|
sa.Column("source", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column("status", sa.String(length=20), nullable=False),
|
||
|
|
sa.Column(
|
||
|
|
"created_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.Column(
|
||
|
|
"updated_at",
|
||
|
|
sa.DateTime(timezone=True),
|
||
|
|
server_default=sa.text("now()"),
|
||
|
|
nullable=False,
|
||
|
|
),
|
||
|
|
sa.PrimaryKeyConstraint("id"),
|
||
|
|
)
|
||
|
|
op.create_index(
|
||
|
|
"ix_memories_owner_type_status",
|
||
|
|
"memories",
|
||
|
|
["owner_id", "memory_type", "status"],
|
||
|
|
unique=False,
|
||
|
|
)
|
||
|
|
op.create_index(
|
||
|
|
"ix_memories_agent_type_status",
|
||
|
|
"memories",
|
||
|
|
["agent_id", "memory_type", "status"],
|
||
|
|
unique=False,
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
"ALTER TABLE memories ADD CONSTRAINT chk_memory_type_agent_id CHECK ((memory_type = 'work' AND agent_id IS NOT NULL) OR (memory_type = 'user' AND agent_id IS NULL))"
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_memories_owner_id",
|
||
|
|
"memories",
|
||
|
|
"users",
|
||
|
|
["owner_id"],
|
||
|
|
["id"],
|
||
|
|
referent_schema="auth",
|
||
|
|
ondelete="CASCADE",
|
||
|
|
)
|
||
|
|
op.create_foreign_key(
|
||
|
|
"fk_memories_agent_id",
|
||
|
|
"memories",
|
||
|
|
"user_agents",
|
||
|
|
["agent_id"],
|
||
|
|
["id"],
|
||
|
|
ondelete="CASCADE",
|
||
|
|
)
|
||
|
|
_enable_rls("memories")
|
||
|
|
|
||
|
|
|
||
|
|
def downgrade() -> None:
|
||
|
|
_drop_rls("memories")
|
||
|
|
op.drop_constraint("fk_memories_agent_id", "memories", type_="foreignkey")
|
||
|
|
op.drop_constraint("fk_memories_owner_id", "memories", type_="foreignkey")
|
||
|
|
op.drop_table("memories")
|
||
|
|
|
||
|
|
_drop_rls("user_agents")
|
||
|
|
op.drop_constraint("fk_user_agents_updated_by", "user_agents", type_="foreignkey")
|
||
|
|
op.drop_constraint("fk_user_agents_created_by", "user_agents", type_="foreignkey")
|
||
|
|
op.drop_constraint("fk_user_agents_llm_id", "user_agents", type_="foreignkey")
|
||
|
|
op.drop_constraint("fk_user_agents_user_id", "user_agents", type_="foreignkey")
|
||
|
|
op.drop_table("user_agents")
|
||
|
|
|
||
|
|
_drop_rls("messages")
|
||
|
|
op.drop_constraint("fk_messages_session_id", "messages", type_="foreignkey")
|
||
|
|
op.drop_table("messages")
|
||
|
|
|
||
|
|
_drop_rls("sessions")
|
||
|
|
op.drop_constraint("fk_sessions_job_id_user_id", "sessions", type_="foreignkey")
|
||
|
|
op.drop_constraint("fk_sessions_user_id", "sessions", type_="foreignkey")
|
||
|
|
op.drop_table("sessions")
|
||
|
|
|
||
|
|
|
||
|
|
def _enable_rls(table_name: str) -> None:
|
||
|
|
for role in ["anon", "authenticated"]:
|
||
|
|
for action in ["select", "insert", "update", "delete"]:
|
||
|
|
op.execute(
|
||
|
|
f"DROP POLICY IF EXISTS {role}_{action}_{table_name} ON {table_name}"
|
||
|
|
)
|
||
|
|
op.execute(f"ALTER TABLE {table_name} ENABLE ROW LEVEL SECURITY")
|
||
|
|
for role in ["anon", "authenticated"]:
|
||
|
|
op.execute(
|
||
|
|
f"CREATE POLICY {role}_select_{table_name} ON {table_name} FOR SELECT TO {role} USING (false)"
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
f"CREATE POLICY {role}_insert_{table_name} ON {table_name} FOR INSERT TO {role} WITH CHECK (false)"
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
f"CREATE POLICY {role}_update_{table_name} ON {table_name} FOR UPDATE TO {role} USING (false) WITH CHECK (false)"
|
||
|
|
)
|
||
|
|
op.execute(
|
||
|
|
f"CREATE POLICY {role}_delete_{table_name} ON {table_name} FOR DELETE TO {role} USING (false)"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _drop_rls(table_name: str) -> None:
|
||
|
|
for role in ["anon", "authenticated"]:
|
||
|
|
op.execute(f"DROP POLICY IF EXISTS {role}_delete_{table_name} ON {table_name}")
|
||
|
|
op.execute(f"DROP POLICY IF EXISTS {role}_update_{table_name} ON {table_name}")
|
||
|
|
op.execute(f"DROP POLICY IF EXISTS {role}_insert_{table_name} ON {table_name}")
|
||
|
|
op.execute(f"DROP POLICY IF EXISTS {role}_select_{table_name} ON {table_name}")
|
||
|
|
op.execute(f"ALTER TABLE {table_name} DISABLE ROW LEVEL SECURITY")
|