From 136fa70c48085320ae7a3e79cf4e5c6347358c83 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Fri, 8 Aug 2025 20:24:44 -0700 Subject: [PATCH] add organization_id to script_blocks (#3150) --- ...70_add_organization_id_to_script_blocks.py | 31 +++++++++++++++++ skyvern/forge/sdk/db/client.py | 34 ++++++++++++++++++- skyvern/forge/sdk/db/models.py | 1 + skyvern/forge/sdk/db/utils.py | 17 +++++++++- skyvern/schemas/scripts.py | 12 +++++++ 5 files changed, 93 insertions(+), 2 deletions(-) create mode 100644 alembic/versions/2025_08_09_0105-a027553be970_add_organization_id_to_script_blocks.py diff --git a/alembic/versions/2025_08_09_0105-a027553be970_add_organization_id_to_script_blocks.py b/alembic/versions/2025_08_09_0105-a027553be970_add_organization_id_to_script_blocks.py new file mode 100644 index 00000000..984799df --- /dev/null +++ b/alembic/versions/2025_08_09_0105-a027553be970_add_organization_id_to_script_blocks.py @@ -0,0 +1,31 @@ +"""add organization_id to script_blocks + +Revision ID: a027553be970 +Revises: 92d665af080e +Create Date: 2025-08-09 01:05:07.564940+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "a027553be970" +down_revision: Union[str, None] = "92d665af080e" +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! ### + op.add_column("script_blocks", sa.Column("organization_id", sa.String(), nullable=False)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("script_blocks", "organization_id") + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index f758bda7..d5f5a194 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -29,6 +29,7 @@ from skyvern.forge.sdk.db.models import ( OrganizationModel, OutputParameterModel, PersistentBrowserSessionModel, + ScriptBlockModel, ScriptFileModel, ScriptModel, StepModel, @@ -55,6 +56,7 @@ from skyvern.forge.sdk.db.utils import ( convert_to_organization_auth_token, convert_to_output_parameter, convert_to_script, + convert_to_script_block, convert_to_script_file, convert_to_step, convert_to_task, @@ -103,7 +105,7 @@ from skyvern.forge.sdk.workflow.models.workflow import ( WorkflowStatus, ) from skyvern.schemas.runs import ProxyLocation, RunEngine, RunType -from skyvern.schemas.scripts import Script, ScriptFile +from skyvern.schemas.scripts import Script, ScriptBlock, ScriptFile from skyvern.webeye.actions.actions import Action from skyvern.webeye.actions.models import AgentStepOutput @@ -3775,3 +3777,33 @@ class AgentDB: ) ).all() return [convert_to_script_file(script_file) for script_file in script_files] + + async def get_script_block( + self, + script_block_id: str, + organization_id: str, + ) -> ScriptBlock | None: + async with self.Session() as session: + record = ( + await session.scalars( + select(ScriptBlockModel) + .filter_by(script_block_id=script_block_id) + .filter_by(organization_id=organization_id) + ) + ).first() + return convert_to_script_block(record) if record else None + + async def get_script_blocks_by_script_revision_id( + self, + script_revision_id: str, + organization_id: str, + ) -> list[ScriptBlock]: + async with self.Session() as session: + records = ( + await session.scalars( + select(ScriptBlockModel) + .filter_by(script_revision_id=script_revision_id) + .filter_by(organization_id=organization_id) + ) + ).all() + return [convert_to_script_block(record) for record in records] diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index be0c8cce..804c405b 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -875,6 +875,7 @@ class ScriptBlockModel(Base): ) script_block_id = Column(String, primary_key=True, default=generate_script_block_id) + organization_id = Column(String, nullable=False) script_id = Column(String, nullable=False) script_revision_id = Column(String, nullable=False, index=True) script_block_label = Column(String, nullable=False) diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index d6dc100a..1fc5d452 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -15,6 +15,7 @@ from skyvern.forge.sdk.db.models import ( OrganizationAuthTokenModel, OrganizationModel, OutputParameterModel, + ScriptBlockModel, ScriptFileModel, ScriptModel, StepModel, @@ -51,7 +52,7 @@ from skyvern.forge.sdk.workflow.models.workflow import ( WorkflowStatus, ) from skyvern.schemas.runs import ProxyLocation -from skyvern.schemas.scripts import Script, ScriptFile +from skyvern.schemas.scripts import Script, ScriptBlock, ScriptFile from skyvern.webeye.actions.actions import ( Action, ActionType, @@ -538,6 +539,20 @@ def convert_to_script_file(script_file_model: ScriptFileModel) -> ScriptFile: ) +def convert_to_script_block(script_block_model: ScriptBlockModel) -> ScriptBlock: + return ScriptBlock( + script_block_id=script_block_model.script_block_id, + organization_id=script_block_model.organization_id, + script_id=script_block_model.script_id, + script_revision_id=script_block_model.script_revision_id, + script_block_label=script_block_model.script_block_label, + script_file_id=script_block_model.script_file_id, + created_at=script_block_model.created_at, + modified_at=script_block_model.modified_at, + deleted_at=script_block_model.deleted_at, + ) + + def hydrate_action(action_model: ActionModel) -> Action: """ Convert ActionModel to the appropriate Action type based on action_type. diff --git a/skyvern/schemas/scripts.py b/skyvern/schemas/scripts.py index 18513797..702ba110 100644 --- a/skyvern/schemas/scripts.py +++ b/skyvern/schemas/scripts.py @@ -125,3 +125,15 @@ class Script(BaseModel): created_at: datetime = Field(description="Timestamp when the script was created") modified_at: datetime = Field(description="Timestamp when the script was last modified") deleted_at: datetime | None = Field(default=None, description="Timestamp when the script was soft deleted") + + +class ScriptBlock(BaseModel): + script_block_id: str + organization_id: str + script_id: str + script_revision_id: str + script_block_label: str + script_file_id: str + created_at: datetime + modified_at: datetime + deleted_at: datetime | None = None