From 4c124174c9a6fb1fd11b49dc90e044f6c6dd6eb2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= Date: Mon, 7 Sep 2026 03:21:33 +0000 Subject: [PATCH] refactor(api): extract human input file upload persistence (#41716) --- .../web/human_input_file_upload.py | 140 +++--- api/controllers/web/human_input_form.py | 18 +- api/extensions/ext_application_services.py | 16 +- .../human_input_file_upload_repository.py | 166 +++++++ .../human_input_file_upload_service.py | 293 ++++++------ .../test_human_input_delivery_test.py | 16 +- .../test_human_input_file_upload_service.py | 5 +- .../web/test_human_input_file_upload.py | 234 ++++++---- .../controllers/web/test_human_input_form.py | 30 +- .../test_ext_application_services.py | 21 + ...test_human_input_file_upload_repository.py | 251 ++++++++++ .../test_human_input_file_upload_service.py | 436 +++++++----------- 12 files changed, 1012 insertions(+), 614 deletions(-) create mode 100644 api/repositories/human_input_file_upload_repository.py create mode 100644 api/tests/unit_tests/repositories/test_human_input_file_upload_repository.py diff --git a/api/controllers/web/human_input_file_upload.py b/api/controllers/web/human_input_file_upload.py index dcddec8f35b..6b00b60c484 100644 --- a/api/controllers/web/human_input_file_upload.py +++ b/api/controllers/web/human_input_file_upload.py @@ -6,37 +6,50 @@ remote URLs. The caller always submits a multipart form: when a non-empty falls back to the local file upload flow. """ -import httpx +from typing import Any + from flask import request from flask_restx import Resource from pydantic import BaseModel, ConfigDict, Field, HttpUrl -from sqlalchemy.orm import sessionmaker +from werkzeug.datastructures import FileStorage import services from controllers.common.errors import ( BlockedFileExtensionError, + FilenameNotExistsError, FileTooLargeError, NoFileUploadedError, - RemoteFileUploadError, + RemoteFileAccessDeniedError, + RemoteFileInvalidResponseError, + RemoteFileInvalidUrlError, + RemoteFileNotFoundError, + RemoteFileUnavailableError, + RemoteFileUrlBlockedError, TooManyFilesError, UnsupportedFileTypeError, ) -from controllers.common.schema import register_schema_models +from controllers.common.schema import JsonResponseWithStatus, register_schema_models from controllers.web import web_ns -from core.file.remote_file_metadata import guess_file_info_from_response -from core.helper import ssrf_proxy -from extensions.ext_database import db +from extensions.ext_application_services import application_services from fields.file_fields import FileResponse, FileWithSignedUrl -from graphon.file import helpers as file_helpers from libs.exception import BaseHTTPException from libs.helper import dump_response -from repositories.factory import DifyAPIRepositoryFactory -from services.file_service import FileService from services.human_input_file_upload_service import ( HITL_UPLOAD_TOKEN_PREFIX, HumanInputFileUploadService, + HumanInputUploadContext, InvalidUploadTokenError, ) +from services.remote_file_service import ( + RemoteFileAccessDeniedError as RemoteFileAccessDeniedServiceError, +) +from services.remote_file_service import ( + RemoteFileInvalidResponseError as RemoteFileInvalidResponseServiceError, +) +from services.remote_file_service import RemoteFileInvalidUrlError as RemoteFileInvalidUrlServiceError +from services.remote_file_service import RemoteFileNotFoundError as RemoteFileNotFoundServiceError +from services.remote_file_service import RemoteFileUnavailableError as RemoteFileUnavailableServiceError +from services.remote_file_service import RemoteFileUrlBlockedError as RemoteFileUrlBlockedServiceError class InvalidUploadTokenBadRequestError(BaseHTTPException): @@ -68,15 +81,6 @@ class HumanInputFileUploadFormPayload(BaseModel): register_schema_models(web_ns, HumanInputFileUploadFormPayload, FileResponse, FileWithSignedUrl) -def _create_upload_service() -> HumanInputFileUploadService: - session_factory = sessionmaker(bind=db.engine) - workflow_run_repository = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_factory) - return HumanInputFileUploadService( - session_factory=session_factory, - workflow_run_repository=workflow_run_repository, - ) - - def _extract_hitl_upload_token() -> str: """Read HITL upload token from Authorization without invoking other bearer auth chains.""" @@ -98,14 +102,14 @@ def _extract_hitl_upload_token() -> str: return token -def _validate_context(service: HumanInputFileUploadService, token: str): +def _validate_context(service: HumanInputFileUploadService, token: str) -> HumanInputUploadContext: try: return service.validate_upload_token(token) except InvalidUploadTokenError as exc: raise InvalidUploadTokenForbiddenError() from exc -def _parse_local_upload_file(): +def _parse_local_upload_file() -> FileStorage: if "file" not in request.files: raise NoFileUploadedError() if len(request.files) > 1: @@ -113,8 +117,6 @@ def _parse_local_upload_file(): file = request.files["file"] if not file.filename: - from controllers.common.errors import FilenameNotExistsError - raise FilenameNotExistsError() return file @@ -124,79 +126,68 @@ def _parse_upload_form() -> HumanInputFileUploadFormPayload: return HumanInputFileUploadFormPayload.model_validate(request.form.to_dict(flat=True)) -def _upload_local_file(context): +def _upload_local_file( + *, + service: HumanInputFileUploadService, + context: HumanInputUploadContext, +) -> dict[str, Any]: file = _parse_local_upload_file() try: - upload_file = FileService(db.engine).upload_file( + upload_file = service.upload_local_file( + context=context, filename=file.filename or "", content=file.read(), mimetype=file.mimetype, - user=context.owner, - source=None, ) except services.errors.file.FileTooLargeError as file_too_large_error: - raise FileTooLargeError(file_too_large_error.description) - except services.errors.file.UnsupportedFileTypeError: - raise UnsupportedFileTypeError() + raise FileTooLargeError(file_too_large_error.description or "File size exceeded.") from file_too_large_error + except services.errors.file.UnsupportedFileTypeError as error: + raise UnsupportedFileTypeError() from error except services.errors.file.BlockedFileExtensionError as exc: raise BlockedFileExtensionError() from exc - return upload_file.id, dump_response(FileResponse, upload_file) + return dump_response(FileResponse, upload_file) -def _upload_remote_file(context, url: str): +def _upload_remote_file( + *, + service: HumanInputFileUploadService, + context: HumanInputUploadContext, + url: str, +) -> dict[str, Any]: try: - resp = ssrf_proxy.head(url=url) - if resp.status_code != httpx.codes.OK: - resp = ssrf_proxy.get(url=url, timeout=3, follow_redirects=True) - if resp.status_code != httpx.codes.OK: - raise RemoteFileUploadError(f"Failed to fetch file from {url}: {resp.text}") - except httpx.RequestError as exc: - raise RemoteFileUploadError(f"Failed to fetch file from {url}: {str(exc)}") - - file_info = guess_file_info_from_response(resp) - if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size): - raise FileTooLargeError() - - content = resp.content if resp.request.method == "GET" else ssrf_proxy.get(url).content - - try: - upload_file = FileService(db.engine).upload_file( - filename=file_info.filename, - content=content, - mimetype=file_info.mimetype, - user=context.owner, - source_url=url, - ) + upload_file = service.upload_remote_file(context=context, url=url) + except RemoteFileInvalidUrlServiceError as error: + raise RemoteFileInvalidUrlError() from error + except RemoteFileUrlBlockedServiceError as error: + raise RemoteFileUrlBlockedError() from error + except RemoteFileNotFoundServiceError as error: + raise RemoteFileNotFoundError() from error + except RemoteFileAccessDeniedServiceError as error: + raise RemoteFileAccessDeniedError() from error + except RemoteFileUnavailableServiceError as error: + raise RemoteFileUnavailableError() from error + except RemoteFileInvalidResponseServiceError as error: + raise RemoteFileInvalidResponseError() from error except services.errors.file.FileTooLargeError as file_too_large_error: - raise FileTooLargeError(file_too_large_error.description) - except services.errors.file.UnsupportedFileTypeError: - raise UnsupportedFileTypeError() + raise FileTooLargeError(file_too_large_error.description or "File size exceeded.") from file_too_large_error + except services.errors.file.UnsupportedFileTypeError as error: + raise UnsupportedFileTypeError() from error except services.errors.file.BlockedFileExtensionError as exc: raise BlockedFileExtensionError() from exc - response = FileWithSignedUrl( - id=upload_file.id, - name=upload_file.name, - size=upload_file.size, - extension=upload_file.extension, - url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id), - mime_type=upload_file.mime_type, - created_by=upload_file.created_by, - created_at=int(upload_file.created_at.timestamp()), - ) - return upload_file.id, response.model_dump(mode="json") + return dump_response(FileWithSignedUrl, upload_file) @web_ns.route("/human-input-forms/files") @web_ns.response(201, "File uploaded successfully", web_ns.models[FileResponse.__name__]) class HumanInputFileUploadApi(Resource): - def post(self): + def post(self) -> JsonResponseWithStatus: """Upload one local file or remote URL file for a HITL human input form.""" token = _extract_hitl_upload_token() - upload_service = _create_upload_service() + upload_service = application_services().human_input_file_uploads context = _validate_context(upload_service, token) form = _parse_upload_form() @@ -204,10 +195,13 @@ class HumanInputFileUploadApi(Resource): # switches the endpoint into the remote-fetch flow; otherwise the # request must carry a local `file`. if form.url is not None: - file_id, response = _upload_remote_file(context=context, url=str(form.url)) + response = _upload_remote_file( + service=upload_service, + context=context, + url=str(form.url), + ) else: - file_id, response = _upload_local_file(context=context) + response = _upload_local_file(service=upload_service, context=context) - upload_service.record_upload_file(context=context, file_id=file_id) # response-contract:ignore pre-dumped response. See above return response, 201 diff --git a/api/controllers/web/human_input_form.py b/api/controllers/web/human_input_form.py index 5668b97abaa..745b68b7f2d 100644 --- a/api/controllers/web/human_input_form.py +++ b/api/controllers/web/human_input_form.py @@ -9,7 +9,6 @@ from typing import Self from flask import request from flask_restx import Resource from sqlalchemy import select -from sqlalchemy.orm import sessionmaker from werkzeug.exceptions import Forbidden from configs import dify_config @@ -21,14 +20,13 @@ from controllers.web import web_ns from controllers.web.error import WebFormRateLimitExceededError from controllers.web.site import WebAppSiteResponse from core.workflow.nodes.human_input.entities import FormInputConfig, UserActionConfig +from extensions.ext_application_services import application_services from extensions.ext_database import db from fields.base import ResponseModel from libs.helper import RateLimiter, dump_response, extract_remote_ip, to_timestamp from models.account import TenantStatus from models.model import App, AppMode, Site -from repositories.factory import DifyAPIRepositoryFactory from services.feature_service import FeatureService -from services.human_input_file_upload_service import HumanInputFileUploadService from services.human_input_service import Form, FormNotFoundError, HumanInputService logger = logging.getLogger(__name__) @@ -99,15 +97,6 @@ _FORM_UPLOAD_TOKEN_RATE_LIMITER = RateLimiter( ) -def _create_upload_service() -> HumanInputFileUploadService: - session_factory = sessionmaker(bind=db.engine) - workflow_run_repository = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_factory) - return HumanInputFileUploadService( - session_factory=session_factory, - workflow_run_repository=workflow_run_repository, - ) - - @web_ns.route("/form/human_input//upload-token") class HumanInputFormUploadTokenApi(Resource): """API for issuing HITL upload tokens for active human input forms.""" @@ -140,7 +129,7 @@ class HumanInputFormUploadTokenApi(Resource): _FORM_UPLOAD_TOKEN_RATE_LIMITER.increment_rate_limit(ip_address) try: - token = _create_upload_service().issue_upload_token(form_token) + token = application_services().human_input_file_uploads.issue_upload_token(form_token) except FormNotFoundError: raise NotFoundError("Form not found") @@ -155,7 +144,6 @@ class HumanInputFormApi(Resource): # NOTE(QuantumGhost): this endpoint is unauthenticated on purpose for now. - # def get(self, _app_model: App, _end_user: EndUser, form_token: str): @web_ns.doc("get_human_input_form") @web_ns.doc(description="Get a human input form definition by token") @web_ns.doc(params={"form_token": "Human input form token"}) @@ -217,7 +205,6 @@ class HumanInputFormApi(Resource): ), ) - # def post(self, _app_model: App, _end_user: EndUser, form_token: str): @web_ns.expect(web_ns.models[HumanInputFormSubmitPayload.__name__]) @web_ns.doc("submit_human_input_form") @web_ns.doc(description="Submit a human input form by token") @@ -272,7 +259,6 @@ class HumanInputFormApi(Resource): selected_action_id=payload.action, form_data=payload.inputs, submission_end_user_id=None, - # submission_end_user_id=_end_user.id, ) except FormNotFoundError: raise NotFoundError("Form not found") diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index b68354b194d..a8050b74ad9 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -47,6 +47,7 @@ from repositories.data_source_oauth_binding_repository import SQLAlchemyDataSour from repositories.explore_banner_query_repository import ExploreBannerQueryRepository from repositories.factory import DifyAPIRepositoryFactory from repositories.file_grant_repository import FileGrantRepository +from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository from repositories.installation_state_repository import InstallationStateRepository from repositories.oauth_access_token_repository import SQLAlchemyOAuthAccessTokenRepository from repositories.oauth_server_repository import RedisOAuthServerTokenRepository, SQLAlchemyOAuthServerRepository @@ -155,6 +156,7 @@ from services.feature_service_gateway import FeatureServiceGateway from services.file_grant_gateways import FileGrantFileGateway, FileGrantRemoteFileGateway, FileGrantTokenGateway from services.file_grant_service import FileGrantService from services.file_service import FileService +from services.human_input_file_upload_service import HumanInputFileUploadService from services.init_validation_service import InitValidationService from services.inner_mail_service import InnerMailService from services.notification_gateway import BillingNotificationGateway @@ -260,6 +262,7 @@ class ApplicationServices: feature_queries: FeatureQueryService file_grants: FileGrantService files: FileService + human_input_file_uploads: HumanInputFileUploadService oauth_server: OAuthServerService init_validation: InitValidationService notifications: NotificationService @@ -423,6 +426,7 @@ def build_application_services( ) workspace_query_repository = WorkspaceQueryRepository(session_factory=database_client) file_service = FileService(session_factory=database_client) + remote_file_service = RemoteFileService(files=file_service) passwords = DefaultAccountPasswordHasher() invitation_tokens = RedisInvitationTokenStore(redis=redis) activation_accounts = SQLAlchemyAccountActivationRepository(session_factory=database_client) @@ -649,6 +653,14 @@ def build_application_services( ), file_grants=_build_file_grant_service(database_client=database_client), files=file_service, + human_input_file_uploads=HumanInputFileUploadService( + uploads=SQLAlchemyHumanInputFileUploadRepository(session_factory=database_client), + workflow_run_repository=DifyAPIRepositoryFactory.create_api_workflow_run_repository( + session_maker=database_client, + ), + files=file_service, + remote_files=remote_file_service, + ), oauth_server=_build_oauth_server_service(database_client=database_client, redis=redis), init_validation=InitValidationService( state=installation_state, @@ -672,9 +684,7 @@ def build_application_services( trial_apps=TrialAppQueryRepository(session_factory=database_client), trial_enabled=trial_app_enabled, ), - remote_files=RemoteFileService( - files=FileService(session_factory=database_client), - ), + remote_files=remote_file_service, trial_app_usage=TrialAppUsageRepository(session_factory=database_client), workflow_run_archives=WorkflowRunArchiveService( bundles=WorkflowRunArchiveBundleQueryRepository(session_factory=database_client), diff --git a/api/repositories/human_input_file_upload_repository.py b/api/repositories/human_input_file_upload_repository.py new file mode 100644 index 00000000000..7406a5337b3 --- /dev/null +++ b/api/repositories/human_input_file_upload_repository.py @@ -0,0 +1,166 @@ +"""SQLAlchemy persistence adapter for human-input file uploads.""" + +from typing import override + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from models.account import Account, Tenant +from models.enums import CreatorUserRole +from models.human_input import ( + HumanInputForm, + HumanInputFormRecipient, + HumanInputFormUploadFile, + HumanInputFormUploadToken, +) +from models.model import App, EndUser +from services.human_input_file_upload_service import ( + HumanInputFileUploadRepository, + HumanInputUploadFormRecord, + HumanInputUploadGrantRecord, +) + + +class SQLAlchemyHumanInputFileUploadRepository(HumanInputFileUploadRepository): + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def get_form_by_recipient_token(self, form_token: str) -> HumanInputUploadFormRecord | None: + stmt = ( + select(HumanInputFormRecipient, HumanInputForm) + .join(HumanInputForm, HumanInputForm.id == HumanInputFormRecipient.form_id) + .where(HumanInputFormRecipient.access_token == form_token) + .limit(1) + ) + with self._session_factory() as session: + row = session.execute(stmt).one_or_none() + if row is None: + return None + recipient, form = row + return self._to_form_record(form=form, recipient_id=recipient.id) + + @override + def create_upload_token(self, *, form: HumanInputUploadFormRecord, upload_token: str) -> None: + with self._session_factory.begin() as session: + session.add( + HumanInputFormUploadToken( + tenant_id=form.tenant_id, + app_id=form.app_id, + form_id=form.form_id, + recipient_id=form.recipient_id, + token=upload_token, + ) + ) + + @override + def get_upload_grant(self, upload_token: str) -> HumanInputUploadGrantRecord | None: + stmt = ( + select(HumanInputFormUploadToken, HumanInputForm) + .join(HumanInputForm, HumanInputForm.id == HumanInputFormUploadToken.form_id) + .where(HumanInputFormUploadToken.token == upload_token) + .limit(1) + ) + with self._session_factory() as session: + row = session.execute(stmt).one_or_none() + if row is None: + return None + token, form = row + return HumanInputUploadGrantRecord( + upload_token_id=token.id, + form=self._to_form_record(form=form, recipient_id=token.recipient_id), + ) + + @override + def get_upload_owner( + self, + *, + owner_id: str, + owner_role: CreatorUserRole, + tenant_id: str, + app_id: str, + ) -> Account | EndUser | None: + with self._session_factory() as session: + if owner_role == CreatorUserRole.END_USER: + return session.scalar( + select(EndUser) + .where( + EndUser.id == owner_id, + EndUser.tenant_id == tenant_id, + EndUser.app_id == app_id, + ) + .limit(1) + ) + + if owner_role != CreatorUserRole.ACCOUNT: + return None + + account = session.scalar(select(Account).where(Account.id == owner_id).limit(1)) + tenant = session.scalar(select(Tenant).where(Tenant.id == tenant_id).limit(1)) + if account is None or tenant is None: + return None + + account.set_current_tenant_with_session(tenant, session=session) + if account.current_tenant_id != tenant_id: + return None + return account + + @override + def get_delivery_test_upload_owner(self, *, tenant_id: str, app_id: str) -> Account | None: + with self._session_factory() as session: + app = session.scalar( + select(App) + .where( + App.id == app_id, + App.tenant_id == tenant_id, + ) + .limit(1) + ) + if app is None or app.created_by is None: + return None + + account = session.scalar(select(Account).where(Account.id == app.created_by).limit(1)) + tenant = session.scalar(select(Tenant).where(Tenant.id == tenant_id).limit(1)) + if account is None or tenant is None: + return None + + account.set_current_tenant_with_session(tenant, session=session) + if account.current_tenant_id != tenant_id: + return None + return account + + @override + def add_file( + self, + *, + tenant_id: str, + app_id: str, + form_id: str, + upload_token_id: str, + file_id: str, + ) -> None: + with self._session_factory.begin() as session: + session.add( + HumanInputFormUploadFile( + tenant_id=tenant_id, + app_id=app_id, + form_id=form_id, + upload_file_id=file_id, + upload_token_id=upload_token_id, + ) + ) + + @staticmethod + def _to_form_record(*, form: HumanInputForm, recipient_id: str) -> HumanInputUploadFormRecord: + return HumanInputUploadFormRecord( + form_id=form.id, + recipient_id=recipient_id, + tenant_id=form.tenant_id, + app_id=form.app_id, + workflow_run_id=form.workflow_run_id, + form_kind=form.form_kind, + status=form.status, + submitted_at=form.submitted_at, + expiration_time=form.expiration_time, + created_at=form.created_at, + ) diff --git a/api/services/human_input_file_upload_service.py b/api/services/human_input_file_upload_service.py index 00b4230870c..538b5f128a6 100644 --- a/api/services/human_input_file_upload_service.py +++ b/api/services/human_input_file_upload_service.py @@ -3,24 +3,18 @@ from __future__ import annotations import secrets from dataclasses import dataclass from datetime import datetime, timedelta - -from sqlalchemy import Engine, select -from sqlalchemy.orm import Session, selectinload, sessionmaker +from typing import Protocol from configs import dify_config from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus from libs.datetime_utils import ensure_naive_utc, naive_utc_now -from models.account import Account, Tenant +from models.account import Account from models.enums import CreatorUserRole -from models.human_input import ( - HumanInputForm, - HumanInputFormRecipient, - HumanInputFormUploadFile, - HumanInputFormUploadToken, -) -from models.model import App, EndUser +from models.model import EndUser, UploadFile from repositories.api_workflow_run_repository import APIWorkflowRunRepository +from services.file_service import FileService from services.human_input_service import FormExpiredError, FormNotFoundError, FormSubmittedError +from services.remote_file_service import RemoteFileService, RemoteFileUploadResult HITL_UPLOAD_TOKEN_PREFIX = "hitl_upload_" _TOKEN_RANDOM_BYTES = 32 @@ -33,6 +27,26 @@ class HumanInputUploadToken: expires_at: datetime +@dataclass(frozen=True) +class HumanInputUploadFormRecord: + form_id: str + recipient_id: str + tenant_id: str + app_id: str + workflow_run_id: str | None + form_kind: HumanInputFormKind + status: HumanInputFormStatus + submitted_at: datetime | None + expiration_time: datetime + created_at: datetime + + +@dataclass(frozen=True) +class HumanInputUploadGrantRecord: + upload_token_id: str + form: HumanInputUploadFormRecord + + @dataclass(frozen=True) class HumanInputUploadContext: tenant_id: str @@ -47,6 +61,35 @@ class InvalidUploadTokenError(Exception): pass +class HumanInputFileUploadRepository(Protocol): + def get_form_by_recipient_token(self, form_token: str) -> HumanInputUploadFormRecord | None: ... + + def create_upload_token(self, *, form: HumanInputUploadFormRecord, upload_token: str) -> None: ... + + def get_upload_grant(self, upload_token: str) -> HumanInputUploadGrantRecord | None: ... + + def get_upload_owner( + self, + *, + owner_id: str, + owner_role: CreatorUserRole, + tenant_id: str, + app_id: str, + ) -> Account | EndUser | None: ... + + def get_delivery_test_upload_owner(self, *, tenant_id: str, app_id: str) -> Account | None: ... + + def add_file( + self, + *, + tenant_id: str, + app_id: str, + form_id: str, + upload_token_id: str, + file_id: str, + ) -> None: ... + + class HumanInputFileUploadService: """Coordinates HITL upload tokens, workflow-run owners, and form-file links. @@ -56,92 +99,96 @@ class HumanInputFileUploadService: uploads are scoped to the app creator account inside the form tenant. """ - _session_maker: sessionmaker[Session] + _uploads: HumanInputFileUploadRepository _workflow_run_repository: APIWorkflowRunRepository + _files: FileService + _remote_files: RemoteFileService def __init__( self, - session_factory: sessionmaker[Session] | Engine, + *, + uploads: HumanInputFileUploadRepository, workflow_run_repository: APIWorkflowRunRepository, + files: FileService, + remote_files: RemoteFileService, ) -> None: - if isinstance(session_factory, Engine): - session_factory = sessionmaker(bind=session_factory) - self._session_maker = session_factory + self._uploads = uploads self._workflow_run_repository = workflow_run_repository + self._files = files + self._remote_files = remote_files def issue_upload_token(self, form_token: str) -> HumanInputUploadToken: """Create an upload token for an active human input recipient token.""" - with self._session_maker() as session, session.begin(): - recipient_model = session.scalar( - select(HumanInputFormRecipient) - .options(selectinload(HumanInputFormRecipient.form)) - .where(HumanInputFormRecipient.access_token == form_token) - .limit(1) - ) - if recipient_model is None or recipient_model.form is None: - raise FormNotFoundError() + form = self._uploads.get_form_by_recipient_token(form_token) + if form is None: + raise FormNotFoundError() - form = recipient_model.form - self._ensure_form_model_active(form) - upload_token = self._generate_unique_upload_token() - token_model = HumanInputFormUploadToken( - tenant_id=form.tenant_id, - app_id=form.app_id, - form_id=form.id, - recipient_id=recipient_model.id, - token=upload_token, - ) - session.add(token_model) - # Snapshot the expiry before commit so callers do not depend on the - # session factory's expire_on_commit policy. - token = HumanInputUploadToken(upload_token=upload_token, expires_at=form.expiration_time) - - return token + self._ensure_form_active(form) + upload_token = self._generate_unique_upload_token() + self._uploads.create_upload_token(form=form, upload_token=upload_token) + return HumanInputUploadToken(upload_token=upload_token, expires_at=form.expiration_time) def validate_upload_token(self, upload_token: str) -> HumanInputUploadContext: """Resolve an upload token and ensure the bound form is still active.""" - query = ( - select(HumanInputFormUploadToken) - .options(selectinload(HumanInputFormUploadToken.form)) - .where(HumanInputFormUploadToken.token == upload_token) - .limit(1) + grant = self._uploads.get_upload_grant(upload_token) + if grant is None: + raise InvalidUploadTokenError() + + form = grant.form + self._ensure_form_active(form) + owner = self._resolve_upload_owner(form=form) + return HumanInputUploadContext( + tenant_id=form.tenant_id, + app_id=form.app_id, + form_id=form.form_id, + recipient_id=form.recipient_id, + upload_token_id=grant.upload_token_id, + owner=owner, ) - with self._session_maker(expire_on_commit=False) as session: - token_model = session.scalars(query).first() - if token_model is None: - raise InvalidUploadTokenError() - - form_model = token_model.form - if form_model is None: - raise InvalidUploadTokenError() - self._ensure_form_model_active(form_model) - - owner = self._resolve_upload_owner(session=session, form_model=form_model) - - return HumanInputUploadContext( - tenant_id=token_model.tenant_id, - app_id=token_model.app_id, - form_id=token_model.form_id, - recipient_id=token_model.recipient_id, - upload_token_id=token_model.id, - owner=owner, - ) def record_upload_file(self, *, context: HumanInputUploadContext, file_id: str) -> None: """Record that a file was uploaded through a specific form upload token.""" - with self._session_maker() as session, session.begin(): - session.add( - HumanInputFormUploadFile( - tenant_id=context.tenant_id, - app_id=context.app_id, - form_id=context.form_id, - upload_file_id=file_id, - upload_token_id=context.upload_token_id, - ) - ) + self._uploads.add_file( + tenant_id=context.tenant_id, + app_id=context.app_id, + form_id=context.form_id, + upload_token_id=context.upload_token_id, + file_id=file_id, + ) + + def upload_local_file( + self, + *, + context: HumanInputUploadContext, + filename: str, + content: bytes, + mimetype: str, + ) -> UploadFile: + upload_file = self._files.upload_file( + filename=filename, + content=content, + mimetype=mimetype, + user=context.owner, + source=None, + ) + self.record_upload_file(context=context, file_id=upload_file.id) + return upload_file + + def upload_remote_file( + self, + *, + context: HumanInputUploadContext, + url: str, + ) -> RemoteFileUploadResult: + upload_file = self._remote_files.upload_from_url( + url=url, + user=context.owner, + ) + self.record_upload_file(context=context, file_id=upload_file.id) + return upload_file def _generate_unique_upload_token(self) -> str: return f"{HITL_UPLOAD_TOKEN_PREFIX}{secrets.token_urlsafe(_TOKEN_RANDOM_BYTES)}" @@ -149,96 +196,54 @@ class HumanInputFileUploadService: def _resolve_upload_owner( self, *, - session: Session, - form_model: HumanInputForm, + form: HumanInputUploadFormRecord, ) -> Account | EndUser: - if form_model.workflow_run_id is None: - if form_model.form_kind == HumanInputFormKind.DELIVERY_TEST: - return self._resolve_delivery_test_upload_owner(session=session, form_model=form_model) + if form.workflow_run_id is None: + if form.form_kind == HumanInputFormKind.DELIVERY_TEST: + owner = self._uploads.get_delivery_test_upload_owner( + tenant_id=form.tenant_id, + app_id=form.app_id, + ) + if owner is not None: + return owner raise InvalidUploadTokenError() workflow_run = self._workflow_run_repository.get_workflow_run_by_id( - tenant_id=form_model.tenant_id, - app_id=form_model.app_id, - run_id=form_model.workflow_run_id, + tenant_id=form.tenant_id, + app_id=form.app_id, + run_id=form.workflow_run_id, ) if workflow_run is None: raise InvalidUploadTokenError() - if workflow_run.created_by_role == CreatorUserRole.END_USER: - end_user = session.scalar( - select(EndUser) - .where( - EndUser.id == workflow_run.created_by, - EndUser.tenant_id == workflow_run.tenant_id, - EndUser.app_id == workflow_run.app_id, - ) - .limit(1) - ) - if end_user is None: - raise InvalidUploadTokenError() - return end_user - - if workflow_run.created_by_role != CreatorUserRole.ACCOUNT: + owner_role = workflow_run.created_by_role + if owner_role not in {CreatorUserRole.ACCOUNT, CreatorUserRole.END_USER}: raise InvalidUploadTokenError() - account = session.scalar(select(Account).where(Account.id == workflow_run.created_by).limit(1)) - if account is None: - raise InvalidUploadTokenError() - - tenant = session.scalar(select(Tenant).where(Tenant.id == workflow_run.tenant_id).limit(1)) - if tenant is None: - raise InvalidUploadTokenError() - - # HITL upload runs outside the normal account auth flow, so hydrate the - # account tenant context explicitly before delegating to FileService. - account.set_current_tenant_with_session(tenant, session=session) - return account - - def _resolve_delivery_test_upload_owner( - self, - *, - session: Session, - form_model: HumanInputForm, - ) -> Account: - app = session.scalar( - select(App) - .where( - App.id == form_model.app_id, - App.tenant_id == form_model.tenant_id, - ) - .limit(1) + owner = self._uploads.get_upload_owner( + owner_id=workflow_run.created_by, + owner_role=owner_role, + tenant_id=workflow_run.tenant_id, + app_id=workflow_run.app_id, ) - if app is None or app.created_by is None: + if owner is None: raise InvalidUploadTokenError() - - account = session.scalar(select(Account).where(Account.id == app.created_by).limit(1)) - if account is None: - raise InvalidUploadTokenError() - - tenant = session.scalar(select(Tenant).where(Tenant.id == form_model.tenant_id).limit(1)) - if tenant is None: - raise InvalidUploadTokenError() - - account.set_current_tenant_with_session(tenant, session=session) - if account.current_tenant_id != form_model.tenant_id: - raise InvalidUploadTokenError() - return account + return owner @staticmethod - def _ensure_form_model_active(form: HumanInputForm) -> None: + def _ensure_form_active(form: HumanInputUploadFormRecord) -> None: if form.submitted_at is not None or form.status == HumanInputFormStatus.SUBMITTED: - raise FormSubmittedError(form.id) + raise FormSubmittedError(form.form_id) if form.status in {HumanInputFormStatus.TIMEOUT, HumanInputFormStatus.EXPIRED}: - raise FormExpiredError(form.id) + raise FormExpiredError(form.form_id) now = naive_utc_now() if ensure_naive_utc(form.expiration_time) <= now: - raise FormExpiredError(form.id) + raise FormExpiredError(form.form_id) global_timeout_seconds = dify_config.HUMAN_INPUT_GLOBAL_TIMEOUT_SECONDS if global_timeout_seconds <= 0 or form.workflow_run_id is None: return global_deadline = ensure_naive_utc(form.created_at) + timedelta(seconds=global_timeout_seconds) if global_deadline <= now: - raise FormExpiredError(form.id) + raise FormExpiredError(form.form_id) diff --git a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py b/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py index 84a0226ba17..3ce7d9e1b93 100644 --- a/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py +++ b/api/tests/test_containers_integration_tests/services/test_human_input_delivery_test.py @@ -1,7 +1,7 @@ import json import uuid from io import BytesIO -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import httpx import pytest @@ -9,7 +9,7 @@ from flask.testing import FlaskClient from sqlalchemy import select from sqlalchemy.orm import Session -import controllers.web.human_input_file_upload as human_input_file_upload_module +import services.remote_file_service as remote_file_service_module from core.workflow.human_input_adapter import ( EmailDeliveryConfig, EmailDeliveryMethod, @@ -257,10 +257,8 @@ def test_human_input_delivery_test_form_accepts_remote_file_upload( }, content=remote_content, ) - head_mock = MagicMock(return_value=head_response) - get_mock = MagicMock(return_value=get_response) - monkeypatch.setattr(human_input_file_upload_module.ssrf_proxy, "head", head_mock) - monkeypatch.setattr(human_input_file_upload_module.ssrf_proxy, "get", get_mock) + request_mock = MagicMock(side_effect=[head_response, get_response]) + monkeypatch.setattr(remote_file_service_module.remote_fetcher, "make_request", request_mock) upload_response = test_client_with_containers.post( "/api/human-input-forms/files", @@ -272,8 +270,10 @@ def test_human_input_delivery_test_form_accepts_remote_file_upload( assert upload_response.status_code == 201, upload_response.get_data(as_text=True) upload_file_id = upload_response.get_json()["id"] assert upload_response.get_json()["url"] - head_mock.assert_called_once_with(url=remote_url) - get_mock.assert_called_once_with(remote_url) + assert request_mock.call_args_list == [ + call("HEAD", url=remote_url), + call("GET", url=remote_url), + ] db_session_with_containers.expire_all() upload_file = db_session_with_containers.get(UploadFile, upload_file_id) diff --git a/api/tests/test_containers_integration_tests/services/test_human_input_file_upload_service.py b/api/tests/test_containers_integration_tests/services/test_human_input_file_upload_service.py index d88d43fad0b..b6f0da19105 100644 --- a/api/tests/test_containers_integration_tests/services/test_human_input_file_upload_service.py +++ b/api/tests/test_containers_integration_tests/services/test_human_input_file_upload_service.py @@ -16,6 +16,7 @@ from models.human_input import ( HumanInputFormUploadToken, StandaloneWebAppRecipientPayload, ) +from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository from services.human_input_file_upload_service import HITL_UPLOAD_TOKEN_PREFIX, HumanInputFileUploadService @@ -61,8 +62,10 @@ def test_issue_upload_token_returns_expiration_with_default_session_expiry( monkeypatch.setattr(service_module.secrets, "token_urlsafe", lambda _bytes: "random-value") service = HumanInputFileUploadService( - session_factory=sessionmaker(bind=db.engine), + uploads=SQLAlchemyHumanInputFileUploadRepository(session_factory=sessionmaker(bind=db.engine)), workflow_run_repository=MagicMock(), + files=MagicMock(), + remote_files=MagicMock(), ) token = service.issue_upload_token("form-token-1") diff --git a/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py b/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py index 3dfa2be777d..a3d2c8d188a 100644 --- a/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py +++ b/api/tests/unit_tests/controllers/web/test_human_input_file_upload.py @@ -9,10 +9,20 @@ from unittest.mock import MagicMock import pytest from flask import Flask -from sqlalchemy import Engine import controllers.web.human_input_file_upload as upload_module -from controllers.common.errors import NoFileUploadedError +from controllers.common.errors import ( + BlockedFileExtensionError, + FileTooLargeError, + NoFileUploadedError, + RemoteFileAccessDeniedError, + RemoteFileInvalidResponseError, + RemoteFileInvalidUrlError, + RemoteFileNotFoundError, + RemoteFileUnavailableError, + RemoteFileUrlBlockedError, + UnsupportedFileTypeError, +) from controllers.web.human_input_file_upload import ( HumanInputFileUploadApi, InvalidUploadTokenForbiddenError, @@ -23,7 +33,26 @@ from models import Account from models.account import AccountStatus from models.enums import CreatorUserRole from models.model import UploadFile -from services.human_input_file_upload_service import HumanInputUploadContext +from services.errors.file import ( + BlockedFileExtensionError as BlockedFileExtensionServiceError, +) +from services.errors.file import FileTooLargeError as FileTooLargeServiceError +from services.errors.file import UnsupportedFileTypeError as UnsupportedFileTypeServiceError +from services.human_input_file_upload_service import ( + HumanInputUploadContext, + InvalidUploadTokenError, +) +from services.remote_file_service import ( + RemoteFileAccessDeniedError as RemoteFileAccessDeniedServiceError, +) +from services.remote_file_service import ( + RemoteFileInvalidResponseError as RemoteFileInvalidResponseServiceError, +) +from services.remote_file_service import RemoteFileInvalidUrlError as RemoteFileInvalidUrlServiceError +from services.remote_file_service import RemoteFileNotFoundError as RemoteFileNotFoundServiceError +from services.remote_file_service import RemoteFileUnavailableError as RemoteFileUnavailableServiceError +from services.remote_file_service import RemoteFileUploadResult +from services.remote_file_service import RemoteFileUrlBlockedError as RemoteFileUrlBlockedServiceError @pytest.fixture @@ -69,23 +98,25 @@ def _upload_file() -> UploadFile: return upload_file -def _patch_upload_service(monkeypatch: pytest.MonkeyPatch, service: MagicMock) -> tuple[MagicMock, dict[str, object]]: - workflow_run_repository = MagicMock() - repo_factory = MagicMock(return_value=workflow_run_repository) - captured: dict[str, object] = {} - - def _service_factory(session_factory, workflow_run_repository): - captured["session_factory"] = session_factory - captured["workflow_run_repository"] = workflow_run_repository - return service - - monkeypatch.setattr( - upload_module.DifyAPIRepositoryFactory, - "create_api_workflow_run_repository", - repo_factory, +def _remote_upload_file() -> RemoteFileUploadResult: + return RemoteFileUploadResult( + id="file-1", + name="sample.txt", + size=6, + extension="txt", + url="signed:file-1", + mime_type="text/plain", + created_by="owner-1", + created_at=datetime(2024, 1, 1), + ) + + +def _patch_upload_service(monkeypatch: pytest.MonkeyPatch, service: MagicMock) -> None: + monkeypatch.setattr( + upload_module, + "application_services", + lambda: SimpleNamespace(human_input_file_uploads=service), ) - monkeypatch.setattr(upload_module, "HumanInputFileUploadService", _service_factory) - return repo_factory, captured def test_human_input_file_upload_route_uses_unified_path() -> None: @@ -111,18 +142,12 @@ def test_local_upload_requires_authorization_before_reading_files(app: Flask) -> HumanInputFileUploadApi().post() -def test_local_upload_ignores_source_and_records_form_file_link( - monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine -) -> None: +def test_local_upload_delegates_to_human_input_upload_service(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: service = MagicMock() - service.validate_upload_token.return_value = _upload_context() - repo_factory, captured = _patch_upload_service(monkeypatch, service) - - file_service = MagicMock() - file_service.upload_file.return_value = _upload_file() - file_service_cls = MagicMock(return_value=file_service) - monkeypatch.setattr(upload_module, "FileService", file_service_cls) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) + context = _upload_context() + service.validate_upload_token.return_value = context + service.upload_local_file.return_value = _upload_file() + _patch_upload_service(monkeypatch, service) data = { "file": (BytesIO(b"content"), "sample.txt"), @@ -139,24 +164,18 @@ def test_local_upload_ignores_source_and_records_form_file_link( assert status == 201 assert result["id"] == "file-1" - file_service.upload_file.assert_called_once() - assert file_service.upload_file.call_args.kwargs["source"] is None - assert file_service.upload_file.call_args.kwargs["user"].id == "owner-1" - repo_factory.assert_called_once() - assert captured["workflow_run_repository"] is repo_factory.return_value - service.record_upload_file.assert_called_once_with( - context=service.validate_upload_token.return_value, - file_id="file-1", + service.upload_local_file.assert_called_once_with( + context=context, + filename="sample.txt", + content=b"content", + mimetype="text/plain", ) -def test_local_upload_missing_file_raises_after_valid_token( - monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine -) -> None: +def test_local_upload_missing_file_raises_after_valid_token(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: service = MagicMock() service.validate_upload_token.return_value = _upload_context() _patch_upload_service(monkeypatch, service) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) with app.test_request_context( "/api/human-input-forms/files", @@ -170,15 +189,10 @@ def test_local_upload_missing_file_raises_after_valid_token( service.validate_upload_token.assert_called_once_with("hitl_upload_token-1") -def test_remote_upload_validates_token_before_fetching_remote_url( - monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine -) -> None: +def test_remote_upload_validates_token_before_fetching_remote_url(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: service = MagicMock() - service.validate_upload_token.side_effect = InvalidUploadTokenForbiddenError() + service.validate_upload_token.side_effect = InvalidUploadTokenError() _patch_upload_service(monkeypatch, service) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) - ssrf_proxy = MagicMock() - monkeypatch.setattr(upload_module, "ssrf_proxy", ssrf_proxy) with app.test_request_context( "/api/human-input-forms/files", @@ -190,41 +204,15 @@ def test_remote_upload_validates_token_before_fetching_remote_url( with pytest.raises(InvalidUploadTokenForbiddenError): HumanInputFileUploadApi().post() - ssrf_proxy.head.assert_not_called() - ssrf_proxy.get.assert_not_called() + service.upload_remote_file.assert_not_called() -def test_remote_upload_records_form_file_link( - monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine -) -> None: +def test_remote_upload_delegates_to_human_input_upload_service(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: service = MagicMock() - service.validate_upload_token.return_value = _upload_context() + context = _upload_context() + service.validate_upload_token.return_value = context + service.upload_remote_file.return_value = _remote_upload_file() _patch_upload_service(monkeypatch, service) - monkeypatch.setattr(upload_module, "db", SimpleNamespace(engine=sqlite_engine)) - - response = MagicMock() - response.status_code = 200 - response.content = b"remote" - response.request.method = "GET" - ssrf_proxy = MagicMock() - ssrf_proxy.head.return_value = response - monkeypatch.setattr(upload_module, "ssrf_proxy", ssrf_proxy) - monkeypatch.setattr( - upload_module, - "guess_file_info_from_response", - lambda _response: SimpleNamespace(filename="sample.txt", extension="txt", mimetype="text/plain", size=6), - ) - - file_service = MagicMock() - file_service.upload_file.return_value = _upload_file() - file_service_cls = MagicMock(return_value=file_service) - file_service_cls.is_file_size_within_limit.return_value = True - monkeypatch.setattr(upload_module, "FileService", file_service_cls) - monkeypatch.setattr( - upload_module.file_helpers, - "get_signed_file_url", - lambda upload_file_id: f"signed:{upload_file_id}", - ) with app.test_request_context( "/api/human-input-forms/files", @@ -237,10 +225,84 @@ def test_remote_upload_records_form_file_link( assert status == 201 assert result["url"] == "signed:file-1" - file_service.upload_file.assert_called_once() - assert file_service.upload_file.call_args.kwargs["source_url"] == "https://example.com/file.txt" - assert file_service.upload_file.call_args.kwargs["user"].id == "owner-1" - service.record_upload_file.assert_called_once_with( - context=service.validate_upload_token.return_value, - file_id="file-1", + service.upload_remote_file.assert_called_once_with( + context=context, + url="https://example.com/file.txt", ) + + +@pytest.mark.parametrize( + ("service_error", "http_error"), + [ + pytest.param(RemoteFileInvalidUrlServiceError(), RemoteFileInvalidUrlError, id="invalid-url"), + pytest.param(RemoteFileUrlBlockedServiceError(), RemoteFileUrlBlockedError, id="blocked-url"), + pytest.param(RemoteFileNotFoundServiceError(), RemoteFileNotFoundError, id="not-found"), + pytest.param(RemoteFileAccessDeniedServiceError(), RemoteFileAccessDeniedError, id="access-denied"), + pytest.param(RemoteFileUnavailableServiceError(), RemoteFileUnavailableError, id="unavailable"), + pytest.param(RemoteFileInvalidResponseServiceError(), RemoteFileInvalidResponseError, id="invalid-response"), + ], +) +def test_remote_upload_maps_remote_file_errors( + monkeypatch: pytest.MonkeyPatch, + app: Flask, + service_error: Exception, + http_error: type[Exception], +) -> None: + service = MagicMock() + service.validate_upload_token.return_value = _upload_context() + service.upload_remote_file.side_effect = service_error + _patch_upload_service(monkeypatch, service) + + with app.test_request_context( + "/api/human-input-forms/files", + method="POST", + headers={"Authorization": "Bearer hitl_upload_token-1"}, + data={"url": "https://example.com/file.txt"}, + content_type="multipart/form-data", + ): + with pytest.raises(http_error) as raised: + HumanInputFileUploadApi().post() + + assert raised.value.__cause__ is service_error + + +@pytest.mark.parametrize( + ("service_error", "http_error"), + [ + pytest.param(FileTooLargeServiceError(), FileTooLargeError, id="too-large"), + pytest.param(UnsupportedFileTypeServiceError(), UnsupportedFileTypeError, id="unsupported"), + pytest.param(BlockedFileExtensionServiceError("Blocked extension"), BlockedFileExtensionError, id="blocked"), + ], +) +@pytest.mark.parametrize("remote", [False, True], ids=["local", "remote"]) +def test_upload_maps_file_service_errors( + monkeypatch: pytest.MonkeyPatch, + app: Flask, + service_error: Exception, + http_error: type[Exception], + remote: bool, +) -> None: + service = MagicMock() + service.validate_upload_token.return_value = _upload_context() + if remote: + service.upload_remote_file.side_effect = service_error + data = {"url": "https://example.com/file.txt"} + else: + service.upload_local_file.side_effect = service_error + data = {"file": (BytesIO(b"content"), "sample.txt")} + _patch_upload_service(monkeypatch, service) + + with app.test_request_context( + "/api/human-input-forms/files", + method="POST", + headers={"Authorization": "Bearer hitl_upload_token-1"}, + data=data, + content_type="multipart/form-data", + ): + with pytest.raises(http_error) as raised: + HumanInputFileUploadApi().post() + + assert raised.value.__cause__ is service_error + if isinstance(service_error, FileTooLargeServiceError): + assert isinstance(raised.value, FileTooLargeError) + assert raised.value.description == "File size exceeded." diff --git a/api/tests/unit_tests/controllers/web/test_human_input_form.py b/api/tests/unit_tests/controllers/web/test_human_input_form.py index 5d9ba9e42e4..d53ba7d60ca 100644 --- a/api/tests/unit_tests/controllers/web/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/web/test_human_input_form.py @@ -10,7 +10,7 @@ from uuid import uuid4 import pytest from flask import Flask from sqlalchemy import Engine -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden import controllers.web.human_input_form as human_input_module @@ -268,9 +268,7 @@ def test_get_form_uses_runtime_select_options(monkeypatch: pytest.MonkeyPatch, a service_mock.resolve_form_inputs.assert_called_once_with(form) -def test_create_upload_token_returns_token_and_form_expiration( - monkeypatch: pytest.MonkeyPatch, app: Flask, sqlite_engine: Engine -): +def test_create_upload_token_returns_token_and_form_expiration(monkeypatch: pytest.MonkeyPatch, app: Flask) -> None: """POST returns a HITL upload token for an active form token.""" expiration_time = datetime(2099, 1, 1, tzinfo=UTC) @@ -279,26 +277,11 @@ def test_create_upload_token_returns_token_and_form_expiration( upload_token="hitl_upload_token-1", expires_at=expiration_time, ) - workflow_run_repository = MagicMock() - repo_factory = MagicMock(return_value=workflow_run_repository) - captured: dict[str, object] = {} - - def _service_factory(session_factory, workflow_run_repository): - captured["session_factory"] = session_factory - captured["workflow_run_repository"] = workflow_run_repository - return service_mock - - monkeypatch.setattr( - human_input_module.DifyAPIRepositoryFactory, - "create_api_workflow_run_repository", - repo_factory, - ) monkeypatch.setattr( human_input_module, - "HumanInputFileUploadService", - _service_factory, + "application_services", + lambda: SimpleNamespace(human_input_file_uploads=service_mock), ) - monkeypatch.setattr(human_input_module, "db", SimpleNamespace(engine=sqlite_engine)) limiter_mock = MagicMock() limiter_mock.is_rate_limited.return_value = False @@ -313,11 +296,6 @@ def test_create_upload_token_returns_token_and_form_expiration( "upload_token": "hitl_upload_token-1", "expires_at": int(expiration_time.timestamp()), } - repo_factory.assert_called_once() - assert captured["workflow_run_repository"] is workflow_run_repository - session_factory = captured["session_factory"] - assert isinstance(session_factory, sessionmaker) - assert session_factory.kw["bind"] is sqlite_engine service_mock.issue_upload_token.assert_called_once_with("token-1") limiter_mock.increment_rate_limit.assert_called_once_with("203.0.113.10") diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index 2be0f0c358e..f0ba570a48c 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -31,6 +31,7 @@ from repositories.account_repository import SQLAlchemyAccountRepository from repositories.app_site_command_repository import AppSiteCommandRepository from repositories.app_statistic_query_repository import AppStatisticQueryRepository from repositories.app_tracing_config_repository import SQLAlchemyAppTracingConfigRepository +from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository from repositories.workflow_app_log_query_repository import WorkflowAppLogQueryRepository from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository @@ -68,6 +69,7 @@ from services.compliance_download_service import ComplianceDownloadService from services.enterprise.enterprise_service import WebAppSettings from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFoundError from services.file_service import FileService +from services.human_input_file_upload_service import HumanInputFileUploadService from services.init_validation_service import InvalidInitializationPasswordError from services.partner_tenant_binding_service import PartnerTenantBindingService from services.retention.workflow_run.archive_download_task_cache import WorkflowRunArchiveDownloadTaskCache @@ -254,6 +256,25 @@ def test_build_application_services_wires_workflow_run_archives( assert workflow_run_archives._sign_download_url is ext_application_services.sign_workflow_run_archive_download_url +def test_build_application_services_wires_human_input_file_uploads( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + human_input_file_uploads = services.human_input_file_uploads + assert isinstance(human_input_file_uploads, HumanInputFileUploadService) + assert isinstance(human_input_file_uploads._uploads, SQLAlchemyHumanInputFileUploadRepository) + assert human_input_file_uploads._uploads._session_factory is sqlite_session_factory + assert human_input_file_uploads._remote_files is services.remote_files + assert human_input_file_uploads._files is services.files + assert services.remote_files._files is services.files + + def test_build_application_services_wires_app_site_boundary( sqlite_session_factory: sessionmaker[Session], ) -> None: diff --git a/api/tests/unit_tests/repositories/test_human_input_file_upload_repository.py b/api/tests/unit_tests/repositories/test_human_input_file_upload_repository.py new file mode 100644 index 00000000000..35639bc8aad --- /dev/null +++ b/api/tests/unit_tests/repositories/test_human_input_file_upload_repository.py @@ -0,0 +1,251 @@ +from datetime import datetime + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus +from models.account import Account, Tenant, TenantAccountJoin +from models.enums import CreatorUserRole, EndUserType +from models.human_input import ( + HumanInputForm, + HumanInputFormRecipient, + HumanInputFormUploadFile, + HumanInputFormUploadToken, + RecipientType, + StandaloneWebAppRecipientPayload, +) +from models.model import App, AppMode, EndUser, IconType +from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository +from services.human_input_file_upload_service import HumanInputUploadFormRecord + +_TENANT_ID = "11111111-1111-1111-1111-111111111111" +_OTHER_TENANT_ID = "22222222-2222-2222-2222-222222222222" +_APP_ID = "33333333-3333-3333-3333-333333333333" +_OTHER_APP_ID = "44444444-4444-4444-4444-444444444444" +_ACCOUNT_ID = "55555555-5555-5555-5555-555555555555" +_END_USER_ID = "66666666-6666-6666-6666-666666666666" +_FORM_ID = "77777777-7777-7777-7777-777777777777" +_RECIPIENT_ID = "88888888-8888-8888-8888-888888888888" +_FILE_ID = "99999999-9999-9999-9999-999999999999" +_WORKFLOW_RUN_ID = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa" +_DELIVERY_ID = "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb" +_FORM_TOKEN = "form-token" +_UPLOAD_TOKEN = "hitl_upload_token" +_CREATED_AT = datetime(2026, 1, 1) +_EXPIRATION_TIME = datetime(2099, 1, 1) + + +def _repository(session_factory: sessionmaker[Session]) -> SQLAlchemyHumanInputFileUploadRepository: + return SQLAlchemyHumanInputFileUploadRepository(session_factory=session_factory) + + +def _persist_form_and_recipient(session: Session) -> None: + session.add( + HumanInputForm( + id=_FORM_ID, + tenant_id=_TENANT_ID, + app_id=_APP_ID, + workflow_run_id=_WORKFLOW_RUN_ID, + form_kind=HumanInputFormKind.RUNTIME, + node_id="human-input", + form_definition="{}", + rendered_content="content", + expiration_time=_EXPIRATION_TIME, + created_at=_CREATED_AT, + ) + ) + session.add( + HumanInputFormRecipient( + id=_RECIPIENT_ID, + form_id=_FORM_ID, + delivery_id=_DELIVERY_ID, + recipient_type=RecipientType.STANDALONE_WEB_APP, + recipient_payload=StandaloneWebAppRecipientPayload().model_dump_json(), + access_token=_FORM_TOKEN, + ) + ) + session.commit() + + +def _persist_app_owner(session: Session) -> Account: + tenant = Tenant(name="Workspace") + tenant.id = _TENANT_ID + account = Account(name="Owner", email="owner@example.com") + account.id = _ACCOUNT_ID + session.add_all( + [ + tenant, + account, + TenantAccountJoin(tenant_id=_TENANT_ID, account_id=_ACCOUNT_ID, current=True), + App( + id=_APP_ID, + tenant_id=_TENANT_ID, + name="App", + description="", + mode=AppMode.WORKFLOW, + icon_type=IconType.EMOJI, + icon="app", + icon_background="#FFFFFF", + enable_site=True, + enable_api=True, + created_by=_ACCOUNT_ID, + updated_by=_ACCOUNT_ID, + ), + ] + ) + session.commit() + return account + + +def test_form_token_upload_token_grant_and_file_link_round_trip( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _persist_form_and_recipient(sqlite_session) + repository = _repository(sqlite_session_factory) + + form = repository.get_form_by_recipient_token(_FORM_TOKEN) + + assert form is not None + assert form == HumanInputUploadFormRecord( + form_id=_FORM_ID, + recipient_id=_RECIPIENT_ID, + tenant_id=_TENANT_ID, + app_id=_APP_ID, + workflow_run_id=_WORKFLOW_RUN_ID, + form_kind=HumanInputFormKind.RUNTIME, + status=HumanInputFormStatus.WAITING, + submitted_at=None, + expiration_time=_EXPIRATION_TIME, + created_at=_CREATED_AT, + ) + assert repository.get_form_by_recipient_token("missing-token") is None + + repository.create_upload_token(form=form, upload_token=_UPLOAD_TOKEN) + grant = repository.get_upload_grant(_UPLOAD_TOKEN) + + assert grant is not None + assert grant.form == form + assert repository.get_upload_grant("missing-token") is None + + repository.add_file( + tenant_id=_TENANT_ID, + app_id=_APP_ID, + form_id=_FORM_ID, + upload_token_id=grant.upload_token_id, + file_id=_FILE_ID, + ) + + with sqlite_session_factory() as session: + token = session.scalar( + select(HumanInputFormUploadToken).where(HumanInputFormUploadToken.token == _UPLOAD_TOKEN) + ) + link = session.scalar( + select(HumanInputFormUploadFile).where(HumanInputFormUploadFile.upload_file_id == _FILE_ID) + ) + + assert token is not None + assert token.tenant_id == _TENANT_ID + assert token.app_id == _APP_ID + assert token.form_id == _FORM_ID + assert token.recipient_id == _RECIPIENT_ID + assert link is not None + assert link.tenant_id == _TENANT_ID + assert link.app_id == _APP_ID + assert link.form_id == _FORM_ID + assert link.upload_token_id == grant.upload_token_id + assert link.upload_file_id == _FILE_ID + + +def test_get_upload_owner_hydrates_account_for_tenant( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _persist_app_owner(sqlite_session) + other_tenant = Tenant(name="Other Workspace") + other_tenant.id = _OTHER_TENANT_ID + sqlite_session.add(other_tenant) + sqlite_session.commit() + repository = _repository(sqlite_session_factory) + + owner = repository.get_upload_owner( + owner_id=_ACCOUNT_ID, + owner_role=CreatorUserRole.ACCOUNT, + tenant_id=_TENANT_ID, + app_id=_APP_ID, + ) + + assert isinstance(owner, Account) + assert owner.id == _ACCOUNT_ID + assert owner.current_tenant_id == _TENANT_ID + assert ( + repository.get_upload_owner( + owner_id=_ACCOUNT_ID, + owner_role=CreatorUserRole.ACCOUNT, + tenant_id=_OTHER_TENANT_ID, + app_id=_APP_ID, + ) + is None + ) + + +def test_get_upload_owner_scopes_end_user_to_tenant_and_app( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _persist_app_owner(sqlite_session) + end_user = EndUser( + tenant_id=_TENANT_ID, + app_id=_APP_ID, + type=EndUserType.BROWSER, + is_anonymous=False, + session_id="session", + external_user_id="external-user", + ) + end_user.id = _END_USER_ID + sqlite_session.add(end_user) + sqlite_session.commit() + repository = _repository(sqlite_session_factory) + + owner = repository.get_upload_owner( + owner_id=_END_USER_ID, + owner_role=CreatorUserRole.END_USER, + tenant_id=_TENANT_ID, + app_id=_APP_ID, + ) + + assert isinstance(owner, EndUser) + assert owner.id == _END_USER_ID + assert ( + repository.get_upload_owner( + owner_id=_END_USER_ID, + owner_role=CreatorUserRole.END_USER, + tenant_id=_OTHER_TENANT_ID, + app_id=_APP_ID, + ) + is None + ) + assert ( + repository.get_upload_owner( + owner_id=_END_USER_ID, + owner_role=CreatorUserRole.END_USER, + tenant_id=_TENANT_ID, + app_id=_OTHER_APP_ID, + ) + is None + ) + + +def test_get_delivery_test_upload_owner_requires_matching_tenant( + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + _persist_app_owner(sqlite_session) + repository = _repository(sqlite_session_factory) + + owner = repository.get_delivery_test_upload_owner(tenant_id=_TENANT_ID, app_id=_APP_ID) + + assert isinstance(owner, Account) + assert owner.id == _ACCOUNT_ID + assert owner.current_tenant_id == _TENANT_ID + assert repository.get_delivery_test_upload_owner(tenant_id=_OTHER_TENANT_ID, app_id=_APP_ID) is None diff --git a/api/tests/unit_tests/services/test_human_input_file_upload_service.py b/api/tests/unit_tests/services/test_human_input_file_upload_service.py index 3447d866c63..2f1b15cf472 100644 --- a/api/tests/unit_tests/services/test_human_input_file_upload_service.py +++ b/api/tests/unit_tests/services/test_human_input_file_upload_service.py @@ -5,297 +5,219 @@ from types import SimpleNamespace from unittest.mock import MagicMock import pytest -from sqlalchemy import create_engine, select -from sqlalchemy.orm import sessionmaker -import models.account as account_module import services.human_input_file_upload_service as service_module from core.workflow.nodes.human_input.enums import HumanInputFormKind, HumanInputFormStatus -from graphon.enums import WorkflowExecutionStatus from libs.datetime_utils import naive_utc_now -from models.account import Account, Tenant, TenantAccountJoin -from models.base import Base -from models.enums import CreatorUserRole, EndUserType, WorkflowRunTriggeredFrom -from models.human_input import ( - HumanInputForm, - HumanInputFormRecipient, - HumanInputFormUploadFile, - HumanInputFormUploadToken, +from models.account import Account +from models.enums import CreatorUserRole +from services.human_input_file_upload_service import ( + HITL_UPLOAD_TOKEN_PREFIX, + HumanInputFileUploadRepository, + HumanInputFileUploadService, + HumanInputUploadContext, + HumanInputUploadFormRecord, + HumanInputUploadGrantRecord, + InvalidUploadTokenError, ) -from models.model import App, AppMode, EndUser -from models.workflow import WorkflowRun, WorkflowType -from services.human_input_file_upload_service import HITL_UPLOAD_TOKEN_PREFIX, HumanInputFileUploadService -from services.human_input_service import FormSubmittedError +from services.human_input_service import FormNotFoundError, FormSubmittedError -@pytest.fixture -def session_maker(monkeypatch: pytest.MonkeyPatch): - engine = create_engine("sqlite:///:memory:") - monkeypatch.setattr(account_module, "db", SimpleNamespace(engine=engine)) - Base.metadata.create_all( - engine, - tables=[ - Tenant.__table__, - Account.__table__, - TenantAccountJoin.__table__, - App.__table__, - EndUser.__table__, - WorkflowRun.__table__, - HumanInputForm.__table__, - HumanInputFormRecipient.__table__, - HumanInputFormUploadToken.__table__, - HumanInputFormUploadFile.__table__, - ], - ) - try: - yield sessionmaker(bind=engine, expire_on_commit=False) - finally: - Base.metadata.drop_all( - engine, - tables=[ - HumanInputFormUploadFile.__table__, - HumanInputFormUploadToken.__table__, - HumanInputFormRecipient.__table__, - HumanInputForm.__table__, - WorkflowRun.__table__, - EndUser.__table__, - App.__table__, - TenantAccountJoin.__table__, - Account.__table__, - Tenant.__table__, - ], - ) - engine.dispose() - - -def _create_waiting_form( - session_maker, +def _active_form( *, - created_by_role: CreatorUserRole = CreatorUserRole.ACCOUNT, + workflow_run_id: str | None = "run-1", form_kind: HumanInputFormKind = HumanInputFormKind.RUNTIME, -) -> tuple[str, str, str]: - form_id = "00000000-0000-0000-0000-000000000001" - recipient_id = "00000000-0000-0000-0000-000000000002" - workflow_run_id = None - if form_kind == HumanInputFormKind.RUNTIME: - workflow_run_id = "00000000-0000-0000-0000-000000000012" - tenant_id = "00000000-0000-0000-0000-000000000010" - app_id = "00000000-0000-0000-0000-000000000011" + status: HumanInputFormStatus = HumanInputFormStatus.WAITING, +) -> HumanInputUploadFormRecord: now = naive_utc_now() - created_by = ( - "00000000-0000-0000-0000-000000000020" - if created_by_role == CreatorUserRole.ACCOUNT - else "00000000-0000-0000-0000-000000000021" + return HumanInputUploadFormRecord( + form_id="form-1", + recipient_id="recipient-1", + tenant_id="tenant-1", + app_id="app-1", + workflow_run_id=workflow_run_id, + form_kind=form_kind, + status=status, + submitted_at=None, + expiration_time=now + timedelta(hours=1), + created_at=now, + ) + + +def _upload_context() -> HumanInputUploadContext: + return HumanInputUploadContext( + tenant_id="tenant-1", + app_id="app-1", + form_id="form-1", + recipient_id="recipient-1", + upload_token_id="token-1", + owner=MagicMock(spec=Account), ) - with session_maker.begin() as session: - tenant = Tenant(name="tenant-1") - tenant.id = tenant_id - session.add(tenant) - if created_by_role == CreatorUserRole.ACCOUNT: - account = Account(name="owner", email="owner@example.com") - account.id = created_by - session.add(account) - session.add( - TenantAccountJoin( - tenant_id=tenant_id, - account_id=created_by, - current=True, - ) - ) - app_creator = created_by - else: - end_user = EndUser( - tenant_id=tenant_id, - app_id=app_id, - type=EndUserType.BROWSER, - is_anonymous=False, - session_id="session-1", - external_user_id="external-1", - ) - end_user.id = created_by - session.add(end_user) - app_creator = "00000000-0000-0000-0000-000000000020" - account = Account(name="owner", email="owner@example.com") - account.id = app_creator - session.add(account) - session.add( - TenantAccountJoin( - tenant_id=tenant_id, - account_id=app_creator, - current=True, - ) - ) - app = App( - tenant_id=tenant_id, - name="app-1", - description="", - mode=AppMode.WORKFLOW, - icon_type="emoji", - icon="app", - icon_background="#ffffff", - enable_site=True, - enable_api=True, - created_by=app_creator, - updated_by=app_creator, - ) - app.id = app_id - session.add(app) - if workflow_run_id is not None: - workflow_run = WorkflowRun( - tenant_id=tenant_id, - app_id=app_id, - workflow_id="00000000-0000-0000-0000-000000000013", - type=WorkflowType.WORKFLOW, - triggered_from=WorkflowRunTriggeredFrom.APP_RUN, - version="1", - graph="{}", - inputs="{}", - status=WorkflowExecutionStatus.RUNNING, - created_by_role=created_by_role, - created_by=created_by, - created_at=now, - ) - workflow_run.id = workflow_run_id - session.add(workflow_run) - session.add( - HumanInputForm( - id=form_id, - tenant_id=tenant_id, - app_id=app_id, - workflow_run_id=workflow_run_id, - form_kind=form_kind, - node_id="node-1", - form_definition="{}", - rendered_content="content", - expiration_time=now + timedelta(hours=1), - created_at=now, - ) - ) - session.add( - HumanInputFormRecipient( - id=recipient_id, - form_id=form_id, - delivery_id="00000000-0000-0000-0000-000000000003", - recipient_type="standalone_web_app", - recipient_payload='{"TYPE": "standalone_web_app"}', - access_token="form-token-1", - ) - ) - return form_id, recipient_id, created_by def _create_service( - session_maker, - workflow_run_repository: MagicMock | None = None, + *, + uploads: MagicMock | None = None, + workflow_runs: MagicMock | None = None, + files: MagicMock | None = None, + remote_files: MagicMock | None = None, ) -> HumanInputFileUploadService: return HumanInputFileUploadService( - session_maker, - workflow_run_repository=workflow_run_repository or MagicMock(), + uploads=uploads if uploads is not None else MagicMock(spec=HumanInputFileUploadRepository), + workflow_run_repository=workflow_runs if workflow_runs is not None else MagicMock(), + files=files if files is not None else MagicMock(), + remote_files=remote_files if remote_files is not None else MagicMock(), ) -def _get_workflow_run(session_maker) -> WorkflowRun: - with session_maker() as session: - workflow_run = session.get(WorkflowRun, "00000000-0000-0000-0000-000000000012") - assert workflow_run is not None - return workflow_run - - -def test_issue_upload_token_persists_token_without_technical_end_user( - monkeypatch: pytest.MonkeyPatch, - session_maker, -) -> None: - form_id, recipient_id, _created_by = _create_waiting_form(session_maker) +def test_issue_upload_token_persists_repository_record(monkeypatch: pytest.MonkeyPatch) -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + form = _active_form() + uploads.get_form_by_recipient_token.return_value = form monkeypatch.setattr(service_module.secrets, "token_urlsafe", lambda _bytes: "random-value") - token = _create_service(session_maker).issue_upload_token("form-token-1") + token = _create_service(uploads=uploads).issue_upload_token("form-token-1") assert token.upload_token == f"{HITL_UPLOAD_TOKEN_PREFIX}random-value" - with session_maker() as session: - token_model = session.scalar(select(HumanInputFormUploadToken)) - assert token_model is not None - assert token_model.form_id == form_id - assert token_model.recipient_id == recipient_id - assert token_model.token == token.upload_token - assert session.scalar(select(EndUser).limit(1)) is None + assert token.expires_at == form.expiration_time + uploads.get_form_by_recipient_token.assert_called_once_with("form-token-1") + uploads.create_upload_token.assert_called_once_with(form=form, upload_token=token.upload_token) -def test_validate_upload_token_returns_account_owner_and_record_file_link(session_maker) -> None: - form_id, recipient_id, created_by = _create_waiting_form(session_maker, created_by_role=CreatorUserRole.ACCOUNT) - token = _create_service(session_maker).issue_upload_token("form-token-1") - workflow_run_repository = MagicMock() - workflow_run_repository.get_workflow_run_by_id.return_value = _get_workflow_run(session_maker) +def test_issue_upload_token_rejects_unknown_form_token() -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + uploads.get_form_by_recipient_token.return_value = None - context = HumanInputFileUploadService( - session_maker, - workflow_run_repository=workflow_run_repository, - ).validate_upload_token(token.upload_token) - assert context.form_id == form_id - assert context.recipient_id == recipient_id - assert isinstance(context.owner, Account) - assert context.owner.id == created_by - assert context.owner.current_tenant_id == "00000000-0000-0000-0000-000000000010" - workflow_run_repository.get_workflow_run_by_id.assert_called_once_with( - tenant_id="00000000-0000-0000-0000-000000000010", - app_id="00000000-0000-0000-0000-000000000011", - run_id="00000000-0000-0000-0000-000000000012", + with pytest.raises(FormNotFoundError): + _create_service(uploads=uploads).issue_upload_token("missing-form-token") + + uploads.create_upload_token.assert_not_called() + + +@pytest.mark.parametrize("owner_role", [CreatorUserRole.ACCOUNT, CreatorUserRole.END_USER]) +def test_validate_upload_token_resolves_workflow_run_owner(owner_role: CreatorUserRole) -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + form = _active_form() + uploads.get_upload_grant.return_value = HumanInputUploadGrantRecord(upload_token_id="token-1", form=form) + owner = MagicMock() + uploads.get_upload_owner.return_value = owner + workflow_runs = MagicMock() + workflow_runs.get_workflow_run_by_id.return_value = SimpleNamespace( + created_by="owner-1", + created_by_role=owner_role, + tenant_id="tenant-1", + app_id="app-1", ) - _create_service(session_maker).record_upload_file( - context=context, - file_id="00000000-0000-0000-0000-000000000099", + context = _create_service(uploads=uploads, workflow_runs=workflow_runs).validate_upload_token("upload-token-1") + + assert context == HumanInputUploadContext( + tenant_id="tenant-1", + app_id="app-1", + form_id="form-1", + recipient_id="recipient-1", + upload_token_id="token-1", + owner=owner, + ) + workflow_runs.get_workflow_run_by_id.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + run_id="run-1", + ) + uploads.get_upload_owner.assert_called_once_with( + owner_id="owner-1", + owner_role=owner_role, + tenant_id="tenant-1", + app_id="app-1", ) - with session_maker() as session: - link = session.scalar(select(HumanInputFormUploadFile)) - assert link is not None - assert link.tenant_id == context.tenant_id - assert link.app_id == context.app_id - assert link.form_id == form_id - assert link.upload_token_id == context.upload_token_id + +def test_validate_upload_token_resolves_delivery_test_owner() -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + form = _active_form(workflow_run_id=None, form_kind=HumanInputFormKind.DELIVERY_TEST) + uploads.get_upload_grant.return_value = HumanInputUploadGrantRecord(upload_token_id="token-1", form=form) + owner = MagicMock(spec=Account) + uploads.get_delivery_test_upload_owner.return_value = owner + workflow_runs = MagicMock() + + context = _create_service(uploads=uploads, workflow_runs=workflow_runs).validate_upload_token("upload-token-1") + + assert context.owner is owner + uploads.get_delivery_test_upload_owner.assert_called_once_with(tenant_id="tenant-1", app_id="app-1") + workflow_runs.get_workflow_run_by_id.assert_not_called() -def test_validate_upload_token_returns_end_user_owner(session_maker) -> None: - form_id, recipient_id, created_by = _create_waiting_form(session_maker, created_by_role=CreatorUserRole.END_USER) - token = _create_service(session_maker).issue_upload_token("form-token-1") - workflow_run_repository = MagicMock() - workflow_run_repository.get_workflow_run_by_id.return_value = _get_workflow_run(session_maker) +def test_validate_upload_token_rejects_unknown_upload_token() -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + uploads.get_upload_grant.return_value = None - context = HumanInputFileUploadService( - session_maker, - workflow_run_repository=workflow_run_repository, - ).validate_upload_token(token.upload_token) - - assert context.form_id == form_id - assert context.recipient_id == recipient_id - assert isinstance(context.owner, EndUser) - assert context.owner.id == created_by + with pytest.raises(InvalidUploadTokenError): + _create_service(uploads=uploads).validate_upload_token("missing-upload-token") -def test_validate_upload_token_allows_delivery_test_form(session_maker) -> None: - form_id, recipient_id, _created_by = _create_waiting_form( - session_maker, - form_kind=HumanInputFormKind.DELIVERY_TEST, - ) - token = _create_service(session_maker).issue_upload_token("form-token-1") - - context = _create_service(session_maker).validate_upload_token(token.upload_token) - - assert context.form_id == form_id - assert context.recipient_id == recipient_id - assert isinstance(context.owner, Account) - assert context.owner.id == "00000000-0000-0000-0000-000000000020" - assert context.owner.current_tenant_id == "00000000-0000-0000-0000-000000000010" - - -def test_validate_upload_token_rejects_submitted_form(session_maker) -> None: - form_id, _recipient_id, _created_by = _create_waiting_form(session_maker) - token = _create_service(session_maker).issue_upload_token("form-token-1") - with session_maker.begin() as session: - form = session.get(HumanInputForm, form_id) - assert form is not None - form.status = HumanInputFormStatus.SUBMITTED - form.submitted_at = naive_utc_now() +def test_validate_upload_token_rejects_submitted_form() -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + form = _active_form(status=HumanInputFormStatus.SUBMITTED) + uploads.get_upload_grant.return_value = HumanInputUploadGrantRecord(upload_token_id="token-1", form=form) with pytest.raises(FormSubmittedError): - _create_service(session_maker).validate_upload_token(token.upload_token) + _create_service(uploads=uploads).validate_upload_token("upload-token-1") + + uploads.get_upload_owner.assert_not_called() + + +def test_upload_local_file_records_the_form_file_link() -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + files = MagicMock() + upload_file = MagicMock(id="file-1") + files.upload_file.return_value = upload_file + context = _upload_context() + + result = _create_service(uploads=uploads, files=files).upload_local_file( + context=context, + filename="sample.txt", + content=b"content", + mimetype="text/plain", + ) + + assert result is upload_file + files.upload_file.assert_called_once_with( + filename="sample.txt", + content=b"content", + mimetype="text/plain", + user=context.owner, + source=None, + ) + uploads.add_file.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + form_id="form-1", + upload_token_id="token-1", + file_id="file-1", + ) + + +def test_upload_remote_file_records_the_form_file_link() -> None: + uploads = MagicMock(spec=HumanInputFileUploadRepository) + remote_files = MagicMock() + upload_file = MagicMock(id="file-1") + remote_files.upload_from_url.return_value = upload_file + context = _upload_context() + + result = _create_service(uploads=uploads, remote_files=remote_files).upload_remote_file( + context=context, + url="https://example.com/sample.txt", + ) + + assert result is upload_file + remote_files.upload_from_url.assert_called_once_with( + url="https://example.com/sample.txt", + user=context.owner, + ) + uploads.add_file.assert_called_once_with( + tenant_id="tenant-1", + app_id="app-1", + form_id="form-1", + upload_token_id="token-1", + file_id="file-1", + )