src/backend/base/langflow/alembic/versions/6e7b581b5648_fix_nullable.py

"""Fix nullable

Revision ID: 6e7b581b5648
Revises: 58b28437a398
Create Date: 2024-04-30 09:17:45.024688

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.engine.reflection import Inspector

# revision identifiers, used by Alembic.
revision: str = "6e7b581b5648"
down_revision: Union[str, None] = "58b28437a398"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    conn = op.get_bind()
    inspector = Inspector.from_engine(conn)  # type: ignore
    # table_names = inspector.get_table_names()
    # ### commands auto generated by Alembic - please adjust! ###
    columns = inspector.get_columns("apikey")
    column_names = {column["name"]: column for column in columns}
    with op.batch_alter_table("apikey", schema=None) as batch_op:
        created_at_column = [column for column in columns if column["name"] == "created_at"][0]
        if "created_at" in column_names and created_at_column.get("nullable"):
            batch_op.alter_column(
                "created_at",
                existing_type=sa.DATETIME(),
                nullable=False,
                existing_server_default=sa.text("(CURRENT_TIMESTAMP)"),  # type: ignore
            )

    # ### end Alembic commands ###


def downgrade() -> None:
    conn = op.get_bind()
    inspector = Inspector.from_engine(conn)  # type: ignore
    # table_names = inspector.get_table_names()
    columns = inspector.get_columns("apikey")
    column_names = {column["name"]: column for column in columns}
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("apikey", schema=None) as batch_op:
        created_at_column = [column for column in columns if column["name"] == "created_at"][0]
        if "created_at" in column_names and not created_at_column.get("nullable"):
            batch_op.alter_column(
                "created_at",
                existing_type=sa.DATETIME(),
                nullable=True,
                existing_server_default=sa.text("(CURRENT_TIMESTAMP)"),  # type: ignore
            )

    # ### end Alembic commands ###