"""Application ports and policies for datasource credential resolution.""" import logging import time from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from typing import Protocol from services.data_source.entities.credential import DatasourceCredentialRecord, DatasourceOAuthClientConfigRecord logger = logging.getLogger(__name__) class NotionCredentialUnavailableError(ValueError): pass def resolve_stored_notion_access_token( *, credential_id: str | None, load_token: Callable[[str], str], integration_token: str | None, ) -> str: """Use the saved credential, then the configured legacy integration. Stored documents never select a tenant default credential: changing that default must not change the integration used to read an existing document. Preserve indexing's environment fallback when credential resolution fails, including refresh failures. Actor-selected imports do not use this policy. """ credential_error: Exception | None = None if credential_id: try: token = load_token(credential_id) if isinstance(token, str) and token: return token except Exception as error: credential_error = error if not integration_token: raise NotionCredentialUnavailableError( "Must specify `integration_token` or set environment variable `NOTION_INTEGRATION_TOKEN`." ) from credential_error logger.warning("Stored Notion credential %s unavailable; using NOTION_INTEGRATION_TOKEN", credential_id) return integration_token class DatasourceProviderCredentialStore(Protocol): """Persisted provider credentials for already-authorized execution paths.""" def get_provider_credential( self, *, workspace_id: str, provider: str, plugin_id: str, credential_id: str | None ) -> DatasourceCredentialRecord | None: ... def list_provider_credentials( self, *, workspace_id: str, provider: str, plugin_id: str ) -> Sequence[DatasourceCredentialRecord]: ... def get_oauth_client_config( self, *, workspace_id: str, provider: str, plugin_id: str ) -> DatasourceOAuthClientConfigRecord: ... def update_if_unchanged( self, *, record: DatasourceCredentialRecord, encrypted_credentials: Mapping[str, object], expires_at: int ) -> bool: ... class DatasourceCredentialStore(Protocol): def get_visible( self, *, workspace_id: str, actor_id: str, credential_id: str, provider: str, plugin_id: str, ) -> DatasourceCredentialRecord | None: ... def update_if_unchanged( self, *, record: DatasourceCredentialRecord, encrypted_credentials: Mapping[str, object], expires_at: int, ) -> bool: ... class TrustedDatasourceCredentialStore(Protocol): def get_for_stored_document( self, *, workspace_id: str, dataset_id: str, document_id: str, credential_id: str | None, provider: str, plugin_id: str, ) -> DatasourceCredentialRecord | None: ... def update_if_unchanged( self, *, record: DatasourceCredentialRecord, encrypted_credentials: Mapping[str, object], expires_at: int, ) -> bool: ... class DatasourceCredentialCodec(Protocol): def decrypt(self, record: DatasourceCredentialRecord) -> dict[str, object]: ... def encrypt( self, record: DatasourceCredentialRecord, credentials: Mapping[str, object], ) -> dict[str, object]: ... @dataclass(frozen=True, slots=True) class RefreshedDatasourceCredential: credentials: Mapping[str, object] expires_at: int class DatasourceCredentialRefresher(Protocol): def refresh( self, *, workspace_id: str, actor_id: str, record: DatasourceCredentialRecord, credentials: Mapping[str, object], ) -> RefreshedDatasourceCredential: ... class ActorDatasourceCredentialResolver(Protocol): """Resolve a credential selected by an authenticated actor.""" def resolve( self, *, workspace_id: str, actor_id: str, credential_id: str, provider: str, plugin_id: str, ) -> dict[str, object]: ... class StoredDatasourceCredentialResolver(Protocol): """Resolve a credential proven by an existing document owner chain.""" def resolve_for_document( self, *, workspace_id: str, dataset_id: str, document_id: str, credential_id: str | None, provider: str, plugin_id: str, ) -> dict[str, object]: ... class DatasourceCredentialError(Exception): """Base class for credential resolution failures.""" class DatasourceCredentialNotFoundError(DatasourceCredentialError): def __init__(self) -> None: super().__init__("Credential not found") class DatasourceCredentialConcurrentUpdateError(DatasourceCredentialError): def __init__(self) -> None: super().__init__("Credential changed while it was being refreshed") class DatasourceCredentialRefreshError(DatasourceCredentialError): def __init__(self, credential_id: str) -> None: super().__init__(f"Failed to refresh datasource credential: {credential_id}") class ActorAwareDatasourceCredentialGateway: def __init__( self, *, credentials: DatasourceCredentialStore, codec: DatasourceCredentialCodec, refresher: DatasourceCredentialRefresher, now: Callable[[], float] = time.time, ) -> None: self._credentials = credentials self._codec = codec self._refresher = refresher self._now = now def resolve( self, *, workspace_id: str, actor_id: str, credential_id: str, provider: str, plugin_id: str, ) -> dict[str, object]: def load_record() -> DatasourceCredentialRecord: record = self._credentials.get_visible( workspace_id=workspace_id, actor_id=actor_id, credential_id=credential_id, provider=provider, plugin_id=plugin_id, ) if record is None: raise DatasourceCredentialNotFoundError() return record return resolve_credential_snapshot( load_record(), actor_id=actor_id, codec=self._codec, refresher=self._refresher, update=self._credentials.update_if_unchanged, load_latest=load_record, now=self._now, ) class TrustedStoredDatasourceCredentialGateway: """Resolve a credential only when a stored document proves the full owner chain.""" def __init__( self, *, credentials: TrustedDatasourceCredentialStore, codec: DatasourceCredentialCodec, refresher: DatasourceCredentialRefresher, now: Callable[[], float] = time.time, ) -> None: self._credentials = credentials self._codec = codec self._refresher = refresher self._now = now def resolve_for_document( self, *, workspace_id: str, dataset_id: str, document_id: str, credential_id: str | None, provider: str, plugin_id: str, ) -> dict[str, object]: def load_record() -> DatasourceCredentialRecord: record = self._credentials.get_for_stored_document( workspace_id=workspace_id, dataset_id=dataset_id, document_id=document_id, credential_id=credential_id, provider=provider, plugin_id=plugin_id, ) if record is None: raise DatasourceCredentialNotFoundError() return record record = load_record() return resolve_credential_snapshot( record, actor_id=record.owner_id, codec=self._codec, refresher=self._refresher, update=self._credentials.update_if_unchanged, load_latest=load_record, now=self._now, ) def credential_needs_refresh(expires_at: int, *, now: float) -> bool: return expires_at != -1 and expires_at - 60 < int(now) class CredentialSnapshotUpdate(Protocol): def __call__( self, *, record: DatasourceCredentialRecord, encrypted_credentials: Mapping[str, object], expires_at: int ) -> bool: ... def resolve_credential_snapshot( record: DatasourceCredentialRecord, *, actor_id: str | None, codec: DatasourceCredentialCodec, refresher: DatasourceCredentialRefresher, update: CredentialSnapshotUpdate, load_latest: Callable[[], DatasourceCredentialRecord], now: Callable[[], float], ) -> dict[str, object]: decrypted = codec.decrypt(record) if not credential_needs_refresh(record.expires_at, now=now()): return decrypted if actor_id is None: raise DatasourceCredentialRefreshError(record.id) refreshed = refresher.refresh( workspace_id=record.workspace_id, actor_id=actor_id, record=record, credentials=decrypted ) if update( record=record, encrypted_credentials=codec.encrypt(record, refreshed.credentials), expires_at=refreshed.expires_at, ): return dict(refreshed.credentials) latest = load_latest() if credential_needs_refresh(latest.expires_at, now=now()): raise DatasourceCredentialConcurrentUpdateError() return codec.decrypt(latest)