diff --git a/api/core/repositories/celery_workflow_execution_repository.py b/api/core/repositories/celery_workflow_execution_repository.py index e87d1cd6b2..d65c71abc8 100644 --- a/api/core/repositories/celery_workflow_execution_repository.py +++ b/api/core/repositories/celery_workflow_execution_repository.py @@ -6,6 +6,7 @@ providing improved performance by offloading database operations to background w """ import logging +from typing import override from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker @@ -92,6 +93,7 @@ class CeleryWorkflowExecutionRepository(WorkflowExecutionRepository): self._triggered_from, ) + @override def save(self, execution: WorkflowExecution): """ Save or update a WorkflowExecution instance asynchronously using Celery. diff --git a/api/core/repositories/celery_workflow_node_execution_repository.py b/api/core/repositories/celery_workflow_node_execution_repository.py index 2451563317..dc2588b489 100644 --- a/api/core/repositories/celery_workflow_node_execution_repository.py +++ b/api/core/repositories/celery_workflow_node_execution_repository.py @@ -7,6 +7,7 @@ providing improved performance by offloading database operations to background w import logging from collections.abc import Sequence +from typing import override from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker @@ -105,6 +106,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): self._triggered_from, ) + @override def save(self, execution: WorkflowNodeExecution): """ Save or update a WorkflowNodeExecution instance to cache and asynchronously to database. @@ -147,6 +149,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): # For now, we'll re-raise the exception raise + @override def get_by_workflow_execution( self, workflow_execution_id: str, diff --git a/api/core/repositories/human_input_repository.py b/api/core/repositories/human_input_repository.py index 740d727e26..7a97b32838 100644 --- a/api/core/repositories/human_input_repository.py +++ b/api/core/repositories/human_input_repository.py @@ -2,7 +2,7 @@ import dataclasses import json from collections.abc import Mapping, Sequence from datetime import datetime -from typing import Any, Protocol +from typing import Any, Protocol, override from sqlalchemy import select from sqlalchemy.orm import Session, selectinload @@ -113,10 +113,12 @@ class _HumanInputFormRecipientEntityImpl(HumanInputFormRecipientEntity): self._recipient_model = recipient_model @property + @override def id(self) -> str: return self._recipient_model.id @property + @override def token(self) -> str: if self._recipient_model.access_token is None: raise AssertionError(f"access_token should not be None for recipient {self._recipient_model.id}") @@ -144,10 +146,12 @@ class _HumanInputFormEntityImpl(HumanInputFormEntity): ) @property + @override def id(self) -> str: return self._form_model.id @property + @override def submission_token(self) -> str | None: if self._console_recipient is not None: return self._console_recipient.access_token @@ -156,30 +160,37 @@ class _HumanInputFormEntityImpl(HumanInputFormEntity): return self._interactive_surface_recipient.access_token @property + @override def recipients(self) -> list[HumanInputFormRecipientEntity]: return list(self._recipients) @property + @override def rendered_content(self) -> str: return self._form_model.rendered_content @property + @override def selected_action_id(self) -> str | None: return self._form_model.selected_action_id @property + @override def submitted_data(self) -> Mapping[str, Any] | None: return self._submitted_data @property + @override def submitted(self) -> bool: return self._form_model.submitted_at is not None @property + @override def status(self) -> HumanInputFormStatus: return self._form_model.status @property + @override def expiration_time(self) -> datetime: return self._form_model.expiration_time diff --git a/api/core/repositories/sqlalchemy_workflow_execution_repository.py b/api/core/repositories/sqlalchemy_workflow_execution_repository.py index 6be3902317..4bca8b34e8 100644 --- a/api/core/repositories/sqlalchemy_workflow_execution_repository.py +++ b/api/core/repositories/sqlalchemy_workflow_execution_repository.py @@ -4,6 +4,7 @@ SQLAlchemy implementation of the WorkflowExecutionRepository. import json import logging +from typing import override from sqlalchemy.engine import Engine from sqlalchemy.orm import sessionmaker @@ -174,6 +175,7 @@ class SQLAlchemyWorkflowExecutionRepository(WorkflowExecutionRepository): return db_model + @override def save(self, execution: WorkflowExecution): """ Save or update a WorkflowExecution domain entity to the database. diff --git a/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py b/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py index b036687bc9..7eda458f85 100644 --- a/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py +++ b/api/core/repositories/sqlalchemy_workflow_node_execution_repository.py @@ -7,7 +7,7 @@ import json import logging from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor -from typing import Any +from typing import Any, override import psycopg2.errors from sqlalchemy import UnaryExpression, asc, desc, select @@ -313,6 +313,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) offload=offload, ) + @override def save(self, execution: WorkflowNodeExecution) -> None: """ Save or update a NodeExecution domain entity to the database. @@ -399,6 +400,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) if db_model.node_execution_id: self._node_execution_cache[db_model.node_execution_id] = db_model + @override def save_execution_data(self, execution: WorkflowNodeExecution): domain_model = execution with self._session_factory(expire_on_commit=False) as session: @@ -518,6 +520,7 @@ class SQLAlchemyWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository) return db_models + @override def get_by_workflow_execution( self, workflow_execution_id: str,