188 lines
8.1 KiB
Python
188 lines
8.1 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, select
|
|
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 UserNotFoundError, UserRepository
|
|
from deerflow.persistence.user.model import 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,
|
|
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:
|
|
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
|
|
|
|
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_user_by_email(self, email: str) -> User | None:
|
|
stmt = select(UserRow).where(UserRow.email == email)
|
|
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
|
|
|
|
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
|
|
# Escape SQL LIKE wildcards so usernames like ``foo_bar`` or ``a%b`` do
|
|
# not match unrelated rows. ``\`` is the chosen escape char below.
|
|
escaped = local.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
pattern = f"{escaped}@%"
|
|
# Case-insensitive prefix match on the email field, oldest first so a
|
|
# subsequent registration cannot shadow an existing account.
|
|
stmt = (
|
|
select(UserRow)
|
|
.where(func.lower(UserRow.email).like(pattern, escape="\\"))
|
|
.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()
|
|
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:
|
|
row = await session.get(UserRow, str(user.id))
|
|
if row is None:
|
|
# Hard fail on concurrent delete: callers (reset_admin,
|
|
# password change handlers, _ensure_admin_user) all
|
|
# fetched the user just before this call, so a missing
|
|
# row here means the row vanished underneath us. Silent
|
|
# success would let the caller log "password reset" for
|
|
# a row that no longer exists.
|
|
raise UserNotFoundError(f"User {user.id} no longer exists")
|
|
row.email = user.email
|
|
row.password_hash = user.password_hash
|
|
row.system_role = user.system_role
|
|
row.oauth_provider = user.oauth_provider
|
|
row.oauth_id = user.oauth_id
|
|
row.needs_setup = user.needs_setup
|
|
row.token_version = user.token_version
|
|
row.approved = user.approved
|
|
row.login_password = user.login_password
|
|
await session.commit()
|
|
return user
|
|
|
|
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()]
|
|
|
|
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
|