diff --git a/api/repositories/installed_app_connection_repository.py b/api/repositories/installed_app_connection_repository.py index 73eb746b1de..92539533625 100644 --- a/api/repositories/installed_app_connection_repository.py +++ b/api/repositories/installed_app_connection_repository.py @@ -1,14 +1,23 @@ """Bounded persistence for Installed App member choices and non-secret credential metadata.""" import json +from collections.abc import Mapping +from dataclasses import asdict from datetime import UTC, datetime +from typing import Any -from sqlalchemy import select +from sqlalchemy import select, update +from sqlalchemy.dialects.mysql import insert as mysql_insert +from sqlalchemy.dialects.postgresql import insert as postgresql_insert +from sqlalchemy.dialects.sqlite import insert as sqlite_insert +from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.sql.dml import Insert from constants.end_user_connections import ConnectionCredentialType from context.connection_execution import ConnectionExecutionContext, WorkspaceConnectionScope from core.app.entities.app_invoke_entities import InvokeFrom +from libs.datetime_utils import naive_utc_now from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantStatus from models.credential_permission import CredentialType from models.enums import ConversationFromSource @@ -32,15 +41,21 @@ class SQLAlchemyInstalledConnectionRepository: def __init__(self, *, session_factory: sessionmaker[Session]) -> None: self._sessions = session_factory - @staticmethod + @classmethod def _graph( + cls, session: Session, scope: WorkspaceConnectionScope, *, workflow_id: str | None = None, conversation_id: str | None = None, - lock: bool = False, ) -> InstalledConnectionGraph: + """Read the current owner chain and publication without acquiring row locks. + + Accept an admitted member scope and optional expected publication/session. + Return a detached graph or raise the existing identity, revision, or + conversation error without writing any state. + """ statement = ( select(InstalledApp, App, Workflow) .join(App, App.id == InstalledApp.app_id) @@ -65,8 +80,6 @@ class SQLAlchemyInstalledConnectionRepository: Tenant.status == TenantStatus.NORMAL, ) ) - if lock: - statement = statement.with_for_update() row = session.execute(statement).one_or_none() if row is None: raise ConnectionIdentityError() @@ -76,20 +89,7 @@ class SQLAlchemyInstalledConnectionRepository: if conversation_id is not None: if ( app.mode != AppMode.ADVANCED_CHAT - or session.scalar( - select(Conversation.id).where( - Conversation.id == conversation_id, - Conversation.app_id == app.id, - Conversation.mode == AppMode.ADVANCED_CHAT, - Conversation.from_account_id == scope.account_id, - Conversation.from_end_user_id.is_(None), - Conversation.from_source == ConversationFromSource.CONSOLE, - Conversation.invoke_from == InvokeFrom.EXPLORE, - Conversation.is_deleted.is_(False), - Conversation.status == "normal", - ) - ) - is None + or session.scalar(cls._conversation_statement(scope, app.id, conversation_id)) is None ): raise ConversationNotExistsError() return InstalledConnectionGraph( @@ -104,6 +104,97 @@ class SQLAlchemyInstalledConnectionRepository: workflow.graph_dict, ) + @staticmethod + def _conversation_statement(scope: WorkspaceConnectionScope, app_id: str, conversation_id: str): + """Select only a live Explore conversation owned by the admitted member.""" + return select(Conversation.id).where( + Conversation.id == conversation_id, + Conversation.app_id == app_id, + Conversation.mode == AppMode.ADVANCED_CHAT, + Conversation.from_account_id == scope.account_id, + Conversation.from_end_user_id.is_(None), + Conversation.from_source == ConversationFromSource.CONSOLE, + Conversation.invoke_from == InvokeFrom.EXPLORE, + Conversation.is_deleted.is_(False), + Conversation.status == "normal", + ) + + @classmethod + def _protect_graph( + cls, + session: Session, + scope: WorkspaceConnectionScope, + graph: InstalledConnectionGraph, + conversation_id: str | None, + ) -> None: + """Protect a previously read owner chain and graph until the caller commits. + + Require an active write transaction. Acquire single-table shared locks in + tenant/account/member/installation/app/workflow/conversation order, so + independent configuration writers remain compatible. Locking projections + read current values even with an existing ORM identity or MySQL snapshot. + Raise an identity, revision, or conversation error before any write if the + owner chain or publication changed. The caller owns rollback on failure. + """ + installation = graph.installation + if ( + installation.id != scope.installed_app_id + or installation.tenant_id != scope.tenant_id + or installation.app_owner_tenant_id != scope.tenant_id + ): + raise ConnectionIdentityError() + identity_statements = ( + select(Tenant.id).where(Tenant.id == scope.tenant_id, Tenant.status == TenantStatus.NORMAL), + select(Account.id).where(Account.id == scope.account_id, Account.status == AccountStatus.ACTIVE), + select(TenantAccountJoin.id).where( + TenantAccountJoin.tenant_id == scope.tenant_id, + TenantAccountJoin.account_id == scope.account_id, + ), + select(InstalledApp.id).where( + InstalledApp.id == scope.installed_app_id, + InstalledApp.tenant_id == scope.tenant_id, + InstalledApp.app_id == installation.app_id, + ), + ) + for statement in identity_statements: + if session.scalar(statement.with_for_update(read=True)) is None: + raise ConnectionIdentityError() + app = session.execute( + select(App.mode, App.workflow_id) + .where( + App.id == installation.app_id, + App.tenant_id == scope.tenant_id, + App.mode.in_((AppMode.ADVANCED_CHAT, AppMode.WORKFLOW)), + ) + .with_for_update(read=True) + ).one_or_none() + if app is None: + raise ConnectionIdentityError() + if app.workflow_id != graph.workflow_id or app.mode != installation.app_mode: + raise ConnectionRevisionChangedError() + workflow_graph = session.scalar( + select(Workflow.graph) + .where( + Workflow.id == graph.workflow_id, + Workflow.tenant_id == scope.tenant_id, + Workflow.app_id == installation.app_id, + Workflow.version != Workflow.VERSION_DRAFT, + ) + .with_for_update(read=True) + ) + if workflow_graph is None: + raise ConnectionIdentityError() + if json.loads(workflow_graph) != graph.graph: + raise ConnectionRevisionChangedError() + if conversation_id is not None and ( + app.mode != AppMode.ADVANCED_CHAT + or session.scalar( + cls._conversation_statement(scope, installation.app_id, conversation_id).with_for_update(read=True) + ) + is None + ): + raise ConversationNotExistsError() + def graph( self, scope: WorkspaceConnectionScope, *, workflow_id: str | None = None, conversation_id: str | None = None ) -> InstalledConnectionGraph: @@ -112,8 +203,17 @@ class SQLAlchemyInstalledConnectionRepository: @staticmethod def _connections( - session: Session, scope: WorkspaceConnectionScope, providers: tuple[str, ...] + session: Session, + scope: WorkspaceConnectionScope, + providers: tuple[str, ...], + *, + connection_ids: frozenset[str] | None, ) -> tuple[ConnectionSummary, ...]: + """Read visible credential metadata without secrets or row locks. + + None selects all credentials for the providers; a set restricts the read + to explicit selections, including an empty set for an empty configuration. + """ statement = select( BuiltinToolProvider.id, BuiltinToolProvider.name, @@ -123,6 +223,8 @@ class SQLAlchemyInstalledConnectionRepository: BuiltinToolProvider.updated_at, BuiltinToolProvider.expires_at, ).where(BuiltinToolProvider.tenant_id == scope.tenant_id, BuiltinToolProvider.provider.in_(providers)) + if connection_ids is not None: + statement = statement.where(BuiltinToolProvider.id.in_(connection_ids)) statement = apply_credential_visibility_filter_for_actor( statement, tenant_id=scope.tenant_id, @@ -153,7 +255,7 @@ class SQLAlchemyInstalledConnectionRepository: def connections(self, scope: WorkspaceConnectionScope, providers: tuple[str, ...]) -> tuple[ConnectionSummary, ...]: with self._sessions() as session: self._graph(session, scope) - return self._connections(session, scope, providers) + return self._connections(session, scope, providers, connection_ids=None) @staticmethod def _configuration_statement(scope: WorkspaceConnectionScope, conversation_id: str | None): @@ -183,16 +285,36 @@ class SQLAlchemyInstalledConnectionRepository: configuration: InstalledConnectionConfiguration, conversation_id: str | None, ) -> None: - current = cls._graph(session, scope, workflow_id=graph.workflow_id, conversation_id=conversation_id, lock=True) - if current != graph: - raise ConnectionRevisionChangedError() + """Stage one complete member configuration in the caller's transaction. + + Require a previously read graph. Parse before locking, then protect its + owner chain/publication and recheck selected credential visibility. Only + the configuration's unique member/session scope is exclusively written. + Failure leaves commit/rollback to the caller and never grants access. + """ from services.workflow.tool_credentials import get_tool_connection_requirements requirements = get_tool_connection_requirements(graph.graph) if configuration.connections.keys() - {item.node_id for item in requirements}: raise ConnectionConfigurationError() + saved = InstalledAppConnectionConfiguration( + tenant_id=scope.tenant_id, + installed_app_id=scope.installed_app_id, + app_id=graph.installation.app_id, + account_id=scope.account_id, + scope_id=conversation_id or "", + connection_revision=configuration.connection_revision, + connections=json.dumps(dict(configuration.connections)), + ) + cls._protect_graph(session, scope, graph, conversation_id) records = { - item.id: item for item in cls._connections(session, scope, tuple(item.provider for item in requirements)) + item.id: item + for item in cls._connections( + session, + scope, + tuple(item.provider for item in requirements), + connection_ids=frozenset(value for value in configuration.connections.values() if value is not None), + ) } for requirement in requirements: connection_id = configuration.connections.get(requirement.node_id) @@ -204,20 +326,55 @@ class SQLAlchemyInstalledConnectionRepository: or record.credential_type not in requirement.allowed_credential_types ): raise ConnectionConfigurationError() - saved = session.scalar(cls._configuration_statement(scope, conversation_id)) - if saved is None: - saved = InstalledAppConnectionConfiguration( - tenant_id=scope.tenant_id, - installed_app_id=scope.installed_app_id, - app_id=graph.installation.app_id, - account_id=scope.account_id, - scope_id=conversation_id or "", - connection_revision=configuration.connection_revision, - connections="{}", + session.execute(cls._configuration_insert_statement(session.get_bind().dialect.name, asdict(saved))) + result = session.execute( + update(InstalledAppConnectionConfiguration) + .where( + InstalledAppConnectionConfiguration.tenant_id == scope.tenant_id, + InstalledAppConnectionConfiguration.installed_app_id == scope.installed_app_id, + InstalledAppConnectionConfiguration.app_id == graph.installation.app_id, + InstalledAppConnectionConfiguration.account_id == scope.account_id, + InstalledAppConnectionConfiguration.scope_id == (conversation_id or ""), ) - session.add(saved) - saved.connection_revision = configuration.connection_revision - saved.connections = json.dumps(dict(configuration.connections)) + .values( + connection_revision=saved.connection_revision, + connections=saved.connections, + updated_at=naive_utc_now(), + ) + ) + assert isinstance(result, CursorResult) + if result.rowcount != 1: + raise ConnectionIdentityError() + + @staticmethod + def _configuration_insert_statement(dialect_name: str, values: Mapping[str, Any]) -> Insert: + """Ensure a configuration exists without overwriting conflicting ownership. + + Accept complete model values and a supported database dialect. The owner + unique key arbitrates concurrent first saves; the caller must follow with + a fully scoped update in the same transaction. Unsupported dialects fail + explicitly, and existing rows are left unchanged. + """ + owner_key = ["installed_app_id", "account_id", "scope_id"] + if dialect_name == "postgresql": + return ( + postgresql_insert(InstalledAppConnectionConfiguration) + .values(**values) + .on_conflict_do_nothing(index_elements=owner_key) + ) + if dialect_name == "mysql": + return ( + mysql_insert(InstalledAppConnectionConfiguration) + .values(**values) + .on_duplicate_key_update(id=InstalledAppConnectionConfiguration.id) + ) + if dialect_name == "sqlite": + return ( + sqlite_insert(InstalledAppConnectionConfiguration) + .values(**values) + .on_conflict_do_nothing(index_elements=owner_key) + ) + raise ValueError(f"Unsupported configuration database dialect: {dialect_name}") def save( self, @@ -226,6 +383,12 @@ class SQLAlchemyInstalledConnectionRepository: configuration: InstalledConnectionConfiguration, conversation_id: str | None, ) -> None: + """Commit a validated complete configuration, rolling back on any failure. + + Require a graph obtained for this member scope. Shared identity/publication + locks and the configuration owner unique key serialize only conflicting + operations; concurrent saves of one scope retain last-writer-wins semantics. + """ with self._sessions.begin() as session: self._save(session, scope, graph, configuration, conversation_id) @@ -233,12 +396,16 @@ class SQLAlchemyInstalledConnectionRepository: def initialize_conversation( cls, session: Session, execution: ConnectionExecutionContext, conversation: Conversation ) -> None: + """Stage initial selections with a new, flushed member-owned conversation. + + Validate the execution publication and owner chain once under shared locks. + Never commit the supplied session: the conversation, message, and choices + remain atomic, and the caller rolls all of them back on failure. + """ scope = execution.workspace if scope is None or conversation.from_account_id != scope.account_id or conversation.app_id != execution.app_id: raise ConnectionIdentityError() - graph = cls._graph( - session, scope, workflow_id=execution.workflow_id, conversation_id=conversation.id, lock=True - ) + graph = cls._graph(session, scope, workflow_id=execution.workflow_id, conversation_id=conversation.id) cls._save( session, scope, diff --git a/api/services/installed_app_connections.py b/api/services/installed_app_connections.py index 6e4a97b7d44..1344d8304d9 100644 --- a/api/services/installed_app_connections.py +++ b/api/services/installed_app_connections.py @@ -85,7 +85,20 @@ class InstalledConnectionStore(Protocol): graph: InstalledConnectionGraph, configuration: InstalledConnectionConfiguration, conversation_id: str | None, - ) -> None: ... + ) -> None: + """Atomically replace selections for the admitted member/session scope. + + Require a previously read graph and validate the current tenant, account, + membership, installation, publication, conversation ownership, and selected + credential visibility before writing. Protect mutable identity/publication + facts with shared locks through commit; the configuration owner unique key + arbitrates first creation and complete replacements serialize per scope. + Independent scopes must not acquire mutually exclusive owner-chain locks. + Preserve last-writer-wins semantics, never persist secrets or grant access, + and roll back on identity, revision, conversation, or configuration errors. + None selects the member's app defaults, not a conversation. + """ + ... class InstalledAppConnectionService: diff --git a/api/tests/e2e_tests/middleware/tools/test_installed_app_connection_repository.py b/api/tests/e2e_tests/middleware/tools/test_installed_app_connection_repository.py new file mode 100644 index 00000000000..e66528bbf22 --- /dev/null +++ b/api/tests/e2e_tests/middleware/tools/test_installed_app_connection_repository.py @@ -0,0 +1,188 @@ +"""Real PostgreSQL contention tests for member-scoped Installed App choices.""" + +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from uuid import uuid4 + +import pytest +from sqlalchemy import delete, event, select, text, update +from sqlalchemy.engine import Connection +from sqlalchemy.exc import OperationalError +from sqlalchemy.orm import ORMExecuteState, Session, SessionTransaction, sessionmaker + +from extensions.ext_database import db +from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantStatus +from models.installed_app_connection import InstalledAppConnectionConfiguration +from models.model import App, Conversation, InstalledApp +from models.tools import BuiltinToolProvider +from models.workflow import Workflow +from repositories.installed_app_connection_repository import SQLAlchemyInstalledConnectionRepository +from services.end_user_connections.errors import ConnectionIdentityError, ConnectionRevisionChangedError +from services.installed_app_connections import InstalledConnectionConfiguration +from tests.helpers.installed_app_connections import ( + InstalledConnectionData, + member_conversation, + seed_installed_connections, +) + + +@pytest.fixture +def cases( + flask_req_ctx: object, +) -> Iterator[tuple[sessionmaker[Session], InstalledConnectionData, InstalledConnectionData]]: + """Seed isolated real records and bound every transaction's lock wait.""" + if db.engine.dialect.name != "postgresql": + pytest.skip("requires the CI PostgreSQL integration database") + sessions = sessionmaker(bind=db.engine, expire_on_commit=False) + + def bound_lock_wait(_session: Session, _transaction: SessionTransaction, connection: Connection) -> None: + connection.execute(text("SET LOCAL lock_timeout = '2s'")) + connection.execute(text("SET LOCAL statement_timeout = '5s'")) + + event.listen(sessions, "after_begin", bound_lock_wait) + first = seed_installed_connections(sessions) + second = seed_installed_connections(sessions, tenant=first.tenant, account=first.account) + try: + yield sessions, first, second + finally: + with sessions.begin() as session: + tenant_id = first.tenant.id + account_ids = [first.account.id, first.other.id, second.other.id] + session.execute( + delete(InstalledAppConnectionConfiguration).where( + InstalledAppConnectionConfiguration.tenant_id == tenant_id + ) + ) + session.execute(delete(Conversation).where(Conversation.app_id.in_([first.app.id, second.app.id]))) + session.execute(delete(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant_id)) + session.execute(delete(Workflow).where(Workflow.tenant_id == tenant_id)) + session.execute(delete(InstalledApp).where(InstalledApp.tenant_id == tenant_id)) + session.execute(delete(App).where(App.tenant_id == tenant_id)) + session.execute(delete(TenantAccountJoin).where(TenantAccountJoin.tenant_id == tenant_id)) + session.execute(delete(Account).where(Account.id.in_(account_ids))) + session.execute(delete(Tenant).where(Tenant.id == tenant_id)) + event.remove(sessions, "after_begin", bound_lock_wait) + + +@pytest.mark.parametrize("independent_scope", ["app", "member", "conversation", "defaults_and_conversation"]) +def test_independent_scope_commits_while_first_save_remains_open(cases, independent_scope): + sessions, first, second = cases + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sessions) + first_scope, second_scope = first.scope, first.scope + first_graph = repository.graph(first_scope) + second_graph = first_graph + first_conversation_id, second_conversation_id = None, None + configuration = first.configuration + if independent_scope == "app": + second_scope = second.scope + second_graph = repository.graph(second_scope) + configuration = second.configuration + elif independent_scope == "member": + second_scope = first.scope.model_copy(update={"account_id": first.other.id}) + else: + with sessions.begin() as session: + conversations = [member_conversation(first), member_conversation(first)] + session.add_all(conversations) + first_conversation_id = conversations[0].id if independent_scope == "conversation" else None + second_conversation_id = conversations[1].id + with ThreadPoolExecutor(max_workers=1) as executor: + with sessions.begin() as session: + repository._save(session, first_scope, first_graph, first.configuration, first_conversation_id) + future = executor.submit(repository.save, second_scope, second_graph, configuration, second_conversation_id) + future.result(timeout=8) + assert repository.configuration(second_scope, second_conversation_id) == configuration + + +def test_new_conversation_initializations_do_not_serialize_on_workspace(cases): + sessions, first, _second = cases + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sessions) + + def initialize() -> str: + with sessions.begin() as session: + conversation = member_conversation(first) + session.add(conversation) + session.flush() + repository.initialize_conversation(session, first.execution, conversation) + return conversation.id + + with ThreadPoolExecutor(max_workers=1) as executor: + with sessions.begin() as session: + conversation = member_conversation(first) + session.add(conversation) + session.flush() + repository.initialize_conversation(session, first.execution, conversation) + second_id = executor.submit(initialize).result(timeout=8) + assert repository.configuration(first.scope, second_id) == first.configuration + + +@pytest.mark.parametrize("existing", [False, True]) +def test_concurrent_first_or_existing_saves_commit_one_complete_configuration(cases, existing): + sessions, first, _second = cases + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sessions) + graph = repository.graph(first.scope) + if existing: + repository.save(first.scope, graph, first.configuration, None) + configurations = ( + InstalledConnectionConfiguration(str(uuid4()), {"root": None}), + InstalledConnectionConfiguration(str(uuid4()), {}), + ) + barrier = Barrier(2, timeout=8) + + def synchronize_insert(state: ORMExecuteState) -> None: + if state.is_insert: + barrier.wait() + + event.listen(sessions, "do_orm_execute", synchronize_insert) + try: + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(repository.save, first.scope, graph, config, None) for config in configurations] + for future in futures: + future.result(timeout=10) + finally: + event.remove(sessions, "do_orm_execute", synchronize_insert) + assert repository.configuration(first.scope, None) in configurations + with sessions() as session: + assert len(session.scalars(select(InstalledAppConnectionConfiguration)).all()) == 1 + + +def test_same_scope_waits_for_its_configuration_writer(cases): + sessions, first, _second = cases + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sessions) + graph = repository.graph(first.scope) + with ThreadPoolExecutor(max_workers=1) as executor: + with sessions.begin() as session: + repository._save(session, first.scope, graph, first.configuration, None) + future = executor.submit(repository.save, first.scope, graph, first.configuration, None) + with pytest.raises(OperationalError, match="lock timeout"): + future.result(timeout=8) + assert repository.configuration(first.scope, None) == first.configuration + + +@pytest.mark.parametrize("change", ["tenant", "account", "member", "installation", "publication", "workflow"]) +def test_invalidation_waits_for_configuration_commit(cases, change): + sessions, first, _second = cases + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sessions) + graph = repository.graph(first.scope) + statements = { + "tenant": update(Tenant).where(Tenant.id == first.tenant.id).values(status=TenantStatus.ARCHIVE), + "account": update(Account).where(Account.id == first.account.id).values(status=AccountStatus.BANNED), + "member": delete(TenantAccountJoin).where(TenantAccountJoin.account_id == first.account.id), + "installation": delete(InstalledApp).where(InstalledApp.id == first.installation.id), + "publication": update(App).where(App.id == first.app.id).values(workflow_id=None), + "workflow": update(Workflow).where(Workflow.id == first.workflow.id).values(graph='{"nodes":[],"edges":[]}'), + } + + def invalidate() -> None: + with sessions.begin() as session: + session.execute(statements[change]) + + with ThreadPoolExecutor(max_workers=1) as executor: + with sessions.begin() as session: + repository._save(session, first.scope, graph, first.configuration, None) + with pytest.raises(OperationalError, match="lock timeout"): + executor.submit(invalidate).result(timeout=8) + assert repository.configuration(first.scope, None) == first.configuration + invalidate() + with pytest.raises((ConnectionIdentityError, ConnectionRevisionChangedError)): + repository.save(first.scope, graph, first.configuration, None) diff --git a/api/tests/helpers/installed_app_connections.py b/api/tests/helpers/installed_app_connections.py new file mode 100644 index 00000000000..df8751ac898 --- /dev/null +++ b/api/tests/helpers/installed_app_connections.py @@ -0,0 +1,122 @@ +"""Real persistence fixtures shared by member-configuration repository tests.""" + +import json +from dataclasses import dataclass +from uuid import uuid4 + +from sqlalchemy import select +from sqlalchemy.orm import Session, sessionmaker + +from context.connection_execution import ( + ConnectionExecutionContext, + ExecutionConnectionSelection, + WorkspaceConnectionScope, +) +from core.plugin.entities.plugin_daemon import CredentialType +from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole +from models.enums import ConversationFromSource, ConversationStatus, InvokeFrom +from models.model import App, AppMode, Conversation, InstalledApp +from models.tools import BuiltinToolProvider +from models.workflow import Workflow, WorkflowType +from services.installed_app_connections import InstalledConnectionConfiguration +from tests.helpers.workflow_test_credentials import tool_node + + +@dataclass(frozen=True) +class InstalledConnectionData: + tenant: Tenant + account: Account + other: Account + app: App + workflow: Workflow + installation: InstalledApp + credential: BuiltinToolProvider + revision: str + + @property + def scope(self) -> WorkspaceConnectionScope: + return WorkspaceConnectionScope( + tenant_id=self.tenant.id, installed_app_id=self.installation.id, account_id=self.account.id + ) + + @property + def configuration(self) -> InstalledConnectionConfiguration: + return InstalledConnectionConfiguration(self.revision, {"root": self.credential.id}) + + @property + def execution(self) -> ConnectionExecutionContext: + return ConnectionExecutionContext( + tenant_id=self.tenant.id, + app_id=self.app.id, + end_user_id="", + account_id=self.account.id, + workspace=self.scope, + workflow_id=self.workflow.id, + connection_revision=self.revision, + selections=( + ExecutionConnectionSelection( + node_id="root", provider=self.credential.provider, connection_id=self.credential.id + ), + ), + ) + + +def seed_installed_connections( + sessions: sessionmaker[Session], *, tenant: Tenant | None = None, account: Account | None = None +) -> InstalledConnectionData: + """Commit a published installation and team credential using a real database. + + None creates a tenant/account; supplying them seeds another app in the same + scope. Return detached rows usable with expire_on_commit=False sessions. + """ + tenant = tenant if tenant is not None else Tenant(name="Workspace") + account = account if account is not None else Account(name="Member", email=f"{uuid4()}@example.com") + other = Account(name="Other member", email=f"{uuid4()}@example.com") + app = App(tenant_id=tenant.id, name="App", mode=AppMode.ADVANCED_CHAT, enable_site=True, enable_api=True) + workflow = Workflow( + id=str(uuid4()), + tenant_id=tenant.id, + app_id=app.id, + type=WorkflowType.CHAT, + version="published", + graph=json.dumps({"nodes": [tool_node("root", None)], "edges": []}), + _features="{}", + created_by=account.id, + ) + app.workflow_id = workflow.id + installation = InstalledApp(tenant_id=tenant.id, app_id=app.id, app_owner_tenant_id=tenant.id) + credential = BuiltinToolProvider( + tenant_id=tenant.id, + user_id=account.id, + provider="vendor/calendar/calendar", + name="Team OAuth", + credential_type=CredentialType.OAUTH2, + encrypted_credentials='{"token":"encrypted"}', + ) + with sessions.begin() as session: + session.merge(tenant) + session.merge(account) + session.add_all([other, app, workflow, installation, credential]) + membership = session.scalar( + select(TenantAccountJoin.id).where( + TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id + ) + ) + if membership is None: + session.add(TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=TenantAccountRole.NORMAL)) + session.add(TenantAccountJoin(tenant_id=tenant.id, account_id=other.id, role=TenantAccountRole.NORMAL)) + return InstalledConnectionData(tenant, account, other, app, workflow, installation, credential, str(uuid4())) + + +def member_conversation(data: InstalledConnectionData) -> Conversation: + """Construct an unsaved, member-owned Explore conversation for the fixture app.""" + return Conversation( + app_id=data.app.id, + mode=AppMode.ADVANCED_CHAT, + from_account_id=data.account.id, + from_source=ConversationFromSource.CONSOLE, + invoke_from=InvokeFrom.EXPLORE, + name="Conversation", + _inputs={}, + status=ConversationStatus.NORMAL, + ) diff --git a/api/tests/unit_tests/repositories/test_installed_app_connection_repository.py b/api/tests/unit_tests/repositories/test_installed_app_connection_repository.py new file mode 100644 index 00000000000..6bc6e1b5b42 --- /dev/null +++ b/api/tests/unit_tests/repositories/test_installed_app_connection_repository.py @@ -0,0 +1,188 @@ +"""Owner-chain protection and scope-local configuration replacement contracts.""" + +from dataclasses import asdict +from uuid import uuid4 + +import pytest +from sqlalchemy import Table, event, select, update +from sqlalchemy.dialects import mysql, postgresql, sqlite +from sqlalchemy.orm import ORMExecuteState, Session, sessionmaker +from sqlalchemy.sql import Select + +from models.installed_app_connection import InstalledAppConnectionConfiguration +from models.model import App, Conversation, Message +from models.workflow import Workflow +from repositories.installed_app_connection_repository import SQLAlchemyInstalledConnectionRepository +from services.end_user_connections.errors import ( + ConnectionConfigurationError, + ConnectionIdentityError, + ConnectionRevisionChangedError, +) +from services.installed_app_connections import InstalledConnectionConfiguration +from tests.helpers.installed_app_connections import ( + InstalledConnectionData, + member_conversation, + seed_installed_connections, +) + + +@pytest.fixture +def data(sqlite_session_factory: sessionmaker[Session]) -> InstalledConnectionData: + return seed_installed_connections(sqlite_session_factory) + + +def test_replacements_preserve_row_identity_and_allow_idempotent_saves(sqlite_session_factory, data): + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sqlite_session_factory) + graph = repository.graph(data.scope) + repository.save(data.scope, graph, data.configuration, None) + with sqlite_session_factory() as session: + original_id = session.scalars(select(InstalledAppConnectionConfiguration.id)).one() + replacement = InstalledConnectionConfiguration(str(uuid4()), {}) + repository.save(data.scope, graph, replacement, None) + repository.save(data.scope, graph, replacement, None) + assert repository.configuration(data.scope, None) == replacement + with sqlite_session_factory() as session: + assert session.scalars(select(InstalledAppConnectionConfiguration.id)).one() == original_id + + +@pytest.mark.parametrize("owner_field", ["tenant_id", "app_id"]) +def test_conflicting_owner_is_not_overwritten(sqlite_session_factory, data, owner_field): + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sqlite_session_factory) + graph = repository.graph(data.scope) + saved = InstalledAppConnectionConfiguration( + tenant_id=data.tenant.id, + installed_app_id=data.installation.id, + app_id=data.app.id, + account_id=data.account.id, + scope_id="", + connection_revision=data.revision, + connections="{}", + ) + foreign_owner = str(uuid4()) + setattr(saved, owner_field, foreign_owner) + with sqlite_session_factory.begin() as session: + session.add(saved) + with pytest.raises(ConnectionIdentityError): + repository.save(data.scope, graph, data.configuration, None) + with sqlite_session_factory() as session: + unchanged = session.get(InstalledAppConnectionConfiguration, saved.id) + assert unchanged is not None + assert getattr(unchanged, owner_field) == foreign_owner + assert unchanged.connections == "{}" + + +@pytest.mark.parametrize("change", ["publication", "graph"]) +def test_protection_uses_current_projection_not_stale_orm_identity(sqlite_session_factory, data, change): + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sqlite_session_factory) + graph = repository.graph(data.scope) + with sqlite_session_factory.begin() as session: + app = session.get(App, data.app.id) + workflow = session.get(Workflow, data.workflow.id) + assert app is not None + assert workflow is not None + statement = ( + update(App).where(App.id == app.id).values(workflow_id=str(uuid4())) + if change == "publication" + else update(Workflow).where(Workflow.id == workflow.id).values(graph='{"nodes":[],"edges":[]}') + ) + session.execute(statement.execution_options(synchronize_session=False)) + assert app.workflow_id == graph.workflow_id + assert workflow.graph_dict == graph.graph + with pytest.raises(ConnectionRevisionChangedError): + repository._save(session, data.scope, graph, data.configuration, None) + assert session.scalar(select(InstalledAppConnectionConfiguration.id)) is None + + +@pytest.mark.parametrize("initialize", [False, True]) +def test_write_paths_use_one_shared_single_table_protection_pass(sqlite_session_factory, data, initialize): + repository = SQLAlchemyInstalledConnectionRepository(session_factory=sqlite_session_factory) + statements: list[Select] = [] + + def observe(state: ORMExecuteState) -> None: + if isinstance(state.statement, Select): + statements.append(state.statement) + + event.listen(sqlite_session_factory, "do_orm_execute", observe) + try: + if initialize: + with sqlite_session_factory.begin() as session: + conversation = member_conversation(data) + session.add(conversation) + session.flush() + repository.initialize_conversation(session, data.execution, conversation) + else: + repository.save(data.scope, repository.graph(data.scope), data.configuration, None) + finally: + event.remove(sqlite_session_factory, "do_orm_execute", observe) + for dialect, shared_clause in [(postgresql.dialect(), "FOR SHARE"), (mysql.dialect(), "LOCK IN SHARE MODE")]: + protected_tables = [] + for statement in statements: + sql = str(statement.compile(dialect=dialect)) + assert "FOR UPDATE" not in sql + if shared_clause in sql: + tables = statement.get_final_froms() + assert len(tables) == 1 + assert isinstance(tables[0], Table) + protected_tables.append(tables[0].name) + expected = ["tenants", "accounts", "tenant_account_joins", "installed_apps", "apps", "workflows"] + assert protected_tables == expected + (["conversations"] if initialize else []) + credential_reads = [statement for statement in statements if "tool_builtin_providers" in str(statement)] + assert len(credential_reads) == 1 + credential_sql = str(credential_reads[0].compile(compile_kwargs={"literal_binds": True})) + assert f"tool_builtin_providers.id IN ('{data.credential.id}')" in credential_sql + assert "encrypted_credentials" not in credential_sql + + +@pytest.mark.parametrize("dialect", [postgresql.dialect(), mysql.dialect(), sqlite.dialect()]) +def test_first_save_insert_preserves_conflicting_row(dialect, data): + row = InstalledAppConnectionConfiguration( + tenant_id=data.tenant.id, + installed_app_id=data.installation.id, + app_id=data.app.id, + account_id=data.account.id, + scope_id="", + connection_revision=data.revision, + connections="{}", + ) + statement = SQLAlchemyInstalledConnectionRepository._configuration_insert_statement(dialect.name, asdict(row)) + sql = str(statement.compile(dialect=dialect)) + if dialect.name == "mysql": + assert "ON DUPLICATE KEY UPDATE id = installed_app_connection_configurations.id" in sql + else: + assert "ON CONFLICT (installed_app_id, account_id, scope_id) DO NOTHING" in sql + + +@pytest.mark.parametrize("failure", ["configuration", "outer_transaction"]) +def test_initialization_does_not_commit_conversation_or_message(sqlite_session_factory, data, failure): + execution = data.execution + if failure == "configuration": + execution = execution.model_copy( + update={"selections": (execution.selections[0].model_copy(update={"connection_id": str(uuid4())}),)} + ) + error = ConnectionConfigurationError if failure == "configuration" else RuntimeError + conversation = member_conversation(data) + message = Message( + app_id=data.app.id, + conversation_id=conversation.id, + from_account_id=data.account.id, + from_source=conversation.from_source, + query="Hello", + answer="", + message="", + _inputs={}, + ) + + def initialize_and_abort() -> None: + with sqlite_session_factory.begin() as session: + session.add_all([conversation, message]) + session.flush() + SQLAlchemyInstalledConnectionRepository.initialize_conversation(session, execution, conversation) + assert session.scalar(select(InstalledAppConnectionConfiguration.id)) is not None + raise RuntimeError("Abort the caller's transaction") + + with pytest.raises(error): + initialize_and_abort() + with sqlite_session_factory() as session: + assert session.get(Conversation, conversation.id) is None + assert session.get(Message, message.id) is None + assert session.scalar(select(InstalledAppConnectionConfiguration.id)) is None