add organization_id to script_blocks (#3150)

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]

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)

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.