ObserverThought reproduce migration script

Skyvern-AI · Dec 6, 2024 · 2d726bf · 2d726bf
1 parent 95bef7f
commit 2d726bf
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 21 deletions.
diff --git a/...6217e9d9_introduce_observercruise_and_.py → ...ed4719d5_introduce_observercruise_and_.py b/...6217e9d9_introduce_observercruise_and_.py → ...ed4719d5_introduce_observercruise_and_.py
@@ -1,8 +1,8 @@
 """Introduce ObserverCruise and ObserverThought. Add workflow_run_block_id and observer_cruise_id to artifacts
 
-Revision ID: 5e266217e9d9
+Revision ID: 4d51ed4719d5
 Revises: de0254717601
-Create Date: 2024-12-06 07:25:50.080125+00:00
+Create Date: 2024-12-06 08:52:52.111448+00:00
 
 """
 
@@ -13,7 +13,7 @@
 from alembic import op
 
 # revision identifiers, used by Alembic.
-revision: str = "5e266217e9d9"
+revision: str = "4d51ed4719d5"
 down_revision: Union[str, None] = "de0254717601"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -48,6 +48,7 @@ def upgrade() -> None:
         sa.Column("organization_id", sa.String(), nullable=True),
         sa.Column("observer_cruise_id", sa.String(), nullable=False),
         sa.Column("workflow_run_id", sa.String(), nullable=True),
+        sa.Column("workflow_run_block_id", sa.String(), nullable=True),
         sa.Column("workflow_id", sa.String(), nullable=True),
         sa.Column("thought", sa.String(), nullable=True),
         sa.Column("answer", sa.String(), nullable=True),
@@ -63,6 +64,10 @@ def upgrade() -> None:
             ["workflow_id"],
             ["workflows.workflow_id"],
         ),
+        sa.ForeignKeyConstraint(
+            ["workflow_run_block_id"],
+            ["workflow_run_blocks.workflow_run_block_id"],
+        ),
         sa.ForeignKeyConstraint(
             ["workflow_run_id"],
             ["workflow_runs.workflow_run_id"],
@@ -73,31 +78,19 @@ def upgrade() -> None:
     op.add_column("artifacts", sa.Column("workflow_run_block_id", sa.String(), nullable=True))
     op.add_column("artifacts", sa.Column("observer_cruise_id", sa.String(), nullable=True))
     op.create_index("org_workflow_run_index", "artifacts", ["organization_id", "workflow_run_id"], unique=False)
+    op.create_foreign_key(None, "artifacts", "observer_cruises", ["observer_cruise_id"], ["observer_cruise_id"])
     op.create_foreign_key(
-        "artifacts_workflow_run_block_id_fkey",
-        "artifacts",
-        "workflow_run_blocks",
-        ["workflow_run_block_id"],
-        ["workflow_run_block_id"],
-    )
-    op.create_foreign_key(
-        "artifacts_observer_cruise_id_fkey",
-        "artifacts",
-        "observer_cruises",
-        ["observer_cruise_id"],
-        ["observer_cruise_id"],
-    )
-    op.create_foreign_key(
-        "artifacts_workflow_run_id_fkey", "artifacts", "workflow_runs", ["workflow_run_id"], ["workflow_run_id"]
+        None, "artifacts", "workflow_run_blocks", ["workflow_run_block_id"], ["workflow_run_block_id"]
     )
+    op.create_foreign_key(None, "artifacts", "workflow_runs", ["workflow_run_id"], ["workflow_run_id"])
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_constraint("artifacts_workflow_run_block_id_fkey", "artifacts", type_="foreignkey")
-    op.drop_constraint("artifacts_observer_cruise_id_fkey", "artifacts", type_="foreignkey")
-    op.drop_constraint("artifacts_workflow_run_id_fkey", "artifacts", type_="foreignkey")
+    op.drop_constraint(None, "artifacts", type_="foreignkey")
+    op.drop_constraint(None, "artifacts", type_="foreignkey")
+    op.drop_constraint(None, "artifacts", type_="foreignkey")
     op.drop_index("org_workflow_run_index", table_name="artifacts")
     op.drop_column("artifacts", "observer_cruise_id")
     op.drop_column("artifacts", "workflow_run_block_id")

diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py
@@ -522,6 +522,7 @@ class ObserverThought(Base):
     organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True)
     observer_cruise_id = Column(String, ForeignKey("observer_cruises.observer_cruise_id"), nullable=False)
     workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True)
+    workflow_run_block_id = Column(String, ForeignKey("workflow_run_blocks.workflow_run_block_id"), nullable=True)
     workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=True)
     thought = Column(String, nullable=True)
     answer = Column(String, nullable=True)