SurfSense/surfsense_backend/alembic/versions/4_add_linkup_api_enum.py

60 lines
1.6 KiB
Python

"""Add LINKUP_API to SearchSourceConnectorType enum
Revision ID: 4
Revises: 3
"""
from typing import Sequence, Union
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "4"
down_revision: Union[str, None] = "3"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
ENUM_NAME = "searchsourceconnectortype"
NEW_VALUE = "LINKUP_API"
op.execute(
f"""
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_enum
WHERE enumlabel = '{NEW_VALUE}'
AND enumtypid = (
SELECT oid FROM pg_type WHERE typname = '{ENUM_NAME}'
)
) THEN
ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}';
END IF;
END$$;
"""
)
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
# Downgrading removal of an enum value requires recreating the type
op.execute(
"ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old"
)
op.execute(
"CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR', 'LINEAR_CONNECTOR')"
)
op.execute(
(
"ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING "
"connector_type::text::searchsourceconnectortype"
)
)
op.execute("DROP TYPE searchsourceconnectortype_old")
pass
# ### end Alembic commands ###