122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
"""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,
|
|
*,
|
|
category_id: str | None = None,
|
|
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)
|
|
if category_id:
|
|
conditions.append(Document.category_id == category_id)
|
|
|
|
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
|