import mimetypes import os import threading import uuid from dataclasses import dataclass from typing import Any try: # pragma: no cover - optional in some test paths import psycopg from psycopg.rows import dict_row except Exception: # pragma: no cover psycopg = None dict_row = None MAX_TEXT_SIZE = 8 * 1024 * 1024 PREVIEW_TEXT_SIZE = 2 * 1024 * 1024 TEXT_EXTENSIONS = { "md", "markdown", "txt", "json", "js", "jsx", "ts", "tsx", "css", "scss", "less", "html", "htm", "py", "vue", "xml", "yaml", "yml", "csv", "log", "sql", "toml", "ini", "cfg", "conf", "sh", "bat", "ps1", "java", "c", "cpp", "h", "hpp", "go", "rs", "swift", "kt", "rb", "php", "pl", "r", "scala", "gradle", "properties", "env", "gitignore", "dockerfile", } BINARY_EXTENSIONS = { "exe", "dll", "so", "dylib", "bin", "dat", "obj", "o", "a", "doc", "docx", "xls", "xlsx", "ppt", "pptx", "odt", "ods", "odp", "pdf", "zip", "rar", "7z", "tar", "gz", "bz2", "xz", "png", "jpg", "jpeg", "gif", "bmp", "ico", "webp", "svg", "mp3", "mp4", "wav", "avi", "mov", "mkv", "flv", "wmv", "ttf", "otf", "woff", "woff2", "eot", "class", "pyc", "pyo", "jar", "war", "ear", "db", "sqlite", "mdb", "accdb", "pem", "key", "crt", "cer", "p12", "pfx", "jks", "msg", "eml", "pst", "ost", "dwg", "dxf", "step", "stl", "fbx", "3ds", "blend", } DEFAULT_MIME_TYPES = { "avi": "video/x-msvideo", "md": "text/markdown", "markdown": "text/markdown", "mkv": "video/x-matroska", "mov": "video/quicktime", "mp4": "video/mp4", "txt": "text/plain", "json": "application/json", "js": "text/javascript", "jsx": "text/javascript", "ts": "text/typescript", "tsx": "text/typescript", "css": "text/css", "html": "text/html", "htm": "text/html", "py": "text/x-python", "vue": "text/plain", "xml": "application/xml", "yaml": "text/yaml", "yml": "text/yaml", "csv": "text/csv", "log": "text/plain", "sql": "text/plain", "toml": "text/plain", "ini": "text/plain", "cfg": "text/plain", "conf": "text/plain", "sh": "text/plain", "bat": "text/plain", "ps1": "text/plain", "jpg": "image/jpeg", "jpeg": "image/jpeg", "flv": "video/x-flv", "m4v": "video/x-m4v", "png": "image/png", "gif": "image/gif", "webp": "image/webp", "svg": "image/svg+xml", "pdf": "application/pdf", "ogg": "video/ogg", "ogv": "video/ogg", "webm": "video/webm", "wmv": "video/x-ms-wmv", } def _now_sql() -> str: return "CURRENT_TIMESTAMP" def get_extension(name: str = "") -> str: value = str(name or "") if "." not in value: return value.lower() if value.lower() == "dockerfile" else "" return value.rsplit(".", 1)[-1].lower() def infer_mime_type(name: str, fallback: str = "") -> str: if fallback: return fallback ext = get_extension(name) guessed = DEFAULT_MIME_TYPES.get(ext) if guessed: return guessed guessed, _ = mimetypes.guess_type(name) return guessed or "application/octet-stream" def is_text_file(name: str, mime_type: str = "") -> bool: ext = get_extension(name) if ext in BINARY_EXTENSIONS: return False mime_value = str(mime_type or "").lower() return ext in TEXT_EXTENSIONS or mime_value.startswith("text/") or "json" in mime_value or "xml" in mime_value def _text_payload(text: str) -> dict[str, Any]: encoded = text.encode("utf-8") preview = encoded[:PREVIEW_TEXT_SIZE].decode("utf-8", errors="ignore") return { "storage_kind": "text", "size": len(encoded), "content_text": text, "preview_text": preview if len(encoded) > PREVIEW_TEXT_SIZE else text, "is_truncated_preview": len(encoded) > PREVIEW_TEXT_SIZE, "blob_data": None, } def prepare_file_payload(name: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: resolved_mime = infer_mime_type(name, mime_type) if is_text_file(name, resolved_mime): text = raw_bytes.decode("utf-8", errors="ignore") return { "mime_type": resolved_mime, **_text_payload(text), } return { "mime_type": resolved_mime, "storage_kind": "blob", "size": len(raw_bytes), "content_text": "", "preview_text": "", "is_truncated_preview": False, "blob_data": raw_bytes, } @dataclass class BlobPayload: content: bytes mime_type: str filename: str class BaseDocumentStore: _UNSET = object() def list_nodes(self) -> list[dict[str, Any]]: raise NotImplementedError def create_folder(self, name: str, parent_id: str | None) -> dict[str, Any]: raise NotImplementedError def create_text_file(self, name: str, parent_id: str | None, content: str = "") -> dict[str, Any]: raise NotImplementedError def upload_file(self, name: str, parent_id: str | None, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: raise NotImplementedError def update_node(self, node_id: str, *, name: str | None | object = _UNSET, parent_id: str | None | object = _UNSET, content: str | None | object = _UNSET) -> dict[str, Any]: raise NotImplementedError def replace_blob(self, node_id: str, filename: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: raise NotImplementedError def delete_node(self, node_id: str) -> None: raise NotImplementedError def get_blob(self, node_id: str) -> BlobPayload: raise NotImplementedError class InMemoryDocumentStore(BaseDocumentStore): def __init__(self) -> None: self.nodes: dict[str, dict[str, Any]] = {} def _serialize(self, node: dict[str, Any]) -> dict[str, Any]: return { "id": node["id"], "parentId": node["parent_id"], "type": node["type"], "name": node["name"], "mimeType": node.get("mime_type") or "", "storageKind": node.get("storage_kind") or "text", "size": int(node.get("size") or 0), "content": node.get("content_text") or "", "previewText": node.get("preview_text") or "", "isTruncatedPreview": bool(node.get("is_truncated_preview")), "createdAt": int(node["created_at"]), "updatedAt": int(node["updated_at"]), } def _timestamp(self) -> int: import time return int(time.time() * 1000) def _get(self, node_id: str) -> dict[str, Any]: node = self.nodes.get(node_id) if node is None: raise KeyError(node_id) return node def list_nodes(self) -> list[dict[str, Any]]: ordered = sorted( self.nodes.values(), key=lambda item: (item["type"] != "folder", item["name"].lower(), item["created_at"]), ) return [self._serialize(node) for node in ordered] def create_folder(self, name: str, parent_id: str | None) -> dict[str, Any]: now = self._timestamp() node = { "id": str(uuid.uuid4()), "parent_id": parent_id, "type": "folder", "name": name, "mime_type": "", "storage_kind": "text", "size": 0, "content_text": "", "preview_text": "", "is_truncated_preview": False, "blob_data": None, "created_at": now, "updated_at": now, } self.nodes[node["id"]] = node return self._serialize(node) def create_text_file(self, name: str, parent_id: str | None, content: str = "") -> dict[str, Any]: now = self._timestamp() payload = _text_payload(content) node = { "id": str(uuid.uuid4()), "parent_id": parent_id, "type": "file", "name": name, "mime_type": infer_mime_type(name), **payload, "created_at": now, "updated_at": now, } self.nodes[node["id"]] = node return self._serialize(node) def upload_file(self, name: str, parent_id: str | None, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: now = self._timestamp() payload = prepare_file_payload(name, raw_bytes, mime_type) node = { "id": str(uuid.uuid4()), "parent_id": parent_id, "type": "file", "name": name, **payload, "created_at": now, "updated_at": now, } self.nodes[node["id"]] = node return self._serialize(node) def update_node(self, node_id: str, *, name: str | None | object = BaseDocumentStore._UNSET, parent_id: str | None | object = BaseDocumentStore._UNSET, content: str | None | object = BaseDocumentStore._UNSET) -> dict[str, Any]: node = self._get(node_id) if name is not BaseDocumentStore._UNSET: node["name"] = name if node["type"] == "file": node["mime_type"] = infer_mime_type(name, node.get("mime_type") or "") if parent_id is not BaseDocumentStore._UNSET: node["parent_id"] = parent_id if content is not BaseDocumentStore._UNSET: payload = _text_payload(content) node.update(payload) node["updated_at"] = self._timestamp() return self._serialize(node) def replace_blob(self, node_id: str, filename: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: node = self._get(node_id) payload = prepare_file_payload(filename, raw_bytes, mime_type) node.update(payload) node["name"] = filename node["mime_type"] = payload["mime_type"] node["updated_at"] = self._timestamp() return self._serialize(node) def delete_node(self, node_id: str) -> None: descendants = {node_id} changed = True while changed: changed = False for node in list(self.nodes.values()): if node.get("parent_id") in descendants and node["id"] not in descendants: descendants.add(node["id"]) changed = True for current_id in descendants: self.nodes.pop(current_id, None) def get_blob(self, node_id: str) -> BlobPayload: node = self._get(node_id) if node["type"] != "file": raise FileNotFoundError(node_id) if node.get("blob_data") is not None: content = node["blob_data"] else: content = (node.get("content_text") or "").encode("utf-8") return BlobPayload( content=content, mime_type=node.get("mime_type") or infer_mime_type(node.get("name") or ""), filename=node["name"], ) class PostgresDocumentStore(BaseDocumentStore): def __init__(self, database_url: str) -> None: if psycopg is None or dict_row is None: raise RuntimeError("psycopg 未安装，无法使用 PostgreSQL 文档存储") self.database_url = database_url self._init_lock = threading.Lock() self._initialized = False def _connect(self): return psycopg.connect(self.database_url, autocommit=True, row_factory=dict_row) def _ensure_initialized(self) -> None: if self._initialized: return with self._init_lock: if self._initialized: return with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ CREATE TABLE IF NOT EXISTS document_nodes ( id TEXT PRIMARY KEY, parent_id TEXT REFERENCES document_nodes(id) ON DELETE CASCADE, type TEXT NOT NULL CHECK (type IN ('folder', 'file')), name TEXT NOT NULL, mime_type TEXT NOT NULL DEFAULT '', storage_kind TEXT NOT NULL DEFAULT 'text', size BIGINT NOT NULL DEFAULT 0, content_text TEXT NOT NULL DEFAULT '', preview_text TEXT NOT NULL DEFAULT '', is_truncated_preview BOOLEAN NOT NULL DEFAULT FALSE, blob_data BYTEA, created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP ) """ ) cur.execute( "CREATE INDEX IF NOT EXISTS document_nodes_parent_idx ON document_nodes(parent_id)" ) self._initialized = True def _serialize(self, row: dict[str, Any]) -> dict[str, Any]: created = row["created_at"] updated = row["updated_at"] return { "id": row["id"], "parentId": row["parent_id"], "type": row["type"], "name": row["name"], "mimeType": row.get("mime_type") or "", "storageKind": row.get("storage_kind") or "text", "size": int(row.get("size") or 0), "content": row.get("content_text") or "", "previewText": row.get("preview_text") or "", "isTruncatedPreview": bool(row.get("is_truncated_preview")), "createdAt": int(created.timestamp() * 1000), "updatedAt": int(updated.timestamp() * 1000), } def _fetch_node(self, node_id: str) -> dict[str, Any]: self._ensure_initialized() with self._connect() as conn: with conn.cursor() as cur: cur.execute("SELECT * FROM document_nodes WHERE id = %s", (node_id,)) row = cur.fetchone() if row is None: raise KeyError(node_id) return row def list_nodes(self) -> list[dict[str, Any]]: self._ensure_initialized() with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ SELECT id, parent_id, type, name, mime_type, storage_kind, size, CASE WHEN storage_kind = 'text' AND size <= %s THEN content_text ELSE '' END AS content_text, preview_text, is_truncated_preview, created_at, updated_at FROM document_nodes ORDER BY CASE WHEN type = 'folder' THEN 0 ELSE 1 END, LOWER(name), created_at """, (MAX_TEXT_SIZE,), ) return [self._serialize(row) for row in cur.fetchall()] def create_folder(self, name: str, parent_id: str | None) -> dict[str, Any]: self._ensure_initialized() node_id = str(uuid.uuid4()) with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ INSERT INTO document_nodes (id, parent_id, type, name, updated_at) VALUES (%s, %s, 'folder', %s, CURRENT_TIMESTAMP) RETURNING * """, (node_id, parent_id, name), ) return self._serialize(cur.fetchone()) def create_text_file(self, name: str, parent_id: str | None, content: str = "") -> dict[str, Any]: self._ensure_initialized() node_id = str(uuid.uuid4()) payload = _text_payload(content) with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ INSERT INTO document_nodes ( id, parent_id, type, name, mime_type, storage_kind, size, content_text, preview_text, is_truncated_preview, blob_data, updated_at ) VALUES (%s, %s, 'file', %s, %s, %s, %s, %s, %s, %s, %s, CURRENT_TIMESTAMP) RETURNING * """, ( node_id, parent_id, name, infer_mime_type(name), payload["storage_kind"], payload["size"], payload["content_text"], payload["preview_text"], payload["is_truncated_preview"], None, ), ) return self._serialize(cur.fetchone()) def upload_file(self, name: str, parent_id: str | None, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: self._ensure_initialized() node_id = str(uuid.uuid4()) payload = prepare_file_payload(name, raw_bytes, mime_type) with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ INSERT INTO document_nodes ( id, parent_id, type, name, mime_type, storage_kind, size, content_text, preview_text, is_truncated_preview, blob_data, updated_at ) VALUES (%s, %s, 'file', %s, %s, %s, %s, %s, %s, %s, %s, CURRENT_TIMESTAMP) RETURNING * """, ( node_id, parent_id, name, payload["mime_type"], payload["storage_kind"], payload["size"], payload["content_text"], payload["preview_text"], payload["is_truncated_preview"], payload["blob_data"], ), ) return self._serialize(cur.fetchone()) def update_node(self, node_id: str, *, name: str | None | object = BaseDocumentStore._UNSET, parent_id: str | None | object = BaseDocumentStore._UNSET, content: str | None | object = BaseDocumentStore._UNSET) -> dict[str, Any]: self._ensure_initialized() row = self._fetch_node(node_id) next_name = row["name"] if name is BaseDocumentStore._UNSET else name next_parent_id = row["parent_id"] if parent_id is BaseDocumentStore._UNSET else parent_id next_mime_type = row.get("mime_type") or "" next_storage_kind = row.get("storage_kind") or "text" next_size = row.get("size") or 0 next_content_text = row.get("content_text") or "" next_preview_text = row.get("preview_text") or "" next_is_truncated = bool(row.get("is_truncated_preview")) next_blob_data = row.get("blob_data") if row["type"] == "file" and name is not BaseDocumentStore._UNSET: next_mime_type = infer_mime_type(next_name, next_mime_type) if content is not BaseDocumentStore._UNSET: payload = _text_payload(content) next_storage_kind = payload["storage_kind"] next_size = payload["size"] next_content_text = payload["content_text"] next_preview_text = payload["preview_text"] next_is_truncated = payload["is_truncated_preview"] next_blob_data = None with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ UPDATE document_nodes SET name = %s, parent_id = %s, mime_type = %s, storage_kind = %s, size = %s, content_text = %s, preview_text = %s, is_truncated_preview = %s, blob_data = %s, updated_at = CURRENT_TIMESTAMP WHERE id = %s RETURNING * """, ( next_name, next_parent_id, next_mime_type, next_storage_kind, next_size, next_content_text, next_preview_text, next_is_truncated, next_blob_data, node_id, ), ) updated = cur.fetchone() if updated is None: raise KeyError(node_id) return self._serialize(updated) def replace_blob(self, node_id: str, filename: str, raw_bytes: bytes, mime_type: str = "") -> dict[str, Any]: self._ensure_initialized() payload = prepare_file_payload(filename, raw_bytes, mime_type) with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ UPDATE document_nodes SET name = %s, mime_type = %s, storage_kind = %s, size = %s, content_text = %s, preview_text = %s, is_truncated_preview = %s, blob_data = %s, updated_at = CURRENT_TIMESTAMP WHERE id = %s AND type = 'file' RETURNING * """, ( filename, payload["mime_type"], payload["storage_kind"], payload["size"], payload["content_text"], payload["preview_text"], payload["is_truncated_preview"], payload["blob_data"], node_id, ), ) updated = cur.fetchone() if updated is None: raise KeyError(node_id) return self._serialize(updated) def delete_node(self, node_id: str) -> None: self._ensure_initialized() with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ WITH RECURSIVE descendants AS ( SELECT id FROM document_nodes WHERE id = %s UNION ALL SELECT child.id FROM document_nodes child INNER JOIN descendants parent ON child.parent_id = parent.id ) DELETE FROM document_nodes WHERE id IN (SELECT id FROM descendants) """, (node_id,), ) def get_blob(self, node_id: str) -> BlobPayload: self._ensure_initialized() with self._connect() as conn: with conn.cursor() as cur: cur.execute( """ SELECT id, name, mime_type, content_text, blob_data FROM document_nodes WHERE id = %s AND type = 'file' """, (node_id,), ) row = cur.fetchone() if row is None: raise FileNotFoundError(node_id) blob_data = row.get("blob_data") if blob_data is None: content = (row.get("content_text") or "").encode("utf-8") else: content = bytes(blob_data) return BlobPayload( content=content, mime_type=row.get("mime_type") or infer_mime_type(row.get("name") or ""), filename=row["name"], ) _document_store: BaseDocumentStore | None = None def get_document_store() -> BaseDocumentStore: global _document_store if _document_store is not None: return _document_store backend = (os.getenv("DOCS_BACKEND") or "postgres").strip().lower() if backend == "memory": _document_store = InMemoryDocumentStore() return _document_store database_url = os.getenv("DATABASE_URL", "").strip() if not database_url: raise RuntimeError("缺少 DATABASE_URL，无法初始化 PostgreSQL 文档存储") _document_store = PostgresDocumentStore(database_url) return _document_store def reset_document_store() -> None: global _document_store _document_store = None