🔧 chore(alembic): add store updates migration script

🔧 chore(alembic): add user id can be null in flow migration script
🔧 chore(flow.py): make is_component field nullable in Flow model
🔧 chore(flow.py): make user_id field nullable in Flow model
🔧 chore(service.py): add schema health check after running migrations and handle autogenerate diffs detected case
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-10-18 23:23:10 -03:00
commit d66944ae7e
4 changed files with 75 additions and 16 deletions

View file

@ -1,8 +1,8 @@
"""Changes for the store
"""Store updates
Revision ID: 595c9c2a2ad4
Revision ID: 7843803a87b5
Revises: eb5866d51fd2
Create Date: 2023-10-18 16:27:37.781613
Create Date: 2023-10-18 23:08:57.744906
"""
from typing import Sequence, Union
@ -13,7 +13,7 @@ import sqlmodel
# revision identifiers, used by Alembic.
revision: str = "595c9c2a2ad4"
revision: str = "7843803a87b5"
down_revision: Union[str, None] = "eb5866d51fd2"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
@ -23,7 +23,7 @@ def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
try:
with op.batch_alter_table("flow", schema=None) as batch_op:
batch_op.add_column(sa.Column("is_component", sa.Boolean(), nullable=False))
batch_op.add_column(sa.Column("is_component", sa.Boolean(), nullable=True))
with op.batch_alter_table("user", schema=None) as batch_op:
batch_op.add_column(
@ -33,17 +33,16 @@ def upgrade() -> None:
)
except Exception:
pass
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
try:
with op.batch_alter_table("user", schema=None) as batch_op:
batch_op.drop_column("store_api_key")
with op.batch_alter_table("user", schema=None) as batch_op:
batch_op.drop_column("store_api_key")
with op.batch_alter_table("flow", schema=None) as batch_op:
batch_op.drop_column("is_component")
with op.batch_alter_table("flow", schema=None) as batch_op:
batch_op.drop_column("is_component")
except Exception:
pass
# ### end Alembic commands ###

View file

@ -0,0 +1,42 @@
"""User id can be null in Flow
Revision ID: f5ee9749d1a6
Revises: 7843803a87b5
Create Date: 2023-10-18 23:12:27.297016
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
import sqlmodel
# revision identifiers, used by Alembic.
revision: str = "f5ee9749d1a6"
down_revision: Union[str, None] = "7843803a87b5"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
try:
with op.batch_alter_table("flow", schema=None) as batch_op:
batch_op.alter_column(
"user_id", existing_type=sa.CHAR(length=32), nullable=True
)
except Exception:
pass
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table("flow", schema=None) as batch_op:
batch_op.alter_column(
"user_id", existing_type=sa.CHAR(length=32), nullable=False
)
# ### end Alembic commands ###

View file

@ -15,7 +15,7 @@ class FlowBase(SQLModelSerializable):
name: str = Field(index=True)
description: Optional[str] = Field(index=True)
data: Optional[Dict] = Field(default=None, nullable=True)
is_component: bool = Field(default=False)
is_component: bool = Field(default=False, nullable=True)
@validator("data")
def validate_json(v):
@ -36,7 +36,7 @@ class FlowBase(SQLModelSerializable):
class Flow(FlowBase, table=True):
id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True)
data: Optional[Dict] = Field(default=None, sa_column=Column(JSON))
user_id: UUID = Field(index=True, foreign_key="user.id")
user_id: UUID = Field(index=True, foreign_key="user.id", nullable=True)
user: "User" = Relationship(back_populates="flows")

View file

@ -10,8 +10,9 @@ from sqlalchemy.exc import OperationalError
from sqlmodel import SQLModel, Session, create_engine
from loguru import logger
from alembic.config import Config
from alembic import command
from alembic import command, util
from langflow.services.database import models # noqa
import time
if TYPE_CHECKING:
from sqlalchemy.engine import Engine
@ -120,7 +121,24 @@ class DatabaseService(Service):
alembic_cfg = Config()
alembic_cfg.set_main_option("script_location", str(self.script_location))
alembic_cfg.set_main_option("sqlalchemy.url", self.database_url)
command.upgrade(alembic_cfg, "head")
try:
command.check(alembic_cfg)
except Exception as exc:
if isinstance(exc, util.exc.CommandError) or isinstance(
exc, util.exc.AutogenerateDiffsDetected
):
command.upgrade(alembic_cfg, "head")
# We should check the schema health after running migrations
try:
command.check(alembic_cfg)
except util.exc.AutogenerateDiffsDetected:
# downgrade to base and upgrade again
logger.warning("Autogenerate diffs detected, downgrading and upgrading")
command.downgrade(alembic_cfg, "-1")
# wait for the database to be ready
time.sleep(5)
command.upgrade(alembic_cfg, "head")
def run_migrations_test(self):
# This method is used for testing purposes only