347 lines
13 KiB
Python
347 lines
13 KiB
Python
"""KnowledgeFS object operations backed by Dify's configured storage provider.
|
|
|
|
KnowledgeFS is a separate Node process, so it reaches this service through a
|
|
trusted inner API rather than importing Python storage implementations or
|
|
holding a second set of provider credentials. Logical KnowledgeFS keys are
|
|
confined below a dedicated physical namespace. Metadata is stored as a small
|
|
sidecar because Dify's common ``BaseStorage`` contract intentionally exposes
|
|
portable byte operations rather than provider-specific object metadata.
|
|
|
|
Object keys are immutable/versioned in KnowledgeFS product flows. A backend
|
|
must also support ``scan`` for bounded logical pagination and cleanup; health
|
|
fails closed when that capability is unavailable.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from base64 import b64decode, b64encode
|
|
from collections.abc import Generator, Mapping
|
|
from dataclasses import dataclass
|
|
from hashlib import sha256
|
|
from typing import Protocol, TypedDict
|
|
|
|
from extensions.ext_storage import storage
|
|
|
|
KNOWLEDGE_FS_OBJECT_MAX_BYTES = 64 * 1024 * 1024
|
|
_MAX_KEY_BYTES = 1024
|
|
_MAX_LIST_LIMIT = 100
|
|
_MAX_METADATA_ENTRIES = 64
|
|
_MAX_METADATA_KEY_BYTES = 64
|
|
_MAX_METADATA_VALUE_BYTES = 1024
|
|
_DATA_ROOT = "knowledge-fs/objects"
|
|
_METADATA_ROOT = "knowledge-fs/object-metadata"
|
|
_HEALTH_SCAN_ROOT = "knowledge-fs/health-capability"
|
|
|
|
|
|
class KnowledgeFSObjectStorageError(Exception):
|
|
"""Base error for the trusted KnowledgeFS storage gateway."""
|
|
|
|
|
|
class KnowledgeFSObjectStorageInvalidInputError(KnowledgeFSObjectStorageError):
|
|
"""The caller supplied an invalid key, cursor, metadata value, or limit."""
|
|
|
|
|
|
class KnowledgeFSObjectStorageTooLargeError(KnowledgeFSObjectStorageInvalidInputError):
|
|
"""The caller supplied an object body above the portable gateway limit."""
|
|
|
|
|
|
class KnowledgeFSObjectStorageChecksumError(KnowledgeFSObjectStorageError):
|
|
"""The supplied checksum does not match the uploaded object body."""
|
|
|
|
|
|
class KnowledgeFSObjectStorageCorruptError(KnowledgeFSObjectStorageError):
|
|
"""The physical object and its portable metadata sidecar disagree."""
|
|
|
|
|
|
class KnowledgeFSObjectStorageUnavailableError(KnowledgeFSObjectStorageError):
|
|
"""The selected Dify storage backend lacks a required portable capability."""
|
|
|
|
|
|
class _StorageBackend(Protocol):
|
|
def save(self, filename: str, data: bytes) -> None: ...
|
|
|
|
def load_once(self, filename: str) -> bytes: ...
|
|
|
|
def load_stream(self, filename: str) -> Generator[bytes, None, None]: ...
|
|
|
|
def exists(self, filename: str) -> bool: ...
|
|
|
|
def delete(self, filename: str) -> None: ...
|
|
|
|
def scan(self, path: str, files: bool = True, directories: bool = False) -> list[str]: ...
|
|
|
|
|
|
class _StoredMetadataPayload(TypedDict):
|
|
checksum_sha256_base64: str
|
|
content_type: str | None
|
|
key: str
|
|
metadata: dict[str, str]
|
|
size_bytes: int
|
|
version: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class KnowledgeFSObjectMetadata:
|
|
"""Portable metadata returned to the KnowledgeFS adapter."""
|
|
|
|
checksum_sha256_base64: str
|
|
content_type: str | None
|
|
key: str
|
|
metadata: Mapping[str, str]
|
|
size_bytes: int
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class KnowledgeFSObjectList:
|
|
"""One keyset-paginated page of logical KnowledgeFS objects."""
|
|
|
|
objects: tuple[KnowledgeFSObjectMetadata, ...]
|
|
next_cursor: str | None = None
|
|
|
|
|
|
class KnowledgeFSObjectStorageService:
|
|
"""Confine KnowledgeFS object I/O to Dify's unified storage namespace."""
|
|
|
|
_backend: _StorageBackend
|
|
|
|
def __init__(self, *, backend: _StorageBackend = storage) -> None:
|
|
self._backend = backend
|
|
|
|
def put_object(
|
|
self,
|
|
*,
|
|
key: str,
|
|
body: bytes,
|
|
metadata: Mapping[str, str],
|
|
checksum_sha256_base64: str | None = None,
|
|
content_type: str | None = None,
|
|
) -> KnowledgeFSObjectMetadata:
|
|
"""Validate and persist one immutable logical object and its metadata sidecar."""
|
|
normalized_key = _normalize_key(key)
|
|
normalized_metadata = _normalize_metadata(metadata)
|
|
normalized_content_type = _normalize_content_type(content_type)
|
|
if len(body) > KNOWLEDGE_FS_OBJECT_MAX_BYTES:
|
|
raise KnowledgeFSObjectStorageTooLargeError(f"object exceeds max bytes {KNOWLEDGE_FS_OBJECT_MAX_BYTES}")
|
|
|
|
actual_checksum = b64encode(sha256(body).digest()).decode()
|
|
if checksum_sha256_base64 is not None:
|
|
expected_checksum = _normalize_checksum(checksum_sha256_base64)
|
|
if expected_checksum != actual_checksum:
|
|
raise KnowledgeFSObjectStorageChecksumError("object checksum does not match body")
|
|
|
|
object_metadata = KnowledgeFSObjectMetadata(
|
|
checksum_sha256_base64=actual_checksum,
|
|
content_type=normalized_content_type,
|
|
key=normalized_key,
|
|
metadata=normalized_metadata,
|
|
size_bytes=len(body),
|
|
)
|
|
data_path = _data_path(normalized_key)
|
|
self._backend.save(data_path, body)
|
|
try:
|
|
self._backend.save(_metadata_path(normalized_key), _encode_metadata(object_metadata))
|
|
except Exception:
|
|
self._backend.delete(data_path)
|
|
raise
|
|
return object_metadata
|
|
|
|
def head_object(self, *, key: str) -> KnowledgeFSObjectMetadata | None:
|
|
"""Return portable metadata without loading the object body."""
|
|
normalized_key = _normalize_key(key)
|
|
if not self._backend.exists(_data_path(normalized_key)):
|
|
return None
|
|
try:
|
|
encoded = self._backend.load_once(_metadata_path(normalized_key))
|
|
except FileNotFoundError as exc:
|
|
raise KnowledgeFSObjectStorageCorruptError("object metadata sidecar is missing") from exc
|
|
metadata = _decode_metadata(encoded)
|
|
if metadata.key != normalized_key:
|
|
raise KnowledgeFSObjectStorageCorruptError("object metadata key does not match")
|
|
return metadata
|
|
|
|
def load_stream(self, *, key: str) -> Generator[bytes, None, None] | None:
|
|
"""Return the configured backend's stream for one logical object."""
|
|
normalized_key = _normalize_key(key)
|
|
data_path = _data_path(normalized_key)
|
|
if not self._backend.exists(data_path):
|
|
return None
|
|
return self._backend.load_stream(data_path)
|
|
|
|
def delete_object(self, *, key: str) -> None:
|
|
"""Idempotently delete an object and its portable metadata sidecar."""
|
|
normalized_key = _normalize_key(key)
|
|
self._backend.delete(_data_path(normalized_key))
|
|
self._backend.delete(_metadata_path(normalized_key))
|
|
|
|
def list_objects(
|
|
self,
|
|
*,
|
|
prefix: str,
|
|
limit: int,
|
|
cursor: str | None = None,
|
|
) -> KnowledgeFSObjectList:
|
|
"""List logical keys using a stable lexical cursor and portable metadata."""
|
|
normalized_prefix = _normalize_prefix(prefix)
|
|
if not isinstance(limit, int) or isinstance(limit, bool) or not 1 <= limit <= _MAX_LIST_LIMIT:
|
|
raise KnowledgeFSObjectStorageInvalidInputError("list limit must be between 1 and 100")
|
|
normalized_cursor = _normalize_key(cursor) if cursor is not None else None
|
|
if normalized_cursor is not None and not normalized_cursor.startswith(normalized_prefix):
|
|
raise KnowledgeFSObjectStorageInvalidInputError("list cursor must be within prefix")
|
|
|
|
try:
|
|
physical_paths = self._backend.scan(
|
|
_physical_scan_prefix(normalized_prefix),
|
|
files=True,
|
|
directories=False,
|
|
)
|
|
except FileNotFoundError:
|
|
physical_paths = []
|
|
except NotImplementedError as exc:
|
|
raise KnowledgeFSObjectStorageUnavailableError(
|
|
"configured Dify storage backend does not support object listing"
|
|
) from exc
|
|
except Exception as exc:
|
|
raise KnowledgeFSObjectStorageUnavailableError("Dify storage listing failed") from exc
|
|
|
|
data_prefix = f"{_DATA_ROOT}/"
|
|
logical_keys = sorted(
|
|
{
|
|
physical_path[len(data_prefix) :]
|
|
for physical_path in physical_paths
|
|
if physical_path.startswith(data_prefix)
|
|
and physical_path[len(data_prefix) :].startswith(normalized_prefix)
|
|
and (normalized_cursor is None or physical_path[len(data_prefix) :] > normalized_cursor)
|
|
}
|
|
)
|
|
page_keys = logical_keys[: limit + 1]
|
|
has_more = len(page_keys) > limit
|
|
objects: list[KnowledgeFSObjectMetadata] = []
|
|
for logical_key in page_keys[:limit]:
|
|
metadata = self.head_object(key=logical_key)
|
|
if metadata is not None:
|
|
objects.append(metadata)
|
|
next_cursor = objects[-1].key if has_more and objects else None
|
|
return KnowledgeFSObjectList(objects=tuple(objects), next_cursor=next_cursor)
|
|
|
|
def health(self) -> bool:
|
|
"""Check the portable list capability without scanning the object namespace."""
|
|
try:
|
|
self._backend.scan(_HEALTH_SCAN_ROOT, files=True, directories=False)
|
|
except FileNotFoundError:
|
|
return True
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
|
|
def _normalize_key(value: str) -> str:
|
|
if not isinstance(value, str):
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object key must be a string")
|
|
normalized = value.strip()
|
|
segments = normalized.split("/")
|
|
if (
|
|
not normalized
|
|
or normalized.startswith("/")
|
|
or "\\" in normalized
|
|
or "\x00" in normalized
|
|
or any(segment in {"", ".", ".."} for segment in segments)
|
|
or len(normalized.encode()) > _MAX_KEY_BYTES
|
|
):
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object key is invalid")
|
|
return normalized
|
|
|
|
|
|
def _normalize_prefix(value: str) -> str:
|
|
if value == "":
|
|
return value
|
|
normalized = value.strip()
|
|
candidate = normalized.removesuffix("/")
|
|
_normalize_key(candidate)
|
|
return normalized
|
|
|
|
|
|
def _normalize_metadata(metadata: Mapping[str, str]) -> dict[str, str]:
|
|
if len(metadata) > _MAX_METADATA_ENTRIES:
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object metadata has too many entries")
|
|
normalized: dict[str, str] = {}
|
|
for key, value in metadata.items():
|
|
if (
|
|
not isinstance(key, str)
|
|
or not isinstance(value, str)
|
|
or not key
|
|
or len(key.encode()) > _MAX_METADATA_KEY_BYTES
|
|
or len(value.encode()) > _MAX_METADATA_VALUE_BYTES
|
|
):
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object metadata is invalid")
|
|
normalized[key] = value
|
|
return normalized
|
|
|
|
|
|
def _normalize_content_type(value: str | None) -> str | None:
|
|
if value is None:
|
|
return None
|
|
normalized = value.strip()
|
|
if not normalized or len(normalized) > 255 or "\r" in normalized or "\n" in normalized:
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object content type is invalid")
|
|
return normalized
|
|
|
|
|
|
def _normalize_checksum(value: str) -> str:
|
|
try:
|
|
decoded = b64decode(value, validate=True)
|
|
except ValueError as exc:
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object checksum is invalid") from exc
|
|
if len(decoded) != 32 or b64encode(decoded).decode() != value:
|
|
raise KnowledgeFSObjectStorageInvalidInputError("object checksum is invalid")
|
|
return value
|
|
|
|
|
|
def _data_path(key: str) -> str:
|
|
return f"{_DATA_ROOT}/{key}"
|
|
|
|
|
|
def _metadata_path(key: str) -> str:
|
|
digest = sha256(key.encode()).hexdigest()
|
|
return f"{_METADATA_ROOT}/{digest[:2]}/{digest[2:4]}/{digest}.json"
|
|
|
|
|
|
def _physical_scan_prefix(prefix: str) -> str:
|
|
parent = prefix.rstrip("/").rpartition("/")[0]
|
|
return f"{_DATA_ROOT}/{parent}" if parent else _DATA_ROOT
|
|
|
|
|
|
def _encode_metadata(metadata: KnowledgeFSObjectMetadata) -> bytes:
|
|
payload: _StoredMetadataPayload = {
|
|
"checksum_sha256_base64": metadata.checksum_sha256_base64,
|
|
"content_type": metadata.content_type,
|
|
"key": metadata.key,
|
|
"metadata": dict(metadata.metadata),
|
|
"size_bytes": metadata.size_bytes,
|
|
"version": 1,
|
|
}
|
|
return json.dumps(payload, separators=(",", ":"), sort_keys=True).encode()
|
|
|
|
|
|
def _decode_metadata(encoded: bytes) -> KnowledgeFSObjectMetadata:
|
|
try:
|
|
payload = json.loads(encoded)
|
|
if not isinstance(payload, dict) or payload.get("version") != 1:
|
|
raise ValueError
|
|
key = _normalize_key(payload["key"])
|
|
checksum = _normalize_checksum(payload["checksum_sha256_base64"])
|
|
content_type = _normalize_content_type(payload["content_type"])
|
|
metadata = _normalize_metadata(payload["metadata"])
|
|
size_bytes = payload["size_bytes"]
|
|
if not isinstance(size_bytes, int) or isinstance(size_bytes, bool) or size_bytes < 0:
|
|
raise ValueError
|
|
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
|
raise KnowledgeFSObjectStorageCorruptError("object metadata sidecar is invalid") from exc
|
|
return KnowledgeFSObjectMetadata(
|
|
checksum_sha256_base64=checksum,
|
|
content_type=content_type,
|
|
key=key,
|
|
metadata=metadata,
|
|
size_bytes=size_bytes,
|
|
)
|