"""Add cascading deletes for sources to agents

Revision ID: e78b4e82db30
Revises: d6632deac81d
Create Date: 2024-12-20 16:30:17.095888

"""

from typing import Sequence, Union

from alembic import op
from letta.settings import settings

# revision identifiers, used by Alembic.
revision: str = "e78b4e82db30"
down_revision: Union[str, None] = "d6632deac81d"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint("sources_agents_agent_id_fkey", "sources_agents", type_="foreignkey")
    op.drop_constraint("sources_agents_source_id_fkey", "sources_agents", type_="foreignkey")
    op.create_foreign_key(None, "sources_agents", "sources", ["source_id"], ["id"], ondelete="CASCADE")
    op.create_foreign_key(None, "sources_agents", "agents", ["agent_id"], ["id"], ondelete="CASCADE")
    # ### end Alembic commands ###


def downgrade() -> None:
    # Skip this migration for SQLite
    if not settings.letta_pg_uri_no_default:
        return

    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_constraint(None, "sources_agents", type_="foreignkey")
    op.drop_constraint(None, "sources_agents", type_="foreignkey")
    op.create_foreign_key("sources_agents_source_id_fkey", "sources_agents", "sources", ["source_id"], ["id"])
    op.create_foreign_key("sources_agents_agent_id_fkey", "sources_agents", "agents", ["agent_id"], ["id"])
    # ### end Alembic commands ###
