diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index dbf0903b1f67a4..0a522453757e1b 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -112,8 +112,7 @@ class _ProviderModelCacheEntry: encrypted_config: str | None @classmethod - def from_record(cls, record: ProviderModel) -> _ProviderModelCacheEntry: - credential = record.__dict__.get("credential") + def from_record(cls, record: ProviderModel, credential: ProviderModelCredential | None) -> _ProviderModelCacheEntry: return cls( id=record.id, provider_name=record.provider_name, @@ -440,33 +439,57 @@ def _get_cached_or_load_records[T: _CacheEntry]( return records -def _attach_active_credentials( +def _load_active_credentials[T: ProviderCredential | ProviderModelCredential]( *, session: Any, records: Sequence[Provider | ProviderModel], - credential_model_cls: type[ProviderCredential | ProviderModelCredential], -) -> None: - credential_ids = [record.credential_id for record in records if getattr(record, "credential_id", None)] + credential_model_cls: type[T], +) -> dict[str, T]: + """Load the credentials referenced by ``records`` through the caller's session. + + Precondition: ``session`` is open and ``credential_model_cls`` matches the record type. + Postcondition: the result maps each credential ID that still exists to its row. + """ + credential_ids = [record.credential_id for record in records if record.credential_id] if not credential_ids: - return + return {} credentials = session.scalars(select(credential_model_cls).where(credential_model_cls.id.in_(credential_ids))).all() - credential_by_id = {credential.id: credential for credential in credentials} - for record in records: - if getattr(record, "credential_id", None): - record.__dict__["credential"] = credential_by_id.get(record.credential_id) + return {credential.id: credential for credential in credentials} + + +def _attach_provider_credentials(*, session: Any, providers: Sequence[Provider]) -> None: + """Attach active credentials so ``providers`` can be read after ``session`` closes. + + Precondition: ``session`` is open. + Postcondition: every provider with a ``credential_id`` has its row attached, or None when the row is missing. + """ + credentials = _load_active_credentials( + session=session, + records=providers, + credential_model_cls=ProviderCredential, + ) + for provider in providers: + if provider.credential_id: + provider.attach_credential(credentials.get(provider.credential_id)) def _load_provider_model_cache_entries(tenant_id: str) -> list[_ProviderModelCacheEntry]: with session_factory.create_session() as session: stmt = select(ProviderModel).where(ProviderModel.tenant_id == tenant_id, ProviderModel.is_valid.is_(True)) provider_models = list(session.scalars(stmt)) - _attach_active_credentials( + credentials = _load_active_credentials( session=session, records=provider_models, credential_model_cls=ProviderModelCredential, ) - return [_ProviderModelCacheEntry.from_record(provider_model) for provider_model in provider_models] + return [ + _ProviderModelCacheEntry.from_record( + provider_model, + credentials.get(provider_model.credential_id) if provider_model.credential_id else None, + ) + for provider_model in provider_models + ] def _load_preferred_model_provider_cache_entries(tenant_id: str) -> list[_TenantPreferredModelProviderCacheEntry]: @@ -982,11 +1005,7 @@ def _get_all_providers(tenant_id: str) -> dict[str, list[Provider]]: with session_factory.create_session() as session: stmt = select(Provider).where(Provider.tenant_id == tenant_id, Provider.is_valid.is_(True)) providers = list(session.scalars(stmt)) - _attach_active_credentials( - session=session, - records=providers, - credential_model_cls=ProviderCredential, - ) + _attach_provider_credentials(session=session, providers=providers) for provider in providers: # Use provider name with prefix after the data migration provider_name_to_provider_records_dict[str(ModelProviderID(provider.provider_name))].append(provider) @@ -1235,6 +1254,7 @@ def _init_trial_provider_records( existed_provider_record.is_valid = True db.session.commit() + _attach_provider_credentials(session=db.session, providers=[existed_provider_record]) provider_name_to_provider_records_dict[provider_name].append(existed_provider_record) return provider_name_to_provider_records_dict diff --git a/api/models/model.py b/api/models/model.py index 88b1909b03ec4f..18e820ace49c24 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -31,7 +31,6 @@ from .account import Account, Tenant from .base import Base, TypeBase, gen_uuidv4_string -from .engine import db from .enums import ( ApiTokenBindingResourceType, ApiTokenType, @@ -1740,19 +1739,6 @@ def set_extra_contents(self, contents: Sequence[dict[str, Any]]) -> None: def extra_contents(self) -> list[ExtraContentDict]: return getattr(self, "_extra_contents", []) - @property - def workflow_run(self): - if self.workflow_run_id: - from sqlalchemy.orm import sessionmaker - - from repositories.factory import DifyAPIRepositoryFactory - - session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) - repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker) - return repo.get_workflow_run_by_id_without_tenant(run_id=self.workflow_run_id) - - return None - def to_dict(self, *, session: Session) -> MessageDict: return { "id": self.id, diff --git a/api/models/provider.py b/api/models/provider.py index b9b68b93e27630..40cf66cf86a87f 100644 --- a/api/models/provider.py +++ b/api/models/provider.py @@ -2,15 +2,13 @@ from datetime import datetime from enum import StrEnum, auto -from functools import cached_property from typing import override from uuid import uuid4 import sqlalchemy as sa -from sqlalchemy import DateTime, String, func, select, text +from sqlalchemy import DateTime, String, func, text from sqlalchemy.orm import Mapped, mapped_column -from core.db.session_factory import session_factory from graphon.model_runtime.entities.model_entities import ModelType from libs.uuid_utils import uuidv7 @@ -81,11 +79,26 @@ def __repr__(self): f" provider_type='{self.provider_type}')>" ) - @cached_property - def credential(self): - if self.credential_id: - with session_factory.create_session() as session: - return session.scalar(select(ProviderCredential).where(ProviderCredential.id == self.credential_id)) + def attach_credential(self, credential: ProviderCredential | None) -> None: + """Keep the active credential loaded through the caller's session. + + Precondition: ``credential`` is the row for ``credential_id``, or None when that row is missing. + Postcondition: reading ``credential`` returns this value without database access. + """ + self.__dict__["_active_credential"] = credential + + @property + def credential(self) -> ProviderCredential | None: + """Return the attached credential without database access. + + Precondition: ``attach_credential`` ran when ``credential_id`` is set. + Raises RuntimeError when it did not, so a missing load never reads as an unset credential. + """ + if not self.credential_id: + return None + if "_active_credential" not in self.__dict__: + raise RuntimeError(f"Credential of provider {self.id} was not attached") + return self.__dict__["_active_credential"] @property def credential_name(self): @@ -145,24 +158,6 @@ class ProviderModel(TypeBase): DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False ) - @cached_property - def credential(self): - if self.credential_id: - with session_factory.create_session() as session: - return session.scalar( - select(ProviderModelCredential).where(ProviderModelCredential.id == self.credential_id) - ) - - @property - def credential_name(self): - credential = self.credential - return credential.credential_name if credential else None - - @property - def encrypted_config(self): - credential = self.credential - return credential.encrypted_config if credential else None - class TenantDefaultModel(TypeBase): __tablename__ = "tenant_default_models" diff --git a/api/models/workflow.py b/api/models/workflow.py index 9ee6420131d5dc..efb0bb93254239 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -1402,19 +1402,6 @@ class WorkflowAppLog(TypeBase): DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) - @property - def workflow_run(self): - if self.workflow_run_id: - from sqlalchemy.orm import sessionmaker - - from repositories.factory import DifyAPIRepositoryFactory - - session_maker = sessionmaker(bind=db.engine, expire_on_commit=False) - repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker) - return repo.get_workflow_run_by_id_without_tenant(run_id=self.workflow_run_id) - - return None - def created_by_account(self, session: orm.Session) -> Account | None: created_by_role = CreatorUserRole(self.created_by_role) return session.get(Account, self.created_by) if created_by_role == CreatorUserRole.ACCOUNT else None 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 408660b62f1042..71898e282129d9 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 @@ -57,11 +57,10 @@ def database_session( sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch, ) -> Session: - """Bind model/controller database access to the shared SQLite session.""" + """Bind controller database access to the shared SQLite session.""" database = SimpleNamespace(engine=sqlite_engine, session=sqlite_session) monkeypatch.setattr(human_input_module, "db", database) - monkeypatch.setattr("models.model.db", database) return sqlite_session diff --git a/api/tests/unit_tests/core/test_provider_manager.py b/api/tests/unit_tests/core/test_provider_manager.py index b0dabc0346f4ee..1678e75317a910 100644 --- a/api/tests/unit_tests/core/test_provider_manager.py +++ b/api/tests/unit_tests/core/test_provider_manager.py @@ -1,3 +1,4 @@ +from collections import defaultdict from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace @@ -1202,6 +1203,28 @@ def test_get_all_providers_attaches_active_credentials(provider_db: Session) -> assert str(ModelProviderID("anthropic")) not in result +def test_init_trial_provider_records_attaches_credential_of_concurrently_created_record(provider_db: Session) -> None: + existing_provider = _build_trial_provider_record() + existing_provider.credential_id = "credential-id" + credential = ProviderCredential( + tenant_id="tenant-id", + provider_name="openai", + credential_name="primary", + encrypted_config='{"api_key": "secret"}', + ) + credential.id = "credential-id" + provider_db.add_all([existing_provider, credential]) + provider_db.commit() + + with patch( + "core.provider_manager.ext_hosting_provider.hosting_configuration.provider_map", + {"openai": _build_hosting_provider()}, + ): + result = ProviderManager._init_trial_provider_records("tenant-id", defaultdict(list)) + + assert result["openai"][0].encrypted_config == '{"api_key": "secret"}' + + def test_invalidate_configurations_cache_bumps_selected_source_version() -> None: fake_redis = _FakeRedis() diff --git a/api/tests/unit_tests/models/test_provider.py b/api/tests/unit_tests/models/test_provider.py index f582423cacb938..e3d08289bb8ec8 100644 --- a/api/tests/unit_tests/models/test_provider.py +++ b/api/tests/unit_tests/models/test_provider.py @@ -260,6 +260,59 @@ def test_provider_token_is_set_false_when_no_credential(self) -> None: # Act & Assert assert provider.token_is_set is False + def test_provider_credential_reads_attached_value(self) -> None: + """Test credential accessors read the attached row without database access.""" + # Arrange + provider = Provider( + tenant_id=str(uuid4()), + provider_name="openai", + credential_id=str(uuid4()), + ) + credential = ProviderCredential( + tenant_id=provider.tenant_id, + provider_name="openai", + credential_name="primary", + encrypted_config='{"api_key": "secret"}', + ) + + # Act + provider.attach_credential(credential) + + # Assert + assert provider.credential is credential + assert provider.credential_name == "primary" + assert provider.encrypted_config == '{"api_key": "secret"}' + assert provider.token_is_set is True + + def test_provider_credential_missing_row_reads_as_unset(self) -> None: + """Test a credential row deleted before loading reads as unset.""" + # Arrange + provider = Provider( + tenant_id=str(uuid4()), + provider_name="openai", + credential_id=str(uuid4()), + ) + + # Act + provider.attach_credential(None) + + # Assert + assert provider.credential is None + assert provider.token_is_set is False + + def test_provider_credential_raises_when_not_attached(self) -> None: + """Test an unloaded credential fails instead of reading as unset.""" + # Arrange + provider = Provider( + tenant_id=str(uuid4()), + provider_name="openai", + credential_id=str(uuid4()), + ) + + # Act & Assert + with pytest.raises(RuntimeError, match="was not attached"): + _ = provider.token_is_set + def test_provider_is_enabled_false_when_not_valid(self) -> None: """Test is_enabled returns False when provider is not valid.""" # Arrange