feat: add Docker support and update backend dependencies
- Introduced `requirements.docker.txt` for Docker-specific dependencies. - Updated `requirements.txt` to include `psycopg[binary]` and `python-multipart`. - Enhanced test suite in `test_main_endpoints.py` to cover document CRUD operations. - Modified `docker-compose.yml` to include PostgreSQL and frontend services. - Added Nginx configuration for reverse proxying API requests. - Refactored file handling in Vue components to support new document storage backend. - Created new utility functions in `docsApi.js` for document management. - Updated configuration to support new API endpoints for document operations. - Adjusted Vite configuration to proxy API requests to the local backend.
This commit is contained in:
“ydy0615”
@@ -0,0 +1,668 @@
|
||||
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
|
||||
Reference in New Issue
Block a user