"""Add subgraph support for Enterprise and Premium tiers

Revision ID: 2bcbf00568eb
Revises: c2f378adcdb4
Create Date: 2025-08-22 14:42:30.714567

"""

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

# revision identifiers, used by Alembic.
revision = "2bcbf00568eb"
down_revision = "c2f378adcdb4"
branch_labels = None
depends_on = None


def upgrade() -> None:
  # ### commands auto generated by Alembic - please adjust! ###
  op.add_column("graphs", sa.Column("parent_graph_id", sa.String(), nullable=True))
  op.add_column("graphs", sa.Column("subgraph_index", sa.Integer(), nullable=True))
  op.add_column("graphs", sa.Column("subgraph_name", sa.String(), nullable=True))
  op.add_column(
    "graphs",
    sa.Column("is_subgraph", sa.Boolean(), nullable=False, server_default="false"),
  )
  op.add_column(
    "graphs",
    sa.Column(
      "subgraph_metadata", postgresql.JSONB(astext_type=sa.Text()), nullable=True
    ),
  )
  op.create_index("idx_graphs_is_subgraph", "graphs", ["is_subgraph"], unique=False)
  op.create_index("idx_graphs_parent", "graphs", ["parent_graph_id"], unique=False)
  op.create_unique_constraint(
    "unique_subgraph_index", "graphs", ["parent_graph_id", "subgraph_index"]
  )
  op.create_check_constraint(
    "check_subgraph_consistency",
    "graphs",
    "(is_subgraph = false AND parent_graph_id IS NULL AND subgraph_index IS NULL AND subgraph_name IS NULL) OR "
    "(is_subgraph = true AND parent_graph_id IS NOT NULL AND subgraph_index IS NOT NULL AND subgraph_name IS NOT NULL)",
  )
  # ### end Alembic commands ###


def downgrade() -> None:
  # ### commands auto generated by Alembic - please adjust! ###
  op.drop_constraint("check_subgraph_consistency", "graphs", type_="check")
  op.drop_constraint("unique_subgraph_index", "graphs", type_="unique")
  op.drop_index("idx_graphs_parent", table_name="graphs")
  op.drop_index("idx_graphs_is_subgraph", table_name="graphs")
  op.drop_column("graphs", "subgraph_metadata")
  op.drop_column("graphs", "is_subgraph")
  op.drop_column("graphs", "subgraph_name")
  op.drop_column("graphs", "subgraph_index")
  op.drop_column("graphs", "parent_graph_id")
  # ### end Alembic commands ###
