345 lines
16 KiB
Python
345 lines
16 KiB
Python
"""SQLAlchemy-backed UserRepository implementation.
|
|
|
|
Uses the shared async session factory from
|
|
``deerflow.persistence.engine`` — the ``users`` table lives in the
|
|
same database as ``threads_meta``, ``runs``, ``run_events``, and
|
|
``feedback``.
|
|
|
|
Constructor takes the session factory directly (same pattern as the
|
|
other four repositories in ``deerflow.persistence.*``). Callers
|
|
construct this after ``init_engine_from_config()`` has run.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import func, or_, select, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from app.gateway.auth.models import User
|
|
from app.gateway.auth.repositories.base import CredentialConflictError, UserNotFoundError, UserRepository
|
|
from deerflow.persistence.user.model import AuthInitializationRow, AuthUsernameRow, UserRow
|
|
|
|
|
|
class SQLiteUserRepository(UserRepository):
|
|
"""Async user repository backed by the shared SQLAlchemy engine."""
|
|
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
# ── Converters ────────────────────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _row_to_user(row: UserRow) -> User:
|
|
return User(
|
|
id=UUID(row.id),
|
|
email=row.email,
|
|
password_hash=row.password_hash,
|
|
system_role=row.system_role, # type: ignore[arg-type]
|
|
# SQLite loses tzinfo on read; reattach UTC so downstream
|
|
# code can compare timestamps reliably.
|
|
created_at=row.created_at if row.created_at.tzinfo else row.created_at.replace(tzinfo=UTC),
|
|
oauth_provider=row.oauth_provider,
|
|
oauth_id=row.oauth_id,
|
|
needs_setup=row.needs_setup,
|
|
token_version=row.token_version,
|
|
approved=row.approved,
|
|
login_password=row.login_password,
|
|
)
|
|
|
|
@staticmethod
|
|
def _user_to_row(user: User) -> UserRow:
|
|
return UserRow(
|
|
id=str(user.id),
|
|
email=user.email,
|
|
username_key=(user.email or "").split("@", 1)[0].strip().lower()[:128],
|
|
password_hash=user.password_hash,
|
|
system_role=user.system_role,
|
|
created_at=user.created_at,
|
|
oauth_provider=user.oauth_provider,
|
|
oauth_id=user.oauth_id,
|
|
needs_setup=user.needs_setup,
|
|
token_version=user.token_version,
|
|
approved=user.approved,
|
|
login_password=user.login_password,
|
|
)
|
|
|
|
# ── CRUD ──────────────────────────────────────────────────────────
|
|
|
|
async def create_user(self, user: User) -> User:
|
|
"""Insert a new user. Raises ``ValueError`` on duplicate email."""
|
|
row = self._user_to_row(user)
|
|
async with self._sf() as session:
|
|
await self._reserve_username(session, user.email)
|
|
session.add(row)
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError as exc:
|
|
await session.rollback()
|
|
raise ValueError(f"Email already registered: {user.email}") from exc
|
|
return user
|
|
|
|
@staticmethod
|
|
async def _reserve_username(session: AsyncSession, email: str) -> None:
|
|
"""Fence new registrations without rewriting historical user rows."""
|
|
local = email.split("@", 1)[0].lower()
|
|
escaped = local.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
existing = await session.scalar(select(UserRow.id).where(UserRow.username_key == local).limit(1))
|
|
if existing is None:
|
|
existing = await session.scalar(
|
|
select(UserRow.id)
|
|
.where(func.lower(UserRow.email).like(f"{escaped}@%", escape="\\"))
|
|
.limit(1)
|
|
)
|
|
if existing is not None:
|
|
raise ValueError("Username already registered")
|
|
session.add(AuthUsernameRow(username=local))
|
|
|
|
async def create_initial_admin(self, user: User) -> User:
|
|
"""Claim a singleton and create the first admin in one transaction.
|
|
|
|
A unique primary key fences concurrent requests across Gateway workers,
|
|
including requests using different email addresses.
|
|
"""
|
|
async with self._sf() as session:
|
|
count = await session.scalar(select(func.count()).select_from(UserRow).where(UserRow.system_role == "admin"))
|
|
if count:
|
|
raise ValueError("System already initialized")
|
|
await self._reserve_username(session, user.email)
|
|
session.add(AuthInitializationRow(id="initial-admin"))
|
|
session.add(self._user_to_row(user))
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError as exc:
|
|
await session.rollback()
|
|
raise ValueError("System already initialized") from exc
|
|
return user
|
|
|
|
async def get_user_by_id(self, user_id: str) -> User | None:
|
|
async with self._sf() as session:
|
|
row = await session.get(UserRow, user_id)
|
|
return self._row_to_user(row) if row is not None else None
|
|
|
|
async def get_users_by_ids(self, user_ids: list[str]) -> list[User]:
|
|
wanted = list(dict.fromkeys(user_id for user_id in user_ids if user_id))
|
|
if not wanted:
|
|
return []
|
|
async with self._sf() as session:
|
|
result = await session.execute(select(UserRow).where(UserRow.id.in_(wanted)))
|
|
return [self._row_to_user(row) for row in result.scalars()]
|
|
|
|
async def get_user_by_email(self, email: str) -> User | None:
|
|
normalized = email.strip().lower()
|
|
stmt = select(UserRow).where(UserRow.email == normalized).order_by(UserRow.created_at.asc()).limit(1)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
row = result.scalar_one_or_none()
|
|
if row is None:
|
|
result = await session.execute(
|
|
select(UserRow)
|
|
.where(func.lower(UserRow.email) == normalized)
|
|
.order_by(UserRow.created_at.asc())
|
|
.limit(1)
|
|
)
|
|
row = result.scalar_one_or_none()
|
|
return self._row_to_user(row) if row is not None else None
|
|
|
|
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
|
|
local = (local_part or "").strip().lower()
|
|
if not local:
|
|
return None
|
|
# Compatibility fallback for rows predating the indexed username key.
|
|
# not match unrelated rows. ``\`` is the chosen escape char below.
|
|
escaped = local.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
pattern = f"{escaped}@%"
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
select(UserRow)
|
|
.where(UserRow.username_key == local)
|
|
.order_by(UserRow.created_at.asc())
|
|
.limit(1)
|
|
)
|
|
row = result.scalar_one_or_none()
|
|
if row is None:
|
|
result = await session.execute(
|
|
select(UserRow)
|
|
.where(func.lower(UserRow.email).like(pattern, escape="\\"))
|
|
.order_by(UserRow.created_at.asc())
|
|
.limit(1)
|
|
)
|
|
row = result.scalar_one_or_none()
|
|
if row is not None and not row.username_key:
|
|
await session.execute(
|
|
update(UserRow)
|
|
.where(UserRow.id == row.id, UserRow.username_key.is_(None))
|
|
.values(username_key=local)
|
|
)
|
|
await session.commit()
|
|
return self._row_to_user(row) if row is not None else None
|
|
|
|
async def find_users_by_email_local_part_like(self, substring: str, *, limit: int = 500) -> list[User]:
|
|
term = (substring or "").strip().lower()
|
|
if not term:
|
|
return []
|
|
# Escape SQL LIKE wildcards so a literal ``_``/``%`` in the query does
|
|
# not turn into a wildcard. ``\`` is the chosen escape char.
|
|
escaped = term.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
pattern = f"%{escaped}%"
|
|
# Prefilter in SQL on the whole email (portable across sqlite/postgres/
|
|
# mysql), then keep only rows whose *local part* contains the term so a
|
|
# domain-only hit (e.g. searching "gmail") doesn't pull unrelated users.
|
|
stmt = select(UserRow).where(func.lower(UserRow.email).like(pattern, escape="\\")).order_by(UserRow.created_at.asc()).limit(max(1, limit))
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
rows = result.scalars().all()
|
|
out: list[User] = []
|
|
for row in rows:
|
|
local = (row.email or "").split("@", 1)[0].lower()
|
|
if term in local:
|
|
out.append(self._row_to_user(row))
|
|
return out
|
|
|
|
async def update_user(self, user: User) -> User:
|
|
async with self._sf() as session:
|
|
# Profile/approval writes cannot restore credentials from a stale
|
|
# User object after a concurrent password reset.
|
|
result = await session.execute(
|
|
update(UserRow)
|
|
.where(UserRow.id == str(user.id), UserRow.token_version == user.token_version)
|
|
.values(
|
|
email=user.email,
|
|
username_key=(user.email or "").split("@", 1)[0].strip().lower()[:128],
|
|
system_role=user.system_role,
|
|
oauth_provider=user.oauth_provider,
|
|
oauth_id=user.oauth_id,
|
|
needs_setup=user.needs_setup,
|
|
approved=user.approved,
|
|
)
|
|
)
|
|
if result.rowcount != 1:
|
|
row = await session.get(UserRow, str(user.id))
|
|
await session.rollback()
|
|
if row is None:
|
|
raise UserNotFoundError(f"User {user.id} no longer exists")
|
|
raise CredentialConflictError("账号密码已更新，请刷新后重试")
|
|
await session.commit()
|
|
return user
|
|
|
|
async def update_credentials(self, user: User, password_hash: str | None, *, email: str | None = None) -> User:
|
|
values = {
|
|
"password_hash": password_hash,
|
|
"login_password": password_hash,
|
|
"token_version": UserRow.token_version + 1,
|
|
"needs_setup": False,
|
|
}
|
|
if email is not None:
|
|
values["email"] = email
|
|
values["username_key"] = email.split("@", 1)[0].strip().lower()[:128]
|
|
async with self._sf() as session:
|
|
if email is not None and email.split("@", 1)[0].lower() != user.email.split("@", 1)[0].lower():
|
|
await self._reserve_username(session, email)
|
|
with session.no_autoflush:
|
|
result = await session.execute(update(UserRow).where(UserRow.id == str(user.id), UserRow.token_version == user.token_version).values(**values))
|
|
if result.rowcount != 1:
|
|
await session.rollback()
|
|
raise CredentialConflictError("账号密码已更新，请重新登录后操作")
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError as exc:
|
|
await session.rollback()
|
|
raise ValueError("Email already registered") from exc
|
|
user.password_hash = password_hash
|
|
user.login_password = password_hash
|
|
user.token_version += 1
|
|
user.needs_setup = False
|
|
if email is not None:
|
|
user.email = email
|
|
return user
|
|
|
|
async def start_login_session(self, user: User) -> User:
|
|
"""Atomically fence older logins without accepting a concurrently reset password.
|
|
|
|
Ordinary login races may advance the version again; the last issued
|
|
generation alone remains usable. Credential/approval changes after the
|
|
original check require the caller to authenticate again instead.
|
|
"""
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
update(UserRow)
|
|
.where(
|
|
UserRow.id == str(user.id),
|
|
UserRow.password_hash == user.password_hash,
|
|
UserRow.login_password == user.login_password,
|
|
UserRow.approved == user.approved,
|
|
UserRow.system_role == user.system_role,
|
|
)
|
|
.values(token_version=UserRow.token_version + 1)
|
|
)
|
|
if result.rowcount != 1:
|
|
await session.rollback()
|
|
raise CredentialConflictError("账号凭据或权限已变更，请重新登录")
|
|
row = (await session.execute(select(UserRow).where(UserRow.id == str(user.id)))).scalar_one()
|
|
saved = self._row_to_user(row)
|
|
await session.commit()
|
|
return saved
|
|
|
|
async def count_users(self) -> int:
|
|
stmt = select(func.count()).select_from(UserRow)
|
|
async with self._sf() as session:
|
|
return await session.scalar(stmt) or 0
|
|
|
|
async def count_admin_users(self) -> int:
|
|
stmt = select(func.count()).select_from(UserRow).where(UserRow.system_role == "admin")
|
|
async with self._sf() as session:
|
|
return await session.scalar(stmt) or 0
|
|
|
|
async def list_users(self) -> list[User]:
|
|
stmt = select(UserRow).order_by(UserRow.created_at.asc())
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_user(row) for row in result.scalars()]
|
|
|
|
@staticmethod
|
|
def _user_search_filter(q: str | None):
|
|
"""Case-insensitive substring match on email or user id (None → no filter)."""
|
|
if not q or not q.strip():
|
|
return None
|
|
needle = f"%{q.strip().lower()}%"
|
|
return or_(
|
|
func.lower(UserRow.email).like(needle),
|
|
func.lower(UserRow.id).like(needle),
|
|
)
|
|
|
|
async def list_users_page(
|
|
self, *, q: str | None = None, limit: int = 20, offset: int = 0
|
|
) -> list[User]:
|
|
stmt = select(UserRow).order_by(UserRow.created_at.asc()).limit(limit).offset(offset)
|
|
search = self._user_search_filter(q)
|
|
if search is not None:
|
|
stmt = stmt.where(search)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_user(row) for row in result.scalars()]
|
|
|
|
async def count_users_filtered(
|
|
self, *, q: str | None = None, approved_or_admin: bool = False
|
|
) -> int:
|
|
stmt = select(func.count()).select_from(UserRow)
|
|
search = self._user_search_filter(q)
|
|
if search is not None:
|
|
stmt = stmt.where(search)
|
|
if approved_or_admin:
|
|
stmt = stmt.where(or_(UserRow.approved.is_(True), UserRow.system_role == "admin"))
|
|
async with self._sf() as session:
|
|
return await session.scalar(stmt) or 0
|
|
|
|
async def get_user_by_oauth(self, provider: str, oauth_id: str) -> User | None:
|
|
stmt = select(UserRow).where(UserRow.oauth_provider == provider, UserRow.oauth_id == oauth_id)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
row = result.scalar_one_or_none()
|
|
return self._row_to_user(row) if row is not None else None
|