3

2026-09-01 13:00:36 +08:00
parent 1d8621717a
commit dfd38c99a0
35 changed files with 3141 additions and 21 deletions
@@ -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
@@ -0,0 +1,104 @@
+"""KnowledgeBase 知识库 Repository。"""
+
+from sqlalchemy import func, select
+from sqlalchemy.orm import Session
+
+from app.models.document import Document
+from app.models.knowledge_base import KnowledgeBase
+
+
+class KnowledgeBaseRepository:
+    def __init__(self, session: Session) -> None:
+        self._session = session
+
+    def get_by_id(self, kb_id: str) -> KnowledgeBase | None:
+        return self._session.get(KnowledgeBase, kb_id)
+
+    def get_by_token_hash(self, token_hash: str) -> KnowledgeBase | None:
+        stmt = select(KnowledgeBase).where(KnowledgeBase.token_hash == token_hash)
+        return self._session.scalars(stmt).first()
+
+    def list_by_user(
+        self, user_id: str, *, page: int = 1, page_size: int = 20
+    ) -> tuple[list[KnowledgeBase], int]:
+        """分页获取用户的知识库列表。返回 (items, total)。"""
+        count_stmt = (
+            select(func.count())
+            .select_from(KnowledgeBase)
+            .where(
+                KnowledgeBase.user_id == user_id,
+                KnowledgeBase.status != "DELETED",
+            )
+        )
+        total = self._session.scalar(count_stmt) or 0
+
+        stmt = (
+            select(KnowledgeBase)
+            .where(
+                KnowledgeBase.user_id == user_id,
+                KnowledgeBase.status != "DELETED",
+            )
+            .order_by(KnowledgeBase.created_at.desc())
+            .offset((page - 1) * page_size)
+            .limit(page_size)
+        )
+        items = list(self._session.scalars(stmt).all())
+        return items, total
+
+    def create(
+        self,
+        *,
+        user_id: str,
+        name: str,
+        description: str | None,
+        token_hash: str,
+        token_encrypted: str,
+        token_hint: str,
+    ) -> KnowledgeBase:
+        kb = KnowledgeBase(
+            user_id=user_id,
+            name=name,
+            description=description,
+            enabled=True,
+            status="active",
+            token_hash=token_hash,
+            token_encrypted=token_encrypted,
+            token_hint=token_hint,
+        )
+        self._session.add(kb)
+        self._session.flush()
+        return kb
+
+    def update(self, kb: KnowledgeBase, **fields) -> None:
+        for key, value in fields.items():
+            if value is not None:
+                setattr(kb, key, value)
+        self._session.flush()
+
+    def delete(self, kb: KnowledgeBase) -> None:
+        """软删除。"""
+        kb.status = "DELETED"
+        self._session.flush()
+
+    def count_documents(self, kb_id: str) -> int:
+        stmt = (
+            select(func.count())
+            .select_from(Document)
+            .where(
+                Document.knowledge_base_id == kb_id,
+                Document.status != "DELETED",
+            )
+        )
+        return self._session.scalar(stmt) or 0
+
+    def get_document_count_map(self, user_id: str) -> dict[str, int]:
+        """批量获取用户所有知识库的文档数量（避免 N+1）。"""
+        stmt = (
+            select(Document.knowledge_base_id, func.count())
+            .where(
+                Document.user_id == user_id,
+                Document.status != "DELETED",
+            )
+            .group_by(Document.knowledge_base_id)
+        )
+        return dict(self._session.execute(stmt).all())
@@ -0,0 +1,33 @@
+"""Plan 套餐 Repository。"""
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from app.models.plan import Plan
+
+
+class PlanRepository:
+    def __init__(self, session: Session) -> None:
+        self._session = session
+
+    def get_by_id(self, plan_id: str) -> Plan | None:
+        return self._session.get(Plan, plan_id)
+
+    def get_by_code(self, code: str) -> Plan | None:
+        stmt = select(Plan).where(Plan.code == code)
+        return self._session.scalars(stmt).first()
+
+    def get_or_create_free(self) -> Plan:
+        """获取 free plan，不存在则创建（seed 逻辑）。"""
+        plan = self.get_by_code("free")
+        if plan is None:
+            plan = Plan(
+                code="free",
+                name="免费版",
+                storage_quota=104_857_600,  # 100 MB
+                max_file_size=20_971_520,   # 20 MB
+                is_active=True,
+            )
+            self._session.add(plan)
+            self._session.flush()
+        return plan
@@ -0,0 +1,64 @@
+"""User 用户 Repository。"""
+
+from sqlalchemy import select, or_
+from sqlalchemy.orm import Session
+
+from app.models.user import User
+
+
+class UserRepository:
+    def __init__(self, session: Session) -> None:
+        self._session = session
+
+    def get_by_id(self, user_id: str) -> User | None:
+        return self._session.get(User, user_id)
+
+    def get_by_username(self, username: str) -> User | None:
+        stmt = select(User).where(User.username == username)
+        return self._session.scalars(stmt).first()
+
+    def get_by_email(self, email: str) -> User | None:
+        stmt = select(User).where(User.email == email)
+        return self._session.scalars(stmt).first()
+
+    def get_by_username_or_email(self, value: str) -> User | None:
+        """登录用：按用户名或邮箱查找。"""
+        stmt = select(User).where(
+            or_(User.username == value, User.email == value.lower())
+        )
+        return self._session.scalars(stmt).first()
+
+    def exists_username_or_email(self, username: str, email: str) -> tuple[bool, bool]:
+        """检查用户名/邮箱是否已存在。返回 (username_taken, email_taken)。"""
+        stmt = select(User).where(
+            or_(User.username == username, User.email == email.lower())
+        )
+        existing = list(self._session.scalars(stmt).all())
+        return (
+            any(u.username == username for u in existing),
+            any(u.email == email.lower() for u in existing),
+        )
+
+    def create(
+        self,
+        *,
+        username: str,
+        email: str,
+        password_hash: str,
+        plan_id: str,
+    ) -> User:
+        user = User(
+            username=username,
+            email=email.lower(),
+            password_hash=password_hash,
+            status="active",
+            plan_id=plan_id,
+            storage_used=0,
+        )
+        self._session.add(user)
+        self._session.flush()
+        return user
+
+    def update_password(self, user: User, password_hash: str) -> None:
+        user.password_hash = password_hash
+        self._session.flush()