dograh/api/alembic/versions/0c1223cc266f_make_json_not_nullable.py

"""make json not nullable

Revision ID: 0c1223cc266f
Revises: 0fe708f2acb9
Create Date: 2025-07-21 14:36:31.182969

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

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


def upgrade() -> None:
    # Update existing NULL values to empty dict for workflows table
    op.execute("""
        UPDATE workflows 
        SET template_context_variables = '{}'::jsonb 
        WHERE template_context_variables IS NULL
    """)
    op.execute("""
        UPDATE workflows 
        SET call_disposition_codes = '{}'::jsonb 
        WHERE call_disposition_codes IS NULL
    """)

    # Update existing NULL values to empty dict for workflow_runs table
    op.execute("""
        UPDATE workflow_runs 
        SET usage_info = '{}'::jsonb 
        WHERE usage_info IS NULL
    """)
    op.execute("""
        UPDATE workflow_runs 
        SET cost_info = '{}'::jsonb 
        WHERE cost_info IS NULL
    """)
    op.execute("""
        UPDATE workflow_runs 
        SET initial_context = '{}'::jsonb 
        WHERE initial_context IS NULL
    """)
    op.execute("""
        UPDATE workflow_runs 
        SET gathered_context = '{}'::jsonb 
        WHERE gathered_context IS NULL
    """)
    op.execute("""
        UPDATE workflow_runs 
        SET annotations = '{}'::jsonb 
        WHERE annotations IS NULL
    """)

    # Update existing NULL values to empty dict for looptalk_test_sessions table
    op.execute("""
        UPDATE looptalk_test_sessions 
        SET results = '{}'::jsonb 
        WHERE results IS NULL
    """)

    # Update existing NULL values to empty dict for looptalk_conversations table
    op.execute("""
        UPDATE looptalk_conversations 
        SET transcript = '{}'::jsonb 
        WHERE transcript IS NULL
    """)
    op.execute("""
        UPDATE looptalk_conversations 
        SET metrics = '{}'::jsonb 
        WHERE metrics IS NULL
    """)

    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column(
        "looptalk_conversations",
        "transcript",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "looptalk_conversations",
        "metrics",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "looptalk_test_sessions",
        "results",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflow_runs",
        "usage_info",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflow_runs",
        "cost_info",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflow_runs",
        "initial_context",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflow_runs",
        "gathered_context",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflow_runs",
        "annotations",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflows",
        "template_context_variables",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    op.alter_column(
        "workflows",
        "call_disposition_codes",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=False,
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.alter_column(
        "workflows",
        "call_disposition_codes",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "workflows",
        "template_context_variables",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "workflow_runs",
        "annotations",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "workflow_runs",
        "gathered_context",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "workflow_runs",
        "initial_context",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "workflow_runs",
        "cost_info",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "workflow_runs",
        "usage_info",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "looptalk_test_sessions",
        "results",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "looptalk_conversations",
        "metrics",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    op.alter_column(
        "looptalk_conversations",
        "transcript",
        existing_type=postgresql.JSON(astext_type=sa.Text()),
        nullable=True,
    )
    # ### end Alembic commands ###