mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-09 21:49:37 +00:00
fix(auth): correct OAuth conflict error message + validate multi-worker Postgres claim with real concurrency benchmark (#5026)
* fix(auth): correct OAuth uniqueness error and index parity on Postgres
create_user() caught any IntegrityError on commit and always reported it
as a duplicate email. The email pre-check already rules out a real email
collision in the common case, so any IntegrityError reaching that handler
is actually idx_users_oauth_identity firing instead -- confirmed against
both backends: SQLite reports "UNIQUE constraint failed:
users.oauth_provider, users.oauth_id", Postgres reports a
UniqueViolationError naming the same index. The caller saw "Email already
registered" for an OAuth account conflict, which is wrong and would send
API consumers debugging the wrong field.
Distinguish the two cases via a substring check on the driver error text
(both backends name the oauth columns) and raise an accurate message for
each.
Also add postgresql_where to the same index, alongside the existing
sqlite_where. This is not a correctness fix -- verified empirically that
Postgres already enforces the same practical uniqueness without it
(NULL is never equal to NULL in either backends unique index, so real
duplicate (provider, id) pairs are already rejected and NULL/NULL rows
are already unconstrained). postgresql_where makes the index genuinely
partial on Postgres too, matching the stated intent in the surrounding
comment and keeping the index smaller as the common case (plain-password
accounts, both columns NULL) accumulates.
* test(bench): add multi-process SQLite vs Postgres concurrency benchmark
CONFIGURATION.md documents that multi-worker deployments must use Postgres
because "SQLite silently ignores row-level locks", but nothing in the repo
exercised that claim against real separate worker processes -- the existing
checkpoint benchmarks (scripts/benchmark/checkpoint/) measure single-process
read/write latency, and the existing Postgres tests
(test_pg_schema_integration.py, test_multi_worker_postgres_gate.py) cover
schema placement and config validation, not throughput or lock behavior
under concurrent load.
run_concurrency_bench.py spawns N real OS processes (subprocess.Popen, not
asyncio tasks or threads within one process) against the shared users
table, mixing reads (get_user_by_email) and writes (create_user) at a
configurable ratio, and reports throughput, error counts by exception
type, and p50/p95/p99/max latency per run.
Measured locally (2/4/8/16 workers, 100 ops/worker, 70/30 read/write):
SQLite completed all operations with zero errors at every worker count
(busy_timeout absorbs contention rather than raising), but total
throughput stayed flat around 28-34 ops/s regardless of worker count, and
p99 latency grew from ~400ms at 2 workers to ~5.9s at 16, with a 22s max.
Postgres throughput scaled with worker count (41 to 66 ops/s) and p99
stayed under 500ms at every worker count tested. Raw JSON output from
both runs is available on request; exact numbers will vary by machine and
are not asserted in the test suite.
test_bench_concurrency.py unit-tests the pure aggregation logic
(percentile math, error grouping, crashed-worker handling) the same way
test_bench_checkpoint_channels.py does for the existing benchmarks --
fast, no DB required, not the full multi-process sweep in CI.
* fix(auth): inspect the driver exception for OAuth conflict detection
str(exc) embeds the full failed INSERT statement, whose column list
always names oauth_provider/oauth_id, so a substring check on it
misclassified every commit-time IntegrityError on the users table as
an OAuth conflict (reproduced on SQLite: a duplicate primary key with
a different email raised "OAuth account already linked: None/None").
_is_oauth_identity_violation now inspects exc.orig instead: constraint_name
on Postgres, both violated column names present (not a bare "oauth"
substring) on SQLite.
Also ships the alembic revision idx_users_oauth_identity's postgresql_where
predicate needed: 0001_baseline created it as a full index on Postgres,
and ORM metadata changes only affect fresh create_all databases, never an
already-versioned deployment.
Addresses review feedback from willem-bd.
* fix(bench): run the concurrency benchmark in an isolated schema and derive paths from the checkout
--pg-url accepted an arbitrary database URL while the code pinned
postgres_schema="public" and unconditionally ran DELETE FROM users --
against any non-disposable database that permanently destroyed every
auth account. Each run now generates a unique throwaway schema
(bench_<uuid>), points both the seeder and every worker subprocess at
it via postgres_schema, and drops only that schema (DROP SCHEMA ...
CASCADE) once the full worker-count sweep finishes.
Also stopped hard-coding /opt/deer-flow/backend as the checkout path
and .venv/bin/python3 as the interpreter: BACKEND_DIR is now derived
from Path(__file__), and workers are spawned with sys.executable (the
orchestrator's own interpreter) instead, so the documented
uv run python scripts/benchmark/concurrency/run_concurrency_bench.py
command works from any checkout.
Addresses review feedback from willem-bd.
* fix: shorten oauth-index revision id, repin migration-head assertions, fix bench read/write mix
- 0017_users_oauth_identity_partial_pg (36 chars) exceeded
alembic_version.version_num's VARCHAR(32) limit, which would fail
stamping/upgrading on both fresh and existing Postgres deployments.
Renamed to 0017_oauth_identity_pg_partial (30 chars).
- Repinned every test asserting 0016_subagent_batches as the migration
head (test_persistence_bootstrap[.py|_concurrency.py|_regression.py],
test_migration_0004/0007/0015) to the new 0017 revision id.
- worker.py's `(i % 100) < int(read_ratio * 100)` assumed n_ops >= 100;
at the documented default (50 ops/worker, 0.7 read ratio) it produced
either all-reads or all-writes, never the claimed mixed workload.
Replaced with read_count()/is_read_op(), which distribute an exact
round(n_ops * read_ratio) reads evenly across the sequence via modular
spacing, and added test_bench_worker.py covering the default values
plus small op counts.
* fix(bench): establish a real physical connection before timing ops
async with sf(): pass entered an empty AsyncSession without checking out
a physical connection -- SQLAlchemy stays lazy until the first statement
executes. That pushed connection-establishment cost onto each worker's
first timed operation instead of conn_time_s, and at 16 workers those 16
cold first-ops (1% of a 1600-op sample) could skew the reported p99.
Execute a real `SELECT 1` before starting the timer instead.
Verified with a real end-to-end run (uv sync + sqlite backend, 2
workers/10 ops, 0 errors) plus the full auth/bench/migration-bootstrap
suites (135 tests) and ruff check/format, all clean.
* fix(bench): synchronize workers before timing, fix percentile off-by-one
Two remaining measurement issues from review:
- run_workers() started the wall clock before spawning any worker, so
throughput/wall_time absorbed N processes' staggered Python-startup and
connection-establishment cost, and early workers could run ahead of ones
still starting. Workers now print READY right before their timed loop
and block on stdin for a GO signal; the orchestrator waits for every
READY, then starts the timer and releases all workers together.
- summarize()'s pct() used int(len(latencies) * p) directly as a
zero-based index -- a one-based-rank-as-index bug that put p95 and p99
at the same slot (the max) for any 20-or-fewer-sample run, and for the
documented 100-sample default. Now delegates to
checkpoint_bench_common.percentile(), the already-correct nearest-rank
implementation used elsewhere in the same benchmark family, instead of
a second, broken one.
Verified: 14/14 unit tests pass (2 new pinned-value regression tests for
the percentile bug, using the reviewer's own 20-sample repro), ruff
clean, and a real 2/4-worker SQLite multi-process smoke run completes
with distinct p95/p99/max latencies and no hang.
* fix(bench): absolute SQLite bench path, surface crash diagnostics, exit nonzero on failure; share OAuth index constant + cover Postgres branch
Three more findings from review at 5fd25a7:
- seed_baseline() cleaned an absolute .deer-flow/bench_data path but
handed DatabaseConfig a relative one, which resolves against the
CALLER's CWD -- not BACKEND_DIR. Invoking the documented command from
anywhere other than backend/ silently pointed the seeder and the
(cwd=BACKEND_DIR) workers at two different directories: workers crashed
with 'unable to open database file' while the run still printed a
well-formed summary and exited 0. Both seed_baseline() and worker.py's
make_session_factory() now use the same absolute path.
- Crashed workers' stderr was captured then discarded, and main() always
exited 0 -- an all-crashed sweep was indistinguishable from a real
(uneventful) measurement to anything checking the exit code or
--out. run_workers() now prints each crash immediately and tags it with
the real worker_id (previously always None); summarize() exposes
crashed_worker_errors alongside the existing crashed_workers count;
main() exits 1 via the new summary_indicates_failure() whenever any
sweep crashed or fell short of expected_total_ops.
- idx_users_oauth_identity was hardcoded separately in the ORM Index and
in _is_oauth_identity_violation's Postgres branch, with no test to
catch drift, and that branch had zero non-skipped coverage (its only
guard needs a live Postgres CI never configures). Exported
OAUTH_IDENTITY_INDEX_NAME from user/model.py as the shared source of
truth (migrations intentionally keep their own frozen literal, matching
every other revision in that package) and added stub-exception unit
tests pinning both the asyncpg constraint_name path and the sqlite
message-substring path, positive and negative.
Verified: 107 passed locally (auth + bench-unit suites), ruff clean, and
two real reproductions -- invoking run_concurrency_bench.py from a
scratch directory outside backend/ (the reviewer's exact repro) now
completes 8/8 ops with crashed_workers: 0 instead of crashing, and the
new crashed_worker_errors/exit-code logic is exercised directly by the
new unit tests against the real summarize()/summary_indicates_failure().
* fix(auth): attribute create_user IntegrityErrors to the right constraint
Two coupled review findings on the classification helpers:
P3 (fall-through) -- after ruling out the OAuth-identity index, create_user
raised "Email already registered: {email}" for every remaining
IntegrityError, including the duplicate-primary-key case the new
regression test exercises, whose address is not registered. Added
_is_email_violation() so the email message is used only for an actual
users.email collision that raced past the pre-check; anything else (in
practice a duplicate id) now raises a neutral
"User already exists (constraint: <name>)".
P2 (unreachable asyncpg branch) -- exc.orig is not the asyncpg error.
SQLAlchemy's asyncpg dialect re-raises its own DBAPI IntegrityError
(pgcode/sqlstate only) 'from' the real asyncpg error, so constraint_name
lives on exc.orig.__cause__. getattr(exc.orig, "constraint_name", None)
was always None on Postgres; the helpers only worked there by accident,
matching asyncpg's DETAIL line in the message fallback. Added
_driver_constraint_name() which walks orig then orig.__cause__, and the
stub tests now model that real shape (orig wrapper + __cause__) instead of
a constraint_name that no driver puts on orig directly.
Tests: 76 passed. New coverage for the email-race path, both new helpers
on each backend, the neutral fallback message, and the cause-chain walk.
* fix(bench): match app SQLite PRAGMAs in workers; fail a sweep on any op error
Two review follow-ups:
- worker.py opened its SQLite engine with only connect_args timeout=30.
synchronous and foreign_keys are per-connection PRAGMAs, so workers ran
at SQLite's synchronous=FULL / foreign_keys=OFF while a real Gateway
worker runs synchronous=NORMAL (persistence/engine.py::_enable_sqlite_wal)
-- an extra fsync per commit on the measured 30%-write path, overstating
SQLite's cost in the direction that flatters the "use Postgres"
conclusion. Added a connect listener applying the same four PRAGMAs, with
a test asserting synchronous/foreign_keys/journal_mode on a real worker
connection.
- summary_indicates_failure() only looked at crashes and the completed vs
expected op counts, so a sweep where every op completed but raised
(e.g. writes hitting OperationalError) passed as a clean measurement:
completed_ops == expected, 0 crashes. Added an "errors > 0" clause; the
error breakdown stays in the JSON, only the exit code changes. Test added.
test_bench_concurrency.py + test_bench_worker.py green (20), plus a real
2-worker sqlite smoke run (6/6 ops, 0 errors, exit 0).
* fix(auth): match the real email index name; only claim "exists" for uniqueness
Review follow-ups on the classification helpers:
- email is mapped_column(unique=True, index=True), which SQLAlchemy and
0001_baseline realise as a single UNIQUE INDEX (ix_users_email), not a
named UNIQUE constraint. _is_email_violation compared the driver
constraint name against "users_email_key", which Postgres never emits,
so that arm was dead on Postgres (SQLite matched via the message). Fixed
to ix_users_email.
- the residual IntegrityError fallback raised "User already exists" for
every remaining IntegrityError -- a NOT NULL / CHECK / foreign-key
violation is not a "user already exists" condition and is not part of
create_user's ValueError contract. Added _is_uniqueness_violation
(sqlstate 23505, or the SQLite "UNIQUE/PRIMARY KEY constraint failed"
message); only that raises the "already exists" ValueError, everything
else propagates unchanged.
- documented scripts/benchmark/concurrency/ in backend/AGENTS.md alongside
the other benchmark family.
Tests: 78 auth + 20 bench-unit pass, ruff clean. New coverage for
_is_uniqueness_violation on both backends and for a non-uniqueness
IntegrityError propagating out of create_user.
* fix(bench): don't pre-close worker stdin (breaks communicate); require --pg-url for postgres
* fix(bench): ruff format; time throughput on the op phase, not teardown
- lint-backend: ruff format the files touched in this PR.
- Throughput window (P2): the orchestrator sampled its wall clock after
every worker's communicate() returned, so it also covered each worker's
engine.dispose(), result serialization and stdout transfer. Each worker
now times just its operation phase (GO -> last op) and reports
ops_elapsed_s; summarize() uses max(ops_elapsed_s) over the workers -- all
released by the same GO -- as the throughput window (ops_window_s).
- Exercise migration 0018 (P2): test_user_oauth_partial_index.py goes
through bootstrap create_all(), which builds the partial index from ORM
metadata and never runs 0018.upgrade(). New Postgres-gated
test_migration_0018_oauth_identity_pg_partial.py alembic-upgrades to 0017
(full index), then 0018 (asserts the predicate appears), then downgrades
(asserts the full index is restored) and re-upgrades.
* docs(middlewares): tighten SandboxAudit and Clarification entries in AGENTS.md
PR #5134 grew agents/middlewares/AGENTS.md ~1.8 KB, pushing the effective
AGENTS.md chain for that directory over the 96 KiB hard limit once this
branch also documents scripts/benchmark/concurrency/ in backend/AGENTS.md.
Condense the two longest middleware entries (SandboxAuditMiddleware,
ClarificationMiddleware) without dropping any identifier, example, issue
reference, ordering constraint, or documented gap; chain back to ~96.8 KiB.
---------
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
8ee3c83508
commit
27cb73659d
@ -127,6 +127,22 @@ The offline test suite must not require network access, provider credentials,
|
||||
or the LongMemEval dataset. Small LongMemEval-shaped fixtures must be synthetic
|
||||
and generated by tests.
|
||||
|
||||
`scripts/benchmark/concurrency/` measures multi-process contention on the
|
||||
`users` table (N separate OS processes, not asyncio tasks) for SQLite vs
|
||||
Postgres -- the scenario `CONFIGURATION.md` requires Postgres for. `worker.py`
|
||||
connects directly via SQLAlchemy (skipping the ~8.5s Alembic bootstrap the
|
||||
orchestrator already ran once) and mirrors the app's per-connection SQLite
|
||||
PRAGMAs; `run_concurrency_bench.py` seeds a disposable per-run Postgres schema,
|
||||
synchronises workers on a READY/GO barrier before timing, and exits non-zero on
|
||||
any crash, short op count, or `errors > 0`. Postgres runs need a throwaway
|
||||
database via `--pg-url`; nothing here touches `public`. Run from `backend/`:
|
||||
|
||||
```bash
|
||||
uv run python scripts/benchmark/concurrency/run_concurrency_bench.py \
|
||||
--backend sqlite --workers 2,4,8,16 --ops-per-worker 50 --read-ratio 0.7
|
||||
uv run pytest tests/test_bench_concurrency.py tests/test_bench_worker.py -q
|
||||
```
|
||||
|
||||
## Commands
|
||||
|
||||
**Root directory** (for full application):
|
||||
|
||||
@ -21,7 +21,96 @@ 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
|
||||
from deerflow.persistence.user.model import OAUTH_IDENTITY_INDEX_NAME, UserRow
|
||||
|
||||
# ``email`` is ``mapped_column(unique=True, index=True)``, which SQLAlchemy
|
||||
# (and 0001_baseline) realise as a single UNIQUE INDEX -- not a named UNIQUE
|
||||
# constraint -- so a Postgres duplicate reports the index name here.
|
||||
_EMAIL_UNIQUE_INDEX_NAME = "ix_users_email"
|
||||
|
||||
|
||||
def _driver_constraint_name(exc: IntegrityError) -> str | None:
|
||||
"""The violated constraint's name from the driver exception, or ``None``
|
||||
when the driver does not expose one.
|
||||
|
||||
``exc.orig`` is NOT the raw driver error. SQLAlchemy's asyncpg dialect
|
||||
re-raises a plain ``AsyncAdapt_asyncpg_dbapi.IntegrityError`` built from a
|
||||
rendered string and carrying only ``pgcode``/``sqlstate``
|
||||
(``sqlalchemy/dialects/postgresql/asyncpg.py::_handle_exception``); the
|
||||
real ``asyncpg.UniqueViolationError`` — the one with ``constraint_name`` —
|
||||
survives as ``exc.orig.__cause__`` (``raise translated_error from error``).
|
||||
aiosqlite exposes no constraint name at all. Check the wrapper, then its
|
||||
cause.
|
||||
"""
|
||||
for obj in (exc.orig, getattr(exc.orig, "__cause__", None)):
|
||||
name = getattr(obj, "constraint_name", None)
|
||||
if name:
|
||||
return str(name)
|
||||
return None
|
||||
|
||||
|
||||
def _is_oauth_identity_violation(exc: IntegrityError) -> bool:
|
||||
"""Distinguish the ``idx_users_oauth_identity`` (oauth_provider, oauth_id)
|
||||
unique-index violation from any OTHER ``IntegrityError`` reaching this
|
||||
commit (a duplicate primary key, or a duplicate-email race that slipped
|
||||
past the pre-check above).
|
||||
|
||||
Never uses the SQLAlchemy wrapper's ``str(exc)``: it embeds the full
|
||||
failed INSERT statement, whose column list names
|
||||
``oauth_provider``/``oauth_id`` on every call regardless of which
|
||||
constraint fired, so a substring check on it misclassifies every
|
||||
commit-time ``IntegrityError`` on this table as an OAuth conflict
|
||||
(reproduced on SQLite: a duplicate ``id`` raised "OAuth account already
|
||||
linked: None/None").
|
||||
|
||||
Postgres: match :data:`OAUTH_IDENTITY_INDEX_NAME` against the driver
|
||||
constraint name (see :func:`_driver_constraint_name`). SQLite: no
|
||||
structured name, only a message naming the columns (``"UNIQUE constraint
|
||||
failed: users.oauth_provider, users.oauth_id"``) — require BOTH oauth
|
||||
column names, not a bare "oauth" substring.
|
||||
"""
|
||||
name = _driver_constraint_name(exc)
|
||||
if name is not None:
|
||||
return name == OAUTH_IDENTITY_INDEX_NAME
|
||||
message = str(exc.orig).lower()
|
||||
return "oauth_provider" in message and "oauth_id" in message
|
||||
|
||||
|
||||
def _is_email_violation(exc: IntegrityError) -> bool:
|
||||
"""True when ``exc`` is the ``ix_users_email`` uniqueness violation, using
|
||||
the same driver-exception inspection as
|
||||
:func:`_is_oauth_identity_violation` (never ``str(exc)``, whose INSERT text
|
||||
names every column)."""
|
||||
name = _driver_constraint_name(exc)
|
||||
if name is not None:
|
||||
return name == _EMAIL_UNIQUE_INDEX_NAME
|
||||
return "users.email" in str(exc.orig).lower()
|
||||
|
||||
|
||||
def _is_uniqueness_violation(exc: IntegrityError) -> bool:
|
||||
"""True for a unique-index / primary-key violation specifically, as
|
||||
opposed to a NOT NULL / CHECK / foreign-key ``IntegrityError`` on the same
|
||||
INSERT -- only the former means "a user like this already exists"."""
|
||||
sqlstate = getattr(exc.orig, "sqlstate", None) or getattr(exc.orig, "pgcode", None)
|
||||
if sqlstate is not None:
|
||||
return sqlstate == "23505" # unique_violation
|
||||
message = str(exc.orig).lower()
|
||||
return "unique constraint failed" in message or "primary key constraint failed" in message
|
||||
|
||||
|
||||
def _violated_constraint(exc: IntegrityError) -> str | None:
|
||||
"""Best-effort name of the constraint behind ``exc``, for a diagnostic that
|
||||
does not pin an unattributed violation on a specific column. Uses the
|
||||
driver constraint name where available, else the column(s) SQLite names in
|
||||
its message (``"UNIQUE constraint failed: users.id"``)."""
|
||||
name = _driver_constraint_name(exc)
|
||||
if name:
|
||||
return name
|
||||
marker = "constraint failed: "
|
||||
message = str(exc.orig)
|
||||
if marker in message:
|
||||
return message.split(marker, 1)[1].splitlines()[0].strip() or None
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_email(email: str) -> str:
|
||||
@ -84,7 +173,11 @@ class SQLiteUserRepository(UserRepository):
|
||||
# ── CRUD ──────────────────────────────────────────────────────────
|
||||
|
||||
async def create_user(self, user: User) -> User:
|
||||
"""Insert a new user. Raises ``ValueError`` on duplicate email.
|
||||
"""Insert a new user. Raises ``ValueError`` on any uniqueness
|
||||
violation -- duplicate email, a duplicate (provider, oauth_id) pair
|
||||
for an OAuth-linked account, or a duplicate id -- with a message
|
||||
naming the specific conflict. Other ``IntegrityError``\\ s (NOT NULL,
|
||||
CHECK, foreign key) propagate unchanged.
|
||||
|
||||
The email is canonicalised to lowercase before insert so the existing
|
||||
unique constraint enforces case-insensitive uniqueness for new rows and
|
||||
@ -103,7 +196,30 @@ class SQLiteUserRepository(UserRepository):
|
||||
await session.commit()
|
||||
except IntegrityError as exc:
|
||||
await session.rollback()
|
||||
raise ValueError(f"Email already registered: {user.email}") from exc
|
||||
# The email pre-check above already ruled out an email
|
||||
# collision under normal (non-racing) conditions, so
|
||||
# IntegrityErrors reaching here are usually
|
||||
# idx_users_oauth_identity -- but not always (a duplicate
|
||||
# primary key, or an email collision that raced past the
|
||||
# pre-check). Attribute the failure to the constraint that
|
||||
# actually fired instead of assuming any one of them.
|
||||
if _is_oauth_identity_violation(exc):
|
||||
raise ValueError(f"OAuth account already linked: {user.oauth_provider}/{user.oauth_id}") from exc
|
||||
if _is_email_violation(exc):
|
||||
# A duplicate address that got past the pre-check: a
|
||||
# concurrent insert of the same email.
|
||||
raise ValueError(f"Email already registered: {user.email}") from exc
|
||||
if _is_uniqueness_violation(exc):
|
||||
# Some other unique index / primary key (in practice a
|
||||
# duplicate id). "Already exists" fits, but don't dress it
|
||||
# up as an email conflict for an address that isn't
|
||||
# registered.
|
||||
constraint = _violated_constraint(exc)
|
||||
raise ValueError(f"User already exists (constraint: {constraint})" if constraint else "User already exists") from exc
|
||||
# A NOT NULL / CHECK / foreign-key IntegrityError is not a
|
||||
# "user already exists" condition and not part of this
|
||||
# method's ValueError contract -- let it propagate.
|
||||
raise
|
||||
return user
|
||||
|
||||
async def get_user_by_id(self, user_id: str) -> User | None:
|
||||
|
||||
@ -64,7 +64,7 @@ it to that middleware's declaration in the same change.
|
||||
`__authorization_outcome` key (so `build_run_config` strips caller-supplied
|
||||
forgeries). Consumers pop it; the publisher and the consumer share only that
|
||||
contract module.
|
||||
10. **SandboxAuditMiddleware** - Audits sandboxed shell/file operations for security logging before tool execution. Command classification is **defense-in-depth and audit, not a security boundary** — the sandbox itself is the isolation boundary. Command substitution is judged by *position*, not by the presence of `$(`: a substitution in **command position** (`$(curl url)`, `` `curl url` ``, the word after a `|`/`&&`/`;`, or any `eval`/`source` argument) executes fetched or interpreted content and is blocked, while **value position** (`x=$(curl url)`, `echo $(curl url)`, an argument, a `for` word list) only captures output and passes (#4611). `_HIGH_RISK_COMMAND_POSITION_PATTERNS` is therefore matched anchored against each split sub-command, never against the whole compound string, and `_split_compound_command(split_pipes=True)` supplies those sub-commands; rules that span a pipe (`| sh`, `base64 -d | ...`) still rely on `_classify_command`'s whole-command Pass 1. `_COMMAND_POSITION_PREFIX` extends the anchor over leading variable assignments and exec wrappers (`FOO=1 $(curl url)`, `env`/`command`/`builtin`/`exec`/`nohup`/`time`/`sudo`/`doas`), which are still command position; its assignment branch requires whitespace before the substitution, which is exactly what keeps `x=$(curl url)` in value position. Two execution contexts are deliberately **position-blind** and matched against the whole command in Pass 1, because they execute what they receive wherever they appear (including as an argument to something else, e.g. `xargs sh -c "$(curl url)"`): an `eval`/`source` argument, and an interpreter's **code-string flag** — `-c` (shells, `python`), `-e` (`perl`/`ruby`/`node`), `-p` (`perl`/`node`), `-r` (`php`) — plus the here-string (`<<<`) that reaches the same place through stdin. All three substitution spellings (`$(cmd`, `<(cmd`, `` `cmd ``) share one `_RISKY_SUBSTITUTION` opener so a rule cannot cover one spelling and miss another. An unquoted newline splits like `;`, because it separates statements the same way: leaving it joined let `echo hi\n$(curl url)` evade the anchored rules that its `;` spelling triggers. A heredoc body is data rather than statements, so `_split_compound_command` records headers (`<<EOF`, `<<-EOF`, `<<'EOF'`) and consumes their bodies verbatim at the newline that starts them — otherwise a body line beginning with `$(curl url)` would be promoted to a command position the shell never creates. Two things that look like headers must not open one, or a body that never terminates swallows every following statement: `<<<` is a here-string (both a lookahead and a lookbehind are needed, or the trailing `<<` of `<<< "text"` reads as a heredoc with delimiter `text`), and a `<<` inside `$(( ... ))` / `(( ... ))` is a bit shift, so arithmetic depth is tracked alongside the quote flags. That is a heuristic, not shell parsing: it exists only to avoid manufacturing command positions *and* to avoid destroying real ones. An unterminated body consumes the rest of the string; an unclosed `((` only disables heredoc detection, so newlines keep splitting and the failure direction stays towards seeing more command positions rather than fewer. Known, deliberate gaps: process substitution outside `eval`/`source` (`. <(curl u)`) is not detected — closing it would require real shell parsing, which is out of scope for this layer. Two-step forms (`x=$(curl u); eval "$x"`) are inherent rather than incidental: any rule that allows output capture allows the first statement, and connecting it to the later `eval` needs dataflow analysis, not pattern matching. There is currently no config gate: the middleware is appended unconditionally in `_build_runtime_middlewares`, so it applies to both the lead agent and subagents.
|
||||
10. **SandboxAuditMiddleware** - Audits sandboxed shell/file operations before tool execution; command classification is **defense-in-depth and audit, not a security boundary** (the sandbox is the isolation boundary). Command substitution is judged by *position*, not the presence of `$(`: **command position** (`$(curl url)`, `` `curl url` ``, the word after `|`/`&&`/`;`, an `eval`/`source` argument) executes fetched content and is blocked; **value position** (`x=$(curl url)`, `echo $(curl url)`, an argument, a `for` word list) only captures output and passes (#4611). So `_HIGH_RISK_COMMAND_POSITION_PATTERNS` is matched anchored against each sub-command from `_split_compound_command(split_pipes=True)`, never the whole string; pipe-spanning rules (`| sh`, `base64 -d | ...`) still use `_classify_command`'s whole-command Pass 1. `_COMMAND_POSITION_PREFIX` extends the anchor over leading assignments and exec wrappers (`FOO=1 $(curl url)`, `env`/`command`/`builtin`/`exec`/`nohup`/`time`/`sudo`/`doas`); its assignment branch requires whitespace before the substitution, which keeps `x=$(curl url)` in value position. Two contexts are deliberately **position-blind** (matched whole-command in Pass 1, since they execute their input anywhere, e.g. `xargs sh -c "$(curl url)"`): an `eval`/`source` argument, and an interpreter **code-string flag** — `-c` (shells, `python`), `-e` (`perl`/`ruby`/`node`), `-p` (`perl`/`node`), `-r` (`php`) — plus the here-string (`<<<`) reaching the same place via stdin. All three substitution spellings (`$(`, `<(`, `` ` ``) share one `_RISKY_SUBSTITUTION` opener. An unquoted newline splits like `;` (else `echo hi\n$(curl url)` evades the anchored rules). Heredoc bodies are data: `_split_compound_command` records headers (`<<EOF`, `<<-EOF`, `<<'EOF'`) and consumes their bodies verbatim, so a body line starting `$(curl url)` isn't promoted to a command position; `<<<` (here-string, needs look-ahead + look-behind) and a `<<` inside `$(( ))`/`(( ))` (bit shift, arithmetic depth tracked with the quote flags) must not open one. This is a heuristic, not shell parsing — an unterminated body consumes the rest of the string, an unclosed `((` only disables heredoc detection, and the failure direction is always toward *more* command positions, not fewer. Known gaps: process substitution outside `eval`/`source` (`. <(curl u)`) is undetected, and two-step forms (`x=$(curl u); eval "$x"`) need dataflow analysis. No config gate — appended unconditionally in `_build_runtime_middlewares`, for both lead and subagents.
|
||||
11. **ReadBeforeWriteMiddleware** - *(optional, if `read_before_write.enabled`, default on)* Outermost write gate (issue #3857): `read_file` stamps a content hash onto its ToolMessage; `write_file` (append/overwrite-existing) and `str_replace` are blocked unless the newest mark for that path matches the file's current hash. Sits outside ToolProgressMiddleware and ToolErrorHandlingMiddleware so a blocked write returns immediately without consuming a ToolProgress slot. Blocked results call `normalize_tool_result` directly to stamp `deerflow_tool_meta` (`recoverable_by_model=True`) before returning, keeping the result well-formed for any outer consumer. Marks live on messages, so summarization dropping the read result invalidates the gate automatically; writes never refresh marks, forcing a re-read between consecutive edits. Gate check + tool execution are serialized per (thread, path) so same-turn parallel writes cannot reuse one stale mark; on sandboxes whose `read_file` reports failures as `"Error: ..."` strings instead of raising (AIO/E2B), uninspectable targets fail open (creation proceeds, no mark stamped). The middleware also owns the sandbox authorization scope for these composed calls: pre-write inspection, the tool body, and post-read hashing share one sync/async provider decision, while `SandboxAuthorizationError` bypasses the generic inspection fail-open paths and becomes an error ToolMessage.
|
||||
12. **ToolProgressMiddleware** - *(optional, if `tool_progress.enabled`)* State-machine-based stagnation guard (RFC #3177). Outer wrapper around ToolErrorHandlingMiddleware so its `wrap_tool_call` receives results already stamped with `deerflow_tool_meta`. Tracks per-(thread, tool) consecutive "no-new-info" calls across three error categories: (a) `recoverable_by_model=True` (no_results, not_found, permission, Jaccard-duplicate success): ACTIVE → WARNED (terminal — hint re-injected on each subsequent problem); (b) `recoverable_by_model=False, action≠stop` (rate_limited, transient): ACTIVE → WARNED → BLOCKED after `warn_escalation_count` more problems; (c) `recoverable_by_model=False, action=stop` (auth, config, internal): immediately BLOCKED on first occurrence. **Division of labor with LoopDetectionMiddleware:** ToolProgressMiddleware is a result-quality guard — fires after tool execution and blocks specific tools that stop producing new information; LoopDetectionMiddleware is a call-pattern guard — fires after the model responds and hard-stops the whole turn when the model repeatedly issues identical tool_calls. Both can inject HumanMessage hints in the same model call without conflict; neither reads the other's internal state.
|
||||
13. **ToolReceiptMiddleware + ToolErrorHandlingMiddleware** - `ToolReceiptMiddleware` is *(optional, if `verification.receipts_enabled`, default on)*. It is the **outermost `wrap_tool_call` layer** — registered ahead of entries 9-12 — because Guardrail/SandboxAudit/ReadBeforeWrite/ToolProgress can short-circuit a call with their own ToolMessage (and SandboxAudit rebuilds medium-risk results); an inner receipt layer would silently gap the ledger on those results (ordering constraints in `deerflow.extensions.ordering`). Normal results still carry the `deerflow_tool_meta` status ToolErrorHandlingMiddleware stamps on the inner return path; short-circuit messages self-stamp meta or fall back to `message.status`. It stamps deterministic provenance (tool name, status, args/output hashes, byte count, timestamp) onto direct `ToolMessage` results and every matching `ToolMessage` carried in `Command.update.messages`, including delegated `task`, `present_file`, `view_image`, and `tool_search` results; before model calls it derives a hidden receipt ledger (display ids r1..rN) from message state, and when the 2,000-character budget is exceeded the newest receipts are retained in chronological order with their original ids plus an older-receipts omission marker. Rendering returns both the text and its retained receipt subset; every response that received a ledger carries only that exact server-owned subset, never omitted receipts. Snapshot validation accepts a strictly consecutive positive original-id range (for example `r24`–`r30`) rather than requiring `r1`, so subagent terminal citation verification resolves ids against evidence present in the citing turn even when later summarization drops and renumbers tool messages. Model-generated citation IDs are digit-bounded before integer conversion; oversized IDs are ignored as malformed input rather than raising through task write-back. Citation parsing deduplicates exact `(id, anchor)` pairs, not IDs alone, so repeated identical references stay compact while every distinct anchor claim is verified. Gateway strips delegated receipts/verdicts from external messages. `ToolErrorHandlingMiddleware` receives `AppConfig`, converts tool exceptions into error `ToolMessage`s so the run can continue instead of aborting, stamps every result with `deerflow_tool_meta` (status / error_type / recoverable_by_model / recommended_next_action / source) via `tool_result_meta.normalize_tool_result`, stamps structured metadata for task exception wrappers, and stamps skill-read metadata for downstream durable-context capture. Task tool result text is generated from the same status/result/error inputs as the structured metadata so callers do not hand-write a second protocol string.
|
||||
@ -102,4 +102,4 @@ Before changing a later authorization phase, read the [authorization RFC](../../
|
||||
32. **TerminalResponseMiddleware** - When a provider returns an empty terminal `AIMessage` after tool execution, injects a hidden recovery prompt and retries the model once; a second empty response is replaced in checkpoint state by a visible error fallback marked for the run worker, so the run finishes as an error instead of a silent success
|
||||
33. **ModelLengthFinishReasonMiddleware** - Records `stop_reason=model_length_capped` when provider-specific length detectors match a terminal `AIMessage` without tool-call intent (`finish_reason=length` / `MAX_TOKENS`, or `stop_reason=max_tokens`), preserving the original assistant content and never reparsing textual tool-call-like envelopes
|
||||
34. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first
|
||||
35. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). In `after_model` it drops sibling tool calls from the same turn so they cannot execute before the user answers; a malformed `ask_clarification` that LangChain parked on `invalid_tool_calls` is the same stop signal (the valid sibling would otherwise still run). `disable_clarification` runs keep the siblings. Payloads are versioned: legacy modes (`free_text` / `choice_with_other`) keep `version: 1` unchanged, while the v2 `form` mode (from `fields`) carries `version: 2` so older frontends reject the payload and degrade to the plain-text fallback. Field normalization is deterministic and lives in the middleware, not the tool schema — the middleware short-circuits before tool execution, so tool-arg typing alone provides no runtime validation. Validation is atomic: any structurally broken entry (non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member like `__proto__`/`constructor`, exceeding the caps of 16 fields / 24 options per field / 200 chars per text, or the whole normalized definition exceeding `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8 — the per-item caps alone admit forms whose IM text fallback would blow channel delivery limits and truncate away trailing fields) degrades the whole form to the legacy option/free-text modes, so a card can never render "complete" while silently missing a business field; benign issues keep local degradation (unknown types — including unhashable JSON like `type: []`, which must never raise from the membership probe — and option-less selects become `text`), and options are trimmed/deduped with blanks dropped (both form-level and top-level) because the frontend parser rejects blank option labels. Model-produced XML-to-dict option payloads are recursively flattened from dict/list containers in source order, scalar string/number leaves are retained, and residual XML tags are removed before the same trimming and deduplication. Checkbox fields are booleans that default to an explicit "no"; `required` on a checkbox means must-agree/consent semantics. The response protocol is deliberately unchanged (v1 `text`/`option` only): form cards submit a readable text summary as `response_kind: "text"`, so journal persistence and answered-card recovery need no new allowlist entries. Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. That reconciliation is **not** limited to `ask_clarification`: every middleware that answers a tool call itself has the same gap, and a result the user saw during the run must not disappear on reload (#4666 — `ReadBeforeWriteMiddleware`'s blocked-write errors were reaching the UI but never the event store). Its scope is bounded by three independent conditions instead of a tool-name allowlist — the message must be user-visible, the call must belong to this run's **lead agent** (`_remember_current_run_tool_calls` records lead-agent calls only, so subagent results stay in their own `subagent.step` feed), and it must not already be persisted. Keep those three; they are what makes a name allowlist unnecessary. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context.
|
||||
35. **ClarificationMiddleware** - Intercepts `ask_clarification`, writes a readable `ToolMessage.content` fallback plus a structured `ToolMessage.artifact.human_input` payload, and interrupts via `Command(goto=END)` (must be last). `after_model` drops same-turn sibling tool calls so they cannot run before the user answers; a malformed `ask_clarification` parked on `invalid_tool_calls` is the same stop signal. `disable_clarification` runs keep the siblings. Payloads are versioned — legacy `free_text`/`choice_with_other` stay `version: 1`; the v2 `form` mode (from `fields`) is `version: 2` so older frontends reject it and fall back to plain text. Field normalization is deterministic and lives in the middleware (it short-circuits before tool execution, so tool-arg typing gives no runtime validation), and it is atomic: any structurally broken entry — non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member (`__proto__`/`constructor`), or exceeding the caps (16 fields / 24 options per field / 200 chars per text / `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8, the per-item caps alone admitting forms whose IM text fallback overruns channel limits) — degrades the whole form to the legacy option/free-text modes, so a card never renders "complete" while missing a field. Benign issues degrade locally (unknown types — incl. unhashable JSON like `type: []`, which must not raise from the membership probe — and option-less selects become `text`); options are trimmed/deduped with blanks dropped (form- and top-level) since the frontend rejects blank labels. XML-to-dict option payloads are recursively flattened from dict/list containers in source order, scalar leaves kept, residual XML tags stripped before that trimming. Checkboxes are booleans defaulting to "no"; `required` on one means consent semantics. The response protocol is unchanged (v1 `text`/`option`): form cards submit a text summary as `response_kind: "text"`, so journal persistence needs no new allowlist entries. Because this middleware can short-circuit before `on_tool_end`, `RunJournal` does a root-run reconciliation for `ToolMessage`s whose `tool_call_id` came from the current run, so cards survive checkpoint compaction. That reconciliation is **not** `ask_clarification`-only — any middleware that answers a tool call has the same gap, and a result the user saw must not vanish on reload (#4666 — `ReadBeforeWriteMiddleware` blocked-write errors reached the UI but not the event store). It is bounded by three conditions, not a name allowlist: the message is user-visible, the call belongs to this run's **lead agent** (`_remember_current_run_tool_calls` records lead-agent calls only; subagent results stay in `subagent.step`), and it is not already persisted. Human Input Card replies are `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden sources (currently `ask_clarification`) as `llm.human.input`.
|
||||
|
||||
@ -86,6 +86,8 @@ on installs that never enabled it. The convention is:
|
||||
- `migrations/versions/0014_managed_subagents.py` — creates the deployment-level managed Subagent catalog table
|
||||
- `migrations/versions/0015_scheduled_task_enqueue.py` — interrupts legacy transient queued rows, adds durable scheduled-run launch leases and attempt counts, expands the one-active-occurrence index to `queued`/`launching`/`running`, and migrates the overlap policy from `skip` to `enqueue`; chains after `0014_managed_subagents`
|
||||
- `migrations/versions/0016_subagent_batches.py` — creates durable native-subagent batch and item tables, including owner/submission idempotency, item identity, lease/recovery state, and result fields
|
||||
- `migrations/versions/0017_personal_access_tokens.py` — creates the personal access token table for programmatic API access
|
||||
- `migrations/versions/0018_oauth_identity_pg_partial.py` — converts `idx_users_oauth_identity` to a partial index on Postgres (`postgresql_where`), matching what `UserRow.__table_args__` already builds via `create_all`; `0001_baseline` never applied the predicate on Postgres, so every `alembic upgrade head`-provisioned deployment carried a full index until this revision. Postgres-only, idempotent (checks `pg_index.indpred` directly), no-op on SQLite (already partial via `sqlite_where`) and on a DB where the index doesn't exist yet. Originally generated as 0017 and renumbered to 0018 after 0017_personal_access_tokens merged first and kept that slot
|
||||
- `persistence/bootstrap.py` — `bootstrap_schema(engine, backend=...)`, the three-branch decision + locking
|
||||
- `extensions/loader.py::load_extensions` — registers each spec's `table_prefix` with `register_extension_table_prefix()`
|
||||
- Tests: `tests/test_persistence_bootstrap.py` (branches), `tests/test_persistence_bootstrap_concurrency.py` (concurrency), `tests/test_persistence_bootstrap_regression.py` (issue #3682), `tests/test_persistence_migrations_env.py` (filter, including extension-owned tables), `tests/test_extension_loader.py::TestTablePrefixRegistration` (spec-to-filter wiring), `tests/blocking_io/test_persistence_bootstrap.py` (asyncio.to_thread anchor), `tests/test_migration_0004_run_ownership_dedupe.py` + `tests/test_migration_0007_scheduled_run_active_dedupe.py` (dedupe-before-unique-index pre-steps)
|
||||
|
||||
@ -0,0 +1,72 @@
|
||||
"""partial postgres predicate for idx_users_oauth_identity.
|
||||
|
||||
Revision ID: 0018_oauth_identity_pg_partial
|
||||
Revises: 0017_personal_access_tokens
|
||||
Create Date: 2026-08-29
|
||||
|
||||
alembic_version.version_num is VARCHAR(32); revision ids in this chain must
|
||||
stay at or under that length or stamping/upgrading a database fails outright.
|
||||
|
||||
Numbering note: originally generated as 0017 against the then-current main
|
||||
head (0016), same as 0017_personal_access_tokens (#5041). That one merged
|
||||
first and kept the slot per this package's own renumber-on-rebase
|
||||
convention, so this revision moved to 0018 and re-parented onto it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision: str = "0018_oauth_identity_pg_partial"
|
||||
down_revision: str | Sequence[str] | None = "0017_personal_access_tokens"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_INDEX_NAME = "idx_users_oauth_identity"
|
||||
_WHERE = sa.text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL")
|
||||
|
||||
# 0001_baseline's create_index call for this index passed sqlite_where but
|
||||
# not postgresql_where, so every deployment provisioned via
|
||||
# `alembic upgrade head` (the only path production actually runs, per this
|
||||
# package's own AGENTS.md) has idx_users_oauth_identity as a FULL
|
||||
# (non-partial) unique index on Postgres, even after UserRow.__table_args__
|
||||
# gained postgresql_where -- ORM metadata only affects fresh create_all
|
||||
# databases, never an already-versioned one. Not a correctness bug (NULL is
|
||||
# never equal to NULL in a unique index on either backend, so real
|
||||
# duplicates are already rejected and unlimited NULL/NULL rows already
|
||||
# coexist) -- see UserRow.__table_args__'s own comment -- but the index
|
||||
# stays full-table-sized instead of covering only the OAuth-linked rows.
|
||||
|
||||
|
||||
def _pg_index_missing_predicate(bind) -> bool:
|
||||
"""True only if the index exists on Postgres as a full (non-partial)
|
||||
index. False if it doesn't exist yet (a fresh create_all-provisioned DB
|
||||
already built the partial form from ORM metadata -- nothing to do) or
|
||||
already has a predicate (this revision already ran, or a legacy DB was
|
||||
stamped past it)."""
|
||||
row = bind.execute(sa.text(f"SELECT indpred FROM pg_index WHERE indexrelid = to_regclass('{_INDEX_NAME}')")).first()
|
||||
return row is not None and row[0] is None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
if bind.dialect.name != "postgresql":
|
||||
return # SQLite already got sqlite_where from 0001_baseline; nothing for this revision to do there.
|
||||
if not _pg_index_missing_predicate(bind):
|
||||
return
|
||||
op.drop_index(_INDEX_NAME, table_name="users")
|
||||
op.create_index(_INDEX_NAME, "users", ["oauth_provider", "oauth_id"], unique=True, postgresql_where=_WHERE)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
if bind.dialect.name != "postgresql":
|
||||
return
|
||||
row = bind.execute(sa.text(f"SELECT indpred FROM pg_index WHERE indexrelid = to_regclass('{_INDEX_NAME}')")).first()
|
||||
if row is None or row[0] is None:
|
||||
return # doesn't exist, or already the pre-migration full form -- nothing to revert
|
||||
op.drop_index(_INDEX_NAME, table_name="users")
|
||||
op.create_index(_INDEX_NAME, "users", ["oauth_provider", "oauth_id"], unique=True)
|
||||
@ -18,6 +18,17 @@ from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from deerflow.persistence.base import Base
|
||||
|
||||
# Single source of truth for the index name, shared with
|
||||
# app.gateway.auth.repositories.sqlite._is_oauth_identity_violation (which
|
||||
# has to match a Postgres driver error's constraint_name against exactly
|
||||
# this string). Previously that name was a separate hardcoded literal
|
||||
# there, with no test catching the two drifting apart. Migration files
|
||||
# under persistence/migrations/versions/ intentionally do NOT import this
|
||||
# -- migrations are frozen historical DDL, not a live view of the model --
|
||||
# so 0018_oauth_identity_pg_partial.py keeps its own literal by
|
||||
# convention (consistent with every other revision in that package).
|
||||
OAUTH_IDENTITY_INDEX_NAME = "idx_users_oauth_identity"
|
||||
|
||||
|
||||
class UserRow(Base):
|
||||
__tablename__ = "users"
|
||||
@ -49,11 +60,28 @@ class UserRow(Base):
|
||||
token_version: Mapped[int] = mapped_column(nullable=False, default=0)
|
||||
|
||||
__table_args__ = (
|
||||
# sqlite_where alone is a SQLAlchemy dialect-specific kwarg -- it
|
||||
# does not apply on the postgresql dialect, so a table created
|
||||
# against Postgres with only sqlite_where builds a FULL
|
||||
# (non-partial) unique index instead of a partial one. This is
|
||||
# NOT a correctness bug: verified empirically against a live
|
||||
# Postgres instance that a full index already enforces the
|
||||
# intended semantics on its own -- Postgres (like SQLite) treats
|
||||
# NULL as never-equal-to-NULL in a unique index, so real
|
||||
# (provider, id) duplicates are already rejected and unlimited
|
||||
# (NULL, NULL) rows are already allowed, per standard SQL NULL
|
||||
# handling, independent of the WHERE predicate. postgresql_where
|
||||
# is added for two smaller, real reasons instead: (1) it matches
|
||||
# the comment above literally (a partial index, on both
|
||||
# backends), and (2) a partial index only indexes non-NULL rows,
|
||||
# so it stays smaller and cheaper to maintain as the
|
||||
# plain-password-account rows (the common case) accumulate.
|
||||
Index(
|
||||
"idx_users_oauth_identity",
|
||||
OAUTH_IDENTITY_INDEX_NAME,
|
||||
"oauth_provider",
|
||||
"oauth_id",
|
||||
unique=True,
|
||||
sqlite_where=text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL"),
|
||||
postgresql_where=text("oauth_provider IS NOT NULL AND oauth_id IS NOT NULL"),
|
||||
),
|
||||
)
|
||||
|
||||
330
backend/scripts/benchmark/concurrency/run_concurrency_bench.py
Executable file
330
backend/scripts/benchmark/concurrency/run_concurrency_bench.py
Executable file
@ -0,0 +1,330 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Real concurrency benchmark: N SEPARATE OS processes (subprocess.Popen,
|
||||
not asyncio.gather, not threading) hitting the SAME users table at the same
|
||||
time, comparing SQLite vs Postgres at 2/4/8/16 workers.
|
||||
|
||||
This tests exactly the scenario DeerFlow's own docs describe
|
||||
(CONFIGURATION.md line 325): "Multi-worker deployments (GATEWAY_WORKERS > 1)
|
||||
must use the Postgres database backend... SQLite silently ignores row-level
|
||||
locks" -- multiple Gateway PROCESSES, each with its own connection, not
|
||||
multiple async tasks inside ONE process (which a prior single-process
|
||||
benchmark already showed has no problem).
|
||||
|
||||
Usage:
|
||||
uv run python scripts/benchmark/concurrency/run_concurrency_bench.py \
|
||||
--backend sqlite --workers 2,4,8,16 --ops-per-worker 50 --read-ratio 0.7
|
||||
|
||||
uv run python scripts/benchmark/concurrency/run_concurrency_bench.py \
|
||||
--backend postgres --workers 2,4,8,16 --ops-per-worker 50 --read-ratio 0.7 \
|
||||
--pg-url postgresql+asyncpg://deerflow_test:deerflow_test_pw@localhost/deerflow_test
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
# scripts/benchmark/concurrency/run_concurrency_bench.py -> backend/ is 3
|
||||
# levels up (concurrency -> benchmark -> scripts -> backend). Derived from
|
||||
# this file's own location, not hard-coded, so the documented
|
||||
# `uv run python scripts/benchmark/concurrency/run_concurrency_bench.py`
|
||||
# command works from any checkout, not just one at a specific fixed path.
|
||||
BACKEND_DIR = Path(__file__).resolve().parents[3]
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
# checkpoint_bench_common.py is a sibling script folder, not a package (see
|
||||
# its own docstring) -- same sys.path-insert-then-import pattern
|
||||
# bench_channels.py/bench_production.py already use for it, reused here so
|
||||
# percentile() has one correct implementation instead of two.
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "checkpoint"))
|
||||
from checkpoint_bench_common import percentile # noqa: E402
|
||||
|
||||
WORKER_SCRIPT = Path(__file__).parent / "worker.py"
|
||||
# One absolute path, shared by the seeder (this file) and every worker
|
||||
# process (worker.py's make_session_factory). DatabaseConfig.sqlite_dir
|
||||
# resolves relative strings against the CALLER's CWD, not this file's
|
||||
# location -- passing the literal ".deer-flow/bench_data" meant the
|
||||
# orchestrator (running from wherever it was invoked) and the workers
|
||||
# (spawned with cwd=BACKEND_DIR) could silently resolve to two different
|
||||
# directories whenever this script is invoked from outside backend/,
|
||||
# leaving workers pointed at a DB the seeder never created (or already
|
||||
# removed).
|
||||
SQLITE_BENCH_DIR = str(BACKEND_DIR / ".deer-flow" / "bench_data")
|
||||
# The orchestrator itself is already running under the correct interpreter
|
||||
# (`uv run python ...`, per this file's own usage docstring above) -- reuse
|
||||
# it for workers instead of a second hard-coded venv path that silently
|
||||
# assumes deer-flow is checked out at /opt/deer-flow.
|
||||
PYTHON = [sys.executable]
|
||||
|
||||
|
||||
async def seed_baseline(backend: str, pg_url: str, pg_schema: str, n_users: int = 100) -> list[str]:
|
||||
"""Populate a known baseline BEFORE the concurrent run starts -- worker
|
||||
reads target these known emails (not ones the workers themselves are
|
||||
creating), so read and write paths don't depend on each other within
|
||||
the same run."""
|
||||
from app.gateway.auth.models import User
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
||||
|
||||
if backend == "sqlite":
|
||||
sqlite_dir = Path(SQLITE_BENCH_DIR)
|
||||
if sqlite_dir.exists():
|
||||
shutil.rmtree(sqlite_dir)
|
||||
cfg = DatabaseConfig(backend="sqlite", sqlite_dir=SQLITE_BENCH_DIR)
|
||||
else:
|
||||
# pg_schema is a unique, disposable schema for this benchmark run
|
||||
# (see main()) -- never "public" or any schema a real deployment
|
||||
# might already be using. init_engine_from_config creates it
|
||||
# automatically and pins search_path to it, so every statement
|
||||
# below (including the DELETE re-seed on repeat worker-count
|
||||
# sweeps) is scoped to this run's own throwaway namespace.
|
||||
cfg = DatabaseConfig(backend="postgres", postgres_url=pg_url, postgres_schema=pg_schema)
|
||||
|
||||
await init_engine_from_config(cfg)
|
||||
sf = get_session_factory()
|
||||
repo = SQLiteUserRepository(sf)
|
||||
if backend == "postgres":
|
||||
# clean prior worker-count sweep's rows so they don't accumulate
|
||||
# across iterations WITHIN this run -- safe here specifically
|
||||
# because it's scoped (via search_path) to this run's own isolated
|
||||
# schema, never a shared/production one.
|
||||
from sqlalchemy import text
|
||||
|
||||
from deerflow.persistence.engine import get_engine
|
||||
|
||||
engine = get_engine()
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text("DELETE FROM users"))
|
||||
emails = []
|
||||
for i in range(n_users):
|
||||
email = f"baseline_{i}@conc-bench-teste.com"
|
||||
u = User(id=uuid4(), email=email, password_hash="h", system_role="user", created_at=datetime.now(UTC), oauth_provider=None, oauth_id=None, needs_setup=False, token_version=0)
|
||||
await repo.create_user(u)
|
||||
emails.append(email)
|
||||
await close_engine()
|
||||
return emails
|
||||
|
||||
|
||||
async def drop_isolated_schema(pg_url: str, pg_schema: str) -> None:
|
||||
"""Drop this run's disposable Postgres schema (and everything in it) once
|
||||
every worker-count sweep has finished. Only ever targets the unique
|
||||
per-run schema main() generated -- never "public" or a caller-supplied
|
||||
name, so there's nothing here that can reach into a real deployment's
|
||||
namespace even if --pg-url points at one."""
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
|
||||
cfg = DatabaseConfig(backend="postgres", postgres_url=pg_url, postgres_schema=pg_schema)
|
||||
engine = create_async_engine(cfg.app_sqlalchemy_url)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{pg_schema}" CASCADE'))
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
def run_workers(backend: str, n_workers: int, ops_per_worker: int, read_ratio: float, known_emails: list[str], pg_url: str, pg_schema: str):
|
||||
emails_arg = ",".join(known_emails)
|
||||
procs = [] # list of (worker_id, Popen) -- worker_id kept alongside so a
|
||||
# crashed worker's diagnostics can be attributed to the right id below,
|
||||
# instead of the placeholder "worker_id": None every crash used to get.
|
||||
for wid in range(n_workers):
|
||||
cmd = PYTHON + [str(WORKER_SCRIPT), backend, str(wid), str(ops_per_worker), str(read_ratio), emails_arg, pg_url, pg_schema]
|
||||
p = subprocess.Popen(cmd, cwd=str(BACKEND_DIR), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
procs.append((wid, p))
|
||||
|
||||
# Start barrier: wait for every worker to print READY (connection
|
||||
# established, right before its own timed loop -- see worker.py) before
|
||||
# releasing any of them. Otherwise early workers run ahead of ones still
|
||||
# starting -- not a controlled N-worker contention measurement. Each
|
||||
# worker times its own operation phase (from GO to its last op) and the
|
||||
# throughput window is the max of those (see summarize); the staggered
|
||||
# Python-startup + connection cost stays out of it (conn_time_s measures
|
||||
# that per-worker). A worker that crashes before printing READY closes
|
||||
# its stdout, so readline() returns "" rather than hanging; the loop
|
||||
# below reports that same worker as crashed via its nonzero returncode.
|
||||
for _wid, p in procs:
|
||||
p.stdout.readline()
|
||||
|
||||
for _wid, p in procs:
|
||||
try:
|
||||
p.stdin.write("GO\n")
|
||||
p.stdin.flush()
|
||||
except (BrokenPipeError, ValueError):
|
||||
pass # worker already exited -- nothing to release
|
||||
# Do NOT close p.stdin here: p.communicate() below flushes and closes
|
||||
# it, and a second close turns that flush into an uncaught ValueError
|
||||
# ("I/O operation on closed file"). The worker reads exactly one line
|
||||
# (the GO above), so the flush is all the release it needs.
|
||||
|
||||
worker_outputs = []
|
||||
for wid, p in procs:
|
||||
stdout, stderr = p.communicate()
|
||||
if p.returncode != 0:
|
||||
# Surfaced immediately (not just embedded in the summary JSON)
|
||||
# so a crash is visible in real time, not just discoverable by
|
||||
# someone reading crashed_workers back out of the final report.
|
||||
print(f"--- worker {wid} crashed (exit {p.returncode}): {stderr[-2000:]} ---", file=sys.stderr)
|
||||
worker_outputs.append({"worker_id": wid, "crashed": True, "stderr": stderr[-2000:], "results": []})
|
||||
continue
|
||||
try:
|
||||
worker_outputs.append(json.loads(stdout.strip().splitlines()[-1]))
|
||||
except Exception as e:
|
||||
msg = f"parse error: {e}; stdout={stdout[-500:]}; stderr={stderr[-500:]}"
|
||||
print(f"--- worker {wid} produced unparseable output: {msg} ---", file=sys.stderr)
|
||||
worker_outputs.append({"worker_id": wid, "crashed": True, "stderr": msg, "results": []})
|
||||
return worker_outputs
|
||||
|
||||
|
||||
def summarize(worker_outputs, n_workers: int, ops_per_worker: int) -> dict:
|
||||
all_results = []
|
||||
crashed = 0
|
||||
crashed_worker_errors = []
|
||||
ops_windows = []
|
||||
for w in worker_outputs:
|
||||
if w.get("crashed"):
|
||||
crashed += 1
|
||||
crashed_worker_errors.append({"worker_id": w.get("worker_id"), "stderr": w.get("stderr")})
|
||||
continue
|
||||
all_results.extend(w["results"])
|
||||
if "ops_elapsed_s" in w:
|
||||
ops_windows.append(w["ops_elapsed_s"])
|
||||
|
||||
# All non-crashed workers are released by the same GO, so the slowest
|
||||
# worker's operation-phase elapsed is the window during which every
|
||||
# worker was contending. Using it (not the orchestrator's post-
|
||||
# communicate() wall clock) keeps per-worker engine.dispose() +
|
||||
# result serialization + stdout transfer out of the throughput figure.
|
||||
ops_window = max(ops_windows) if ops_windows else 0.0
|
||||
|
||||
total_ops = len(all_results)
|
||||
errors = [r for r in all_results if not r["ok"]]
|
||||
latencies = sorted(r["latency_s"] for r in all_results)
|
||||
err_types = {}
|
||||
for r in errors:
|
||||
key = r["err"].split(":")[0] if r["err"] else "unknown"
|
||||
err_types[key] = err_types.get(key, 0) + 1
|
||||
|
||||
def pct(p):
|
||||
# p is a 0..1 fraction here (0.50/0.95/0.99); checkpoint_bench_common's
|
||||
# percentile() takes 0..100 and already does correct nearest-rank
|
||||
# interpolation ((n-1)*percentile/100, not int(n*p) used directly as
|
||||
# an index -- that off-by-one made p95/p99 both resolve to the max
|
||||
# for any sample of 20 or fewer values, and for the documented
|
||||
# 100-sample default).
|
||||
if not latencies:
|
||||
return None
|
||||
return percentile(latencies, p * 100)
|
||||
|
||||
return {
|
||||
"n_workers": n_workers,
|
||||
"ops_per_worker": ops_per_worker,
|
||||
"expected_total_ops": n_workers * ops_per_worker,
|
||||
"completed_ops": total_ops,
|
||||
"crashed_workers": crashed,
|
||||
"crashed_worker_errors": crashed_worker_errors,
|
||||
"errors": len(errors),
|
||||
"error_types": err_types,
|
||||
"ops_window_s": round(ops_window, 3),
|
||||
"throughput_ops_per_s": round(total_ops / ops_window, 2) if ops_window > 0 else None,
|
||||
"latency_p50_ms": round(pct(0.50) * 1000, 3) if pct(0.50) is not None else None,
|
||||
"latency_p95_ms": round(pct(0.95) * 1000, 3) if pct(0.95) is not None else None,
|
||||
"latency_p99_ms": round(pct(0.99) * 1000, 3) if pct(0.99) is not None else None,
|
||||
"latency_max_ms": round(latencies[-1] * 1000, 3) if latencies else None,
|
||||
}
|
||||
|
||||
|
||||
def summary_indicates_failure(summary: dict) -> bool:
|
||||
"""True if this worker-count sweep's summary represents a broken run,
|
||||
not a real measurement:
|
||||
|
||||
- a crashed worker, or fewer completed results than expected without a
|
||||
crash -- an all-crashed sweep still produces a well-formed-looking
|
||||
summary (crashed_workers: N, completed_ops: 0,
|
||||
throughput_ops_per_s: 0.0);
|
||||
- any failed op (errors > 0). This benchmark's conclusions rest on
|
||||
"0 errors on both backends"; a sweep where every op completed but
|
||||
raised (e.g. writes hitting OperationalError) has completed_ops ==
|
||||
expected and 0 crashes, so it would otherwise pass as a clean
|
||||
measurement. The error breakdown stays in the printed JSON either
|
||||
way -- this only stops the exit code from calling it clean."""
|
||||
return summary["crashed_workers"] > 0 or summary["completed_ops"] != summary["expected_total_ops"] or summary.get("errors", 0) > 0
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--backend", required=True, choices=["sqlite", "postgres"])
|
||||
ap.add_argument("--workers", required=True, help="comma-separated list, e.g. 2,4,8,16")
|
||||
ap.add_argument("--ops-per-worker", type=int, default=50)
|
||||
ap.add_argument("--read-ratio", type=float, default=0.7)
|
||||
ap.add_argument("--pg-url", default="")
|
||||
ap.add_argument("--baseline-users", type=int, default=100)
|
||||
ap.add_argument("--out", default=None)
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.backend == "postgres" and not args.pg_url:
|
||||
ap.error("--pg-url is required for --backend postgres")
|
||||
|
||||
worker_counts = [int(x) for x in args.workers.split(",")]
|
||||
all_summaries = []
|
||||
# A sweep where every worker crashed still produces a well-formed
|
||||
# summary (crashed_workers, completed_ops: 0, throughput_ops_per_s:
|
||||
# 0.0) -- exiting 0 for that made a broken run indistinguishable from
|
||||
# a real (if uneventful) measurement to anything checking the exit
|
||||
# code, and let a garbage --out file sit next to a real one the same
|
||||
# way. Tracked across the whole worker-count sweep, not just the last
|
||||
# iteration, so one bad n_workers value in the middle doesn't get
|
||||
# masked by later ones succeeding.
|
||||
had_failure = False
|
||||
|
||||
# One disposable schema for this ENTIRE invocation (reused across the
|
||||
# worker-count sweep below, dropped once at the very end) -- --pg-url
|
||||
# accepts an arbitrary database URL, so this must never touch "public"
|
||||
# or any namespace a real deployment might be using. Unset for sqlite;
|
||||
# seed_baseline/run_workers/worker.py ignore it on that backend.
|
||||
pg_schema = f"bench_{uuid4().hex[:12]}" if args.backend == "postgres" else ""
|
||||
|
||||
try:
|
||||
for n_workers in worker_counts:
|
||||
print(f"--- seeding baseline ({args.backend}, {args.baseline_users} users{f', schema={pg_schema}' if pg_schema else ''}) ---", file=sys.stderr)
|
||||
emails = asyncio.run(seed_baseline(args.backend, args.pg_url, pg_schema, args.baseline_users))
|
||||
|
||||
print(f"--- running {n_workers} workers ({args.backend}, {args.ops_per_worker} ops/worker) ---", file=sys.stderr)
|
||||
worker_outputs = run_workers(args.backend, n_workers, args.ops_per_worker, args.read_ratio, emails, args.pg_url, pg_schema)
|
||||
summary = summarize(worker_outputs, n_workers, args.ops_per_worker)
|
||||
summary["backend"] = args.backend
|
||||
all_summaries.append(summary)
|
||||
print(json.dumps(summary, indent=2), file=sys.stderr)
|
||||
if summary_indicates_failure(summary):
|
||||
had_failure = True
|
||||
finally:
|
||||
if pg_schema:
|
||||
print(f"--- dropping isolated schema {pg_schema} ---", file=sys.stderr)
|
||||
asyncio.run(drop_isolated_schema(args.pg_url, pg_schema))
|
||||
|
||||
result = {"backend": args.backend, "read_ratio": args.read_ratio, "runs": all_summaries}
|
||||
output = json.dumps(result, indent=2)
|
||||
if args.out:
|
||||
Path(args.out).write_text(output)
|
||||
# Printed unconditionally, failure or not -- a broken run's diagnostics
|
||||
# (crashed_worker_errors, the mismatched op counts) are exactly what's
|
||||
# needed to debug it, so the JSON goes out before the exit code below
|
||||
# can make anything piping/discarding stdout on a nonzero exit lose it.
|
||||
print(output)
|
||||
|
||||
if had_failure:
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
221
backend/scripts/benchmark/concurrency/worker.py
Normal file
221
backend/scripts/benchmark/concurrency/worker.py
Normal file
@ -0,0 +1,221 @@
|
||||
#!/usr/bin/env python3
|
||||
"""A worker running as its own separate OS process -- its own engine/connection,
|
||||
sharing nothing with the parent process or other workers, to simulate real
|
||||
Gateway workers (separate processes, not asyncio tasks inside one process).
|
||||
|
||||
Connects DIRECTLY via SQLAlchemy (bypassing init_engine_from_config's
|
||||
Alembic schema-state bootstrap, which costs ~8.5s per call regardless of
|
||||
backend -- a real, separately-disclosed cost, but not what this benchmark
|
||||
measures. The schema is already bootstrapped once by the orchestrator's
|
||||
seed_baseline() before any worker starts, so a worker attaching directly is
|
||||
exactly what a warm Gateway worker process does after its own one-time
|
||||
startup, and isolates DB lock/throughput behavior from Python/import
|
||||
cold-start cost).
|
||||
|
||||
Each worker does a fixed mix of reads (get_user_by_email) and writes
|
||||
(create_user) against the SAME shared users table, and prints one JSON line
|
||||
to stdout (latency + success/error per op), collected by the orchestrator
|
||||
(run_concurrency_bench.py) afterward.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
# scripts/benchmark/concurrency/worker.py -> backend/ is 3 levels up, same
|
||||
# derivation as run_concurrency_bench.py (which spawns this file as a
|
||||
# subprocess with cwd already set to BACKEND_DIR, but this file is also
|
||||
# runnable/importable on its own, so it derives its own sys.path entry
|
||||
# rather than relying on the parent's cwd).
|
||||
BACKEND_DIR = Path(__file__).resolve().parents[3]
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
from sqlalchemy import event, text # noqa: E402
|
||||
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine # noqa: E402
|
||||
|
||||
from app.gateway.auth.models import User # noqa: E402
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository # noqa: E402
|
||||
from deerflow.config.database_config import DatabaseConfig # noqa: E402
|
||||
|
||||
# Must be the exact same absolute path run_concurrency_bench.py's
|
||||
# seed_baseline() computes (SQLITE_BENCH_DIR there) -- DatabaseConfig
|
||||
# resolves a relative sqlite_dir against the CALLER's CWD, so a shared
|
||||
# relative literal here and there silently pointed the seeder and the
|
||||
# workers at different directories whenever this script is invoked from
|
||||
# outside backend/ (the seeder ran in-process from the invoker's own CWD;
|
||||
# workers are spawned with cwd=BACKEND_DIR, which don't necessarily match).
|
||||
SQLITE_BENCH_DIR = str(BACKEND_DIR / ".deer-flow" / "bench_data")
|
||||
|
||||
# The exact per-connection PRAGMAs the app sets on every SQLite connection
|
||||
# (deerflow/persistence/engine.py::_enable_sqlite_wal). journal_mode is
|
||||
# persistent so WAL would be picked up incidentally from the seeder's engine,
|
||||
# but synchronous and foreign_keys are per-connection: without this a worker
|
||||
# runs at SQLite's synchronous=FULL / foreign_keys=OFF defaults and its write
|
||||
# path pays a different (heavier) per-commit fsync cost than the deployment
|
||||
# being modelled. Kept in sync with that listener by hand -- there are only
|
||||
# these four lines and both sites cite each other.
|
||||
_APP_SQLITE_PRAGMAS = (
|
||||
"PRAGMA journal_mode=WAL;",
|
||||
"PRAGMA synchronous=NORMAL;",
|
||||
"PRAGMA foreign_keys=ON;",
|
||||
"PRAGMA busy_timeout=30000;",
|
||||
)
|
||||
|
||||
|
||||
def read_count(n_ops: int, read_ratio: float) -> int:
|
||||
"""Exact number of reads out of n_ops -- round() rather than int()'s
|
||||
truncation-toward-zero, so the documented default (n_ops=50,
|
||||
read_ratio=0.7) yields 35 reads, not 34."""
|
||||
return round(n_ops * read_ratio)
|
||||
|
||||
|
||||
def is_read_op(i: int, n_ops: int, n_reads: int) -> bool:
|
||||
"""Whether op index i (0-based) is a read, given exactly n_reads reads
|
||||
spread evenly across n_ops slots.
|
||||
|
||||
The previous check, `(i % 100) < int(read_ratio * 100)`, assumed n_ops
|
||||
was always >= 100: for the documented/default n_ops=50 (or any n_ops <=
|
||||
100), i % 100 == i, so every op satisfies i < read_ratio*100 up to i=69
|
||||
and every op after that fails it -- meaning the entire 50-op run was
|
||||
either all reads or all writes depending on read_ratio, never the
|
||||
claimed mixed workload.
|
||||
|
||||
This instead uses modular (Bresenham-style) spacing: stepping i*n_reads
|
||||
through n_ops slots visits exactly n_reads distinct residues below
|
||||
n_reads, evenly distributed rather than clustered at the front, and
|
||||
produces precisely n_reads True values across the n_ops calls for any
|
||||
n_ops/n_reads pair.
|
||||
"""
|
||||
if n_ops <= 0 or n_reads <= 0:
|
||||
return False
|
||||
if n_reads >= n_ops:
|
||||
return True
|
||||
return (i * n_reads) % n_ops < n_reads
|
||||
|
||||
|
||||
def make_session_factory(backend: str, pg_url: str, pg_schema: str):
|
||||
"""Build engine + session factory directly, without the Alembic
|
||||
bootstrap dance -- caller guarantees the schema already exists.
|
||||
|
||||
pg_schema is the SAME disposable per-run schema the orchestrator's
|
||||
seed_baseline() already bootstrapped -- never "public" -- so a worker
|
||||
attaching directly lands in the right namespace instead of falling
|
||||
back to whatever the connection's default search_path happens to be.
|
||||
"""
|
||||
if backend == "sqlite":
|
||||
cfg = DatabaseConfig(backend="sqlite", sqlite_dir=SQLITE_BENCH_DIR)
|
||||
url = cfg.app_sqlalchemy_url
|
||||
engine = create_async_engine(url, connect_args={"timeout": 30})
|
||||
|
||||
@event.listens_for(engine.sync_engine, "connect")
|
||||
def _match_app_sqlite_pragmas(dbapi_conn, _record): # noqa: ARG001 — SQLAlchemy contract
|
||||
cursor = dbapi_conn.cursor()
|
||||
try:
|
||||
for pragma in _APP_SQLITE_PRAGMAS:
|
||||
cursor.execute(pragma)
|
||||
finally:
|
||||
cursor.close()
|
||||
else:
|
||||
cfg = DatabaseConfig(backend="postgres", postgres_url=pg_url, postgres_schema=pg_schema)
|
||||
url = cfg.app_sqlalchemy_url
|
||||
engine = create_async_engine(url, connect_args={"server_settings": {"search_path": pg_schema}})
|
||||
return engine, async_sessionmaker(engine, expire_on_commit=False)
|
||||
|
||||
|
||||
async def run_worker(backend: str, worker_id: int, n_ops: int, read_ratio: float, known_emails: list[str], pg_url: str, pg_schema: str):
|
||||
t_conn0 = time.perf_counter()
|
||||
engine, sf = make_session_factory(backend, pg_url, pg_schema)
|
||||
# Force a real physical connection now (not lazy) so conn_time reflects
|
||||
# the actual cost of a worker's first DB round-trip, same as a real
|
||||
# Gateway worker would pay on its first request. Entering an empty
|
||||
# AsyncSession does NOT check out a connection -- SQLAlchemy stays lazy
|
||||
# until the first statement executes -- so this must run an actual
|
||||
# lightweight query, not just `async with sf(): pass`, or the first
|
||||
# timed op in the loop below silently absorbs connection-establishment
|
||||
# cost instead of conn_time (a real distortion at 16 workers: those 16
|
||||
# cold first-ops are 1% of a 1600-op sample and can skew the reported
|
||||
# p99).
|
||||
async with sf() as session:
|
||||
await session.execute(text("SELECT 1"))
|
||||
conn_time = time.perf_counter() - t_conn0
|
||||
repo = SQLiteUserRepository(sf)
|
||||
|
||||
# Signal ready, then block for the orchestrator's start signal, before
|
||||
# touching the operation loop's timer. Without this, each worker starts
|
||||
# its timed loop as soon as ITS OWN imports+connection finish -- so the
|
||||
# orchestrator's wall_time (started before any worker was even spawned)
|
||||
# includes N staggered process-startup costs, and early workers run
|
||||
# ahead of workers that are still starting. This makes every worker
|
||||
# cross the same starting line together, so wall_time measures actual
|
||||
# concurrent execution instead of startup skew (see run_workers() in
|
||||
# run_concurrency_bench.py for the other half of this handshake).
|
||||
print("READY", flush=True)
|
||||
sys.stdin.readline()
|
||||
# Time the operation phase only. The orchestrator's wall clock is sampled
|
||||
# after communicate() returns, so it also covers this worker's
|
||||
# engine.dispose(), result serialization and stdout transfer -- teardown
|
||||
# that isn't contention. The orchestrator uses max(ops_elapsed_s) over
|
||||
# workers (all released by the same GO) as the throughput window instead.
|
||||
t_ops0 = time.perf_counter()
|
||||
|
||||
results = []
|
||||
n_reads = read_count(n_ops, read_ratio)
|
||||
for i in range(n_ops):
|
||||
is_read = is_read_op(i, n_ops, n_reads)
|
||||
t0 = time.perf_counter()
|
||||
ok = True
|
||||
err = None
|
||||
try:
|
||||
if is_read:
|
||||
email = known_emails[(worker_id * n_ops + i) % len(known_emails)]
|
||||
await repo.get_user_by_email(email)
|
||||
else:
|
||||
u = User(
|
||||
id=uuid4(),
|
||||
email=f"bench_w{worker_id}_{i}_{uuid4().hex[:8]}@conc-bench-teste.com",
|
||||
password_hash="h",
|
||||
system_role="user",
|
||||
created_at=datetime.now(UTC),
|
||||
oauth_provider=None,
|
||||
oauth_id=None,
|
||||
needs_setup=False,
|
||||
token_version=0,
|
||||
)
|
||||
await repo.create_user(u)
|
||||
except Exception as e:
|
||||
ok = False
|
||||
err = f"{type(e).__name__}: {str(e)[:200]}"
|
||||
elapsed = time.perf_counter() - t0
|
||||
results.append({"op": "read" if is_read else "write", "ok": ok, "err": err, "latency_s": elapsed})
|
||||
|
||||
ops_elapsed = time.perf_counter() - t_ops0
|
||||
await engine.dispose()
|
||||
return {
|
||||
"worker_id": worker_id,
|
||||
"conn_time_s": conn_time,
|
||||
"ops_elapsed_s": ops_elapsed,
|
||||
"results": results,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
backend = sys.argv[1]
|
||||
worker_id = int(sys.argv[2])
|
||||
n_ops = int(sys.argv[3])
|
||||
read_ratio = float(sys.argv[4])
|
||||
known_emails = sys.argv[5].split(",")
|
||||
pg_url = sys.argv[6] if len(sys.argv) > 6 else ""
|
||||
pg_schema = sys.argv[7] if len(sys.argv) > 7 else ""
|
||||
|
||||
out = asyncio.run(run_worker(backend, worker_id, n_ops, read_ratio, known_emails, pg_url, pg_schema))
|
||||
print(json.dumps(out))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -425,6 +425,263 @@ def test_sqlite_round_trip_new_fields():
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# ── IntegrityError classification (OAuth conflict vs. everything else) ──────
|
||||
#
|
||||
# Regression coverage for a misclassification bug: create_user's
|
||||
# except IntegrityError handler used to substring-match "oauth" against
|
||||
# str(exc) (the SQLAlchemy wrapper), which always contains the failed
|
||||
# INSERT statement's full column list -- including oauth_provider/oauth_id
|
||||
# -- regardless of which constraint actually fired. Fixed to inspect
|
||||
# exc.orig (the driver exception) instead.
|
||||
|
||||
|
||||
def test_create_user_duplicate_primary_key_is_not_misreported_as_oauth(tmp_path):
|
||||
"""A duplicate `id` (primary-key violation, unrelated to OAuth) must be
|
||||
reported as neither an OAuth conflict nor an email conflict: it names
|
||||
neither `oauth` (the false positive a substring check on str(exc)
|
||||
produces) nor `second@test.com` (an address that is not registered)."""
|
||||
import asyncio
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
|
||||
|
||||
async def _run() -> None:
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
|
||||
url = f"sqlite+aiosqlite:///{tmp_path}/scratch.db"
|
||||
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
||||
try:
|
||||
repo = SQLiteUserRepository(get_session_factory())
|
||||
shared_id = uuid4()
|
||||
first = User(id=shared_id, email="first@test.com", password_hash="h", system_role="user")
|
||||
await repo.create_user(first)
|
||||
|
||||
# Same id, different email: an email-uniqueness collision is
|
||||
# ruled out by construction (different email, and the pre-check
|
||||
# would catch a real email dupe first anyway) -- this can only
|
||||
# be the id primary-key constraint, never idx_users_oauth_identity.
|
||||
duplicate_id = User(id=shared_id, email="second@test.com", password_hash="h", system_role="user")
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await repo.create_user(duplicate_id)
|
||||
message = str(exc_info.value)
|
||||
assert "OAuth" not in message, f"primary-key violation misreported as an OAuth conflict: {message}"
|
||||
assert "second@test.com" not in message, f"primary-key violation misreported as an email conflict: {message}"
|
||||
assert "User already exists" in message
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_create_user_duplicate_email_race_still_reports_email(tmp_path):
|
||||
"""An email collision that reaches the DB (the pre-check bypassed to
|
||||
simulate the concurrent-insert race) must still get the email-specific
|
||||
message, not the neutral fallback."""
|
||||
import asyncio
|
||||
|
||||
from app.gateway.auth.repositories import sqlite as sqlite_repo
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
|
||||
|
||||
async def _run() -> None:
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
|
||||
url = f"sqlite+aiosqlite:///{tmp_path}/scratch.db"
|
||||
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
||||
try:
|
||||
repo = SQLiteUserRepository(get_session_factory())
|
||||
await repo.create_user(User(email="race@test.com", password_hash="h", system_role="user"))
|
||||
|
||||
racing = User(email="race@test.com", password_hash="h", system_role="user")
|
||||
with patch.object(sqlite_repo.AsyncSession, "scalar", return_value=None):
|
||||
with pytest.raises(ValueError, match="Email already registered: race@test.com"):
|
||||
await repo.create_user(racing)
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_create_user_propagates_non_uniqueness_integrity_error(tmp_path):
|
||||
"""A NOT NULL / CHECK / FK IntegrityError is not a "user already exists"
|
||||
condition and is not part of create_user's ValueError contract -- it must
|
||||
propagate as-is, not be relabeled "User already exists"."""
|
||||
import asyncio
|
||||
import sqlite3
|
||||
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from app.gateway.auth.repositories import sqlite as sqlite_repo
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
|
||||
|
||||
async def _run() -> None:
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
|
||||
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path}/scratch.db", sqlite_dir=str(tmp_path))
|
||||
try:
|
||||
repo = SQLiteUserRepository(get_session_factory())
|
||||
not_null = IntegrityError(
|
||||
"INSERT INTO users ...",
|
||||
{},
|
||||
orig=sqlite3.IntegrityError("NOT NULL constraint failed: users.system_role"),
|
||||
)
|
||||
with patch.object(sqlite_repo.AsyncSession, "commit", side_effect=not_null):
|
||||
with pytest.raises(IntegrityError):
|
||||
await repo.create_user(User(email="x@test.com", password_hash="h", system_role="user"))
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
def test_create_user_real_oauth_conflict_still_reported_correctly(tmp_path):
|
||||
"""The actual case _is_oauth_identity_violation exists to detect: two
|
||||
users sharing an (oauth_provider, oauth_id) pair must still raise the
|
||||
OAuth-specific message, not just "any IntegrityError"."""
|
||||
import asyncio
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
|
||||
|
||||
async def _run() -> None:
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
||||
|
||||
url = f"sqlite+aiosqlite:///{tmp_path}/scratch.db"
|
||||
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
||||
try:
|
||||
repo = SQLiteUserRepository(get_session_factory())
|
||||
first = User(email="oauth-a@test.com", password_hash=None, system_role="user", oauth_provider="github", oauth_id="dup-123")
|
||||
await repo.create_user(first)
|
||||
|
||||
duplicate = User(email="oauth-b@test.com", password_hash=None, system_role="user", oauth_provider="github", oauth_id="dup-123")
|
||||
with pytest.raises(ValueError, match="OAuth account already linked"):
|
||||
await repo.create_user(duplicate)
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
# The IntegrityError classification helpers are pure functions of the driver
|
||||
# exception -- the end-to-end tests above cover the SQLite branch (real engine,
|
||||
# real IntegrityError). The Postgres/asyncpg branch needs a real Postgres and
|
||||
# its only e2e guard, test_oauth_identity_uniqueness_enforced_end_to_end, is
|
||||
# skipped in CI (no workflow sets DEERFLOW_TEST_POSTGRES_URL). The stubs below
|
||||
# pin it with no DB of either kind.
|
||||
#
|
||||
# The shape matters: SQLAlchemy's asyncpg dialect does NOT hand us the asyncpg
|
||||
# error as `exc.orig`. It re-raises its own AsyncAdapt_asyncpg_dbapi
|
||||
# .IntegrityError (pgcode/sqlstate only) `from` the real asyncpg error, so
|
||||
# `constraint_name` lives on `exc.orig.__cause__`, not `exc.orig`.
|
||||
def _pg_integrity_error(constraint_name: str, sqlstate: str = "23505"):
|
||||
"""A stub in the shape SQLAlchemy's asyncpg dialect actually produces:
|
||||
the `orig` wrapper carries `pgcode`/`sqlstate` but no constraint_name;
|
||||
the real driver error (which does) is its `__cause__`. Default sqlstate
|
||||
23505 = unique_violation."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
wrapper = SimpleNamespace(pgcode=sqlstate, sqlstate=sqlstate)
|
||||
wrapper.__cause__ = SimpleNamespace(constraint_name=constraint_name)
|
||||
return SimpleNamespace(orig=wrapper)
|
||||
|
||||
|
||||
def test_driver_constraint_name_reads_from_asyncpg_cause_chain():
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _driver_constraint_name
|
||||
|
||||
assert _driver_constraint_name(_pg_integrity_error("users_pkey")) == "users_pkey"
|
||||
# No cause, no constraint_name anywhere -> None (SQLite path).
|
||||
assert _driver_constraint_name(SimpleNamespace(orig=SimpleNamespace())) is None
|
||||
|
||||
|
||||
def test_is_oauth_identity_violation_matches_postgres_constraint_name():
|
||||
from app.gateway.auth.repositories.sqlite import _is_oauth_identity_violation
|
||||
from deerflow.persistence.user.model import OAUTH_IDENTITY_INDEX_NAME
|
||||
|
||||
assert _is_oauth_identity_violation(_pg_integrity_error(OAUTH_IDENTITY_INDEX_NAME)) is True
|
||||
|
||||
|
||||
def test_is_oauth_identity_violation_rejects_other_postgres_constraints():
|
||||
"""A primary-key (or any other) constraint name on the SAME table must
|
||||
not be misclassified as the OAuth index -- the Postgres-side equivalent
|
||||
of the SQLite primary-key-violation regression test above."""
|
||||
from app.gateway.auth.repositories.sqlite import _is_oauth_identity_violation
|
||||
|
||||
assert _is_oauth_identity_violation(_pg_integrity_error("users_pkey")) is False
|
||||
|
||||
|
||||
def test_is_oauth_identity_violation_matches_sqlite_message():
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _is_oauth_identity_violation
|
||||
|
||||
orig = sqlite3.IntegrityError("UNIQUE constraint failed: users.oauth_provider, users.oauth_id")
|
||||
exc = SimpleNamespace(orig=orig)
|
||||
assert _is_oauth_identity_violation(exc) is True
|
||||
|
||||
|
||||
def test_is_oauth_identity_violation_rejects_sqlite_email_violation():
|
||||
"""sqlite3 has no constraint_name -- a different UNIQUE violation on
|
||||
the same table (email) must not match on a bare "oauth" substring."""
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _is_oauth_identity_violation
|
||||
|
||||
orig = sqlite3.IntegrityError("UNIQUE constraint failed: users.email")
|
||||
exc = SimpleNamespace(orig=orig)
|
||||
assert _is_oauth_identity_violation(exc) is False
|
||||
|
||||
|
||||
def test_is_email_violation_matches_both_backends():
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _EMAIL_UNIQUE_INDEX_NAME, _is_email_violation
|
||||
|
||||
# email is unique=True + index=True -> a single UNIQUE INDEX, so Postgres
|
||||
# reports the index name (ix_users_email), not a users_email_key constraint.
|
||||
assert _EMAIL_UNIQUE_INDEX_NAME == "ix_users_email"
|
||||
assert _is_email_violation(_pg_integrity_error(_EMAIL_UNIQUE_INDEX_NAME)) is True
|
||||
sqlite = SimpleNamespace(orig=sqlite3.IntegrityError("UNIQUE constraint failed: users.email"))
|
||||
assert _is_email_violation(sqlite) is True
|
||||
|
||||
|
||||
def test_is_email_violation_rejects_primary_key_and_oauth():
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _is_email_violation
|
||||
|
||||
assert _is_email_violation(_pg_integrity_error("users_pkey")) is False
|
||||
assert _is_email_violation(SimpleNamespace(orig=sqlite3.IntegrityError("UNIQUE constraint failed: users.id"))) is False
|
||||
|
||||
|
||||
def test_is_uniqueness_violation_distinguishes_unique_from_not_null_and_check():
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _is_uniqueness_violation
|
||||
|
||||
# Postgres: 23505 unique_violation vs 23502 not_null_violation.
|
||||
assert _is_uniqueness_violation(_pg_integrity_error("ix_users_email", sqlstate="23505")) is True
|
||||
assert _is_uniqueness_violation(_pg_integrity_error("x", sqlstate="23502")) is False
|
||||
# SQLite message forms.
|
||||
assert _is_uniqueness_violation(SimpleNamespace(orig=sqlite3.IntegrityError("UNIQUE constraint failed: users.id"))) is True
|
||||
assert _is_uniqueness_violation(SimpleNamespace(orig=sqlite3.IntegrityError("PRIMARY KEY constraint failed"))) is True
|
||||
assert _is_uniqueness_violation(SimpleNamespace(orig=sqlite3.IntegrityError("NOT NULL constraint failed: users.system_role"))) is False
|
||||
|
||||
|
||||
def test_violated_constraint_extracts_name_from_both_backends():
|
||||
import sqlite3
|
||||
from types import SimpleNamespace
|
||||
|
||||
from app.gateway.auth.repositories.sqlite import _violated_constraint
|
||||
|
||||
assert _violated_constraint(_pg_integrity_error("users_pkey")) == "users_pkey"
|
||||
assert _violated_constraint(SimpleNamespace(orig=sqlite3.IntegrityError("UNIQUE constraint failed: users.id"))) == "users.id"
|
||||
assert _violated_constraint(SimpleNamespace(orig=RuntimeError("opaque driver error"))) is None
|
||||
|
||||
|
||||
def test_update_user_raises_when_row_concurrently_deleted(tmp_path):
|
||||
"""Concurrent-delete during update_user must hard-fail, not silently no-op.
|
||||
|
||||
|
||||
215
backend/tests/test_bench_concurrency.py
Normal file
215
backend/tests/test_bench_concurrency.py
Normal file
@ -0,0 +1,215 @@
|
||||
"""Unit tests for scripts/benchmark/concurrency/run_concurrency_bench.py's
|
||||
pure aggregation logic (percentile/error/crash accounting) -- fast, no DB
|
||||
required, following the same pattern as test_bench_checkpoint_channels.py
|
||||
(load the script as a module, unit-test its helpers directly rather than
|
||||
running the actual multi-process sweep in CI)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_module():
|
||||
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/concurrency/run_concurrency_bench.py"
|
||||
spec = importlib.util.spec_from_file_location("run_concurrency_bench", path)
|
||||
assert spec is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
bench = _load_module()
|
||||
|
||||
|
||||
def _result(ok: bool, latency_s: float, err: str | None = None, op: str = "read") -> dict:
|
||||
return {"op": op, "ok": ok, "err": err, "latency_s": latency_s}
|
||||
|
||||
|
||||
def _worker(worker_id, results: list[dict], ops_elapsed_s: float = 1.0) -> dict:
|
||||
return {"worker_id": worker_id, "ops_elapsed_s": ops_elapsed_s, "results": results}
|
||||
|
||||
|
||||
def _crashed(worker_id, stderr: str = "boom") -> dict:
|
||||
return {"worker_id": worker_id, "crashed": True, "stderr": stderr, "results": []}
|
||||
|
||||
|
||||
def test_summarize_counts_completed_ops_across_workers() -> None:
|
||||
workers = [
|
||||
_worker(0, [_result(True, 0.001), _result(True, 0.002)]),
|
||||
_worker(1, [_result(True, 0.003)]),
|
||||
]
|
||||
summary = bench.summarize(workers, n_workers=2, ops_per_worker=2)
|
||||
assert summary["completed_ops"] == 3
|
||||
assert summary["crashed_workers"] == 0
|
||||
assert summary["errors"] == 0
|
||||
|
||||
|
||||
def test_summarize_separates_crashed_workers_from_completed_ops() -> None:
|
||||
workers = [
|
||||
_worker(0, [_result(True, 0.001)]),
|
||||
_crashed(None),
|
||||
]
|
||||
summary = bench.summarize(workers, n_workers=2, ops_per_worker=5)
|
||||
assert summary["crashed_workers"] == 1
|
||||
assert summary["completed_ops"] == 1
|
||||
assert summary["expected_total_ops"] == 10
|
||||
|
||||
|
||||
def test_summarize_groups_errors_by_exception_type() -> None:
|
||||
workers = [
|
||||
_worker(
|
||||
0,
|
||||
[
|
||||
_result(False, 0.5, err="OperationalError: database is locked"),
|
||||
_result(False, 0.6, err="OperationalError: database is locked"),
|
||||
_result(False, 0.1, err="IntegrityError: duplicate key"),
|
||||
_result(True, 0.001),
|
||||
],
|
||||
)
|
||||
]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=4)
|
||||
assert summary["errors"] == 3
|
||||
assert summary["error_types"] == {"OperationalError": 2, "IntegrityError": 1}
|
||||
|
||||
|
||||
def test_summarize_percentiles_are_monotonic_and_within_observed_range() -> None:
|
||||
"""p50 <= p95 <= p99 <= max always holds for any nonempty, nonnegative
|
||||
latency distribution -- a basic sanity invariant on the aggregation
|
||||
math itself, independent of what a real run happens to produce."""
|
||||
latencies = [0.001 * i for i in range(1, 101)] # 1ms..100ms, evenly spaced
|
||||
workers = [_worker(0, [_result(True, latency) for latency in latencies])]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=100)
|
||||
|
||||
assert summary["latency_p50_ms"] <= summary["latency_p95_ms"]
|
||||
assert summary["latency_p95_ms"] <= summary["latency_p99_ms"]
|
||||
assert summary["latency_p99_ms"] <= summary["latency_max_ms"]
|
||||
assert summary["latency_max_ms"] == 100.0 # the largest input, in ms
|
||||
|
||||
|
||||
def test_summarize_percentiles_match_hand_computed_nearest_rank_values() -> None:
|
||||
"""Exact p50/p95/p99 for a small, hand-computable distribution -- the
|
||||
monotonicity test above can't catch a bug where p95 and p99 both
|
||||
collapse to the same (wrong) value, since max <= max still holds.
|
||||
|
||||
latencies here are 1ms..20ms. The reviewer's finding: the old
|
||||
`int(len(latencies) * p)` used directly as a zero-based index put both
|
||||
p95 and p99 at index 19 -- the maximum -- for any 20-sample run.
|
||||
checkpoint_bench_common.percentile's (n-1)*p/100 nearest-rank
|
||||
interpolation, reused here, keeps them distinct from the max.
|
||||
"""
|
||||
latencies = [0.001 * i for i in range(1, 21)] # 1ms..20ms
|
||||
workers = [_worker(0, [_result(True, latency) for latency in latencies])]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=20)
|
||||
|
||||
assert summary["latency_p50_ms"] == 10.5
|
||||
assert summary["latency_p95_ms"] == 19.05
|
||||
assert summary["latency_p99_ms"] == 19.81
|
||||
assert summary["latency_max_ms"] == 20.0
|
||||
# the specific bug: p95 and p99 must NOT both equal the max
|
||||
assert summary["latency_p95_ms"] != summary["latency_max_ms"]
|
||||
assert summary["latency_p99_ms"] != summary["latency_max_ms"]
|
||||
|
||||
|
||||
def test_summarize_percentiles_match_hand_computed_values_at_documented_sample_size() -> None:
|
||||
"""Same shape of check at the documented default sample size (100 ops)."""
|
||||
latencies = [0.001 * i for i in range(1, 101)] # 1ms..100ms
|
||||
workers = [_worker(0, [_result(True, latency) for latency in latencies])]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=100)
|
||||
|
||||
assert summary["latency_p50_ms"] == 50.5
|
||||
assert summary["latency_p95_ms"] == 95.05
|
||||
assert summary["latency_p99_ms"] == 99.01
|
||||
|
||||
|
||||
def test_summarize_handles_empty_results_without_crashing() -> None:
|
||||
"""All workers crashed -- no ops completed at all. Percentiles must
|
||||
degrade to None rather than raising (e.g. dividing by zero, or
|
||||
indexing an empty sorted list); with no operation window there is also
|
||||
no throughput to report."""
|
||||
workers = [_crashed(None)]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=10)
|
||||
assert summary["completed_ops"] == 0
|
||||
assert summary["latency_p50_ms"] is None
|
||||
assert summary["latency_p99_ms"] is None
|
||||
assert summary["throughput_ops_per_s"] is None
|
||||
|
||||
|
||||
def test_summarize_surfaces_crashed_worker_diagnostics_not_just_a_count() -> None:
|
||||
"""crashed_workers stays an int count (existing contract) but the
|
||||
stderr each crashed worker actually printed must also be reachable
|
||||
from the summary -- previously captured in run_workers() and then
|
||||
discarded, leaving an all-crashed sweep with zero explanation of why."""
|
||||
workers = [
|
||||
_worker(0, [_result(True, 0.001)]),
|
||||
_crashed(1, "OperationalError: unable to open database file"),
|
||||
_crashed(2, "sqlite3.IntegrityError: UNIQUE constraint failed"),
|
||||
]
|
||||
summary = bench.summarize(workers, n_workers=3, ops_per_worker=1)
|
||||
assert summary["crashed_workers"] == 2
|
||||
assert summary["crashed_worker_errors"] == [
|
||||
{"worker_id": 1, "stderr": "OperationalError: unable to open database file"},
|
||||
{"worker_id": 2, "stderr": "sqlite3.IntegrityError: UNIQUE constraint failed"},
|
||||
]
|
||||
|
||||
|
||||
def test_summary_indicates_failure_for_an_all_crashed_sweep() -> None:
|
||||
"""The exact scenario the reviewer's repro produced: every worker
|
||||
crashed, so the summary looks well-formed (completed_ops: 0,
|
||||
throughput_ops_per_s: None) but represents no real measurement at all.
|
||||
main() must treat this as a failure (nonzero exit), not a quiet 0-op
|
||||
result -- this pins the check that decides that, independent of
|
||||
main()'s argparse/subprocess machinery."""
|
||||
workers = [_crashed(0), _crashed(1)]
|
||||
summary = bench.summarize(workers, n_workers=2, ops_per_worker=4)
|
||||
assert bench.summary_indicates_failure(summary) is True
|
||||
|
||||
|
||||
def test_summary_indicates_failure_when_completed_ops_falls_short_without_a_crash() -> None:
|
||||
"""Defense in depth: even if crashed_workers is 0, fewer completed ops
|
||||
than expected must still count as a failure rather than being silently
|
||||
accepted as a (misleadingly short) real measurement."""
|
||||
summary = {"crashed_workers": 0, "completed_ops": 3, "expected_total_ops": 4}
|
||||
assert bench.summary_indicates_failure(summary) is True
|
||||
|
||||
|
||||
def test_summary_indicates_failure_when_every_op_errored_without_a_crash() -> None:
|
||||
"""A sweep where every op ran to completion but raised (e.g. every write
|
||||
hitting OperationalError) has crashed_workers == 0 and
|
||||
completed_ops == expected_total_ops, so the two checks above accept it.
|
||||
errors > 0 must also disqualify it -- the PR's conclusions rest on
|
||||
'0 errors on both backends'."""
|
||||
workers = [
|
||||
_worker(
|
||||
0,
|
||||
[_result(False, 0.5, err="OperationalError: database is locked", op="write") for _ in range(4)],
|
||||
)
|
||||
]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=4)
|
||||
assert summary["crashed_workers"] == 0
|
||||
assert summary["completed_ops"] == summary["expected_total_ops"] == 4
|
||||
assert summary["errors"] == 4
|
||||
assert bench.summary_indicates_failure(summary) is True
|
||||
|
||||
|
||||
def test_summary_indicates_failure_is_false_for_a_clean_run() -> None:
|
||||
workers = [_worker(0, [_result(True, 0.001) for _ in range(4)])]
|
||||
summary = bench.summarize(workers, n_workers=1, ops_per_worker=4)
|
||||
assert bench.summary_indicates_failure(summary) is False
|
||||
|
||||
|
||||
def test_summarize_throughput_uses_slowest_worker_operation_window() -> None:
|
||||
"""Throughput is completed_ops / max(worker ops_elapsed_s), not over the
|
||||
orchestrator's post-communicate() wall clock -- so a worker that spends
|
||||
extra time in teardown/IPC after its last op does not deflate the
|
||||
number, and the slowest still-contending worker sets the window."""
|
||||
workers = [
|
||||
_worker(0, [_result(True, 0.001) for _ in range(30)], ops_elapsed_s=5.0),
|
||||
_worker(1, [_result(True, 0.001) for _ in range(20)], ops_elapsed_s=2.0),
|
||||
]
|
||||
summary = bench.summarize(workers, n_workers=2, ops_per_worker=25)
|
||||
assert summary["ops_window_s"] == 5.0
|
||||
assert summary["throughput_ops_per_s"] == 10.0 # 50 ops / 5s
|
||||
100
backend/tests/test_bench_worker.py
Normal file
100
backend/tests/test_bench_worker.py
Normal file
@ -0,0 +1,100 @@
|
||||
"""Unit tests for scripts/benchmark/concurrency/worker.py's read/write op
|
||||
scheduling -- fast, no DB required, same load-as-module pattern as
|
||||
test_bench_concurrency.py (which covers run_concurrency_bench.py's
|
||||
aggregation logic)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_module():
|
||||
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/concurrency/worker.py"
|
||||
spec = importlib.util.spec_from_file_location("bench_worker", path)
|
||||
assert spec is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec.loader is not None
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
worker = _load_module()
|
||||
|
||||
|
||||
def _schedule(n_ops: int, read_ratio: float) -> list[bool]:
|
||||
n_reads = worker.read_count(n_ops, read_ratio)
|
||||
return [worker.is_read_op(i, n_ops, n_reads) for i in range(n_ops)]
|
||||
|
||||
|
||||
def test_default_ops_per_worker_and_read_ratio_produce_a_mixed_workload() -> None:
|
||||
"""The documented/default invocation (--ops-per-worker 50 --read-ratio
|
||||
0.7, per run_concurrency_bench.py's own --help and usage docstring) must
|
||||
produce both reads and writes. The prior `(i % 100) < int(read_ratio *
|
||||
100)` check made every op a read here (50 < 100), reproducing the
|
||||
reviewer's finding directly."""
|
||||
schedule = _schedule(n_ops=50, read_ratio=0.7)
|
||||
assert schedule.count(True) == 35
|
||||
assert schedule.count(False) == 15
|
||||
|
||||
|
||||
def test_reads_are_spread_across_the_sequence_not_clustered() -> None:
|
||||
"""A schedule that is technically mixed but front-loaded (e.g. all 35
|
||||
reads before any write) would still misrepresent steady-state concurrent
|
||||
load. Both halves of the sequence must contain each op type."""
|
||||
schedule = _schedule(n_ops=50, read_ratio=0.7)
|
||||
first_half, second_half = schedule[:25], schedule[25:]
|
||||
assert True in first_half and False in first_half
|
||||
assert True in second_half and False in second_half
|
||||
|
||||
|
||||
def test_read_count_rounds_rather_than_truncates() -> None:
|
||||
assert worker.read_count(50, 0.7) == 35 # round(35.0), not int(34.999...)
|
||||
|
||||
|
||||
def test_zero_read_ratio_is_all_writes() -> None:
|
||||
schedule = _schedule(n_ops=20, read_ratio=0.0)
|
||||
assert schedule.count(True) == 0
|
||||
|
||||
|
||||
def test_full_read_ratio_is_all_reads() -> None:
|
||||
schedule = _schedule(n_ops=20, read_ratio=1.0)
|
||||
assert schedule.count(True) == 20
|
||||
|
||||
|
||||
def test_small_op_counts_still_get_an_exact_mix() -> None:
|
||||
"""n_ops well under 100 (e.g. a quick manual smoke run) is exactly the
|
||||
regime the old modulo-100 logic silently broke."""
|
||||
for n_ops in (1, 2, 3, 7, 13, 30, 99):
|
||||
schedule = _schedule(n_ops=n_ops, read_ratio=0.7)
|
||||
assert len(schedule) == n_ops
|
||||
assert schedule.count(True) == round(n_ops * 0.7)
|
||||
|
||||
|
||||
def test_sqlite_worker_engine_matches_the_app_connection_pragmas(tmp_path, monkeypatch) -> None:
|
||||
"""synchronous and foreign_keys are per-connection PRAGMAs: without the
|
||||
worker mirroring the app's connect listener
|
||||
(persistence/engine.py::_enable_sqlite_wal) it runs at SQLite's
|
||||
synchronous=FULL / foreign_keys=OFF defaults, and its measured write
|
||||
path pays a heavier per-commit fsync than the deployment being
|
||||
benchmarked -- overstating SQLite's cost in the direction that flatters
|
||||
the 'use Postgres' conclusion."""
|
||||
import asyncio
|
||||
|
||||
from sqlalchemy import text
|
||||
|
||||
monkeypatch.setattr(worker, "SQLITE_BENCH_DIR", str(tmp_path))
|
||||
|
||||
async def _check() -> None:
|
||||
engine, _sf = worker.make_session_factory("sqlite", pg_url="", pg_schema="")
|
||||
try:
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(text("PRAGMA synchronous"))).scalar() == 1 # NORMAL
|
||||
assert (await conn.execute(text("PRAGMA foreign_keys"))).scalar() == 1 # ON
|
||||
assert (await conn.execute(text("PRAGMA journal_mode"))).scalar().lower() == "wal"
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
asyncio.run(_check())
|
||||
@ -157,7 +157,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p
|
||||
with sqlite3.connect(db_path) as raw:
|
||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||
# Bootstrap upgrades through the later revisions after 0004.
|
||||
assert version_row[0] == "0017_personal_access_tokens"
|
||||
assert version_row[0] == "0018_oauth_identity_pg_partial"
|
||||
|
||||
# Sanity: the invariant the index enforces is now true — at most one
|
||||
# active row per thread.
|
||||
|
||||
@ -173,7 +173,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm
|
||||
|
||||
with sqlite3.connect(db_path) as raw:
|
||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||
assert version_row[0] == "0017_personal_access_tokens"
|
||||
assert version_row[0] == "0018_oauth_identity_pg_partial"
|
||||
|
||||
# Sanity: the invariant the index enforces now holds — at most one
|
||||
# active row per task_id.
|
||||
|
||||
@ -57,7 +57,7 @@ async def test_migration_interrupts_legacy_queue_and_adds_claim_fields(tmp_path:
|
||||
|
||||
# Bootstrap always advances to the repository head after exercising
|
||||
# the 0015 migration behavior below.
|
||||
assert version == "0017_personal_access_tokens"
|
||||
assert version == "0018_oauth_identity_pg_partial"
|
||||
assert {"lease_owner", "lease_expires_at", "attempt_count"} <= columns.keys()
|
||||
assert columns["attempt_count"]["nullable"] is False
|
||||
assert overlap_policy == "enqueue"
|
||||
|
||||
@ -0,0 +1,85 @@
|
||||
"""Migration test for 0018_oauth_identity_pg_partial.
|
||||
|
||||
The other OAuth-index test (test_user_oauth_partial_index.py) provisions the
|
||||
schema through ``init_engine_from_config``, which on an empty database takes
|
||||
the bootstrap ``create_all()`` path -- it builds the partial index straight
|
||||
from current ORM metadata and stamps alembic at head, so
|
||||
``0018.upgrade()`` never actually runs. Existing installations only get the
|
||||
partial predicate via the migration, so this exercises that path directly:
|
||||
alembic-upgrade to 0017 (index created full by 0001_baseline, which passed
|
||||
``sqlite_where`` but not ``postgresql_where``), then to 0018, and assert the
|
||||
predicate appears; the downgrade must restore the full index.
|
||||
|
||||
Postgres-only (the revision is a no-op on SQLite -- 0001_baseline already
|
||||
gave that backend ``sqlite_where``). Opt in with DEERFLOW_TEST_POSTGRES_URL,
|
||||
same as test_pg_schema_integration.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic import command as alembic_command
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from deerflow.persistence.bootstrap import _get_alembic_config
|
||||
|
||||
POSTGRES_URL = os.getenv("DEERFLOW_TEST_POSTGRES_URL")
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.asyncio,
|
||||
pytest.mark.skipif(not POSTGRES_URL, reason="set DEERFLOW_TEST_POSTGRES_URL to run live PostgreSQL tests"),
|
||||
]
|
||||
|
||||
_PREVIOUS = "0017_personal_access_tokens"
|
||||
_REVISION = "0018_oauth_identity_pg_partial"
|
||||
_INDEX = "idx_users_oauth_identity"
|
||||
|
||||
|
||||
async def _index_predicate(engine, schema: str) -> str | None:
|
||||
"""The index's WHERE predicate as SQL text, or None if it is a full
|
||||
(non-partial) index. Raises if the index does not exist."""
|
||||
async with engine.connect() as conn:
|
||||
row = (
|
||||
await conn.execute(
|
||||
sa.text("SELECT pg_get_expr(i.indpred, i.indrelid) FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid JOIN pg_namespace n ON n.oid = c.relnamespace WHERE c.relname = :idx AND n.nspname = :schema"),
|
||||
{"idx": _INDEX, "schema": schema},
|
||||
)
|
||||
).fetchone()
|
||||
assert row is not None, f"{_INDEX} not found in schema {schema}"
|
||||
return row[0]
|
||||
|
||||
|
||||
async def test_0018_adds_partial_predicate_and_downgrade_restores_full_index() -> None:
|
||||
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
|
||||
engine = create_async_engine(POSTGRES_URL or "")
|
||||
cfg = _get_alembic_config(engine, postgres_schema=schema)
|
||||
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
|
||||
|
||||
# env.py drives migrations with its own asyncio.run, so the sync
|
||||
# command API must run off the test loop (same wrapper production uses).
|
||||
await asyncio.to_thread(alembic_command.upgrade, cfg, _PREVIOUS)
|
||||
assert await _index_predicate(engine, schema) is None, "0001_baseline should create a full index on Postgres"
|
||||
|
||||
await asyncio.to_thread(alembic_command.upgrade, cfg, _REVISION)
|
||||
predicate = await _index_predicate(engine, schema)
|
||||
assert predicate is not None, "0018 did not make the index partial"
|
||||
assert "oauth_provider" in predicate and "oauth_id" in predicate
|
||||
|
||||
await asyncio.to_thread(alembic_command.downgrade, cfg, _PREVIOUS)
|
||||
assert await _index_predicate(engine, schema) is None, "downgrade did not restore the full index"
|
||||
|
||||
# Round-trips: re-upgrading is idempotent (0018 also guards on indpred).
|
||||
await asyncio.to_thread(alembic_command.upgrade, cfg, "head")
|
||||
assert await _index_predicate(engine, schema) is not None
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default
|
||||
asyncio_test = pytest.mark.asyncio
|
||||
|
||||
|
||||
HEAD = "0017_personal_access_tokens"
|
||||
HEAD = "0018_oauth_identity_pg_partial"
|
||||
BASELINE = "0001_baseline"
|
||||
|
||||
|
||||
|
||||
@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
HEAD = "0017_personal_access_tokens"
|
||||
HEAD = "0018_oauth_identity_pg_partial"
|
||||
|
||||
|
||||
def _url(tmp_path: Path) -> str:
|
||||
|
||||
@ -76,7 +76,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No
|
||||
cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()}
|
||||
assert "token_usage_by_model" in cols
|
||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||
assert version_row[0] == "0017_personal_access_tokens"
|
||||
assert version_row[0] == "0018_oauth_identity_pg_partial"
|
||||
|
||||
# And the read path that originally 500'd must now succeed.
|
||||
sf = get_session_factory()
|
||||
@ -116,6 +116,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path
|
||||
# No duplicate column -- list, not set, to catch dupes.
|
||||
assert cols.count("token_usage_by_model") == 1
|
||||
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
|
||||
assert version_row[0] == "0017_personal_access_tokens"
|
||||
assert version_row[0] == "0018_oauth_identity_pg_partial"
|
||||
finally:
|
||||
await close_engine()
|
||||
|
||||
136
backend/tests/test_user_oauth_partial_index.py
Normal file
136
backend/tests/test_user_oauth_partial_index.py
Normal file
@ -0,0 +1,136 @@
|
||||
"""Live PostgreSQL regression test for the ``users`` table's OAuth identity
|
||||
index (idx_users_oauth_identity) -- same opt-in pattern as
|
||||
test_pg_schema_integration.py (set DEERFLOW_TEST_POSTGRES_URL to run).
|
||||
|
||||
Confirms the index is created as a genuine PARTIAL index on Postgres
|
||||
(``postgresql_where``), matching the SQLite side (``sqlite_where``) and the
|
||||
docstring on UserRow.__table_args__. Without ``postgresql_where``, Postgres
|
||||
still enforces the same practical uniqueness (NULL is never equal to NULL
|
||||
in a unique index on either backend), so this is a size/maintenance
|
||||
regression test for the partial predicate itself, not a correctness test
|
||||
for the uniqueness constraint -- that is covered separately below.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import close_engine, get_engine, init_engine_from_config
|
||||
|
||||
POSTGRES_URL = os.getenv("DEERFLOW_TEST_POSTGRES_URL")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not POSTGRES_URL,
|
||||
reason="set DEERFLOW_TEST_POSTGRES_URL to run live PostgreSQL tests",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_oauth_identity_index_is_partial_on_postgres():
|
||||
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
|
||||
db_config = DatabaseConfig(backend="postgres", postgres_url=POSTGRES_URL or "", postgres_schema=schema)
|
||||
|
||||
await init_engine_from_config(db_config)
|
||||
engine = get_engine()
|
||||
assert engine is not None
|
||||
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
row = (
|
||||
await conn.execute(
|
||||
text("SELECT indexdef FROM pg_indexes WHERE schemaname = :schema AND indexname = 'idx_users_oauth_identity'"),
|
||||
{"schema": schema},
|
||||
)
|
||||
).fetchone()
|
||||
assert row is not None, "idx_users_oauth_identity was not created in the target schema"
|
||||
indexdef = row[0]
|
||||
assert "WHERE" in indexdef, f"expected a partial index (WHERE clause), got: {indexdef}"
|
||||
assert "oauth_provider" in indexdef and "oauth_id" in indexdef
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await close_engine()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_oauth_identity_uniqueness_enforced_end_to_end():
|
||||
"""Correctness check (independent of whether the index is partial):
|
||||
a genuine duplicate (provider, oauth_id) pair is rejected, and
|
||||
multiple plain-password accounts (both fields NULL) are allowed to
|
||||
coexist -- the two behaviours the index exists to guarantee."""
|
||||
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
|
||||
db_config = DatabaseConfig(backend="postgres", postgres_url=POSTGRES_URL or "", postgres_schema=schema)
|
||||
|
||||
await init_engine_from_config(db_config)
|
||||
engine = get_engine()
|
||||
assert engine is not None
|
||||
|
||||
try:
|
||||
from app.gateway.auth.models import User
|
||||
from app.gateway.auth.repositories.sqlite import SQLiteUserRepository
|
||||
from deerflow.persistence.engine import get_session_factory
|
||||
|
||||
repo = SQLiteUserRepository(get_session_factory())
|
||||
|
||||
first = User(
|
||||
id=uuid.uuid4(),
|
||||
email="oauth-user-1@example.com",
|
||||
password_hash=None,
|
||||
system_role="user",
|
||||
created_at=datetime.now(UTC),
|
||||
oauth_provider="github",
|
||||
oauth_id="dup-check-123",
|
||||
needs_setup=False,
|
||||
token_version=0,
|
||||
)
|
||||
await repo.create_user(first)
|
||||
|
||||
duplicate = User(
|
||||
id=uuid.uuid4(),
|
||||
email="oauth-user-2@example.com",
|
||||
password_hash=None,
|
||||
system_role="user",
|
||||
created_at=datetime.now(UTC),
|
||||
oauth_provider="github",
|
||||
oauth_id="dup-check-123",
|
||||
needs_setup=False,
|
||||
token_version=0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="OAuth account already linked"):
|
||||
await repo.create_user(duplicate)
|
||||
|
||||
# two plain-password accounts (NULL, NULL) must coexist without conflict
|
||||
plain_a = User(
|
||||
id=uuid.uuid4(),
|
||||
email="plain-a@example.com",
|
||||
password_hash="h",
|
||||
system_role="user",
|
||||
created_at=datetime.now(UTC),
|
||||
oauth_provider=None,
|
||||
oauth_id=None,
|
||||
needs_setup=False,
|
||||
token_version=0,
|
||||
)
|
||||
plain_b = User(
|
||||
id=uuid.uuid4(),
|
||||
email="plain-b@example.com",
|
||||
password_hash="h",
|
||||
system_role="user",
|
||||
created_at=datetime.now(UTC),
|
||||
oauth_provider=None,
|
||||
oauth_id=None,
|
||||
needs_setup=False,
|
||||
token_version=0,
|
||||
)
|
||||
await repo.create_user(plain_a)
|
||||
await repo.create_user(plain_b) # must not raise
|
||||
finally:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await close_engine()
|
||||
Loading…
x
Reference in New Issue
Block a user