From bb9c183fb336c5846b38203d2e9e38b7d7212911 Mon Sep 17 00:00:00 2001 From: Aaron Perez Date: Thu, 2 Jul 2026 19:35:17 -0500 Subject: [PATCH] refactor(SKY-11703): remove the AgentDB *args facade entirely (#7033) --- run_streaming.py | 4 +- scripts/check_no_direct_db_delegates.sh | 8 +- skyvern/forge/sdk/db/agent_db.py | 250 ------------------- skyvern/forge/sdk/db/repositories/tasks.py | 3 + tests/unit/forge/sdk/db/test_repositories.py | 14 +- 5 files changed, 24 insertions(+), 255 deletions(-) diff --git a/run_streaming.py b/run_streaming.py index c696a9e85..6048295a5 100644 --- a/run_streaming.py +++ b/run_streaming.py @@ -34,7 +34,7 @@ async def run() -> None: try: if workflow_run_id: - workflow_run = await app.DATABASE.get_workflow_run(workflow_run_id=workflow_run_id) + workflow_run = await app.DATABASE.workflow_runs.get_workflow_run(workflow_run_id=workflow_run_id) if not workflow_run or workflow_run.status in [ WorkflowRunStatus.completed, WorkflowRunStatus.failed, @@ -44,7 +44,7 @@ async def run() -> None: file_name = f"{workflow_run_id}.png" elif task_id: - task = await app.DATABASE.get_task(task_id=task_id, organization_id=organization_id) + task = await app.DATABASE.tasks.get_task(task_id=task_id, organization_id=organization_id) if not task or task.status.is_final(): continue file_name = f"{task_id}.png" diff --git a/scripts/check_no_direct_db_delegates.sh b/scripts/check_no_direct_db_delegates.sh index b2042ed93..7a5ef0917 100755 --- a/scripts/check_no_direct_db_delegates.sh +++ b/scripts/check_no_direct_db_delegates.sh @@ -17,8 +17,12 @@ delegate_methods=$( ) if [ -z "$delegate_methods" ]; then - echo "ERROR: Could not extract delegate methods from $AGENT_DB" >&2 - exit 1 + # The AgentDB *args backward-compat facade was removed entirely (SKY-11703). + # With no delegate methods left, there is nothing to police: a direct + # delegate call is now a plain AttributeError/mypy error, caught elsewhere. + # Nothing to check, so pass cleanly. + echo "No AgentDB delegate methods defined; nothing to check." + exit 0 fi # Build a grep alternation pattern for delegate method names. diff --git a/skyvern/forge/sdk/db/agent_db.py b/skyvern/forge/sdk/db/agent_db.py index 792ce8b2c..1fa842a61 100644 --- a/skyvern/forge/sdk/db/agent_db.py +++ b/skyvern/forge/sdk/db/agent_db.py @@ -38,7 +38,6 @@ from skyvern.forge.sdk.db.repositories.workflows import WorkflowsRepository from skyvern.forge.sdk.db.utils import ( _custom_json_serializer, ) -from skyvern.forge.sdk.trace import traced LOG = structlog.get_logger() @@ -413,252 +412,3 @@ class AgentDB(BaseAlchemyDB): def is_retryable_error(self, error: SQLAlchemyError) -> bool: error_msg = str(error).lower() return "server closed the connection" in error_msg - - # ====================================================================== - # Backward-compatible delegate methods - # ====================================================================== - - # -- Task delegates -- - - async def create_task(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.create_task(*args, **kwargs) - - async def create_step(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.create_step(*args, **kwargs) - - async def get_task(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_task(*args, **kwargs) - - async def get_tasks_by_ids(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_tasks_by_ids(*args, **kwargs) - - async def get_step(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_step(*args, **kwargs) - - async def get_task_steps(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_task_steps(*args, **kwargs) - - async def get_steps_by_task_ids(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_steps_by_task_ids(*args, **kwargs) - - async def get_total_unique_step_order_count_by_task_ids(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_total_unique_step_order_count_by_task_ids(*args, **kwargs) - - async def get_workflow_run_progress_timestamps(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_workflow_run_progress_timestamps(*args, **kwargs) - - async def get_task_step_models(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_task_step_models(*args, **kwargs) - - async def get_task_step_count(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_task_step_count(*args, **kwargs) - - async def get_task_actions(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_task_actions(*args, **kwargs) - - async def get_task_actions_hydrated(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_task_actions_hydrated(*args, **kwargs) - - async def get_tasks_actions(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_tasks_actions(*args, **kwargs) - - async def get_action_count_for_step(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_action_count_for_step(*args, **kwargs) - - async def get_first_step(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_first_step(*args, **kwargs) - - async def get_latest_step(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_latest_step(*args, **kwargs) - - @traced(name="skyvern.db.update_step") - async def update_step(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.update_step(*args, **kwargs) - - async def clear_task_failure_reason(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.clear_task_failure_reason(*args, **kwargs) - - @traced(name="skyvern.db.update_task") - async def update_task(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.update_task(*args, **kwargs) - - async def update_task_2fa_state(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.update_task_2fa_state(*args, **kwargs) - - async def bulk_update_tasks(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.bulk_update_tasks(*args, **kwargs) - - async def get_tasks(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_tasks(*args, **kwargs) - - async def get_tasks_count(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_tasks_count(*args, **kwargs) - - async def get_running_tasks_info_globally(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_running_tasks_info_globally(*args, **kwargs) - - async def get_latest_task_by_workflow_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_latest_task_by_workflow_id(*args, **kwargs) - - async def get_last_task_for_workflow_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_last_task_for_workflow_run(*args, **kwargs) - - async def get_tasks_by_workflow_run_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_tasks_by_workflow_run_id(*args, **kwargs) - - async def delete_task_steps(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.delete_task_steps(*args, **kwargs) - - async def get_previous_actions_for_task(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_previous_actions_for_task(*args, **kwargs) - - async def delete_task_actions(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.delete_task_actions(*args, **kwargs) - - # -- Workflow run delegates -- - - async def get_running_workflow_runs_info_globally(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_running_workflow_runs_info_globally(*args, **kwargs) - - async def create_workflow_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.create_workflow_run(*args, **kwargs) - - async def update_workflow_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.update_workflow_run(*args, **kwargs) - - async def bulk_update_workflow_runs(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.bulk_update_workflow_runs(*args, **kwargs) - - async def clear_workflow_run_failure_reason(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.clear_workflow_run_failure_reason(*args, **kwargs) - - async def get_all_runs(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_all_runs(*args, **kwargs) - - async def get_all_runs_v2(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_all_runs_v2(*args, **kwargs) - - async def get_workflow_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_run(*args, **kwargs) - - async def get_last_queued_workflow_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_last_queued_workflow_run(*args, **kwargs) - - async def get_workflow_runs_by_ids(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_runs_by_ids(*args, **kwargs) - - async def get_last_running_workflow_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_last_running_workflow_run(*args, **kwargs) - - async def get_last_workflow_run_for_browser_session(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_last_workflow_run_for_browser_session(*args, **kwargs) - - async def get_last_workflow_run_for_browser_address(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_last_workflow_run_for_browser_address(*args, **kwargs) - - async def get_workflows_depending_on(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflows_depending_on(*args, **kwargs) - - async def get_workflow_runs(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_runs(*args, **kwargs) - - async def get_workflow_runs_count(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_runs_count(*args, **kwargs) - - async def get_workflow_runs_for_workflow_permanent_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_runs_for_workflow_permanent_id(*args, **kwargs) - - async def get_workflow_runs_by_parent_workflow_run_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_runs_by_parent_workflow_run_id(*args, **kwargs) - - async def get_workflow_run_output_parameters(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_run_output_parameters(*args, **kwargs) - - async def get_workflow_run_output_parameter_by_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_run_output_parameter_by_id(*args, **kwargs) - - async def create_or_update_workflow_run_output_parameter(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.create_or_update_workflow_run_output_parameter(*args, **kwargs) - - async def update_workflow_run_output_parameter(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.update_workflow_run_output_parameter(*args, **kwargs) - - async def create_workflow_run_parameter(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.create_workflow_run_parameter(*args, **kwargs) - - async def create_workflow_run_parameters(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.create_workflow_run_parameters(*args, **kwargs) - - async def get_workflow_run_parameters(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_run_parameters(*args, **kwargs) - - async def get_workflow_run_block_errors(self, *args: Any, **kwargs: Any) -> Any: - return await self.workflow_runs.get_workflow_run_block_errors(*args, **kwargs) - - async def create_task_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.create_task_run(*args, **kwargs) - - async def update_task_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.update_task_run(*args, **kwargs) - - async def sync_task_run_status(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.sync_task_run_status(*args, **kwargs) - - async def update_job_run_compute_cost(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.update_job_run_compute_cost(*args, **kwargs) - - async def cache_task_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.cache_task_run(*args, **kwargs) - - async def get_cached_task_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_cached_task_run(*args, **kwargs) - - async def get_run(self, *args: Any, **kwargs: Any) -> Any: - return await self.tasks.get_run(*args, **kwargs) - - # -- Observer delegates -- - - async def get_task_v2(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_task_v2(*args, **kwargs) - - async def delete_thoughts(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.delete_thoughts(*args, **kwargs) - - async def get_task_v2_by_workflow_run_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_task_v2_by_workflow_run_id(*args, **kwargs) - - async def get_thought(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_thought(*args, **kwargs) - - async def get_thoughts(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_thoughts(*args, **kwargs) - - async def create_task_v2(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.create_task_v2(*args, **kwargs) - - async def create_thought(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.create_thought(*args, **kwargs) - - async def update_thought(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.update_thought(*args, **kwargs) - - async def update_task_v2(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.update_task_v2(*args, **kwargs) - - async def create_workflow_run_block(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.create_workflow_run_block(*args, **kwargs) - - async def delete_workflow_run_blocks(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.delete_workflow_run_blocks(*args, **kwargs) - - async def update_workflow_run_block(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.update_workflow_run_block(*args, **kwargs) - - async def get_workflow_run_block(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_workflow_run_block(*args, **kwargs) - - async def get_workflow_run_block_by_task_id(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_workflow_run_block_by_task_id(*args, **kwargs) - - async def get_workflow_run_blocks(self, *args: Any, **kwargs: Any) -> Any: - return await self.observer.get_workflow_run_blocks(*args, **kwargs) diff --git a/skyvern/forge/sdk/db/repositories/tasks.py b/skyvern/forge/sdk/db/repositories/tasks.py index 69ce1b0ed..e17b9bd96 100644 --- a/skyvern/forge/sdk/db/repositories/tasks.py +++ b/skyvern/forge/sdk/db/repositories/tasks.py @@ -26,6 +26,7 @@ from skyvern.forge.sdk.db.utils import convert_to_step, convert_to_task, hydrate from skyvern.forge.sdk.models import Step, StepStatus from skyvern.forge.sdk.schemas.runs import Run from skyvern.forge.sdk.schemas.tasks import OrderBy, SortDirection, Task, TaskStatus +from skyvern.forge.sdk.trace import traced from skyvern.forge.sdk.utils.sanitization import sanitize_postgres_text from skyvern.schemas.runs import ProxyLocationInput, RunStatus, RunType from skyvern.schemas.steps import AgentStepOutput @@ -457,6 +458,7 @@ class TasksRepository(BaseRepository): ) return None + @traced(name="skyvern.db.update_step") @db_operation("update_step") async def update_step( self, @@ -533,6 +535,7 @@ class TasksRepository(BaseRepository): else: raise NotFoundError("Task not found") + @traced(name="skyvern.db.update_task") @db_operation("update_task") async def update_task( self, diff --git a/tests/unit/forge/sdk/db/test_repositories.py b/tests/unit/forge/sdk/db/test_repositories.py index 0799ff03a..80bfbadcb 100644 --- a/tests/unit/forge/sdk/db/test_repositories.py +++ b/tests/unit/forge/sdk/db/test_repositories.py @@ -1,5 +1,6 @@ """Tests for all OSS repository instantiations + dependency injection.""" +import inspect from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -255,8 +256,19 @@ def test_agent_db_has_typed_repo_attributes(): assert isinstance(db.tasks, TasksRepository) assert isinstance(db.credentials, CredentialRepository) assert isinstance(db.credential_folders, CredentialFoldersRepository) - assert hasattr(db, "get_task") # backward compat delegate # Migrated domains no longer have delegates on AgentDB: assert not hasattr(db, "create_workflow") assert not hasattr(db, "get_organization") assert not hasattr(db, "get_credential") + + +def test_agent_db_defines_no_delegator_methods(): + """All data access goes through typed repository attributes; AgentDB itself defines no forwarding methods.""" + from skyvern.forge.sdk.db.agent_db import AgentDB + + defined = {name for name, member in vars(AgentDB).items() if inspect.isfunction(member)} + assert defined == {"__init__", "is_retryable_error"}, ( + f"Unexpected methods on AgentDB: {sorted(defined - {'__init__', 'is_retryable_error'})}. " + "Add data-access methods to the domain repository and call it via the typed attribute " + "(e.g. db.tasks.get_task) instead of adding delegators to AgentDB." + )