改用mysql

This commit is contained in:
amb
2026-09-02 11:19:49 +08:00
parent 410bd30b2e
commit 73c4d83f42
24 changed files with 831 additions and 258 deletions
@@ -20,7 +20,7 @@ depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# document_categories 新增字段
op.add_column('document_categories', sa.Column('parent_id', sa.String(length=32), nullable=True, comment='父分类 IDNULL = 顶层)'))
op.add_column('document_categories', sa.Column('path', sa.Text(), server_default='/', nullable=False, comment='物化路径'))
op.add_column('document_categories', sa.Column('path', sa.String(length=500), server_default='/', nullable=False, comment='物化路径'))
op.add_column('document_categories', sa.Column('is_folder', sa.Boolean(), server_default=sa.text('1'), nullable=False, comment='True=文件夹'))
op.create_index(op.f('ix_document_categories_parent_id'), 'document_categories', ['parent_id'], unique=False)
@@ -28,7 +28,7 @@ def upgrade() -> None:
op.add_column('documents', sa.Column('content', sa.Text(), nullable=True, comment='直接输入的文本内容'))
op.add_column('documents', sa.Column('content_format', sa.String(length=16), server_default='markdown', nullable=False, comment='内容格式'))
# SQLite 不支持 ALTER COLUMN,用 batch mode 重建表
# ALTER COLUMNMySQL 直接支持,SQLite 需要 batch mode
with op.batch_alter_table('documents', schema=None) as batch_op:
batch_op.alter_column('storage_path', existing_type=sa.String(1024), nullable=True)
batch_op.alter_column('sha256', existing_type=sa.String(64), nullable=True)
@@ -0,0 +1,125 @@
"""add_user_role
Revision ID: 4e40432cab9f
Revises: 4bd4c7f26818
Create Date: 2026-09-02 10:29:41.369850
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import mysql
# revision identifiers, used by Alembic.
revision: str = '4e40432cab9f'
down_revision: Union[str, None] = '4bd4c7f26818'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.alter_column('document_categories', 'path',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=500),
comment='物化路径,如 /01_公司层/04_岗位AI角色/',
existing_comment='物化路径',
existing_nullable=False,
existing_server_default=sa.text("'/'"))
op.alter_column('document_categories', 'is_folder',
existing_type=mysql.TINYINT(display_width=1),
comment='True=文件夹(可含子项),False=叶子分类',
existing_comment='True=文件夹',
existing_nullable=False,
existing_server_default=sa.text("'1'"))
op.create_foreign_key(None, 'document_categories', 'document_categories', ['parent_id'], ['id'])
op.alter_column('documents', 'original_filename',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=512),
comment='原始文件名或文本标题 (仅展示)',
existing_comment='原始文件名 (仅展示)',
existing_nullable=False)
op.alter_column('documents', 'storage_path',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=1024),
comment='相对于 data/ 的物理存储路径(文本内容文档为 NULL)',
existing_nullable=True)
op.alter_column('documents', 'content',
existing_type=mysql.TEXT(collation='utf8mb4_unicode_ci'),
comment='直接输入的文本内容(非文件上传时使用)',
existing_comment='直接输入的文本内容',
existing_nullable=True)
op.alter_column('documents', 'content_format',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=16),
comment='内容格式:markdown / text',
existing_comment='内容格式',
existing_nullable=False,
existing_server_default=sa.text("'markdown'"))
op.alter_column('documents', 'file_size',
existing_type=mysql.INTEGER(),
comment='文件大小 (字节),文本内容文档为内容长度',
existing_comment='文件大小 (字节)',
existing_nullable=False)
op.alter_column('documents', 'file_ext',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=16),
comment='文件扩展名 (.docx/.pdf/.md)',
existing_comment='文件扩展名 (.docx/.pdf)',
existing_nullable=False)
op.alter_column('documents', 'sha256',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=64),
comment='文件 SHA-256 哈希(文本内容文档为 NULL)',
existing_nullable=True)
op.add_column('users', sa.Column('role', sa.String(length=16), nullable=False, comment='角色 (internal=内部员工 / customer=客户)'))
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column('users', 'role')
op.alter_column('documents', 'sha256',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=64),
comment=None,
existing_comment='文件 SHA-256 哈希(文本内容文档为 NULL)',
existing_nullable=True)
op.alter_column('documents', 'file_ext',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=16),
comment='文件扩展名 (.docx/.pdf)',
existing_comment='文件扩展名 (.docx/.pdf/.md)',
existing_nullable=False)
op.alter_column('documents', 'file_size',
existing_type=mysql.INTEGER(),
comment='文件大小 (字节)',
existing_comment='文件大小 (字节),文本内容文档为内容长度',
existing_nullable=False)
op.alter_column('documents', 'content_format',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=16),
comment='内容格式',
existing_comment='内容格式:markdown / text',
existing_nullable=False,
existing_server_default=sa.text("'markdown'"))
op.alter_column('documents', 'content',
existing_type=mysql.TEXT(collation='utf8mb4_unicode_ci'),
comment='直接输入的文本内容',
existing_comment='直接输入的文本内容(非文件上传时使用)',
existing_nullable=True)
op.alter_column('documents', 'storage_path',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=1024),
comment=None,
existing_comment='相对于 data/ 的物理存储路径(文本内容文档为 NULL)',
existing_nullable=True)
op.alter_column('documents', 'original_filename',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=512),
comment='原始文件名 (仅展示)',
existing_comment='原始文件名或文本标题 (仅展示)',
existing_nullable=False)
op.drop_constraint(None, 'document_categories', type_='foreignkey')
op.alter_column('document_categories', 'is_folder',
existing_type=mysql.TINYINT(display_width=1),
comment='True=文件夹',
existing_comment='True=文件夹(可含子项),False=叶子分类',
existing_nullable=False,
existing_server_default=sa.text("'1'"))
op.alter_column('document_categories', 'path',
existing_type=mysql.VARCHAR(collation='utf8mb4_unicode_ci', length=500),
comment='物化路径',
existing_comment='物化路径,如 /01_公司层/04_岗位AI角色/',
existing_nullable=False,
existing_server_default=sa.text("'/'"))
# ### end Alembic commands ###
+32 -3
View File
@@ -21,22 +21,48 @@ router = APIRouter(prefix="/auth", tags=["auth"])
@router.post("/register", response_model=UserResponse, status_code=201)
def register(body: RegisterRequest, response: Response, db: Session = Depends(get_db)) -> UserResponse:
"""注册新用户并自动登录。"""
"""注册新用户(默认 role=customer并自动登录。"""
auth_service = AuthService(db)
user, token = auth_service.register(body.username, body.email, body.password)
user, token = auth_service.register(body.username, body.email, body.password, role="customer")
response.set_cookie(value=token, **get_cookie_params())
return _user_response(user)
@router.post("/register-internal", response_model=UserResponse, status_code=201)
def register_internal(
body: RegisterRequest,
response: Response,
current_user: User = Depends(get_current_user),
db: Session = Depends(get_db),
) -> UserResponse:
"""注册内部员工账号(仅已登录的内部用户可调用)。"""
if current_user.role != "internal":
from app.core.errors import PermissionDeniedError
raise PermissionDeniedError("仅内部员工可创建内部账号。")
auth_service = AuthService(db)
user, token = auth_service.register(body.username, body.email, body.password, role="internal")
response.set_cookie(value=token, **get_cookie_params())
return _user_response(user)
@router.post("/login", response_model=UserResponse)
def login(body: LoginRequest, response: Response, db: Session = Depends(get_db)) -> UserResponse:
"""登录"""
"""普通登录(所有用户可用)"""
auth_service = AuthService(db)
user, token = auth_service.login(body.username_or_email, body.password)
response.set_cookie(value=token, **get_cookie_params())
return _user_response(user)
@router.post("/internal-login", response_model=UserResponse)
def internal_login(body: LoginRequest, response: Response, db: Session = Depends(get_db)) -> UserResponse:
"""内部登录(仅 role=internal 用户可用)。"""
auth_service = AuthService(db)
user, token = auth_service.internal_login(body.username_or_email, body.password)
response.set_cookie(value=token, **get_cookie_params())
return _user_response(user)
@router.post("/logout", status_code=204)
def logout(
response: Response,
@@ -70,6 +96,7 @@ def get_me(user: User = Depends(get_current_user)) -> MeResponse:
username=user.username,
email=user.email,
status=user.status,
role=user.role,
storage_used=user.storage_used,
storage_quota=settings.default_storage_quota,
created_at=user.created_at,
@@ -94,6 +121,7 @@ def update_me(
username=user.username,
email=user.email,
status=user.status,
role=user.role,
storage_used=user.storage_used,
storage_quota=settings.default_storage_quota,
created_at=user.created_at,
@@ -123,6 +151,7 @@ def _user_response(user: User) -> UserResponse:
username=user.username,
email=user.email,
status=user.status,
role=user.role,
storage_used=user.storage_used,
storage_quota=settings.default_storage_quota,
created_at=user.created_at,
+12 -4
View File
@@ -1,8 +1,8 @@
"""全局配置:pydantic-settings,全部来自环境变量 / .env。
规则docs/technical-review.md §1.3
规则:
- 密钥不得硬编码;SECRET_KEY 缺失或仍为模板值时,生产环境拒绝启动。
- DATABASE_URL 可切换 PostgreSQL(扩展接口 1
- DATABASE_URL 支持 MySQL / SQLite
"""
from functools import lru_cache
@@ -22,7 +22,7 @@ class Settings(BaseSettings):
secret_key: str = Field(min_length=16)
# --- 数据库 ---
database_url: str = "sqlite:///./data/app.db"
database_url: str = "mysql+pymysql://admin:Lzcc6-01@47.109.98.44:33306/amb_rag?charset=utf8mb4"
# --- 文件存储 ---
storage_root: str = "./data"
@@ -42,6 +42,14 @@ class Settings(BaseSettings):
def is_production(self) -> bool:
return self.environment == "production"
@property
def is_mysql(self) -> bool:
return "mysql" in self.database_url
@property
def is_sqlite(self) -> bool:
return "sqlite" in self.database_url
@property
def storage_root_path(self) -> Path:
return Path(self.storage_root).resolve()
@@ -63,4 +71,4 @@ class Settings(BaseSettings):
def get_settings() -> Settings:
s = Settings() # type: ignore[call-arg]
s.validate_secrets()
return s
return s
+42 -19
View File
@@ -1,18 +1,17 @@
"""SQLAlchemy 2.x 同步引擎 + 会话管理MVP: SQLite
"""SQLAlchemy 2.x 同步引擎 + 会话管理。
切换 PostgreSQL 时(扩展接口 1):
1. DATABASE_URL 改为 postgresql+asyncpg://...
2. 引擎换 create_async_engine + async_sessionmaker
3. get_session 改 async generator + yield
4. 业务层 Repository 调用加 await
支持 MySQL / SQLite,通过 DATABASE_URL 自动切换。
"""
from collections.abc import Generator
from sqlalchemy import create_engine
from sqlalchemy import create_engine, event, text
from sqlalchemy.orm import Session, sessionmaker
from app.core.config import get_settings
from app.core.logging import get_logger
logger = get_logger(__name__)
_engine = None
_session_factory: sessionmaker[Session] | None = None
@@ -22,24 +21,48 @@ def get_engine():
global _engine
if _engine is None:
settings = get_settings()
_engine = create_engine(
settings.database_url,
pool_pre_ping=True,
echo=False,
# SQLite 专属:启用 WAL 模式(并发读 + 写串行化)
connect_args={"check_same_thread": False} if "sqlite" in settings.database_url else {},
)
# SQLite: 启用 WAL 模式与外键约束
if "sqlite" in settings.database_url:
from sqlalchemy import event, text
engine_kwargs = {
"pool_pre_ping": True,
"echo": False,
}
if settings.is_mysql:
# MySQL 配置
engine_kwargs.update({
"pool_size": 10,
"max_overflow": 20,
"pool_recycle": 3600, # 1 小时回收连接,防止 MySQL 超时断开
"connect_args": {
"charset": "utf8mb4",
},
})
elif settings.is_sqlite:
# SQLite 配置
engine_kwargs["connect_args"] = {"check_same_thread": False}
_engine = create_engine(settings.database_url, **engine_kwargs)
# SQLite: 启用 WAL 模式 + 外键约束
if settings.is_sqlite:
@event.listens_for(_engine, "connect")
def _set_sqlite_pragma(dbapi_conn, _): # type: ignore[no-untyped-def]
def _set_sqlite_pragma(dbapi_conn, _):
cursor = dbapi_conn.cursor()
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
# MySQL: 设置字符集和 SQL 模式
if settings.is_mysql:
@event.listens_for(_engine, "connect")
def _set_mysql_session(dbapi_conn, _):
cursor = dbapi_conn.cursor()
cursor.execute("SET NAMES utf8mb4")
cursor.execute("SET SESSION sql_mode='STRICT_TRANS_TABLES,NO_ZERO_DATE,NO_ZERO_IN_DATE,ERROR_FOR_DIVISION_BY_ZERO'")
cursor.close()
logger.info("Database engine created: %s", "MySQL" if settings.is_mysql else "SQLite")
return _engine
@@ -73,4 +96,4 @@ def dispose_engine() -> None:
if _engine is not None:
_engine.dispose()
_engine = None
_session_factory = None
_session_factory = None
+1 -1
View File
@@ -35,7 +35,7 @@ class DocumentCategory(UUIDPrimaryKeyMixin, TimestampMixin, Base):
comment="分类名称",
)
path: Mapped[str] = mapped_column(
Text,
String(500),
default="/",
nullable=False,
comment="物化路径,如 /01_公司层/04_岗位AI角色/",
+6
View File
@@ -39,6 +39,12 @@ class User(UUIDPrimaryKeyMixin, TimestampMixin, Base):
nullable=False,
comment="状态 (active/disabled)",
)
role: Mapped[str] = mapped_column(
String(16),
default="customer",
nullable=False,
comment="角色 (internal=内部员工 / customer=客户)",
)
plan_id: Mapped[str] = mapped_column(
String(32),
ForeignKey("plans.id"),
+2
View File
@@ -46,12 +46,14 @@ class UserRepository:
email: str,
password_hash: str,
plan_id: str,
role: str = "customer",
) -> User:
user = User(
username=username,
email=email.lower(),
password_hash=password_hash,
status="active",
role=role,
plan_id=plan_id,
storage_used=0,
)
+2
View File
@@ -45,6 +45,7 @@ class UserResponse(BaseModel):
username: str
email: str
status: str
role: str
storage_used: int
storage_quota: int
created_at: str
@@ -57,6 +58,7 @@ class MeResponse(BaseModel):
username: str
email: str
status: str
role: str
storage_used: int
storage_quota: int
created_at: str
+25 -2
View File
@@ -19,9 +19,12 @@ class AuthService:
self._user_repo = UserRepository(session)
self._plan_repo = PlanRepository(session)
def register(self, username: str, email: str, password: str) -> tuple[User, str]:
def register(self, username: str, email: str, password: str, role: str = "customer") -> tuple[User, str]:
"""注册新用户。返回 (user, session_token)。
Args:
role: "customer"(默认)或 "internal"
Raises:
ConflictError: 用户名或邮箱已存在
"""
@@ -37,6 +40,7 @@ class AuthService:
email=email,
password_hash=hash_password(password),
plan_id=plan.id,
role=role,
)
self._session.commit()
@@ -44,7 +48,7 @@ class AuthService:
return user, token
def login(self, username_or_email: str, password: str) -> tuple[User, str]:
"""登录。返回 (user, session_token)。
"""普通登录(所有用户可用)。返回 (user, session_token)。
Raises:
InvalidCredentialsError: 用户名/邮箱或密码错误
@@ -60,6 +64,25 @@ class AuthService:
token = create_session(user.id)
return user, token
def internal_login(self, username_or_email: str, password: str) -> tuple[User, str]:
"""内部登录(仅 internal 角色可用)。返回 (user, session_token)。
Raises:
InvalidCredentialsError: 用户名/邮箱或密码错误,或非内部用户
"""
user = self._user_repo.get_by_username_or_email(username_or_email)
if user is None:
raise InvalidCredentialsError()
if not verify_password(password, user.password_hash):
raise InvalidCredentialsError()
if user.status != "active":
raise InvalidCredentialsError("账户已被禁用。")
if user.role != "internal":
raise InvalidCredentialsError("此入口仅限内部员工使用。")
token = create_session(user.id)
return user, token
def logout(self, session_token: str) -> None:
"""登出:删除 session。"""
delete_session(session_token)
+154
View File
@@ -0,0 +1,154 @@
"""
管理命令行工具
用法:
python manage.py create-internal-user # 交互式创建内部员工
python manage.py create-internal-user --username admin --email admin@company.com --password 12345678
python manage.py list-users # 列出所有用户
python manage.py list-users --role internal # 只列出内部员工
"""
import argparse
import sys
from pathlib import Path
# 确保能导入 app 模块
sys.path.insert(0, str(Path(__file__).parent))
from app.core.config import get_settings
from app.core.db import get_session_factory, get_engine
from app.core.security import hash_password
from app.models import Base, User, Plan
from sqlalchemy import select
def ensure_tables():
"""确保表存在。"""
engine = get_engine()
Base.metadata.create_all(engine)
def get_or_create_free_plan(session) -> Plan:
"""获取或创建 free plan。"""
plan = session.scalars(select(Plan).where(Plan.code == "free")).first()
if not plan:
plan = Plan(
code="free",
name="免费版",
storage_quota=104_857_600,
max_file_size=20_971_520,
is_active=True,
)
session.add(plan)
session.flush()
return plan
def cmd_create_internal_user(args):
"""创建内部员工账号。"""
ensure_tables()
factory = get_session_factory()
username = args.username
email = args.email
password = args.password
# 交互式输入
if not username:
username = input("用户名: ").strip()
if not email:
email = input("邮箱: ").strip()
if not password:
import getpass
password = getpass.getpass("密码: ").strip()
if not username or not email or not password:
print("错误:用户名、邮箱、密码不能为空。")
sys.exit(1)
if len(password) < 8:
print("错误:密码长度不能少于 8 位。")
sys.exit(1)
with factory() as session:
# 检查用户名/邮箱是否已存在
existing = session.scalars(
select(User).where((User.username == username) | (User.email == email.lower()))
).first()
if existing:
if existing.username == username:
print(f"错误:用户名 '{username}' 已被占用。")
else:
print(f"错误:邮箱 '{email}' 已被注册。")
sys.exit(1)
plan = get_or_create_free_plan(session)
user = User(
username=username,
email=email.lower(),
password_hash=hash_password(password),
status="active",
role="internal",
plan_id=plan.id,
storage_used=0,
)
session.add(user)
session.commit()
print(f"\n✅ 内部员工账号创建成功!")
print(f" 用户名:{username}")
print(f" 邮箱:{email}")
print(f" 角色:internal")
print(f" 登录地址:/internal-login")
def cmd_list_users(args):
"""列出所有用户。"""
ensure_tables()
factory = get_session_factory()
with factory() as session:
stmt = select(User).order_by(User.created_at.desc())
if args.role:
stmt = stmt.where(User.role == args.role)
users = list(session.scalars(stmt).all())
if not users:
print("暂无用户。")
return
print(f"\n{'用户名':<15} {'邮箱':<25} {'角色':<10} {'状态':<8} {'创建时间'}")
print("-" * 80)
for u in users:
print(f"{u.username:<15} {u.email:<25} {u.role:<10} {u.status:<8} {u.created_at[:19]}")
print(f"\n{len(users)} 个用户")
def main():
parser = argparse.ArgumentParser(description="AI Knowledge Link 管理工具")
subparsers = parser.add_subparsers(dest="command", help="可用命令")
# create-internal-user
p_create = subparsers.add_parser("create-internal-user", help="创建内部员工账号")
p_create.add_argument("--username", "-u", help="用户名")
p_create.add_argument("--email", "-e", help="邮箱")
p_create.add_argument("--password", "-p", help="密码")
# list-users
p_list = subparsers.add_parser("list-users", help="列出所有用户")
p_list.add_argument("--role", "-r", choices=["internal", "customer"], help="按角色过滤")
args = parser.parse_args()
if args.command == "create-internal-user":
cmd_create_internal_user(args)
elif args.command == "list-users":
cmd_list_users(args)
else:
parser.print_help()
if __name__ == "__main__":
main()
+1
View File
@@ -11,6 +11,7 @@ jinja2>=3.1
# --- 数据库 ---
sqlalchemy>=2.0.30
alembic>=1.13
pymysql>=1.1
# --- 安全 ---
argon2-cffi>=23.1
+28 -10
View File
@@ -1,7 +1,7 @@
"""测试配置:每个测试用例使用隔离的内存数据库"""
"""测试配置:每个测试用例使用隔离的内存 SQLite(不依赖远程 MySQL"""
import pytest
from sqlalchemy import create_engine
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker
from app.core import db as db_module
@@ -10,16 +10,10 @@ from app.models import Base
@pytest.fixture(autouse=True)
def _isolate_db(tmp_path, monkeypatch):
"""每个测试用例:创建独立内存 SQLite → 建表 → 替换全局引擎 → 测试结束自动清理。
这确保测试之间完全隔离,不共享任何数据。
"""
"""每个测试用例:创建独立内存 SQLite → 建表 → 替换全局引擎 → 测试结束自动清理。"""
test_db_url = f"sqlite:///{tmp_path}/test.db"
engine = create_engine(test_db_url, connect_args={"check_same_thread": False})
# 启用 WAL + 外键
from sqlalchemy import event, text
@event.listens_for(engine, "connect")
def _set_pragma(dbapi_conn, _):
cursor = dbapi_conn.cursor()
@@ -34,9 +28,33 @@ def _isolate_db(tmp_path, monkeypatch):
monkeypatch.setattr(db_module, "_engine", engine)
monkeypatch.setattr(db_module, "_session_factory", test_factory)
# Mock settings 让 is_mysql 返回 False
from app.core.config import get_settings, Settings
original_settings = get_settings()
class TestSettings:
"""测试用 settings mock,保持 SQLite 行为。"""
environment = "local"
database_url = test_db_url
storage_root = str(tmp_path)
default_storage_quota = 104_857_600
default_max_file_size = 20_971_520
rate_limit_per_token_per_min = 60
rate_limit_per_ip_per_min = 30
frontend_origin = "http://localhost:5173"
is_production = False
is_mysql = False
is_sqlite = True
@property
def storage_root_path(self):
from pathlib import Path
return Path(self.storage_root)
monkeypatch.setattr("app.core.config.get_settings", lambda: TestSettings())
# 清空内存 session 存储
from app.core import session as session_module
session_module._store.clear()
yield