"""Add consent indexes to stagedresource

Revision ID: d45dec7e541d
Revises: 3b0f4905378c
Create Date: 2025-07-01 19:50:25.009417

"""

import sqlalchemy as sa
from alembic import op
from loguru import logger

# revision identifiers, used by Alembic.
revision = "d45dec7e541d"
down_revision = "3b0f4905378c"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    connection = op.get_bind()

    # Check stagedresource table size to decide if we should create indexes immediately
    stagedresource_count = connection.execute(
        sa.text("SELECT COUNT(*) FROM stagedresource")
    ).scalar()

    if stagedresource_count < 1000000:
        # Create indexes normally for smaller tables
        op.create_index(
            "ix_stagedresource_monitor_config_resource_type_consent",
            "stagedresource",
            [
                "monitor_config_id",
                "resource_type",
                sa.text("(meta->>'consent_aggregated')"),
            ],
            unique=False,
        )
        op.create_index(
            "ix_stagedresource_system_vendor_consent",
            "stagedresource",
            ["system_id", "vendor_id", sa.text("(meta->>'consent_aggregated')")],
            unique=False,
        )
    else:
        logger.warning(
            "The stagedresource table has more than 1 million rows, "
            "skipping index creation. Indexes will be created during application startup "
            "via post_upgrade_index_creation.py"
        )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # Drop indexes if they exist (they may not exist if they were deferred)
    op.execute("DROP INDEX IF EXISTS ix_stagedresource_system_vendor_consent")
    op.execute(
        "DROP INDEX IF EXISTS ix_stagedresource_monitor_config_resource_type_consent"
    )

    # ### end Alembic commands ###
