3

backend/app/repositories/doc_repo.py

@@ -0,0 +1,118 @@
+"""Document 文档 Repository。"""
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import Session
+
+from app.models.document import Document
+
+
+class DocumentRepository:
+    def __init__(self, session: Session) -> None:
+        self._session = session
+
+    def get_by_id(self, doc_id: str) -> Document | None:
+        return self._session.get(Document, doc_id)
+
+    def get_by_doc_token_hash(self, token_hash: str) -> Document | None:
+        stmt = select(Document).where(Document.doc_token_hash == token_hash)
+        return self._session.scalars(stmt).first()
+
+    def list_by_knowledge_base(
+        self,
+        kb_id: str,
+        *,
+        page: int = 1,
+        page_size: int = 50,
+        status: str | None = None,
+    ) -> tuple[list[Document], int]:
+        """分页获取知识库的文档列表。"""
+        conditions = [
+            Document.knowledge_base_id == kb_id,
+            Document.status != "DELETED",
+        ]
+        if status:
+            conditions.append(Document.status == status)
+
+        count_stmt = select(func.count()).select_from(Document).where(*conditions)
+        total = self._session.scalar(count_stmt) or 0
+
+        stmt = (
+            select(Document)
+            .where(*conditions)
+            .order_by(Document.created_at.desc())
+            .offset((page - 1) * page_size)
+            .limit(page_size)
+        )
+        items = list(self._session.scalars(stmt).all())
+        return items, total
+
+    def list_by_user(
+        self, user_id: str, *, page: int = 1, page_size: int = 20
+    ) -> tuple[list[Document], int]:
+        """分页获取用户的文档列表。"""
+        conditions = [Document.user_id == user_id, Document.status != "DELETED"]
+
+        count_stmt = select(func.count()).select_from(Document).where(*conditions)
+        total = self._session.scalar(count_stmt) or 0
+
+        stmt = (
+            select(Document)
+            .where(*conditions)
+            .order_by(Document.created_at.desc())
+            .offset((page - 1) * page_size)
+            .limit(page_size)
+        )
+        items = list(self._session.scalars(stmt).all())
+        return items, total
+
+    def create(
+        self,
+        *,
+        knowledge_base_id: str,
+        user_id: str,
+        original_filename: str,
+        storage_path: str,
+        file_size: int,
+        mime_type: str,
+        file_ext: str,
+        sha256: str,
+        doc_token_hash: str,
+        doc_token_encrypted: str,
+        doc_token_hint: str,
+        category_id: str | None = None,
+    ) -> Document:
+        doc = Document(
+            knowledge_base_id=knowledge_base_id,
+            user_id=user_id,
+            original_filename=original_filename,
+            storage_path=storage_path,
+            file_size=file_size,
+            mime_type=mime_type,
+            file_ext=file_ext,
+            sha256=sha256,
+            doc_token_hash=doc_token_hash,
+            doc_token_encrypted=doc_token_encrypted,
+            doc_token_hint=doc_token_hint,
+            category_id=category_id,
+            status="PENDING",
+        )
+        self._session.add(doc)
+        self._session.flush()
+        return doc
+
+    def update(self, doc: Document, **fields) -> None:
+        for key, value in fields.items():
+            setattr(doc, key, value)
+        self._session.flush()
+
+    def delete(self, doc: Document) -> None:
+        doc.status = "DELETED"
+        self._session.flush()
+
+    def count_by_user(self, user_id: str) -> int:
+        stmt = (
+            select(func.count())
+            .select_from(Document)
+            .where(Document.user_id == user_id, Document.status != "DELETED")
+        )
+        return self._session.scalar(stmt) or 0