Compare commits

10 Commits

Author SHA1 Message Date
DarkSky
8e8a781ed1 feat(editor): improve html content color mapping (#15516)
Some checks failed
Build & Test / Build AFFiNE native (x86_64-unknown-linux-gnu) (push) Has been cancelled
Build & Test / Analyze (javascript, affine) (push) Has been cancelled
Build & Test / Analyze (javascript, blocksuite) (push) Has been cancelled
Build & Test / Analyze (typescript, affine) (push) Has been cancelled
Build & Test / Analyze (typescript, blocksuite) (push) Has been cancelled
Build & Test / Lint (push) Has been cancelled
Build & Test / Typecheck (push) Has been cancelled
Build & Test / Mobile native build filter (push) Has been cancelled
Build & Test / Rust test filter (push) Has been cancelled
Build & Test / Build AFFiNE native (aarch64-apple-darwin) (push) Has been cancelled
Build & Test / Check yarn binary (push) Has been cancelled
Build & Test / E2E BlockSuite Test (1) (push) Has been cancelled
Build & Test / E2E BlockSuite Test (2) (push) Has been cancelled
Build & Test / Build AFFiNE native (x86_64-apple-darwin) (push) Has been cancelled
Build & Test / E2E BlockSuite Test (3) (push) Has been cancelled
Build & Test / E2E BlockSuite Test (4) (push) Has been cancelled
Build & Test / E2E BlockSuite Test (5) (push) Has been cancelled
Build & Test / E2E BlockSuite Cross Browser Test (push) Has been cancelled
Build & Test / E2E Test (1) (push) Has been cancelled
Build & Test / E2E Test (2) (push) Has been cancelled
Build & Test / E2E Test (3) (push) Has been cancelled
Build & Test / E2E Test (4) (push) Has been cancelled
Build & Test / E2E Test (5) (push) Has been cancelled
Build & Test / E2E Mobile Test (1) (push) Has been cancelled
Build & Test / E2E Mobile Test (2) (push) Has been cancelled
Build & Test / Build AFFiNE native (aarch64-pc-windows-msvc) (push) Has been cancelled
Build & Test / Build AFFiNE native (x86_64-pc-windows-msvc) (push) Has been cancelled
Build & Test / Build Server native (push) Has been cancelled
Build & Test / Build @affine/electron renderer (push) Has been cancelled
Build & Test / Copilot test filter (push) Has been cancelled
Build & Test / Build Android app (push) Has been cancelled
Build & Test / Test Android app (push) Has been cancelled
Build & Test / Build iOS app (push) Has been cancelled
Build & Test / Lint Rust (push) Has been cancelled
Build & Test / Check Git Status (push) Has been cancelled
Build & Test / Desktop Test (ubuntu-latest, linux, x64, x86_64-unknown-linux-gnu, true) (push) Has been cancelled
Build & Test / Desktop Test (windows-latest, windows, x64, x86_64-pc-windows-msvc, true) (push) Has been cancelled
Build & Test / Unit Test (1) (push) Has been cancelled
Build & Test / Unit Test (2) (push) Has been cancelled
Build & Test / Unit Test (3) (push) Has been cancelled
Build & Test / Unit Test (4) (push) Has been cancelled
Build & Test / Unit Test (5) (push) Has been cancelled
Build & Test / Native Unit Test (push) Has been cancelled
Build & Test / Server Test (0, 4) (push) Has been cancelled
Build & Test / Server Test (1, 4) (push) Has been cancelled
Build & Test / Server Test (2, 4) (push) Has been cancelled
Build & Test / Server Test (3, 4) (push) Has been cancelled
Build & Test / Server Test with Elasticsearch (push) Has been cancelled
Build & Test / Server Test with Embedded Search (push) Has been cancelled
Build & Test / Server Test with Manticore Search (push) Has been cancelled
Build & Test / Server E2E Test (push) Has been cancelled
Build & Test / Run native tests (push) Has been cancelled
Build & Test / Server Copilot Api Test (push) Has been cancelled
Build & Test / Frontend Copilot E2E Test (1, 5) (push) Has been cancelled
Build & Test / Frontend Copilot E2E Test (2, 5) (push) Has been cancelled
Build & Test / Frontend Copilot E2E Test (3, 5) (push) Has been cancelled
Build & Test / Frontend Copilot E2E Test (4, 5) (push) Has been cancelled
Build & Test / Frontend Copilot E2E Test (5, 5) (push) Has been cancelled
Build & Test / Cloud Desktop E2E Test (push) Has been cancelled
Build & Test / Cloud E2E Test 1/2 (push) Has been cancelled
Build & Test / Cloud E2E Test 2/2 (push) Has been cancelled
Build & Test / Desktop Test (macos-latest, macos, arm64, aarch64-apple-darwin, true) (push) Has been cancelled
Build & Test / Desktop Test (macos-latest, macos, x64, x86_64-apple-darwin, false) (push) Has been cancelled
Build & Test / 3, 2, 1 Launch (push) Has been cancelled
Copilot Test Automatically / Setup Test (push) Has been cancelled
Release / Prepare (push) Has been cancelled
Release / Canary Gate (push) Has been cancelled
Release / Release Cloud (push) Has been cancelled
Release / Release Docker Image (push) Has been cancelled
Release / Release Desktop (push) Has been cancelled
Release / Release Mobile (push) Has been cancelled
fix #15514

#### PR Dependency Tree


* **PR #15516** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
  * Improved color handling when importing HTML content.
* Added support for hex, RGB/RGBA, percentage, alpha, transparent,
named, and HSL color values.
  * Correctly maps supported colors to the app’s color themes.
* Prevents invalid, translucent, or unsuitable colors from being
applied.
  * Preserves style values containing additional colons during import.
* Improved imported text formatting for supported and unsupported
colors.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-08-24 10:52:02 +08:00
keepClamDown
cd6593c659 fix(ios): ios local access limit for self-hosted workspace (#15338)
## Summary
- Route iOS nbstore worker token reads through the main-thread
MessagePort and skip Capacitor Auth on `/socket.io` polling so self-host
WebSocket/XHR sync no longer hangs on Connect timeout.
- Harden workspace `flavour:id` routing, DocSyncPeer abort/status
handling, and `resetSync` so local selfhost edits push and Mac browsers
can receive them.
- Soften root-doc readiness waits and session-exchange throttling to
keep mobile selfhost sign-in/sync stable under retries.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Summary by CodeRabbit

* **New Features**
* Improved workspace switching across local and remote environments,
preserving workspace type and server context.
  * Added support for page navigation with query parameters.
* Added iOS local-network permission messaging for self-hosted
workspaces.
* **Bug Fixes**
* Improved document synchronization, reset handling, prioritized
document refreshes, and retry behavior.
* Prevented authentication headers and refresh attempts for socket
connection requests.
* Improved workspace reopening and routing when multiple workspace types
share an ID.
  * Fixed handling of unlimited data query limits.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: DarkSky <darksky2048@gmail.com>
2026-08-24 09:33:58 +08:00
DarkSky
b530198a3b feat(server): improve indexer perf (#15512)
#### PR Dependency Tree


* **PR #15512** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Search now supports generation-based indexing with embedded and remote
providers.
* Added automatic search reconciliation and improved handling of
document, workspace, and permission changes.
* Added clearer search status errors for unavailable, syncing, unready,
or failed indexes.
* Added Manticore Search end-to-end support and provider-specific search
behavior.

* **Improvements**
* Search and aggregate pagination now report returned results and
continuation status more accurately.
* Improved permission filtering to prevent inaccessible documents from
appearing in results.
  * Admin provider selection now consistently enables indexing.

* **Documentation**
* Clarified search pagination, aggregation counts, provider
configuration, and end-to-end setup.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-08-23 17:54:49 +08:00
DarkSky
a8eef53966 chore(server): split module 2026-08-22 08:05:46 +08:00
DarkSky
478a1a0f30 fix(server): clean up unnecessary refreshes 2026-08-22 07:29:58 +08:00
DarkSky
47306699c4 chore(server): fix connection 2026-08-22 06:01:41 +08:00
DarkSky
dec1a01449 fix(server): improve self hosted usability (#15510)
fix #15505
fix #15502
fix #15496
fix #15491

#### PR Dependency Tree


* **PR #15510** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added configurable delays for invitations, invite links, and document
publishing by newly created accounts.
- Added workspace action checks that explain blocked actions and retry
timing.
- BYOK setup now verifies model capabilities and saves only validated
options.

- **Bug Fixes**
- Improved BYOK probing for chat, structured responses, tool calls,
embeddings, reranking, and image generation.
  - Preserved probe request order and strengthened response validation.
  - Authentication configuration changes now reload correctly.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
2026-08-22 04:47:21 +08:00
qiaoyanfei
591f874dad fix(android): stabilize IME composition and deletion handling (#15370)
## Summary
Fix Android editor IME corruption around composition, autocorrect
replay, delete, Enter, and old WebView delete behavior.

## What changed
- Add Android WebView InputConnection wrapper for editor IME handling.
- Route Android delete events through BlockSuite editor input.
- Guard against keyboard autocorrect/composition replay after delete or
space.
- Stabilize delete fallback on older Android/WebView versions.
- Gate IME diagnostic logs behind debug builds.
- Add Android IME fix notes and regression coverage.

## Validation
- Manual Android testing passed.
- `git diff --check upstream/canary...HEAD`

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Improved Android text editing for Backspace, Delete, Enter, composing
text, and keyboard events.
  * Prevented input methods from targeting the wrong editor area.
* Improved caret-based text handling, focus synchronization, and
composing-session cleanup.

* **Platform Improvements**
* Added a dedicated Android input bridge for smoother IME interactions
and fallback keyboard behavior.
* Improved editor actions, input recovery, and trusted-page validation
for Android communication.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: DarkSky <darksky2048@gmail.com>
2026-08-21 13:55:14 +08:00
donqu1xotevincent
3feb17cde3 docs: refresh contributing & building guides, surface CLA signing flow (#15503) 2026-08-21 08:15:26 +08:00
keepClamDown
73e2d10973 feat(ios): polish onboarding and native sign-in flow (#15217)
## Summary
- Reworks the iOS onboarding, native sign-in, and paywall flow so
first-install and cold-start entry stay native, adaptive, and aligned
with the current login/subscription gates.
- Removes onboarding swipe paging, refines onboarding/paywall layout
behavior, and keeps AI/paywall entry behavior consistent for logged-in
and subscribed users.
- Adds the new mobile all-docs empty states with localized copy and
dialog entry points, and closes the remaining review follow-ups by
removing the onboarding plan artifact and dropping the iOS AI
subscription bypass.

## Test plan
- Built the iOS app for simulator with `xcodebuild -workspace
App.xcworkspace -scheme App -configuration Debug -sdk iphonesimulator
-destination 'generic/platform=iOS Simulator' ARCHS=arm64
ONLY_ACTIVE_ARCH=YES CODE_SIGNING_ALLOWED=NO CODE_SIGNING_REQUIRED=NO
build` during the native onboarding/sign-in flow work.
- Let the repo pre-commit hooks run on the latest follow-up commit
(`prettier` + `eslint --fix`).
- Checked diagnostics for the edited TS/TSX files after the review
follow-up changes.
- Manually iterated on onboarding/native sign-in/paywall UI states in
simulator during implementation.

---------

Co-authored-by: DarkSky <darksky2048@gmail.com>
2026-08-21 08:12:27 +08:00
295 changed files with 17778 additions and 8159 deletions

View File

@@ -12,7 +12,6 @@ services:
environment:
DATABASE_URL: postgresql://affine:affine@db:5432/affine
REDIS_SERVER_HOST: redis
AFFINE_INDEXER_SEARCH_ENDPOINT: http://indexer:9308
db:
image: pgvector/pgvector:pg16
@@ -26,19 +25,5 @@ services:
redis:
image: redis
indexer:
image: manticoresearch/manticore:${MANTICORE_VERSION:-10.1.0}
ulimits:
nproc: 65535
nofile:
soft: 65535
hard: 65535
memlock:
soft: -1
hard: -1
volumes:
- manticoresearch_data:/var/lib/manticore
volumes:
postgres-data:
manticoresearch_data:

View File

@@ -10,6 +10,3 @@ DB_DATABASE_NAME=affine
# enable for arm64, e.g.: macOS M1+
# ELASTIC_VERSION_ARM64=-arm64
# ELASTIC_PLATFORM=linux/arm64
# manticoresearch
MANTICORE_VERSION=10.1.0

View File

@@ -32,22 +32,6 @@ services:
volumes:
- mailpit_data:/data
# https://manual.manticoresearch.com/Starting_the_server/Docker
manticoresearch:
image: manticoresearch/manticore:${MANTICORE_VERSION:-10.1.0}
ports:
- 9308:9308
ulimits:
nproc: 65535
nofile:
soft: 65535
hard: 65535
memlock:
soft: -1
hard: -1
volumes:
- manticoresearch_data:/var/lib/manticore
# elasticsearch:
# image: docker.elastic.co/elasticsearch/elasticsearch:${ELASTIC_VERSION:-9.0.1}${ELASTIC_VERSION_ARM64}
# platform: ${ELASTIC_PLATFORM}
@@ -94,6 +78,5 @@ networks:
volumes:
postgres_data:
manticoresearch_data:
mailpit_data:
elasticsearch_data:

View File

@@ -86,18 +86,6 @@
"concurrency": 1
}
},
"queues.indexer": {
"type": "object",
"description": "The config for indexer job queue\n@default {\"concurrency\":1}",
"properties": {
"concurrency": {
"type": "number"
}
},
"default": {
"concurrency": 1
}
},
"queues.notification": {
"type": "object",
"description": "The config for notification job queue\n@default {\"concurrency\":10}",
@@ -199,9 +187,9 @@
"description": "Whether require email verification before accessing restricted resources(not implemented).\n@default true",
"default": true
},
"newAccountShareActionDelay": {
"newAccountActionDelay": {
"type": "number",
"description": "Minimum account age in seconds before new accounts can invite members or create share links.\n@default 86400",
"description": "Minimum account age in seconds before new accounts can invite members, create invite links, or publish documents. Set to 0 to disable.\n@default 86400",
"default": 86400
},
"trustedCloudflareHeaders": {
@@ -1488,38 +1476,33 @@
"properties": {
"enabled": {
"type": "boolean",
"description": "Enable indexer plugin\n@default false\n@environment `AFFINE_INDEXER_ENABLED`",
"description": "Enable indexer plugin\n@default false",
"default": false
},
"provider.type": {
"type": "string",
"description": "Indexer search provider. Self-hosted uses the embedded provider by default; remote providers require an endpoint.\n@default \"embedded\"\n@environment `AFFINE_INDEXER_SEARCH_PROVIDER`",
"description": "Indexer search provider. Self-hosted uses the embedded provider by default; remote providers require an endpoint.\n@default \"embedded\"",
"default": "embedded"
},
"provider.endpoint": {
"type": "string",
"description": "Remote indexer endpoint. Not used by the embedded provider.\n@default \"\"\n@environment `AFFINE_INDEXER_SEARCH_ENDPOINT`",
"description": "Remote indexer endpoint. Not used by the embedded provider.\n@default \"\"",
"default": ""
},
"provider.apiKey": {
"type": "string",
"description": "Indexer search service api key. Optional for elasticsearch\n@default \"\"\n@environment `AFFINE_INDEXER_SEARCH_API_KEY`\n@link https://www.elastic.co/guide/server/current/api-key.html",
"description": "Indexer search service api key. Optional for remote providers\n@default \"\"\n@link https://www.elastic.co/guide/server/current/api-key.html",
"default": ""
},
"provider.username": {
"type": "string",
"description": "Indexer search service auth username, if not set, basic auth will be disabled. Optional for elasticsearch\n@default \"\"\n@environment `AFFINE_INDEXER_SEARCH_USERNAME`\n@link https://www.elastic.co/guide/en/elasticsearch/reference/current/http-clients.html",
"description": "Indexer search service auth username, if not set, basic auth will be disabled. Optional for remote providers\n@default \"\"\n@link https://www.elastic.co/guide/en/elasticsearch/reference/current/http-clients.html",
"default": ""
},
"provider.password": {
"type": "string",
"description": "Indexer search service auth password, if not set, basic auth will be disabled. Optional for elasticsearch\n@default \"\"\n@environment `AFFINE_INDEXER_SEARCH_PASSWORD`",
"description": "Indexer search service auth password, if not set, basic auth will be disabled. Optional for remote providers\n@default \"\"",
"default": ""
},
"autoIndex.batchSize": {
"type": "number",
"description": "Number of workspaces automatically indexed per batch\n@default 10",
"default": 10
}
}
},

6
.github/CLA.md vendored
View File

@@ -1,4 +1,4 @@
<!-- To indicate your agreement, simply edit this file and submit a pull request. -->
<!-- Do NOT edit this file to sign. To sign the CLA, visit https://cla-assistant.io/toeverything/AFFiNE and sign in with GitHub. -->
# AFFiNE Contributor License Agreement
@@ -35,4 +35,6 @@ You accept and agree to the following terms and conditions for your past, presen
## How To Sign
Visit https://cla-assistant.io/toeverything/AFFiNE and sign it.
Visit **<https://cla-assistant.io/toeverything/AFFiNE>**, sign in with your GitHub account, and agree.
Once signed, the `license/cla` check on your pull requests passes automatically. If you signed after opening a PR, click the **recheck** link in the CLA bot's comment on the PR (or push a new commit) to refresh the check. Note that every committer on a PR must sign, and each commit's author email must be linked to a GitHub account.

18
.github/PULL_REQUEST_TEMPLATE.md vendored Normal file
View File

@@ -0,0 +1,18 @@
<!--
Thank you for contributing to AFFiNE!
The PR title must follow Conventional Commits (enforced by CI):
type(scope): description e.g. fix(editor): keep selection after paste
Types: feat fix docs style refactor perf test build ci chore revert
-->
## Description
<!-- What does this PR do? Link related issues, e.g. "Closes #1234". Screenshots or recordings are welcome for UI changes. -->
## Checklist
- [ ] I have signed the [AFFiNE Contributor License Agreement](https://cla-assistant.io/toeverything/AFFiNE) — required before merge; the `license/cla` check must be green ([how it works](https://github.com/toeverything/AFFiNE/blob/canary/docs/BUILDING.md#sign-the-cla-first))
- [ ] The PR targets the `canary` branch and its title follows [Conventional Commits](https://www.conventionalcommits.org/)
- [ ] Tests are added or updated where it makes sense
- [ ] `yarn lint` and `yarn typecheck` pass locally

View File

@@ -33,6 +33,9 @@ runs:
shell: bash
run: |
case '${{ inputs.indexer-provider }}' in
disabled)
default_indexer='{"enabled":false,"provider":{"type":"embedded","endpoint":""}}'
;;
embedded)
default_indexer='{"enabled":true,"provider":{"type":"embedded","endpoint":""}}'
;;

View File

@@ -933,8 +933,6 @@ jobs:
NODE_ENV: test
DATABASE_URL: postgresql://affine:affine@localhost:5432/affine
REDIS_SERVER_HOST: localhost
SERVER_CONFIG: >-
{"indexer":{"enabled":true,"provider":{"type":"elasticsearch","endpoint":"http://localhost:9200"}}}
services:
postgres:
image: pgvector/pgvector:pg16
@@ -1005,8 +1003,74 @@ jobs:
name: affine
fail_ci_if_error: false
server-test-manticore:
name: Server Test with Manticore
server-test-embedded:
name: Server Test with Embedded Search
runs-on: ubuntu-latest
needs:
- build-server-native
env:
NODE_ENV: test
DATABASE_URL: postgresql://affine:affine@localhost:5432/affine
REDIS_SERVER_HOST: localhost
services:
postgres:
image: pgvector/pgvector:pg16
env:
POSTGRES_PASSWORD: affine
options: >-
--health-cmd pg_isready
--health-interval 10s
--health-timeout 5s
--health-retries 5
ports:
- 5432:5432
redis:
image: redis
ports:
- 6379:6379
steps:
- uses: actions/checkout@v7
- name: Setup Node.js
uses: ./.github/actions/setup-node
with:
extra-flags: workspaces focus @affine/monorepo @affine/server
electron-install: false
full-cache: true
- name: Download server-native.node
uses: actions/download-artifact@v4
with:
name: server-native.node
path: ./packages/backend/native
- name: Prepare Server Test Environment
with:
indexer-provider: embedded
uses: ./.github/actions/server-test-env
- name: Run embedded Search provider E2E
run: >-
yarn affine @affine/server e2e:coverage
src/__tests__/e2e/indexer/aggregate.spec.ts
src/__tests__/e2e/indexer/search-docs.spec.ts
src/__tests__/e2e/indexer/search.spec.ts
--concurrency 1
--forbid-only
env:
CARGO_TARGET_DIR: '${{ github.workspace }}/target'
- name: Upload server test coverage results
uses: codecov/codecov-action@v5
with:
token: ${{ secrets.CODECOV_TOKEN }}
files: ./packages/backend/server/.coverage/lcov.info
flags: server-test
name: affine
fail_ci_if_error: false
server-test-manticoresearch:
name: Server Test with Manticore Search
runs-on: ubuntu-latest
needs:
- build-server-native
@@ -1014,9 +1078,6 @@ jobs:
NODE_ENV: test
DATABASE_URL: postgresql://affine:affine@localhost:5432/affine
REDIS_SERVER_HOST: localhost
SEARCH_MS_URL: http://localhost:9308
SERVER_CONFIG: >-
{"indexer":{"enabled":true,"provider":{"type":"manticoresearch","endpoint":"http://localhost:9308"}}}
services:
postgres:
image: pgvector/pgvector:pg16
@@ -1038,15 +1099,18 @@ jobs:
ports:
- 1025:1025
- 8025:8025
indexer:
image: manticoresearch/manticore:10.1.0
options: >-
--health-cmd "wget -q -O /dev/null http://localhost:9308/"
--health-interval 5s
--health-timeout 3s
--health-retries 12
manticore:
image: manticoresearch/manticore:29.0.2
env:
'searchd.listen': 'http:9308'
'searchd.network_timeout': '30'
ports:
- 9308:9308
options: >-
--health-cmd "wget -qO- --post-data='SHOW TABLES' 'http://127.0.0.1:9308/sql?mode=raw' >/dev/null"
--health-interval 5s
--health-timeout 5s
--health-retries 30
steps:
- uses: actions/checkout@v7
@@ -1057,13 +1121,6 @@ jobs:
electron-install: false
full-cache: true
- name: Setup Rust
uses: ./.github/actions/build-rust
with:
target: x86_64-unknown-linux-gnu
package: '@affine/server-native'
no-build: 'true'
- name: Download server-native.node
uses: actions/download-artifact@v4
with:
@@ -1071,17 +1128,14 @@ jobs:
path: ./packages/backend/native
- name: Prepare Server Test Environment
uses: ./.github/actions/server-test-env
with:
indexer-provider: manticoresearch
uses: ./.github/actions/server-test-env
- name: Run server tests with Manticore
run: yarn affine @affine/server e2e:coverage src/__tests__/e2e/indexer/search-docs.spec.ts --forbid-only
- name: Run Manticore provider contract
run: >-
cargo test --manifest-path packages/backend/native/Cargo.toml
remote_providers_apply_search_and_delete_the_same_contract --lib
- name: Run Manticore Search provider E2E
run: yarn affine @affine/server e2e:coverage src/__tests__/e2e/indexer/manticore-provider.spec.ts --forbid-only
env:
CARGO_TARGET_DIR: '${{ github.workspace }}/target'
- name: Upload server test coverage results
uses: codecov/codecov-action@v5
@@ -1135,6 +1189,8 @@ jobs:
path: ./packages/backend/native
- name: Prepare Server Test Environment
with:
indexer-provider: disabled
uses: ./.github/actions/server-test-env
- name: Run server tests
@@ -1597,7 +1653,8 @@ jobs:
- native-unit-test
- server-test
- server-test-elasticsearch
- server-test-manticore
- server-test-embedded
- server-test-manticoresearch
- server-e2e-test
- rust-test
- rust-test-filter

4
Cargo.lock generated
View File

@@ -4758,9 +4758,9 @@ dependencies = [
[[package]]
name = "llm_adapter"
version = "0.2.20"
version = "0.2.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5d13be366ea35d2a9966ad5770e3d070e495af4e1e6dd009638dee4b55ab45ed"
checksum = "02c6b4fa5178b8183331a7d51e8f7ece5421ecb966d3431f3d9ce3c724c6fcc0"
dependencies = [
"base64",
"jsonschema",

View File

@@ -97,7 +97,7 @@ Thanks for checking us out, we appreciate your interest and sincerely hope that
Calling all developers, testers, tech writers and more! Contributions of all types are more than welcome, you can read more in [docs/types-of-contributions.md](docs/types-of-contributions.md). If you are interested in contributing code, read our [docs/CONTRIBUTING.md](docs/CONTRIBUTING.md) and feel free to check out our GitHub issues to get stuck in to show us what you’re made of.
**Before you start contributing, please make sure you have read and accepted our [Contributor License Agreement]. To indicate your agreement, simply edit this file and submit a pull request.**
**Before you start contributing, please sign our [Contributor License Agreement] — it takes less than a minute with your GitHub account. Pull requests cannot be merged until every committer has signed (the `license/cla` check on your PR). Full text: [CLA.md](.github/CLA.md).**
For **bug reports**, **feature requests** and other **suggestions** you can also [create a new issue](https://github.com/toeverything/AFFiNE/issues/new/choose) and choose the most appropriate template for your feedback.
@@ -216,10 +216,9 @@ See [LICENSE] for details.
[update page]: https://affine.pro/blog?tag=Release%20Note
[jobs available]: ./docs/jobs.md
[latest packages]: https://github.com/toeverything/AFFiNE/pkgs/container/affine-self-hosted
[contributor license agreement]: https://github.com/toeverything/affine/edit/canary/.github/CLA.md
[contributor license agreement]: https://cla-assistant.io/toeverything/AFFiNE
[stars-icon]: https://img.shields.io/github/stars/toeverything/AFFiNE.svg?style=flat&logo=github&colorB=red&label=stars
[codecov]: https://codecov.io/gh/toeverything/affine/branch/canary/graphs/badge.svg?branch=canary
[node-version-icon]: https://img.shields.io/badge/node-%3E=18.16.1-success
[typescript-version-icon]: https://img.shields.io/github/package-json/dependency-version/toeverything/affine/dev/typescript
[react-version-icon]: https://img.shields.io/github/package-json/dependency-version/toeverything/AFFiNE/react?filename=packages%2Ffrontend%2Fcore%2Fpackage.json&color=rgb(97%2C228%2C251)
[blocksuite-icon]: https://img.shields.io/github/package-json/dependency-version/toeverything/AFFiNE/@blocksuite/store?color=6880ff&filename=packages%2Ffrontend%2Fcore%2Fpackage.json&label=blocksuite

View File

@@ -3929,48 +3929,40 @@ describe('markdown to snapshot', () => {
});
});
test('html inline color span imports to nearest supported text color', async () => {
const markdown = `<span style="color: #00afde;">Hello</span>`;
const blockSnapshot: BlockSnapshot = {
type: 'block',
id: 'matchesReplaceMap[0]',
flavour: 'affine:note',
props: {
xywh: '[0,0,800,95]',
background: DefaultTheme.noteBackgrounColor,
index: 'a0',
hidden: false,
displayMode: NoteDisplayMode.DocAndEdgeless,
},
children: [
{
type: 'block',
id: 'matchesReplaceMap[1]',
flavour: 'affine:paragraph',
props: {
type: 'text',
text: {
'$blocksuite:internal:text$': true,
delta: [
{
insert: 'Hello',
attributes: {
color: 'var(--affine-v2-text-highlight-fg-blue)',
},
},
],
},
},
children: [],
},
],
};
test.each([
['#00afde', 'blue'],
['rgb(0 175 222 / 100%)', 'blue'],
['#c83030', 'red'],
['red', 'red'],
['hsl(0, 100%, 50%)', 'red'],
['#db7123', 'orange'],
['#ac7400', 'yellow'],
['#04b745', 'green'],
['#0e4841', 'teal'],
['#7c3aed', 'purple'],
['#7a7a7a', 'grey'],
['rgb(26, 26, 26)', null],
['#333', null],
['#fff', null],
['rgba(0, 175, 222, 0.5)', null],
])('maps supported HTML color %s conservatively', async (color, mapped) => {
const mdAdapter = new MarkdownAdapter(createJob(), provider);
const rawBlockSnapshot = await mdAdapter.toBlockSnapshot({
file: markdown,
file: `<span style="color: ${color};">Hello</span>`,
});
expect(rawBlockSnapshot.children[0]?.props.text).toEqual({
'$blocksuite:internal:text$': true,
delta: [
mapped
? {
insert: 'Hello',
attributes: {
color: `var(--affine-v2-text-highlight-fg-${mapped})`,
},
}
: { insert: 'Hello' },
],
});
expect(nanoidReplacement(rawBlockSnapshot)).toEqual(blockSnapshot);
});
test('paragraph', async () => {

View File

@@ -19,6 +19,7 @@ import {
getScrollContainer,
matchModels,
} from '@blocksuite/affine-shared/utils';
import { IS_ANDROID } from '@blocksuite/global/env';
import { Point } from '@blocksuite/global/gfx';
import type { PointerEventState } from '@blocksuite/std';
import { BlockComponent, BlockSelection, TextSelection } from '@blocksuite/std';
@@ -413,7 +414,10 @@ export class PageRootBlockComponent extends BlockComponent<RootBlockModel> {
return !(isNote && displayOnEdgeless);
});
this.contentEditable = String(!this.store.readonly$.value);
// Android IMEs can target this outer editable root instead of a block's
// inline editor, leaving composition text in the DOM without committing it
// to the document model. Keep only the block editors editable on Android.
this.contentEditable = String(!this.store.readonly$.value && !IS_ANDROID);
return html`
<div class="affine-page-root-block-container">${children} ${widgets}</div>

View File

@@ -23,6 +23,7 @@
"@blocksuite/affine-shared": "workspace:*",
"@blocksuite/std": "workspace:*",
"@blocksuite/store": "workspace:*",
"@ctrl/tinycolor": "^4.1.0",
"@toeverything/theme": "^1.1.23",
"@types/hast": "^3.0.4",
"collapse-white-space": "^2.1.0",

View File

@@ -1,9 +1,16 @@
import { parseStringToRgba } from '@blocksuite/affine-components/color-picker';
import { TinyColor } from '@ctrl/tinycolor';
import { cssVarV2, darkThemeV2, lightThemeV2 } from '@toeverything/theme/v2';
type Rgb = { r: number; g: number; b: number };
type Oklab = { l: number; a: number; b: number; chroma: number; hue: number };
const ACHROMATIC_CHROMA_THRESHOLD = 0.02;
const DEFAULT_TEXT_LIGHTNESS_MIN = 0.4;
const DEFAULT_TEXT_LIGHTNESS_MAX = 0.9;
const MAX_COLOR_DISTANCE = 0.18;
const MAX_CHROMA_DISTANCE = 0.12;
const MAX_HUE_DISTANCE = 45;
const COLOR_DISTANCE_THRESHOLD = 90;
const supportedTextColorNames = [
'red',
'orange',
@@ -15,68 +22,81 @@ const supportedTextColorNames = [
'grey',
] as const;
const supportedTextColors = supportedTextColorNames.map(name => ({
name,
cssVar: cssVarV2(`text/highlight/fg/${name}`),
light: lightThemeV2[`text/highlight/fg/${name}`],
dark: darkThemeV2[`text/highlight/fg/${name}`],
}));
export const parseCssColor = (value: string) => {
const parsed = new TinyColor(value);
if (!parsed.isValid) return null;
const { r, g, b, a } = parsed.toRgb();
return { r, g, b, alpha: a };
};
const hexToRgb = (value: string): Rgb | null => {
const hex = value.replace('#', '');
if (![3, 4, 6, 8].includes(hex.length)) {
return null;
}
const normalized =
hex.length === 3 || hex.length === 4
? hex
.slice(0, 3)
.split('')
.map(c => c + c)
.join('')
: hex.slice(0, 6);
const intVal = Number.parseInt(normalized, 16);
if (Number.isNaN(intVal)) {
return null;
}
const srgbToLinear = (channel: number) => {
const value = channel / 255;
return value <= 0.04045
? value / 12.92
: Math.pow((value + 0.055) / 1.055, 2.4);
};
const rgbToOklab = ({ r, g, b }: Rgb): Oklab => {
const red = srgbToLinear(r);
const green = srgbToLinear(g);
const blue = srgbToLinear(b);
const l = Math.cbrt(
0.4122214708 * red + 0.5363325363 * green + 0.0514459929 * blue
);
const m = Math.cbrt(
0.2119034982 * red + 0.6806995451 * green + 0.1073969566 * blue
);
const s = Math.cbrt(
0.0883024619 * red + 0.2817188376 * green + 0.6299787005 * blue
);
const result = {
l: 0.2104542553 * l + 0.793617785 * m - 0.0040720468 * s,
a: 1.9779984951 * l - 2.428592205 * m + 0.4505937099 * s,
b: 0.0259040371 * l + 0.7827717662 * m - 0.808675766 * s,
};
return {
r: (intVal >> 16) & 255,
g: (intVal >> 8) & 255,
b: intVal & 255,
...result,
chroma: Math.hypot(result.a, result.b),
hue: (Math.atan2(result.b, result.a) * 180) / Math.PI,
};
};
export const parseCssColor = (value: string): Rgb | null => {
const trimmed = value.trim();
if (!trimmed) {
return null;
}
if (trimmed.startsWith('#')) {
return hexToRgb(trimmed);
}
if (/^rgba?\(/i.test(trimmed)) {
const rgba = parseStringToRgba(trimmed);
return {
r: Math.round(rgba.r * 255),
g: Math.round(rgba.g * 255),
b: Math.round(rgba.b * 255),
};
}
return null;
};
const supportedTextColors = supportedTextColorNames.map(name => ({
name,
cssVar: cssVarV2(`text/highlight/fg/${name}`),
references: [
lightThemeV2[`text/highlight/fg/${name}`],
darkThemeV2[`text/highlight/fg/${name}`],
].flatMap(color => {
const parsed = parseCssColor(color);
return parsed ? [rgbToOklab(parsed)] : [];
}),
}));
const colorDistance = (a: Rgb, b: Rgb) => {
const dr = a.r - b.r;
const dg = a.g - b.g;
const db = a.b - b.b;
return Math.sqrt(dr * dr + dg * dg + db * db);
const colorDistance = (a: Oklab, b: Oklab) =>
Math.hypot(a.l - b.l, a.a - b.a, a.b - b.b);
const hueDistance = (a: number, b: number) => {
const distance = Math.abs(a - b) % 360;
return Math.min(distance, 360 - distance);
};
export const resolveNearestSupportedColor = (color: string): string | null => {
const target = parseCssColor(color);
if (!target) {
const parsed = parseCssColor(color);
if (!parsed || parsed.alpha < 1) {
return null;
}
const target = rgbToOklab(parsed);
const achromatic = target.chroma < ACHROMATIC_CHROMA_THRESHOLD;
if (
achromatic &&
(target.l < DEFAULT_TEXT_LIGHTNESS_MIN ||
target.l > DEFAULT_TEXT_LIGHTNESS_MAX)
) {
return null;
}
let nearest:
| {
cssVar: string;
@@ -85,21 +105,26 @@ export const resolveNearestSupportedColor = (color: string): string | null => {
| undefined;
for (const supported of supportedTextColors) {
const light = parseCssColor(supported.light);
const dark = parseCssColor(supported.dark);
for (const ref of [light, dark]) {
if (!ref) continue;
const distance = colorDistance(target, ref);
if (achromatic !== (supported.name === 'grey')) {
continue;
}
for (const reference of supported.references) {
const distance = colorDistance(target, reference);
if (
distance > MAX_COLOR_DISTANCE ||
(!achromatic &&
(Math.abs(target.chroma - reference.chroma) > MAX_CHROMA_DISTANCE ||
hueDistance(target.hue, reference.hue) > MAX_HUE_DISTANCE))
) {
continue;
}
if (!nearest || distance < nearest.distance) {
nearest = { cssVar: supported.cssVar, distance };
}
}
}
if (nearest && nearest.distance <= COLOR_DISTANCE_THRESHOLD) {
return nearest.cssVar;
}
return null;
return nearest?.cssVar ?? null;
};
export const extractColorFromStyle = (
@@ -110,10 +135,11 @@ export const extractColorFromStyle = (
}
const declarations = style.split(';');
for (const declaration of declarations) {
const [rawKey, rawValue] = declaration.split(':');
if (!rawKey || !rawValue) continue;
if (rawKey.trim().toLowerCase() === 'color') {
return rawValue.trim();
const colon = declaration.indexOf(':');
if (colon === -1) continue;
const key = declaration.slice(0, colon).trim().toLowerCase();
if (key === 'color') {
return declaration.slice(colon + 1).trim();
}
}
return null;

View File

@@ -57,12 +57,14 @@ export function getPrefixText(inlineEditor: InlineEditor) {
const inlineRange = inlineEditor.getInlineRange();
if (!inlineRange || inlineRange.length > 0) return '';
const nearestLineBreakIndex = inlineEditor.yTextString
.slice(0, inlineRange.index)
.lastIndexOf('\n');
const prefixText = inlineEditor.yTextString.slice(
nearestLineBreakIndex + 1,
inlineRange.index
);
return prefixText;
const maxMarkdownPrefixLength = 512;
const prefixStart = Math.max(0, inlineRange.index - maxMarkdownPrefixLength);
const yTextString = inlineEditor.yTextString;
const prefixWindow = yTextString.slice(prefixStart, inlineRange.index);
const nearestLineBreakIndex = prefixWindow.lastIndexOf('\n');
if (nearestLineBreakIndex === -1 && prefixStart > 0) return '';
return nearestLineBreakIndex === -1
? prefixWindow
: prefixWindow.slice(nearestLineBreakIndex + 1);
}

View File

@@ -1,6 +1,7 @@
import { describe, expect, test } from 'vitest';
import { describe, expect, test, vi } from 'vitest';
import { bindKeymap } from '../event/keymap.js';
import { UIEventState, UIEventStateContext } from '../event/base.js';
import { androidBindKeymapPatch, bindKeymap } from '../event/keymap.js';
const createKeyboardEvent = (options: {
key: string;
@@ -117,3 +118,86 @@ describe('bindKeymap', () => {
expect(handled).toBe(false);
});
});
describe('androidBindKeymapPatch', () => {
const beforeInputCtx = (inputType: string) => {
const event = new InputEvent('beforeinput', {
inputType,
cancelable: true,
});
return { ctx: UIEventStateContext.from(new UIEventState(event)), event };
};
test('routes deleteContentBackward to the Backspace binding', () => {
const backspace = vi.fn(() => true);
const handler = androidBindKeymapPatch({ Backspace: backspace });
const { ctx } = beforeInputCtx('deleteContentBackward');
expect(handler(ctx)).toBe(true);
expect(backspace).toHaveBeenCalledOnce();
});
test('routes insertParagraph to the Enter binding', () => {
const enter = vi.fn((ctx: UIEventStateContext) => {
ctx.get('keyboardState').raw.preventDefault();
return true;
});
const handler = androidBindKeymapPatch({ Enter: enter });
const { ctx, event } = beforeInputCtx('insertParagraph');
const preventDefault = vi.spyOn(event, 'preventDefault');
expect(handler(ctx)).toBe(true);
expect(enter).toHaveBeenCalledOnce();
expect(preventDefault).toHaveBeenCalledOnce();
expect(ctx.get('keyboardState').raw.key).toBe('Enter');
expect(ctx.get('keyboardState').composing).toBe(false);
});
test('propagates preventDefault when the binding returns false', () => {
const backspace = vi.fn((ctx: UIEventStateContext) => {
ctx.get('keyboardState').raw.preventDefault();
return false;
});
const handler = androidBindKeymapPatch({ Backspace: backspace });
const { ctx, event } = beforeInputCtx('deleteContentBackward');
expect(handler(ctx)).toBe(false);
expect(event.defaultPrevented).toBe(true);
});
test('does nothing for insertParagraph without an Enter binding', () => {
const handler = androidBindKeymapPatch({ Backspace: vi.fn(() => true) });
const { ctx } = beforeInputCtx('insertParagraph');
expect(handler(ctx)).toBe(false);
expect(ctx.has('keyboardState')).toBe(false);
});
test('ignores non-input events', () => {
const enter = vi.fn(() => true);
const backspace = vi.fn(() => true);
const ctx = UIEventStateContext.from(
new UIEventState(new KeyboardEvent('keydown', { key: 'Enter' }))
);
expect(
androidBindKeymapPatch({ Enter: enter, Backspace: backspace })(ctx)
).toBeUndefined();
expect(enter).not.toHaveBeenCalled();
expect(backspace).not.toHaveBeenCalled();
});
test('ignores unrelated input types', () => {
const enter = vi.fn(() => true);
const backspace = vi.fn(() => true);
const handler = androidBindKeymapPatch({
Enter: enter,
Backspace: backspace,
});
const { ctx } = beforeInputCtx('insertText');
expect(handler(ctx)).toBe(false);
expect(enter).not.toHaveBeenCalled();
expect(backspace).not.toHaveBeenCalled();
});
});

View File

@@ -3,6 +3,7 @@ import { BlockSuiteError, ErrorCode } from '@blocksuite/global/exceptions';
import { base, keyName } from 'w3c-keyname';
import type { UIEventHandler } from './base.js';
import { KeyboardEventState } from './state/index.js';
function normalizeKeyName(name: string) {
const parts = name.split(/-(?!$)/);
@@ -127,13 +128,38 @@ export function androidBindKeymapPatch(
const event = ctx.get('defaultState').event;
if (!(event instanceof InputEvent)) return;
if (
event.inputType === 'deleteContentBackward' &&
'Backspace' in bindings
) {
return bindings['Backspace'](ctx);
const bindingName =
event.inputType === 'deleteContentBackward'
? 'Backspace'
: event.inputType === 'deleteContentForward'
? 'Delete'
: event.inputType === 'insertParagraph'
? 'Enter'
: undefined;
if (!bindingName || !(bindingName in bindings)) return false;
if (!ctx.has('keyboardState')) {
const keyboardEvent = new KeyboardEvent('keydown', {
key: bindingName,
code: bindingName,
cancelable: true,
});
Object.defineProperty(keyboardEvent, 'isComposing', {
configurable: true,
value: event.isComposing,
});
ctx.add(
new KeyboardEventState({
event: keyboardEvent,
composing: event.isComposing,
})
);
}
return false;
const handled = bindings[bindingName](ctx);
if (handled || ctx.get('keyboardState').raw.defaultPrevented) {
event.preventDefault();
}
return handled;
};
}

View File

@@ -13,11 +13,58 @@ import { isMaybeInlineRangeEqual } from '../utils/inline-range.js';
import { transformInput } from '../utils/transform-input.js';
import type { BeforeinputHookCtx, CompositionEndHookCtx } from './hook.js';
type AndroidIMEInputType = 'deleteContentBackward' | 'deleteContentForward';
type AndroidIMEInputDetail = {
inputType?: AndroidIMEInputType;
handled?: boolean;
};
type AndroidIMEBridge = {
getProtocolVersion?: () => number;
finishComposingSession?: () => void;
finishDeleteSession?: () => void;
setEditorFocused?: (focused: boolean) => void;
};
declare global {
interface HTMLElementEventMap {
'affine-android-ime-input': CustomEvent<AndroidIMEInputDetail>;
}
}
export class EventService<TextAttributes extends BaseTextAttributes> {
private _compositionInlineRange: InlineRange | null = null;
private _isComposing = false;
private readonly _androidIMEBridge = () => {
const bridge = (
globalThis as typeof globalThis & {
AffineAndroidIME?: AndroidIMEBridge;
}
).AffineAndroidIME;
return bridge?.getProtocolVersion?.() === 1 ? bridge : undefined;
};
private readonly _finishAndroidComposingSession = (isDelete: boolean) => {
if (!IS_ANDROID) return;
window.setTimeout(() => {
const bridge = this._androidIMEBridge();
if (isDelete) {
bridge?.finishDeleteSession?.();
} else {
bridge?.finishComposingSession?.();
}
}, 0);
};
private readonly _setAndroidEditorFocused = (focused: boolean) => {
if (!IS_ANDROID) return;
this._androidIMEBridge()?.setEditorFocused?.(focused);
};
private readonly _getClosestInlineRoot = (node: Node): Element | null => {
const el = node instanceof Element ? node : node.parentElement;
return el?.closest(`[${INLINE_ROOT_ATTR}]`) ?? null;
@@ -203,6 +250,97 @@ export class EventService<TextAttributes extends BaseTextAttributes> {
);
this.editor.slots.inputting.next(event.data ?? '');
if (
IS_ANDROID &&
(ctx.raw.inputType === 'deleteContentBackward' ||
ctx.raw.inputType === 'deleteContentForward' ||
ctx.raw.inputType === 'insertParagraph' ||
ctx.raw.inputType === 'insertLineBreak' ||
(ctx.raw.inputType === 'insertText' &&
(ctx.data === ' ' || ctx.data === '\n')))
) {
this._finishAndroidComposingSession(
ctx.raw.inputType === 'deleteContentBackward' ||
ctx.raw.inputType === 'deleteContentForward'
);
}
};
private readonly _onAndroidIMEInput = async (
event: CustomEvent<AndroidIMEInputDetail>
) => {
if (!IS_ANDROID) return;
const inputType = event.detail?.inputType;
if (
inputType !== 'deleteContentBackward' &&
inputType !== 'deleteContentForward'
) {
return;
}
const range = this.editor.rangeService.getNativeRange();
if (!range || !this._isRangeCompletelyInRoot(range)) return;
event.detail.handled = true;
event.preventDefault();
event.stopPropagation();
if (this.editor.isReadonly) return;
let inlineRange = this.editor.toInlineRange(range);
if (!inlineRange) {
this.editor.rerenderWholeEditor();
await this.editor.waitForUpdate();
const newRange = this.editor.rangeService.getNativeRange();
inlineRange = newRange ? this.editor.toInlineRange(newRange) : null;
if (!inlineRange) return;
}
if (inlineRange.length === 0) {
if (inputType === 'deleteContentBackward') {
if (inlineRange.index === 0) return;
inlineRange = {
index: inlineRange.index - 1,
length: 1,
};
} else {
if (inlineRange.index >= this.editor.yTextLength) return;
inlineRange = {
index: inlineRange.index,
length: 1,
};
}
}
this._isComposing = false;
this._compositionInlineRange = null;
const raw = new InputEvent('beforeinput', {
inputType,
bubbles: true,
cancelable: true,
composed: true,
});
const ctx: BeforeinputHookCtx<TextAttributes> = {
inlineEditor: this.editor,
raw,
inlineRange,
data: null,
attributes: {} as TextAttributes,
};
this.editor.hooks.beforeinput?.(ctx);
transformInput<TextAttributes>(
ctx.raw.inputType,
ctx.data,
ctx.attributes,
ctx.inlineRange,
this.editor as never
);
this.editor.slots.inputting.next('');
this._finishAndroidComposingSession(true);
};
private readonly _onClick = (event: MouseEvent) => {
@@ -265,6 +403,7 @@ export class EventService<TextAttributes extends BaseTextAttributes> {
}
this.editor.slots.inputting.next(event.data ?? '');
this._finishAndroidComposingSession(false);
};
private readonly _onCompositionStart = (event: CompositionEvent) => {
@@ -432,6 +571,30 @@ export class EventService<TextAttributes extends BaseTextAttributes> {
this.editor.disposables.addFromEvent(eventSource, 'beforeinput', e => {
this._onBeforeInput(e).catch(console.error);
});
this.editor.disposables.addFromEvent(
eventSource,
'affine-android-ime-input',
e => {
this._onAndroidIMEInput(e).catch(console.error);
}
);
this.editor.disposables.addFromEvent(eventSource, 'focusin', () => {
this._setAndroidEditorFocused(true);
});
this.editor.disposables.addFromEvent(
eventSource,
'focusout',
(event: FocusEvent) => {
const relatedTarget = event.relatedTarget;
if (
relatedTarget instanceof Node &&
this.editor.rootElement?.contains(relatedTarget)
) {
return;
}
this._setAndroidEditorFocused(false);
}
);
this.editor.disposables.addFromEvent(
eventSource,
'compositionstart',

View File

@@ -1,68 +1,83 @@
# Building AFFiNE Web
> **Warning**:
>
> **Warning**
> This document is not guaranteed to be up-to-date.
> If you find any outdated information, please feel free to open an issue or submit a PR.
>
> **Note**
> For developing & building desktop client app, please refer to [building-desktop-client-app.md](./building-desktop-client-app.md)
> This guide covers building and developing the **web app**.
> For the desktop client app, see [building-desktop-client-app.md](./building-desktop-client-app.md).
> For running the server (cloud features) locally, see [developing-server.md](./developing-server.md).
## Table of Contents
- [Sign the CLA first](#sign-the-cla-first)
- [Prerequisites](#prerequisites)
- [Setup Environment](#setup-environment)
- [Start Development Server](#start-development-server)
- [Testing](#testing)
- [Linting and Type Checking](#linting-and-type-checking)
- [Submitting a Pull Request](#submitting-a-pull-request)
- [Troubleshooting](#troubleshooting)
- [Related Documents](#related-documents)
## Sign the CLA first
AFFiNE requires every contributor to sign the [Contributor License Agreement](../.github/CLA.md) before a pull request can be **merged**. The `license/cla` status check on your PR stays red — and blocks the merge — until every committer on the PR has signed. Signing takes less than a minute, so do it before (or right after) opening your first PR:
1. Open **<https://cla-assistant.io/toeverything/AFFiNE>**.
2. Sign in with your GitHub account and agree.
If you opened your PR before signing, the CLA-assistant bot will have left a comment on it with a signing link. After signing, click the **recheck** link in that same comment (or push a new commit) to refresh the check.
If `license/cla` is still red after you signed:
- **Every committer must sign.** The bot's comment lists each commit author in the PR — every name marked with :x: still needs to sign, including co-authors.
- **Your commit email must be linked to your GitHub account.** CLA assistant matches commits to GitHub users by the commit author email. Run `git log --format='%an %ae'` and make sure those emails appear in your [GitHub email settings](https://github.com/settings/emails). Otherwise, either add the email to GitHub, or rewrite the commits with the correct identity and force-push.
## Prerequisites
AFFiNE client has both **Node.js** & **Rust** toolchains.
### Install Node.js
### Node.js
We suggest develop our product under node.js LTS(Long-term support) version
Develop with the Node.js version pinned in [`.nvmrc`](../.nvmrc) (currently Node.js **22**; `package.json` requires `>=22.12.0 <23.0.0`).
#### Option 1: Manually install node.js
install [Node LTS version](https://nodejs.org/en/download)
> Up to now, the major node.js version is 20.x
#### Option 2: Use node version manager
install [fnm](https://github.com/Schniz/fnm)
The easiest way is a version manager that reads `.nvmrc`:
```sh
fnm use
# with fnm (https://github.com/Schniz/fnm)
fnm use --install-if-missing
# or with nvm (https://github.com/nvm-sh/nvm)
nvm install && nvm use
```
### Install Rust Tools
Alternatively, install Node.js 22 (LTS) manually from <https://nodejs.org/en/download>.
Please follow the official guide at https://www.rust-lang.org/tools/install.
### Yarn
### Setup Node.js Environment
This setup requires modern yarn (currently `4.x`), run this if your yarn version is `1.x`
Reference: [Yarn installation doc](https://yarnpkg.com/getting-started/install)
We use modern Yarn (currently **4.x**, pinned by the `packageManager` field in `package.json`). Yarn 1 will not work. Enable it via [Corepack](https://yarnpkg.com/corepack), which ships with Node.js:
```sh
corepack enable
corepack prepare yarn@stable --activate
```
```sh
# install dependencies
yarn install
```
After this, `yarn` inside the repository automatically resolves to the pinned version — verify with `yarn -v` (it should print `4.x`, not `1.x`).
### Clone repository
### Rust
#### Linux & MacOS
Install the Rust toolchain via [rustup](https://rustup.rs/). The required version is pinned in [`rust-toolchain.toml`](../rust-toolchain.toml), and rustup installs it automatically the first time you build inside the repository.
## Setup Environment
### Clone the repository
#### Linux & macOS
```sh
git clone https://github.com/toeverything/AFFiNE
cd AFFiNE
```
#### Windows
@@ -80,30 +95,55 @@ git config --global core.symlinks true
git clone https://github.com/toeverything/AFFiNE
```
### Install dependencies
```sh
yarn install
```
This also initializes the workspace (`yarn affine init`) and installs the git hooks through the `postinstall` script.
### Build Native Dependencies
Run the following script. It will build the native module at [`/packages/frontend/native`](/packages/frontend/native) and build Node.js binding using [NAPI.rs](https://napi.rs/).
This could take a while if you build it for the first time.
Note: use `strip` from system instead of `binutils` if you are running MacOS. [see problem here](https://github.com/toeverything/AFFiNE/discussions/2840)
Run the following script. It will build the native module at [`packages/frontend/native`](../packages/frontend/native) and build Node.js binding using [NAPI.rs](https://napi.rs/). This could take a while if you build it for the first time.
```sh
yarn affine @affine/native build
```
> Note: use `strip` from system instead of `binutils` if you are running macOS. [See problem here](https://github.com/toeverything/AFFiNE/discussions/2840)
### Build Server Dependencies
Only needed if you plan to run the local server (cloud features) or the cloud E2E suites:
```sh
yarn affine @affine/server-native build
```
## Start Development Server
```sh
yarn dev
```
You will be prompted to choose which package to run — pick **`@affine/web`** for the web app, then open **<http://localhost:8080>**.
You can skip the prompt with the `-p` (`--package`) flag:
```sh
yarn dev -p @affine/web
```
Other dev targets include `@affine/server`, `@affine/electron`, `@affine/mobile`, `@affine/admin`, `@affine/ios` and `@affine/android`.
Running `@affine/web` alone is enough for most editor and UI work — workspaces are stored locally in the browser. To work on **cloud** features (accounts, sync, collaboration, AI), run the local server as well: follow [developing-server.md](./developing-server.md).
## Testing
Adding test cases is strongly encouraged when you contribute new features and bug fixes.
We use [Playwright](https://playwright.dev/) for E2E test, and [vitest](https://vitest.dev/) for unit test.
To test locally, please make sure browser binaries are already installed via `npx playwright install`.
Start server before tests by following [`docs/developing-server.md`](./developing-server.md) first.
We use [Vitest](https://vitest.dev/) for unit tests and [Playwright](https://playwright.dev/) for E2E tests.
### Unit Test
@@ -113,8 +153,80 @@ yarn test
### E2E Test
```shell
# there are `affine-local`, `affine-migration`, `affine-local`, `affine-prototype` e2e tests,
# which are run under different situations.
yarn workspace @affine-test/affine-local e2e
Install the browser binaries once before the first run:
```sh
npx playwright install
```
The E2E suites live in [`tests`](../tests):
| Suite | Run with | Notes |
| ---------------------- | ------------------------------------------------------ | --------------------------------------------------- |
| `affine-local` | `yarn workspace @affine-test/affine-local e2e` | Web app, no server needed |
| `affine-cloud` | `yarn workspace @affine-test/affine-cloud e2e` | Requires the [local server](./developing-server.md) |
| `affine-cloud-copilot` | `yarn workspace @affine-test/affine-cloud-copilot e2e` | Requires the local server |
| `affine-desktop` | `yarn workspace @affine-test/affine-desktop e2e` | Desktop (Electron) app |
| `affine-desktop-cloud` | `yarn workspace @affine-test/affine-desktop-cloud e2e` | Desktop + local server |
| `affine-mobile` | `yarn workspace @affine-test/affine-mobile e2e` | Mobile UI |
There is also `@affine-test/blocksuite` (`yarn workspace @affine-test/blocksuite test`) for BlockSuite integration tests.
## Linting and Type Checking
CI runs these checks on every PR — running them locally first saves you a review round trip:
```sh
# lint (oxlint) + format check (oxfmt)
yarn lint
# auto-fix lint & format issues
yarn lint:fix
# TypeScript type check
yarn typecheck
```
## Submitting a Pull Request
1. Make sure you have [signed the CLA](#sign-the-cla-first).
2. Fork the repository and create your branch from **`canary`** (the default development branch).
3. Make your changes. Add or update tests where it makes sense, and run `yarn lint`, `yarn typecheck` and the relevant test suites locally.
4. Open the PR against the `canary` branch of `toeverything/AFFiNE`.
5. Give the PR a title that follows [Conventional Commits](https://www.conventionalcommits.org/) — this is enforced by the `PR Title Lint` CI check:
```text
type(scope): short description
# examples
fix(editor): keep text selection after paste
feat(core): add custom icons for folders
docs: update building guide
```
- **Types**: `feat`, `fix`, `docs`, `style`, `refactor`, `perf`, `test`, `build`, `ci`, `chore`, `revert`
- **Scopes** (optional; if present, must be one of): `admin`, `electron`, `server`, `core`, `web`, `mobile`, `ios`, `android`, `mobile-native`, `docs`, `component`, `env`, `graphql`, `hooks`, `i18n`, `native`, `templates`, `debug`, `nbstore`, `infra`, `editor`, `tools`, `y-octo`, `client`
6. Your PR can be merged once:
- the **`license/cla`** check is green — every committer has signed the CLA;
- CI passes — build & tests, PR title lint;
- a maintainer has reviewed and approved it. Reviewers are assigned automatically, and an automated reviewer may also leave comments — please address or answer them.
## Troubleshooting
- **`yarn install` fails or complains about the Node/Yarn version** — check that `node -v` matches [`.nvmrc`](../.nvmrc) (`fnm use` / `nvm use`) and that Corepack is enabled (`yarn -v` should print `4.x`, not `1.x`).
- **Native module fails to build on macOS** — make sure `strip` is the system one, not the one from `binutils` ([details](https://github.com/toeverything/AFFiNE/discussions/2840)).
- **`EPERM: operation not permitted, symlink` on Windows** — enable Developer Mode and symlinks before cloning; see [Windows](#windows).
- **App fails to start after pulling the latest `canary`** — dependencies or Rust bindings may have changed: re-run `yarn install` and `yarn affine @affine/native build`.
- **Playwright can't find browsers** — run `npx playwright install`.
- **The `license/cla` check stays red** — see [Sign the CLA first](#sign-the-cla-first).
## Related Documents
- [developing-server.md](./developing-server.md) — run the AFFiNE server locally (cloud features)
- [building-desktop-client-app.md](./building-desktop-client-app.md) — build the desktop (Electron) client
- [contributing/tutorial.md](./contributing/tutorial.md) — a walkthrough of the codebase
- [types-of-contributions.md](./types-of-contributions.md) — ways to contribute beyond code
- [issue-triaging.md](./issue-triaging.md) — how issues are triaged
- [contributing/releases.md](./contributing/releases.md) — how releases are cut
- [CODE_OF_CONDUCT.md](./CODE_OF_CONDUCT.md)

View File

@@ -1 +1,36 @@
# Please visit https://docs.affine.pro/contributing
# Contributing to AFFiNE
Thank you for your interest in contributing! ❤️ AFFiNE welcomes contributions of all kinds — code, docs, bug reports, feature ideas and translations. See [types-of-contributions.md](./types-of-contributions.md) for the full picture.
This page is the five-minute overview. The full contributor handbook lives at <https://docs.affine.pro/contributing>.
## Before your first pull request: sign the CLA
We can only merge pull requests whose authors have signed the [Contributor License Agreement](../.github/CLA.md). Every PR is checked automatically (the `license/cla` status check), and an unsigned CLA is the single most common reason PRs get stuck before merge.
**How to sign (takes under a minute):**
1. Open <https://cla-assistant.io/toeverything/AFFiNE>.
2. Sign in with your GitHub account and agree.
Already opened a PR? Sign, then click the **recheck** link in the CLA bot's comment on your PR (or push a new commit). Note that **every** committer on the PR must sign, and each commit's author email must be [linked to a GitHub account](https://github.com/settings/emails). See [BUILDING.md — Sign the CLA first](./BUILDING.md#sign-the-cla-first) for troubleshooting.
## Contribution flow
1. **Find something to work on.** Browse [good first issues](https://github.com/toeverything/AFFiNE/contribute) or the [issue tracker](https://github.com/toeverything/AFFiNE/issues). Issues that are still in triage haven't been reviewed yet — better not to start work on those. For bigger changes, open a [discussion](https://github.com/toeverything/AFFiNE/discussions) or talk to us on [Discord](https://affine.pro/redirect/discord) first.
2. **Set up your environment.** Follow [BUILDING.md](./BUILDING.md) for the web app. For the server (cloud features) see [developing-server.md](./developing-server.md); for the desktop client see [building-desktop-client-app.md](./building-desktop-client-app.md).
3. **Make your change** on a branch created from `canary`. Add tests where it makes sense, and run `yarn lint`, `yarn typecheck` and the relevant tests locally.
4. **Open a PR to `canary`** with a [Conventional Commits](https://www.conventionalcommits.org/) title, e.g. `fix(editor): keep selection after paste` — the title format is enforced by CI.
5. **Get it merged.** A PR merges once the `license/cla` check is green, CI passes, and a maintainer approves the review.
## Code of conduct
Please read our [Code of Conduct](./CODE_OF_CONDUCT.md).
## Useful links
- Codebase tour: [contributing/tutorial.md](./contributing/tutorial.md)
- How issues are triaged: [issue-triaging.md](./issue-triaging.md)
- Release process: [contributing/releases.md](./contributing/releases.md)
- Reporting security issues: [SECURITY.md](../SECURITY.md)
- Community: [Discord](https://affine.pro/redirect/discord) · [GitHub Discussions](https://github.com/toeverything/AFFiNE/discussions)

View File

@@ -34,6 +34,7 @@ export declare class BackendRuntime {
claimInviteAbuseAction(actionId: string, workerId: string): Promise<boolean>
claimRetryableInviteAbuseActions(workerId: string, limit: number): Promise<Array<RuntimeInviteAbuseClaimedAction>>
markInviteAbuseAction(actionId: string, workerId: string, status: string, error?: string | undefined | null): Promise<boolean>
evaluateWorkspaceActionV1(actorUserId: string, workspaceId: string): Promise<RuntimeWorkspaceActionDecision>
assertWorkspaceInviteQuotaV1(input: RuntimeWorkspaceInviteQuotaInput): Promise<RuntimeWorkspaceInviteQuotaDecision>
commitWorkspaceInviteQuotaV1(reservationId: string, usage: RuntimeWorkspaceInviteQuotaUsage): Promise<boolean>
releaseWorkspaceInviteQuotaV1(reservationId: string): Promise<boolean>
@@ -55,10 +56,6 @@ export declare class BackendRuntime {
getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean>
cleanupExpiredRuntimeStates(limit: number): Promise<number>
refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRefreshResult>
recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRecalibrationResult>
writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsSnapshotResult>
recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise<RuntimeWorkspaceStatsDailyRecalibrationResult>
constructor(privateKey?: string | undefined | null, configPaths?: Array<string> | undefined | null)
start(): Promise<void>
stop(): Promise<void>
@@ -67,10 +64,7 @@ export declare class BackendRuntime {
runMigrations(): Promise<void>
searchAuthorized(actorUserId: string, workspaceId: string, request: RuntimeSearchRequest): Promise<SearchOperationOutput>
aggregateAuthorized(actorUserId: string, workspaceId: string, request: RuntimeAggregateRequest): Promise<SearchOperationOutput>
indexSearchDocument(workspaceId: string, docId: string): Promise<void>
deleteSearchDocument(workspaceId: string, docId: string): Promise<void>
reconcileSearchWorkspace(workspaceId: string): Promise<void>
deleteSearchWorkspace(workspaceId: string): Promise<void>
reconcileSearchProjection(limit?: number | undefined | null): Promise<number>
filterReadableDocs(actorUserId: string, workspaceId: string, docIds: Array<string>): Promise<Array<string>>
searchStatus(): Promise<any>
embeddingHealth(): Promise<EmbeddingHealth>
@@ -112,7 +106,6 @@ export declare class StorageRuntime {
rebuildWorkspaceDocBlobRefs(workspaceId: string, limit: number): Promise<RuntimeDocBlobRefsResult>
reconcileWorkspaceDocuments(workspaceId: string): Promise<RuntimeDocumentCleanupReconcileResult>
executeDocumentCleanupCandidates(workspaceId: string | undefined | null, gracePeriodDays: number, limit: number): Promise<RuntimeDocumentCleanupExecuteResult>
ackDocumentCleanupEffect(workspaceId: string, docId: string, cleanupVersion: string, effect: string): Promise<RuntimeDocumentCleanupAckResult>
constructor()
start(): Promise<void>
configure(configJson: string): void
@@ -1260,16 +1253,11 @@ export interface RuntimeDocHistoryInput {
historyMaxAgeMs: number
}
export interface RuntimeDocumentCleanupAckResult {
completed: boolean
}
export interface RuntimeDocumentCleanupEffect {
workspaceId: string
docId: string
cleanupVersion: string
commentObjectsDone: boolean
searchDone: boolean
}
export interface RuntimeDocumentCleanupExecuteResult {
@@ -1496,6 +1484,12 @@ export interface RuntimeVerificationTokenRecord {
expiresAtMs: number
}
export interface RuntimeWorkspaceActionDecision {
allowed: boolean
retryAfterSeconds?: number
reason?: string
}
export interface RuntimeWorkspaceArtifact {
id: string
workspaceId: string
@@ -1544,30 +1538,6 @@ export interface RuntimeWorkspaceInviteQuotaUsage {
targetDomains: Array<RuntimeQuotaTargetDomainInput>
}
export interface RuntimeWorkspaceStatsDailyRecalibrationResult {
processed: number
lastSid: number
snapshotted: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsRecalibrationResult {
processed: number
lastSid: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsRefreshResult {
processed: number
backlog: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsSnapshotResult {
snapshotted: number
skipped: boolean
}
export declare function safeFetch(request: SafeFetchRequest): Promise<SafeFetchResponse>
export type SafeFetchMethod = 'get'|

View File

@@ -6,10 +6,10 @@ use llm_adapter::{
AttachmentKind, AttachmentSource, ModelFeature, ModelInput, ModelOutput, ModelRequirements, declared_model_matches,
},
core::{
CoreContent, CoreMessage, CoreRequest, CoreRole, CoreToolDefinition, EmbeddingRequest, ImageOptions,
ImageProviderOptions, ImageRequest, RerankCandidate, RerankRequest, StructuredRequest,
CoreContent, CoreMessage, CoreRequest, CoreRole, CoreToolChoice, CoreToolDefinition, EmbeddingRequest,
ImageOptions, ImageProviderOptions, ImageRequest, RerankCandidate, RerankRequest, StructuredRequest,
},
router::{ExecutablePreparedRoute, ExecutableRequest, dispatch_prepared_route},
router::{ExecutablePreparedRoute, ExecutableRequest, ExecutableResponse, dispatch_prepared_route},
target::{
BackendCredential, BackendEndpoint, BackendOperation, BackendTargetInput, EgressPolicy, compile_backend_target,
},
@@ -30,9 +30,11 @@ pub(super) async fn execute_probe(
checks: Vec<ByokProbeCheckInput>,
) -> RuntimeResult<ByokProbeResultOutput> {
let tested_at_ms = chrono::Utc::now().timestamp_millis();
let mut requested = HashSet::new();
let mut requested = Vec::new();
let mut requested_set = HashSet::new();
for check in checks {
if !requested.insert((check.model_id.clone(), check.operation.clone())) {
let key = (check.model_id.clone(), check.operation.clone());
if !requested_set.insert(key.clone()) {
return Err(RuntimeError::invalid_input("duplicate BYOK probe check"));
}
if !matches!(
@@ -41,6 +43,7 @@ pub(super) async fn execute_probe(
) {
return Err(RuntimeError::invalid_input("unknown BYOK probe operation"));
}
requested.push(key);
}
let mut models = Vec::new();
@@ -160,7 +163,49 @@ fn dispatch_check(
Err(_) => return failed(checked_at, "invalid_probe_request"),
};
match dispatch_prepared_route(&DefaultHttpClient::default(), &route) {
Ok(_) => verified(checked_at),
Ok(ExecutableResponse::Chat(response)) => {
let valid = if operation == "tool_calling" {
response
.message
.content
.iter()
.any(|content| matches!(content, CoreContent::ToolCall { name, .. } if name == "byok_probe"))
} else {
response
.message
.content
.iter()
.any(|content| matches!(content, CoreContent::Text { text } if !text.trim().is_empty()))
};
if valid {
verified(checked_at)
} else {
failed(checked_at, "invalid_response")
}
}
Ok(ExecutableResponse::Structured(response)) => {
let valid = response.output_json.as_ref().is_some_and(|output| {
let ExecutableRequest::Structured(request) = &route.request else {
return false;
};
llm_adapter::schema::validate_json_schema(&request.schema, output).is_ok()
});
if valid {
verified(checked_at)
} else {
failed(checked_at, "invalid_response")
}
}
Ok(ExecutableResponse::Embedding(response)) if operation == "embedding" && !response.embeddings.is_empty() => {
verified(checked_at)
}
Ok(ExecutableResponse::Rerank(response)) if operation == "rerank" && !response.scores.is_empty() => {
verified(checked_at)
}
Ok(ExecutableResponse::Image(response)) if operation == "image" && !response.images.is_empty() => {
verified(checked_at)
}
Ok(_) => failed(checked_at, "invalid_response"),
Err(error) => failed(checked_at, backend_error_kind(&error)),
}
}
@@ -169,7 +214,11 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest {
let message = CoreMessage {
role: CoreRole::User,
content: vec![CoreContent::Text {
text: "Reply with OK.".to_string(),
text: match operation {
"tool_calling" => "Call the byok_probe tool.".to_string(),
"structured" => "Return exactly {\"ok\":true}.".to_string(),
_ => "Reply with OK.".to_string(),
},
}],
};
match operation {
@@ -177,8 +226,8 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest {
model: String::new(),
messages: vec![message],
stream: false,
max_tokens: Some(8),
temperature: Some(0.0),
max_tokens: Some(64),
temperature: None,
tools: if operation == "tool_calling" {
vec![CoreToolDefinition {
name: "byok_probe".to_string(),
@@ -188,7 +237,9 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest {
} else {
vec![]
},
tool_choice: None,
tool_choice: (operation == "tool_calling").then_some(CoreToolChoice::Specific {
name: "byok_probe".to_string(),
}),
include: None,
reasoning: None,
response_schema: None,
@@ -202,8 +253,8 @@ fn probe_request_for_operation(operation: &str) -> ExecutableRequest {
"required": ["ok"],
"additionalProperties": false
}),
max_tokens: Some(16),
temperature: Some(0.0),
max_tokens: Some(128),
temperature: None,
reasoning: None,
strict: Some(true),
response_mime_type: Some("application/json".to_string()),
@@ -426,7 +477,48 @@ mod tests {
let mut stream = stream.unwrap();
let request = read_request(&mut stream);
let responses = request.starts_with("POST /v1/responses ");
let body = if responses {
let embedding = request.starts_with("POST /v1/embeddings ");
let image = request.starts_with("POST /v1/images/generations ");
let rerank = request.contains("\"logprobs\":true");
let tool_calling = request.contains("byok_probe");
let body = if embedding {
json!({
"model": "smoke-model",
"data": [{ "embedding": [0.1], "index": 0 }],
"usage": { "prompt_tokens": 1, "total_tokens": 1 }
})
} else if image {
json!({
"created": 0,
"data": [{ "url": "https://example.com/smoke.png" }]
})
} else if rerank {
json!({
"model": "smoke-model",
"choices": [{
"logprobs": { "content": [{
"top_logprobs": [
{ "token": "Yes", "logprob": 0.0 },
{ "token": "No", "logprob": -1.0 }
]
}] }
}]
})
} else if responses && tool_calling {
json!({
"id": "resp_smoke",
"model": "smoke-model",
"status": "completed",
"output": [{
"type": "function_call",
"id": "fc_smoke",
"call_id": "call_smoke",
"name": "byok_probe",
"arguments": "{}"
}],
"usage": { "input_tokens": 1, "output_tokens": 1, "total_tokens": 2 }
})
} else if responses {
json!({
"id": "resp_smoke",
"model": "smoke-model",
@@ -439,6 +531,25 @@ mod tests {
}],
"usage": { "input_tokens": 1, "output_tokens": 1, "total_tokens": 2 }
})
} else if tool_calling {
json!({
"id": "chat_smoke",
"model": "smoke-model",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": null,
"tool_calls": [{
"id": "call_smoke",
"type": "function",
"function": { "name": "byok_probe", "arguments": "{}" }
}]
},
"finish_reason": "tool_calls"
}],
"usage": { "prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2 }
})
} else {
json!({
"id": "chat_smoke",
@@ -489,7 +600,7 @@ mod tests {
#[test]
fn openai_compatible_probe_smoke_uses_the_selected_dialect() {
let operations = ["chat", "structured", "tool_calling"];
let operations = ["chat", "structured", "tool_calling", "embedding", "rerank", "image"];
let (endpoint, requests, server) = serve_openai_compatible(operations.len() * 2);
for dialect in [OpenAiDialect::Responses, OpenAiDialect::ChatCompletions] {
@@ -520,19 +631,39 @@ mod tests {
.iter()
.filter(|request| request.starts_with("POST /v1/responses "))
.count(),
operations.len()
3
);
assert_eq!(
requests
.iter()
.filter(|request| request.starts_with("POST /v1/chat/completions "))
.count(),
operations.len()
5
);
assert_eq!(
requests
.iter()
.filter(|request| request.starts_with("POST /v1/embeddings "))
.count(),
2
);
assert_eq!(
requests
.iter()
.filter(|request| request.starts_with("POST /v1/images/generations "))
.count(),
2
);
assert!(requests.iter().all(|request| !request.contains("/models")));
assert_eq!(
requests.iter().filter(|request| request.contains("byok_probe")).count(),
2
);
assert!(
requests
.iter()
.filter(|request| !request.contains("\"logprobs\":true"))
.all(|request| !request.contains("\"temperature\""))
);
}
}

View File

@@ -2,6 +2,3 @@ pub(super) const MAGIC_LINK_OTP_PURPOSE: &str = "magic_link_otp";
pub(super) const MAX_MAGIC_LINK_OTP_ATTEMPTS: i32 = 10;
pub(super) const WORKSPACE_INVITE_LINK_ID_PURPOSE: &str = "workspace_invite_link:id";
pub(super) const WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE: &str = "workspace_invite_link:workspace";
pub(super) const WORKSPACE_STATS_LEASE_KEY: &str = "workspace:admin-stats:refresh";
pub(super) const WORKSPACE_STATS_LOCK_NAMESPACE: i64 = 97_301;
pub(super) const WORKSPACE_STATS_REFRESH_LOCK_KEY: i64 = 1;

View File

@@ -16,7 +16,6 @@ mod scope_compiler;
mod search;
#[cfg(test)]
mod tests;
mod workspace_stats;
use std::{
sync::{Arc, RwLock},
time::Duration,
@@ -26,6 +25,8 @@ use byok::LocalLeasePayload;
use copilot::{backend_provider, executable_protocol};
use embedding::register_artifact_source;
use napi::{Result, bindgen_prelude::Buffer};
#[cfg(test)]
pub(crate) use search::SEARCH_TEST_LOCK;
use search::SearchRuntime;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
@@ -41,7 +42,7 @@ pub(crate) use super::types;
pub(super) use super::{
BackendRuntimeConfig, ConfigSource, InviteQuotaConfig, RuntimeError, RuntimeResult,
migrations::{embedding_schema_health, migrate_all_tables},
napi_error, to_napi_error,
napi_error, to_napi_error, webpki_tls_config,
};
use crate::llm::{
ByokLocalLeaseOutput, ByokPolicyOutput, ByokProbeResultOutput, ByokProfileOutput, CreateByokLocalLeaseInput,
@@ -66,10 +67,12 @@ fn search_operation_output(result: RuntimeResult<serde_json::Value>) -> SearchOp
error_code: Some(
match error {
RuntimeError::SearchWorkspaceDenied => "workspace_denied",
RuntimeError::SearchPermissionUnavailable => "permission_unavailable",
RuntimeError::SearchPermissionUnavailable => "permission_syncing",
RuntimeError::SearchIndexNotReady => "index_not_ready",
RuntimeError::SearchPermissionSyncing => "permission_syncing",
RuntimeError::SearchIndexFailed(_) => "index_failed",
RuntimeError::SearchProviderUnavailable => "provider_unavailable",
RuntimeError::SearchUnsupportedQuery => "unsupported_query",
RuntimeError::SearchReplayGap => "provider_unavailable",
RuntimeError::InvalidInput(_) | RuntimeError::Json { .. } => "invalid_request",
_ => "internal",
}
@@ -358,46 +361,14 @@ impl BackendRuntime {
}
#[napi]
pub async fn index_search_document(&self, workspace_id: String, doc_id: String) -> Result<()> {
let search = self.search_runtime().await?;
let result = if self.role.owns_background() {
search.index_document(&workspace_id, &doc_id).await
} else {
search.project_document_only(&workspace_id, &doc_id).await
};
result.map_err(to_napi_error)
}
#[napi]
pub async fn delete_search_document(&self, workspace_id: String, doc_id: String) -> Result<()> {
let search = self.search_runtime().await?;
let result = if self.role.owns_background() {
search.delete_document(&workspace_id, &doc_id).await
} else {
search.delete_document_only(&workspace_id, &doc_id).await
};
result.map_err(to_napi_error)
}
#[napi]
pub async fn reconcile_search_workspace(&self, workspace_id: String) -> Result<()> {
pub async fn reconcile_search_projection(&self, limit: Option<i32>) -> Result<i32> {
self.require_background()?;
self
.search_runtime()
.await?
.reconcile_workspace(permission::SystemSearchCapability::ReconcileIndex, &workspace_id)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn delete_search_workspace(&self, workspace_id: String) -> Result<()> {
self.require_background()?;
self
.search_runtime()
.await?
.delete_workspace(&workspace_id)
.reconcile_pending(limit.unwrap_or(100))
.await
.map(|count| count as i32)
.map_err(to_napi_error)
}

View File

@@ -29,7 +29,7 @@ impl PermissionAuthorizer {
) -> RuntimeResult<AuthorizedSearchScope> {
match actor {
SearchActor::User { user_id } => {
let snapshot = self.store.search_snapshot(workspace_id, user_id).await?;
let snapshot = self.store.permission_snapshot(workspace_id, user_id, &[]).await?;
let owner_or_admin = matches!(snapshot.evaluation.workspace.role.as_deref(), Some("owner" | "admin"))
&& snapshot.evaluation.workspace.member_state.as_deref() == Some("active");
let decision = evaluate_permission(snapshot.evaluation)
@@ -54,17 +54,12 @@ impl PermissionAuthorizer {
};
Ok(AuthorizedSearchScope {
workspace_id: workspace_id.to_string(),
permission_revision: snapshot.revision,
docs,
})
}
}
}
pub(in crate::runtime::backend_runtime) async fn revision(&self, workspace_id: &str) -> RuntimeResult<i64> {
self.store.revision(workspace_id).await
}
pub(in crate::runtime::backend_runtime) async fn filter_readable_docs(
&self,
workspace_id: &str,
@@ -72,6 +67,9 @@ impl PermissionAuthorizer {
doc_ids: Vec<String>,
) -> RuntimeResult<BTreeSet<String>> {
let snapshot = self.store.permission_snapshot(workspace_id, user_id, &doc_ids).await?;
if snapshot.capability == DocAclCapability::Unknown {
return Err(RuntimeError::SearchPermissionUnavailable);
}
let output = evaluate_permission(snapshot.evaluation).map_err(|_| RuntimeError::SearchPermissionUnavailable)?;
Ok(
output

View File

@@ -3,9 +3,7 @@ mod store;
mod types;
pub(super) use authorizer::PermissionAuthorizer;
#[cfg(test)]
pub(super) use types::AclPredicate;
pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor, SystemSearchCapability};
pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor};
#[cfg(test)]
mod tests;

View File

@@ -20,10 +20,6 @@ impl PermissionStore {
Self { pool }
}
pub(super) async fn search_snapshot(&self, workspace_id: &str, user_id: &str) -> RuntimeResult<PermissionSnapshot> {
self.permission_snapshot(workspace_id, user_id, &[]).await
}
pub(super) async fn permission_snapshot(
&self,
workspace_id: &str,
@@ -40,12 +36,10 @@ impl PermissionStore {
.await
.map_err(|error| RuntimeError::database("configure permission snapshot", error))?;
let row = sqlx::query(
r#"SELECT revision.revision,
policy.visibility, coalesce(policy.sharing_enabled, true) AS sharing_enabled,
r#"SELECT policy.visibility, coalesce(policy.sharing_enabled, true) AS sharing_enabled,
coalesce(policy.member_default_doc_role, 'manager') AS member_default_doc_role,
member.role, member.state
FROM workspaces workspace
LEFT JOIN workspace_permission_revisions revision ON revision.workspace_id=workspace.id
LEFT JOIN workspace_access_policies policy ON policy.workspace_id=workspace.id
LEFT JOIN LATERAL (
SELECT role,state FROM workspace_members
@@ -60,10 +54,6 @@ impl PermissionStore {
.await
.map_err(|error| RuntimeError::database("load workspace permission facts", error))?
.ok_or_else(|| RuntimeError::invalid_input("workspace_not_found"))?;
let revision = row
.try_get::<Option<i64>, _>("revision")
.map_err(|error| RuntimeError::database("decode permission revision", error))?
.ok_or_else(|| RuntimeError::invalid_state("permission_state_unavailable"))?;
let sharing_enabled: bool = row
.try_get("sharing_enabled")
.map_err(|error| RuntimeError::database("decode workspace sharing", error))?;
@@ -134,7 +124,6 @@ impl PermissionStore {
member_state.as_deref() == Some("active") && matches!(role.as_deref(), Some("member" | "admin" | "owner"));
Ok(PermissionSnapshot {
revision,
capability,
evaluation: PermissionEvaluationInputV1 {
version: 1,
@@ -163,15 +152,6 @@ impl PermissionStore {
sharing_enabled,
})
}
pub(super) async fn revision(&self, workspace_id: &str) -> RuntimeResult<i64> {
sqlx::query_scalar("SELECT revision FROM workspace_permission_revisions WHERE workspace_id=$1")
.bind(workspace_id)
.fetch_optional(&self.pool)
.await
.map_err(|error| RuntimeError::database("read permission revision", error))?
.ok_or_else(|| RuntimeError::invalid_state("permission_state_unavailable"))
}
}
async fn load_doc_acl_capability(

View File

@@ -7,7 +7,14 @@ static PERMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_
async fn setup() -> Option<(PgPool, String, String)> {
let database_url = std::env::var("DATABASE_URL").ok()?;
let pool = PgPool::connect(&database_url).await.unwrap();
crate::runtime::migrations::migrate_search_tables(&pool).await.unwrap();
let legacy_relations: Vec<Option<String>> = sqlx::query_scalar(
"SELECT to_regclass(name) FROM \
unnest(ARRAY['workspace_permission_revisions','workspace_permission_changes','search_runtime_generations']) name",
)
.fetch_all(&pool)
.await
.unwrap();
assert!(legacy_relations.into_iter().all(|relation| relation.is_none()));
let suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("search-permission-user-{suffix}");
let workspace_id = format!("search-permission-workspace-{suffix}");
@@ -67,7 +74,15 @@ async fn non_team_is_all_and_team_uses_projected_acl() {
};
assert_eq!(predicate.actor_user_id, user_id);
assert!(predicate.active_member);
assert!(team.permission_revision > free.permission_revision);
sqlx::query("UPDATE workspace_members SET role='admin' WHERE workspace_id=$1 AND user_id=$2")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
let admin = authorizer.authorize_search(&actor, &workspace_id).await.unwrap();
assert_eq!(admin.docs, DocReadScope::All);
}
#[tokio::test]
@@ -112,13 +127,20 @@ async fn inactive_member_is_denied_and_unknown_capability_fails_closed() {
}
#[tokio::test]
async fn fact_changes_advance_revision_and_write_ordered_change() {
async fn canonical_doc_acl_facts_are_evaluated_without_search_state() {
let _guard = PERMISSION_TEST_LOCK.lock().await;
let Some((pool, workspace_id, user_id)) = setup().await else {
return;
};
let authorizer = PermissionAuthorizer::new(pool.clone());
let before = authorizer.revision(&workspace_id).await.unwrap();
sqlx::query(
"INSERT INTO doc_access_policies(workspace_id,doc_id,visibility,member_default_role) \
VALUES($1,'doc','private','none'),($1,'hidden','private','none')",
)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO doc_grants(workspace_id,doc_id,principal_type,principal_id,role) VALUES($1,'doc','user',$2,'reader')",
)
@@ -127,39 +149,23 @@ async fn fact_changes_advance_revision_and_write_ordered_change() {
.execute(&pool)
.await
.unwrap();
let after = authorizer.revision(&workspace_id).await.unwrap();
assert_eq!(after, before + 1);
let change: (i64, Option<String>, String) = sqlx::query_as(
"SELECT revision,doc_id,scope FROM workspace_permission_changes WHERE workspace_id=$1 AND revision=$2",
)
.bind(&workspace_id)
.bind(after)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(change, (after, Some("doc".to_string()), "doc_grant".to_string()));
sqlx::query("UPDATE workspace_members SET updated_at=now() WHERE workspace_id=$1 AND user_id=$2")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
let readable = authorizer
.filter_readable_docs(&workspace_id, &user_id, vec!["doc".to_string(), "hidden".to_string()])
.await
.unwrap();
assert_eq!(authorizer.revision(&workspace_id).await.unwrap(), after);
assert_eq!(readable, ["doc".to_string()].into_iter().collect());
let moved_workspace_id = format!("{workspace_id}-moved");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&moved_workspace_id)
.execute(&pool)
.await
.unwrap();
assert_eq!(authorizer.revision(&moved_workspace_id).await.unwrap(), 0);
sqlx::query("UPDATE doc_grants SET workspace_id=$1 WHERE workspace_id=$2 AND doc_id='doc'")
.bind(&moved_workspace_id)
sqlx::query("UPDATE doc_access_policies SET member_default_role='reader' WHERE workspace_id=$1 AND doc_id='hidden'")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
assert_eq!(authorizer.revision(&workspace_id).await.unwrap(), after + 1);
assert_eq!(authorizer.revision(&moved_workspace_id).await.unwrap(), 1);
let readable = authorizer
.filter_readable_docs(&workspace_id, &user_id, vec!["doc".to_string(), "hidden".to_string()])
.await
.unwrap();
assert_eq!(
readable,
["doc".to_string(), "hidden".to_string()].into_iter().collect()
);
}

View File

@@ -1,10 +1,5 @@
use crate::permission::PermissionEvaluationInputV1;
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum SystemSearchCapability {
ReconcileIndex,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum SearchActor {
User { user_id: String },
@@ -33,12 +28,10 @@ pub(in crate::runtime::backend_runtime) enum DocReadScope {
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) struct AuthorizedSearchScope {
pub(in crate::runtime::backend_runtime) workspace_id: String,
pub(in crate::runtime::backend_runtime) permission_revision: i64,
pub(in crate::runtime::backend_runtime) docs: DocReadScope,
}
pub(super) struct PermissionSnapshot {
pub(super) revision: i64,
pub(super) capability: DocAclCapability,
pub(super) evaluation: PermissionEvaluationInputV1,
pub(super) actor_user_id: String,

View File

@@ -5,7 +5,7 @@ use super::{
BackendRuntime, RuntimeError, RuntimeInviteAbuseClaimedAction, RuntimeResult, napi_error, workspace_subject_key,
};
async fn invite_abuse_user_quarantined_or_banned(pool: &PgPool, user_id: &str) -> RuntimeResult<bool> {
pub(super) async fn invite_abuse_user_quarantined_or_banned(pool: &PgPool, user_id: &str) -> RuntimeResult<bool> {
let row: Option<i32> = sqlx::query_scalar(
r#"
SELECT 1
@@ -22,7 +22,7 @@ async fn invite_abuse_user_quarantined_or_banned(pool: &PgPool, user_id: &str) -
Ok(row.is_some())
}
async fn invite_abuse_workspace_quarantined(pool: &PgPool, workspace_id: &str) -> RuntimeResult<bool> {
pub(super) async fn invite_abuse_workspace_quarantined(pool: &PgPool, workspace_id: &str) -> RuntimeResult<bool> {
let row: Option<i32> = sqlx::query_scalar(
r#"
SELECT 1

View File

@@ -6,6 +6,7 @@ mod reservation;
mod workspace_invite;
mod workspace_invite_policy;
use invite_abuse_actions::{invite_abuse_user_quarantined_or_banned, invite_abuse_workspace_quarantined};
use mail_delivery::{build_mail_scopes, decision_from_violation as mail_decision_from_violation, mail_class};
use napi::Result;
use reservation::{
@@ -15,8 +16,8 @@ use reservation::{
use sha2::{Digest, Sha256};
use workspace_invite_policy::{
ActorFacts, InviteAbuseDecision, InviteActivityFacts, QuotaFacts, WorkspaceFacts, build_invite_scopes,
evaluate_projection, high_confidence_invite_abuse, invite_commit_usage_for_scope, source_cohort_subject_key,
subject_hash, sum_domains,
evaluate_projection, high_confidence_invite_abuse, invite_commit_usage_for_scope, new_account_action_retry_after,
source_cohort_subject_key, subject_hash, sum_domains,
};
#[cfg(test)]
@@ -26,7 +27,8 @@ pub(super) use super::{
types::{
RuntimeInviteAbuseActionRequired, RuntimeInviteAbuseClaimedAction, RuntimeMailDeliveryQuotaDecision,
RuntimeMailDeliveryQuotaInput, RuntimeQuotaSourceInput, RuntimeQuotaTargetDomainInput,
RuntimeWorkspaceInviteQuotaDecision, RuntimeWorkspaceInviteQuotaInput, RuntimeWorkspaceInviteQuotaUsage,
RuntimeWorkspaceActionDecision, RuntimeWorkspaceInviteQuotaDecision, RuntimeWorkspaceInviteQuotaInput,
RuntimeWorkspaceInviteQuotaUsage,
},
};

View File

@@ -5,19 +5,32 @@ use sqlx::{PgPool, Row};
use super::{
ActorFacts, BackendRuntime, InviteAbuseDecision, InviteActivityFacts, InviteQuotaConfig, QuotaFacts, QuotaViolation,
RuntimeError, RuntimeInviteAbuseActionRequired, RuntimeResult, RuntimeWorkspaceInviteQuotaDecision,
RuntimeWorkspaceInviteQuotaInput, RuntimeWorkspaceInviteQuotaUsage, WorkspaceFacts, build_invite_scopes,
commit_reservation, evaluate_projection, high_confidence_invite_abuse, invite_commit_usage_for_scope, napi_error,
normalize_domain, release_reservation, reserve_scopes, short_hash, source_cohort_subject_key, source_prefix,
subject_hash, sum_domains, workspace_subject_key,
RuntimeError, RuntimeInviteAbuseActionRequired, RuntimeResult, RuntimeWorkspaceActionDecision,
RuntimeWorkspaceInviteQuotaDecision, RuntimeWorkspaceInviteQuotaInput, RuntimeWorkspaceInviteQuotaUsage,
WorkspaceFacts, build_invite_scopes, commit_reservation, evaluate_projection, high_confidence_invite_abuse,
invite_abuse_user_quarantined_or_banned, invite_abuse_workspace_quarantined, invite_commit_usage_for_scope,
napi_error, new_account_action_retry_after, normalize_domain, release_reservation, reserve_scopes, short_hash,
source_cohort_subject_key, source_prefix, subject_hash, sum_domains, workspace_subject_key,
};
async fn load_actor(pool: &PgPool, user_id: &str) -> RuntimeResult<ActorFacts> {
let row = sqlx::query(
r#"
SELECT email, created_at, registered, email_verified IS NOT NULL AS email_verified, disabled
SELECT
users.email,
users.created_at,
users.registered,
users.email_verified IS NOT NULL AS email_verified,
users.disabled,
CASE
WHEN quota.known
AND NOT quota.stale
AND (quota.stale_after IS NULL OR quota.stale_after > clock_timestamp())
THEN quota.plan
END AS quota_plan
FROM users
WHERE id = $1
LEFT JOIN effective_user_quota_states quota ON quota.user_id = users.id
WHERE users.id = $1
"#,
)
.bind(user_id)
@@ -32,6 +45,7 @@ async fn load_actor(pool: &PgPool, user_id: &str) -> RuntimeResult<ActorFacts> {
registered: row.get("registered"),
email_verified: row.get("email_verified"),
disabled: row.get("disabled"),
quota_plan: row.get("quota_plan"),
})
}
@@ -285,6 +299,49 @@ fn decision_from_violation(violation: QuotaViolation, reason: &str) -> RuntimeWo
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn evaluate_workspace_action_v1(
&self,
actor_user_id: String,
workspace_id: String,
) -> Result<RuntimeWorkspaceActionDecision> {
let runtime_config = self.config()?;
let pool = self.pool().await?;
if invite_abuse_user_quarantined_or_banned(&pool, &actor_user_id).await? {
return Ok(RuntimeWorkspaceActionDecision {
allowed: false,
retry_after_seconds: None,
reason: Some("abuse_subject".to_string()),
});
}
if invite_abuse_workspace_quarantined(&pool, &workspace_id).await? {
return Ok(RuntimeWorkspaceActionDecision {
allowed: false,
retry_after_seconds: None,
reason: Some("abuse_workspace".to_string()),
});
}
let now: DateTime<Utc> = sqlx::query_scalar("SELECT clock_timestamp()")
.fetch_one(&pool)
.await
.map_err(|err| RuntimeError::database("failed to read database clock", err))?;
let actor = load_actor(&pool, &actor_user_id).await?;
let quota = load_quota(&pool, &workspace_id).await?;
let current_quota = quota.as_ref().filter(|quota| evaluate_projection(quota, now).is_none());
let retry_after_seconds = new_account_action_retry_after(
runtime_config.deployment,
&runtime_config.invite_quota,
&actor,
current_quota,
now,
);
Ok(RuntimeWorkspaceActionDecision {
allowed: retry_after_seconds.is_none(),
retry_after_seconds,
reason: retry_after_seconds.map(|_| "new_account_action_delay".to_string()),
})
}
#[napi]
pub async fn assert_workspace_invite_quota_v1(
&self,
@@ -388,6 +445,22 @@ impl BackendRuntime {
action_required: None,
});
}
if let Some(retry_after_seconds) =
new_account_action_retry_after(runtime_config.deployment, config, &actor, Some(&quota), now)
{
return Ok(RuntimeWorkspaceInviteQuotaDecision {
allowed: false,
reservation_id: None,
retry_after_seconds: Some(retry_after_seconds),
reason: Some("new_account_action_delay".to_string()),
scope_key: None,
window_seconds: None,
limit: None,
current: None,
requested: Some(input.target_count),
action_required: None,
});
}
if let Some(abuse_decision) = high_confidence_invite_abuse(&input, &actor, config) {
let reason = abuse_decision.reason;
let scope_key = match abuse_decision.subject_kind {

View File

@@ -10,6 +10,7 @@ use super::{
InviteQuotaConfig, RuntimeQuotaTargetDomainInput, RuntimeWorkspaceInviteQuotaInput, ScopeLimit, bucket_seconds,
high_risk_domain, napi_error, normalize_domain, scope, short_hash, source_prefix, workspace_subject_key,
};
use crate::llm::Deployment;
#[derive(Clone, Debug)]
pub(super) struct ActorFacts {
@@ -18,6 +19,7 @@ pub(super) struct ActorFacts {
pub(super) registered: bool,
pub(super) email_verified: bool,
pub(super) disabled: bool,
pub(super) quota_plan: Option<String>,
}
#[derive(Clone, Debug)]
@@ -108,9 +110,6 @@ fn base_invite_limits(
let mut per_day = 15;
let mut per_week = 30;
if account_age < Duration::hours(24) {
return (0, 0, 0, 0);
}
if !actor.email_verified {
single = 1;
per_hour = 1;
@@ -185,6 +184,31 @@ pub(super) fn evaluate_projection(quota: &QuotaFacts, now: DateTime<Utc>) -> Opt
None
}
pub(super) fn new_account_action_retry_after(
deployment: Deployment,
config: &InviteQuotaConfig,
actor: &ActorFacts,
workspace_quota: Option<&QuotaFacts>,
now: DateTime<Utc>,
) -> Option<i32> {
if deployment == Deployment::SelfHosted || config.new_account_action_delay_seconds <= 0 {
return None;
}
if actor
.quota_plan
.as_deref()
.is_some_and(|plan| matches!(plan, "pro" | "lifetime_pro" | "ai"))
|| workspace_quota.map(|quota| quota.plan.as_str()).is_some_and(|plan| {
matches!(plan, "pro" | "lifetime_pro" | "ai") || plan.contains("team") && !plan.contains("trial")
})
{
return None;
}
let remaining = config.new_account_action_delay_seconds - (now - actor.created_at).num_seconds();
(remaining > 0).then(|| remaining.min(i64::from(i32::MAX)) as i32)
}
pub(super) fn build_invite_scopes(
input: &RuntimeWorkspaceInviteQuotaInput,
actor: &ActorFacts,
@@ -442,6 +466,7 @@ mod tests {
registered: true,
email_verified: true,
disabled: false,
quota_plan: None,
}
}
@@ -467,7 +492,7 @@ mod tests {
}
#[test]
fn seat_based_weekly_limit_binds_paid_team_and_high_risk_domain() {
fn invite_scopes_apply_plan_ceiling_domain_risk_and_graduated_limits() {
let now = Utc.with_ymd_and_hms(2026, 7, 6, 0, 0, 0).single().unwrap();
let input = RuntimeWorkspaceInviteQuotaInput {
actor_user_id: "u1".to_string(),
@@ -500,6 +525,67 @@ mod tests {
.find(|scope| scope.scope_key == "invite:quota_subject_domain:workspace:w1:qq.com")
.unwrap();
assert_eq!(high_risk.limit, 5);
let fresh_actor_scopes = build_invite_scopes(
&input,
&user(now - Duration::hours(1)),
&workspace(now - Duration::hours(1)),
&quota("paid_team", 10),
&InviteActivityFacts::default(),
&invite_config(),
now,
)
.unwrap();
assert_eq!(fresh_actor_scopes[0].limit, 3);
let mut fresh_actor = user(now - Duration::hours(1));
assert_eq!(
new_account_action_retry_after(
Deployment::Cloud,
&invite_config(),
&fresh_actor,
Some(&quota("free", 3)),
now,
),
Some(23 * 60 * 60)
);
assert_eq!(
new_account_action_retry_after(
Deployment::SelfHosted,
&invite_config(),
&fresh_actor,
Some(&quota("free", 3)),
now,
),
None
);
assert_eq!(
new_account_action_retry_after(
Deployment::Cloud,
&invite_config(),
&fresh_actor,
Some(&quota("paid_team", 10)),
now,
),
None
);
let mut no_delay = invite_config();
no_delay.new_account_action_delay_seconds = 0;
assert_eq!(
new_account_action_retry_after(Deployment::Cloud, &no_delay, &fresh_actor, Some(&quota("free", 3)), now,),
None
);
fresh_actor.quota_plan = Some("pro".to_string());
assert_eq!(
new_account_action_retry_after(
Deployment::Cloud,
&invite_config(),
&fresh_actor,
Some(&quota("free", 3)),
now,
),
None
);
}
#[test]

View File

@@ -250,7 +250,15 @@ mod tests {
};
let _guard = crate::runtime::migrations::EMBEDDING_TEST_LOCK.lock().await;
let pool = PgPool::connect(&database_url).await.unwrap();
crate::runtime::migrations::migrate_search_tables(&pool).await.unwrap();
let legacy_relations: Vec<Option<String>> = sqlx::query_scalar(
"SELECT to_regclass(name) FROM \
unnest(ARRAY['workspace_permission_revisions','workspace_permission_changes','search_runtime_generations']) \
name",
)
.fetch_all(&pool)
.await
.unwrap();
assert!(legacy_relations.into_iter().all(|relation| relation.is_none()));
let suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("scope-user-{suffix}");
let collaborator_id = format!("scope-collaborator-{suffix}");

View File

@@ -1,193 +0,0 @@
use napi::bindgen_prelude::Buffer;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row};
use super::{SCHEMA_FINGERPRINT, store::SearchTable};
use crate::{
runtime::{RuntimeError, RuntimeResult},
search_index::EmbeddedSearchIndex,
};
const DIRTY_CHANGE_THRESHOLD: i64 = 1_000;
const RETAINED_CHANGES: i64 = 10_000;
const MAX_CHECKPOINT_AGE_SECONDS: i64 = 300;
pub(super) async fn restore(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
table: SearchTable,
) -> RuntimeResult<Option<i64>> {
let row = sqlx::query(
"SELECT source_cursor, checkpoint_blob, checksum FROM search_runtime_checkpoints WHERE table_key=$1 AND \
schema_fingerprint=$2",
)
.bind(table.as_str())
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load embedded search checkpoint", error))?;
let Some(row) = row else { return Ok(None) };
let cursor: i64 = row
.try_get("source_cursor")
.map_err(|error| RuntimeError::database("decode search checkpoint cursor", error))?;
let bytes: Vec<u8> = row
.try_get("checkpoint_blob")
.map_err(|error| RuntimeError::database("decode search checkpoint blob", error))?;
let checksum: String = row
.try_get("checksum")
.map_err(|error| RuntimeError::database("decode search checkpoint checksum", error))?;
if digest(&bytes) != checksum {
return Ok(None);
}
if embedded
.restore(table.as_str().to_string(), Buffer::from(bytes))
.await
.is_err()
{
return Ok(None);
}
Ok(Some(cursor))
}
pub(super) async fn persist(pool: &PgPool, embedded: &EmbeddedSearchIndex, cursors: [i64; 2]) -> RuntimeResult<()> {
for table in SearchTable::ORDERED {
let cursor = cursors[table.cursor_index()];
let persisted: Option<(i64, bool)> = sqlx::query_as(
"SELECT source_cursor, updated_at < now() - make_interval(secs => $2) AS expired FROM \
search_runtime_checkpoints WHERE table_key=$1",
)
.bind(table.as_str())
.bind(MAX_CHECKPOINT_AGE_SECONDS as f64)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load persisted checkpoint cursor", error))?;
let (persisted_cursor, expired) = persisted.unwrap_or((0, true));
if cursor <= persisted_cursor || (cursor - persisted_cursor < DIRTY_CHANGE_THRESHOLD && !expired) {
continue;
}
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin embedded checkpoint", error))?;
sqlx::query("SET LOCAL synchronous_commit = off")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("configure embedded checkpoint commit", error))?;
let leader: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock(hashtextextended($1, 0))")
.bind(format!("search-checkpoint/{}", table.as_str()))
.fetch_one(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("acquire embedded checkpoint lease", error))?;
if !leader {
continue;
}
let persisted: Option<(i64, bool)> = sqlx::query_as(
"SELECT source_cursor, updated_at < now() - make_interval(secs => $2) AS expired FROM \
search_runtime_checkpoints WHERE table_key=$1",
)
.bind(table.as_str())
.bind(MAX_CHECKPOINT_AGE_SECONDS as f64)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("reload persisted checkpoint cursor", error))?;
let (persisted_cursor, expired) = persisted.unwrap_or((0, true));
if cursor <= persisted_cursor || (cursor - persisted_cursor < DIRTY_CHANGE_THRESHOLD && !expired) {
continue;
}
embedded.optimize(table.as_str().to_string()).await?;
let checkpoint = embedded.checkpoint(table.as_str().to_string()).await?;
let bytes = checkpoint.data.to_vec();
let saved = sqlx::query(
r#"INSERT INTO search_runtime_checkpoints
(table_key,schema_fingerprint,source_cursor,checkpoint_sequence,checkpoint_blob,checksum,blob_size)
VALUES ($1,$2,$3,$4,$5,$6,$7)
ON CONFLICT (table_key) DO UPDATE SET schema_fingerprint=EXCLUDED.schema_fingerprint,
source_cursor=EXCLUDED.source_cursor,checkpoint_sequence=EXCLUDED.checkpoint_sequence,
checkpoint_blob=EXCLUDED.checkpoint_blob,checksum=EXCLUDED.checksum,blob_size=EXCLUDED.blob_size,updated_at=now()
WHERE search_runtime_checkpoints.source_cursor < EXCLUDED.source_cursor"#,
)
.bind(table.as_str())
.bind(SCHEMA_FINGERPRINT)
.bind(cursor)
.bind(checkpoint.sequence)
.bind(&bytes)
.bind(digest(&bytes))
.bind(bytes.len() as i64)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("persist embedded search checkpoint", error))?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit embedded search checkpoint", error))?;
if saved.rows_affected() == 1 {
embedded
.mark_checkpoint_persisted(table.as_str().to_string(), checkpoint.sequence)
.await?;
}
}
gc(pool).await
}
pub(super) async fn gc(pool: &PgPool) -> RuntimeResult<()> {
for table in SearchTable::ORDERED {
let minimum: Option<i64> = sqlx::query_scalar(
r#"SELECT COALESCE(MIN(watermark),0) FROM (
SELECT c.source_cursor AS watermark FROM search_runtime_provider_cursors c
JOIN search_runtime_generations g USING(generation_id)
WHERE c.table_key=$1 AND g.provider<>'embedded' AND g.state IN ('active','pending')
UNION ALL
SELECT checkpoint.source_cursor FROM search_runtime_checkpoints checkpoint
WHERE checkpoint.table_key=$1 AND EXISTS (
SELECT 1 FROM search_runtime_generations generation
WHERE generation.provider='embedded' AND generation.state IN ('active','pending')
)
) retained"#,
)
.bind(table.as_str())
.fetch_one(pool)
.await
.map_err(|error| RuntimeError::database("compute search retention watermark", error))?;
let retained_from = minimum.unwrap_or(0).saturating_sub(RETAINED_CHANGES);
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search change gc", error))?;
sqlx::query("DELETE FROM search_runtime_changes WHERE table_key=$1 AND stream_sequence <= $2")
.bind(table.as_str())
.bind(retained_from)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("gc search changes", error))?;
sqlx::query("UPDATE search_runtime_streams SET retained_from=GREATEST(retained_from,$2) WHERE table_key=$1")
.bind(table.as_str())
.bind(retained_from)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("advance search retention watermark", error))?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit search change gc", error))?;
}
sqlx::query(
r#"DELETE FROM workspace_permission_changes permission_change
USING workspace_permission_revisions head
WHERE permission_change.workspace_id=head.workspace_id
AND permission_change.revision <= (
SELECT COALESCE(MIN(cursor.permission_revision),head.revision)
FROM search_runtime_permission_cursors cursor
JOIN search_runtime_generations generation USING(generation_id)
WHERE cursor.workspace_id=permission_change.workspace_id AND generation.state IN ('active','pending')
) - $1"#,
)
.bind(RETAINED_CHANGES)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("gc search permission changes", error))?;
Ok(())
}
fn digest(bytes: &[u8]) -> String {
Sha256::digest(bytes).iter().map(|byte| format!("{byte:02x}")).collect()
}

View File

@@ -3,7 +3,7 @@ use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{SCHEMA_FINGERPRINT, provider::RemoteProvider, types::SearchTable};
use super::{SCHEMA_FINGERPRINT, SearchProvider, SearchTable};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
#[derive(Clone)]
@@ -22,82 +22,120 @@ impl ActiveGeneration {
}
}
pub(super) async fn prepare(
pub(super) async fn ensure(
pool: &PgPool,
config: &SearchRuntimeConfig,
remote: Option<&RemoteProvider>,
remote: Option<&SearchProvider>,
rebuild_embedded: bool,
) -> RuntimeResult<ActiveGeneration> {
let fingerprint = config_fingerprint(config);
let expected_config_hash = config_hash(config);
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search generation", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-runtime-generation', 0))")
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("lock search generation", error))?;
let existing = sqlx::query(
r#"SELECT generation_id,provider,manifest FROM search_runtime_generations
WHERE state IN ('active','pending') AND provider=$1 AND config_fingerprint=$2 AND schema_fingerprint=$3
ORDER BY (state='active') DESC LIMIT 1"#,
r#"SELECT id,provider,config_hash,schema_version,manifest,state
FROM search_projection.generations
WHERE state='building'
ORDER BY created_at DESC
LIMIT 1"#,
)
.bind(&config.provider)
.bind(&fingerprint)
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load search generation", error))?;
let generation = if let Some(row) = existing {
decode(row)?
} else {
let pending: bool =
sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM search_runtime_generations WHERE state='pending')")
.fetch_one(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("check pending search generation", error))?;
if pending {
return Err(RuntimeError::invalid_state("search_generation_change_in_progress"));
}
let generation_id = Uuid::new_v4();
let suffix = generation_id.simple().to_string();
let manifest = if config.provider == "embedded" {
json!({"doc":"doc","block":"block"})
let provider: String = row
.try_get("provider")
.map_err(|error| RuntimeError::database("decode search generation provider", error))?;
let existing_config_hash: Vec<u8> = row
.try_get("config_hash")
.map_err(|error| RuntimeError::database("decode search generation config hash", error))?;
let schema_version: i32 = row
.try_get("schema_version")
.map_err(|error| RuntimeError::database("decode search generation schema version", error))?;
if provider != config.provider
|| existing_config_hash != expected_config_hash
|| schema_version != SCHEMA_FINGERPRINT
{
let id: Uuid = row
.try_get("id")
.map_err(|error| RuntimeError::database("decode superseded search generation id", error))?;
sqlx::query(
"UPDATE search_projection.generations SET state='failed', last_error='search generation superseded by \
configuration' WHERE id=$1 AND state='building'",
)
.bind(id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("fail superseded search generation", error))?;
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
} else {
json!({
"doc":format!("affine_search_doc_{suffix}"),
"block":format!("affine_search_block_{suffix}"),
})
};
sqlx::query(
r#"INSERT INTO search_runtime_generations
(generation_id,provider,state,config_fingerprint,schema_fingerprint,manifest)
VALUES ($1,$2,'pending',$3,$4,$5)"#,
)
.bind(generation_id)
.bind(&config.provider)
.bind(&fingerprint)
.bind(SCHEMA_FINGERPRINT)
.bind(&manifest)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("create pending search generation", error))?;
for table in [SearchTable::Doc, SearchTable::Block] {
sqlx::query("INSERT INTO search_runtime_provider_cursors(generation_id,table_key) VALUES ($1,$2)")
.bind(generation_id)
.bind(table.as_str())
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("initialize search generation cursor", error))?;
let state: String = row
.try_get("state")
.map_err(|error| RuntimeError::database("decode search generation state", error))?;
if rebuild_embedded {
if state == "building" {
sqlx::query(
"UPDATE search_projection.generations SET state='failed', last_error='embedded generation lost on \
restart' WHERE id=$1 AND state='building'",
)
.bind(
row
.try_get::<Uuid, _>("id")
.map_err(|error| RuntimeError::database("decode search generation id", error))?,
)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("fail stale embedded generation", error))?;
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
} else {
decode(row)?
}
} else {
decode(row)?
}
}
ActiveGeneration {
id: generation_id,
manifest,
} else {
let active = sqlx::query(
r#"SELECT id,provider,config_hash,schema_version,manifest
FROM search_projection.generations
WHERE state='active'
ORDER BY activated_at DESC NULLS LAST LIMIT 1"#,
)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load active search generation", error))?;
if let Some(row) = active {
let provider: String = row
.try_get("provider")
.map_err(|error| RuntimeError::database("decode active search provider", error))?;
let existing_config_hash: Vec<u8> = row
.try_get("config_hash")
.map_err(|error| RuntimeError::database("decode active search config hash", error))?;
let schema_version: i32 = row
.try_get("schema_version")
.map_err(|error| RuntimeError::database("decode active search schema version", error))?;
if !rebuild_embedded
&& provider == config.provider
&& existing_config_hash == expected_config_hash
&& schema_version == SCHEMA_FINGERPRINT
{
decode(row)?
} else {
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
}
} else {
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
}
};
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit pending search generation", error))?;
.map_err(|error| RuntimeError::database("commit search generation", error))?;
if let Some(remote) = remote {
for table in [SearchTable::Doc, SearchTable::Block] {
@@ -114,14 +152,14 @@ pub(super) async fn load_active(
pool: &PgPool,
config: &SearchRuntimeConfig,
) -> RuntimeResult<Option<ActiveGeneration>> {
let fingerprint = config_fingerprint(config);
let row = sqlx::query(
r#"SELECT generation_id,provider,manifest FROM search_runtime_generations
WHERE state='active' AND provider=$1 AND config_fingerprint=$2 AND schema_fingerprint=$3
r#"SELECT id,provider,config_hash,manifest
FROM search_projection.generations
WHERE state='active' AND provider=$1 AND config_hash=$2 AND schema_version=$3
ORDER BY activated_at DESC NULLS LAST LIMIT 1"#,
)
.bind(&config.provider)
.bind(&fingerprint)
.bind(config_hash(config))
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(pool)
.await
@@ -129,28 +167,122 @@ pub(super) async fn load_active(
row.map(decode).transpose()
}
pub(super) async fn fail(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
sqlx::query("UPDATE search_runtime_generations SET state='failed' WHERE generation_id=$1 AND state='pending'")
.bind(generation.id)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("fail pending search generation", error))?;
async fn fail(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
sqlx::query(
"UPDATE search_projection.generations SET state='failed', last_error=$2 WHERE id=$1 AND state NOT IN \
('active','draining')",
)
.bind(generation.id)
.bind("search generation build failed")
.execute(pool)
.await
.map_err(|error| RuntimeError::database("fail search generation", error))?;
Ok(())
}
pub(super) async fn activate(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
pub(super) async fn activate(
pool: &PgPool,
generation: &ActiveGeneration,
config: &SearchRuntimeConfig,
) -> RuntimeResult<()> {
let expected_config_hash = config_hash(config);
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search generation activation", error))?;
sqlx::query("UPDATE search_runtime_generations SET state='draining' WHERE state='active' AND generation_id<>$1")
.bind(generation.id)
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("drain previous search generation", error))?;
.map_err(|error| RuntimeError::database("lock search generation activation", error))?;
let row = sqlx::query(
r#"SELECT generation.state,generation.provider,generation.config_hash,generation.schema_version,
generation.scan_high_water_sid IS NOT NULL
AND generation.scan_cursor_sid IS NOT NULL
AND generation.scan_cursor_sid >= generation.scan_high_water_sid AS scan_complete,
NOT EXISTS (
SELECT 1
FROM search_projection.workspace_states state
WHERE state.generation_id=generation.id
AND (
NOT state.covered
OR state.required_permission_version > state.applied_permission_version
OR state.pending_scope <> 'none'
OR state.last_error IS NOT NULL
)
) AS workspaces_covered,
NOT EXISTS (
SELECT 1
FROM workspaces workspace
WHERE workspace.sid <= COALESCE(generation.scan_high_water_sid,0)
AND NOT EXISTS (
SELECT 1
FROM search_projection.workspace_states state
WHERE state.generation_id=generation.id AND state.workspace_id=workspace.id
)
) AS workspaces_seeded,
NOT EXISTS (
SELECT 1 FROM search_projection.document_states state
WHERE state.generation_id=generation.id
AND (state.target_source_version <> state.published_source_version
OR state.target_source_exists <> state.published_source_exists
OR state.target_permission_version <> state.published_permission_version)
) AS publications_complete
FROM search_projection.generations generation
WHERE generation.id=$1
FOR UPDATE"#,
)
.bind(generation.id)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load search generation activation state", error))?;
let Some(row) = row else {
return Err(RuntimeError::SearchIndexNotReady);
};
let state: String = row
.try_get("state")
.map_err(|error| RuntimeError::database("decode search generation activation state", error))?;
let provider: String = row
.try_get("provider")
.map_err(|error| RuntimeError::database("decode search generation activation provider", error))?;
let row_config_hash: Vec<u8> = row
.try_get("config_hash")
.map_err(|error| RuntimeError::database("decode search generation activation config hash", error))?;
let schema_version: i32 = row
.try_get("schema_version")
.map_err(|error| RuntimeError::database("decode search generation activation schema", error))?;
let scan_complete: bool = row
.try_get("scan_complete")
.map_err(|error| RuntimeError::database("decode search generation scan state", error))?;
let workspaces_covered: bool = row
.try_get("workspaces_covered")
.map_err(|error| RuntimeError::database("decode search generation workspace coverage", error))?;
let workspaces_seeded: bool = row
.try_get("workspaces_seeded")
.map_err(|error| RuntimeError::database("decode search generation workspace seed", error))?;
let publications_complete: bool = row
.try_get("publications_complete")
.map_err(|error| RuntimeError::database("decode search generation publication state", error))?;
if state != "building"
|| provider != config.provider
|| row_config_hash != expected_config_hash
|| schema_version != SCHEMA_FINGERPRINT
|| !scan_complete
|| !workspaces_covered
|| !workspaces_seeded
|| !publications_complete
{
return Err(RuntimeError::SearchIndexNotReady);
}
sqlx::query(
"UPDATE search_runtime_generations SET state='active', activated_at=coalesce(activated_at,now()) WHERE \
generation_id=$1 AND state IN ('pending','active')",
"UPDATE search_projection.generations SET state='draining', drained_at=now() WHERE state='active' AND id<>$1",
)
.bind(generation.id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("drain previous search generation", error))?;
sqlx::query(
"UPDATE search_projection.generations SET state='active', activated_at=coalesce(activated_at,now()) WHERE id=$1 \
AND state='building'",
)
.bind(generation.id)
.execute(&mut *transaction)
@@ -162,7 +294,7 @@ pub(super) async fn activate(pool: &PgPool, generation: &ActiveGeneration) -> Ru
.map_err(|error| RuntimeError::database("commit search generation activation", error))
}
fn config_fingerprint(config: &SearchRuntimeConfig) -> String {
pub(super) fn config_hash(config: &SearchRuntimeConfig) -> Vec<u8> {
let mut hash = Sha256::new();
for value in [
&config.provider,
@@ -174,16 +306,108 @@ fn config_fingerprint(config: &SearchRuntimeConfig) -> String {
hash.update(value.as_bytes());
hash.update([0]);
}
hash.finalize().iter().map(|byte| format!("{byte:02x}")).collect()
hash.finalize().to_vec()
}
fn decode(row: sqlx::postgres::PgRow) -> RuntimeResult<ActiveGeneration> {
Ok(ActiveGeneration {
id: row
.try_get("generation_id")
.try_get("id")
.map_err(|error| RuntimeError::database("decode search generation id", error))?,
manifest: row
.try_get("manifest")
.map_err(|error| RuntimeError::database("decode search generation manifest", error))?,
})
}
async fn create_generation(
transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
config: &SearchRuntimeConfig,
config_hash: Vec<u8>,
) -> RuntimeResult<ActiveGeneration> {
let generation_id = Uuid::new_v4();
let suffix = generation_id.simple().to_string();
let manifest = if config.provider == "embedded" {
json!({"doc":"doc","block":"block"})
} else {
json!({
"doc":format!("affine_search_doc_{suffix}"),
"block":format!("affine_search_block_{suffix}")
})
};
sqlx::query(
r#"INSERT INTO search_projection.generations
(id,provider,state,config_hash,schema_version,manifest)
VALUES ($1,$2,'building',$3,$4,$5)"#,
)
.bind(generation_id)
.bind(&config.provider)
.bind(&config_hash)
.bind(SCHEMA_FINGERPRINT)
.bind(&manifest)
.execute(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("create search generation", error))?;
Ok(ActiveGeneration {
id: generation_id,
manifest,
})
}
#[cfg(test)]
mod tests {
use sqlx::PgPool;
use uuid::Uuid;
use super::ensure;
use crate::runtime::{
SearchRuntimeConfig, backend_runtime::search::SEARCH_TEST_LOCK, migrations::migrate_search_tables,
};
#[tokio::test]
async fn configuration_change_supersedes_an_incomplete_candidate() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await
.unwrap();
let stale_id = Uuid::new_v4();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(stale_id)
.execute(&pool)
.await
.unwrap();
let generation = ensure(&pool, &SearchRuntimeConfig::default(), None, false)
.await
.unwrap();
assert_ne!(generation.id, stale_id);
let stale: (String, Option<String>) =
sqlx::query_as("SELECT state,last_error FROM search_projection.generations WHERE id=$1")
.bind(stale_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stale.0, "failed");
assert_eq!(
stale.1.as_deref(),
Some("search generation superseded by configuration")
);
sqlx::query("DELETE FROM search_projection.generations WHERE id IN ($1,$2)")
.bind(stale_id)
.bind(generation.id)
.execute(&pool)
.await
.unwrap();
}
}

View File

@@ -1,17 +1,32 @@
mod checkpoint;
mod generation;
mod projection;
mod provider;
mod query;
mod query_runtime;
mod result_filter;
mod runtime;
mod store;
mod types;
mod worker;
use generation::{ActiveGeneration, activate, config_hash, ensure, load_active};
use projection::{ProjectionInput, project_document};
use provider::{SearchChange, SearchProvider, projection_external_id};
use query::{compile, compile_aggregate};
use result_filter::{candidates, retain_visible_nodes};
pub(super) use runtime::SearchRuntime;
use types::{AggregateOptions, RuntimeSearchQuery, SearchOptions, SearchTable};
pub(super) use types::{RuntimeAggregateRequest, RuntimeSearchRequest};
use worker::{reconcile_workspace, sweep_generation_orphans};
const SCHEMA_FINGERPRINT: &str = "search-runtime-v5";
use super::{
permission::{AuthorizedSearchScope, DocReadScope, PermissionAuthorizer, SearchActor},
webpki_tls_config,
};
const SCHEMA_FINGERPRINT: i32 = 1;
#[cfg(test)]
pub(crate) static SEARCH_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
fn exact_token(value: &str) -> String {
use sha2::{Digest, Sha256};
@@ -30,4 +45,19 @@ fn provider_payload(payload: &serde_json::Value) -> serde_json::Value {
}
#[cfg(test)]
mod tests;
mod tests {
use super::SCHEMA_FINGERPRINT;
#[test]
fn terminal_schema_uses_one_generation_version() {
assert_eq!(SCHEMA_FINGERPRINT, 1);
}
#[test]
fn external_document_ids_keep_document_and_block_tables_distinct() {
let cases = [("workspace/doc", false), ("workspace/doc/block", true)];
for (external_id, block) in cases {
assert_eq!(external_id.matches('/').count() >= 2, block);
}
}
}

View File

@@ -1,32 +1,42 @@
use serde_json::{Value, json};
use sqlx::{PgPool, Row};
use super::store::ProjectionInput;
use crate::{
permission::doc_role_allows,
runtime::{RuntimeError, RuntimeResult, storage_runtime::load_current_doc},
};
#[derive(Clone, Debug, PartialEq)]
pub(super) struct ProjectionInput {
pub(super) workspace_id: String,
pub(super) doc_id: String,
pub(super) payload: Value,
}
pub(super) async fn project_document(
pool: &PgPool,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<Option<(ProjectionInput, Vec<ProjectionInput>)>> {
let Some(current) = load_current_doc(pool, workspace_id, doc_id).await? else {
let Some(current) = load_current_doc(pool, workspace_id, doc_id)
.await
.map_err(|error| match error {
RuntimeError::InvalidState(message) => RuntimeError::SearchSourceInvalid(message),
error => error,
})?
else {
return Ok(None);
};
let revision = current.updated_at.timestamp_millis();
let projection =
affine_doc_loader::project_document_search(current.blob, doc_id.to_string(), revision.to_string())
.map_err(|error| RuntimeError::invalid_state(format!("search document projection failed: {error}")))?;
let metadata = sqlx::query(
.map_err(|error| RuntimeError::SearchSourceInvalid(format!("document projection failed: {error}")))?;
let Some(metadata) = sqlx::query(
r#"SELECT snapshot.created_at,snapshot.updated_at,snapshot.created_by,snapshot.updated_by,
revision.revision AS acl_revision,
coalesce(doc_policy.visibility,'private') AS visibility,
doc_policy.public_role,
coalesce(doc_policy.member_default_role,workspace_policy.member_default_doc_role,'manager') AS member_default_role
FROM snapshots snapshot
LEFT JOIN workspace_permission_revisions revision ON revision.workspace_id=snapshot.workspace_id
LEFT JOIN workspace_access_policies workspace_policy ON workspace_policy.workspace_id=snapshot.workspace_id
LEFT JOIN doc_access_policies doc_policy
ON doc_policy.workspace_id=snapshot.workspace_id AND doc_policy.doc_id=snapshot.guid
@@ -37,11 +47,9 @@ pub(super) async fn project_document(
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load search projection metadata", error))?
.ok_or_else(|| RuntimeError::invalid_state("search snapshot metadata unavailable"))?;
let acl_revision = metadata
.try_get::<Option<i64>, _>("acl_revision")
.map_err(|error| RuntimeError::database("decode search ACL revision", error))?
.ok_or_else(|| RuntimeError::invalid_state("permission_state_unavailable"))?;
else {
return Ok(None);
};
let visibility: String = metadata
.try_get("visibility")
.map_err(|error| RuntimeError::database("decode search doc visibility", error))?;
@@ -51,14 +59,18 @@ pub(super) async fn project_document(
let member_default_role: String = metadata
.try_get("member_default_role")
.map_err(|error| RuntimeError::database("decode search member default role", error))?;
let acl_public_readable = visibility == "public"
&& public_role
.as_deref()
.is_some_and(|role| doc_role_allows(role, "Doc.Read").unwrap_or(false));
let acl_member_default_readable = doc_role_allows(&member_default_role, "Doc.Read")
.map_err(|_| RuntimeError::invalid_state("permission_state_unavailable"))?;
let acl_public_readable = if visibility == "public" {
match public_role.as_deref() {
Some(role) => readable_role(role, "public")?,
None => false,
}
} else {
false
};
let acl_member_default_readable = readable_role(&member_default_role, "member default")?;
let grants = sqlx::query(
"SELECT principal_id,role FROM doc_grants WHERE workspace_id=$1 AND doc_id=$2 AND principal_type='user'",
"SELECT principal_id,role FROM doc_grants WHERE workspace_id=$1 AND doc_id=$2 AND principal_type='user' ORDER BY \
principal_id,role",
)
.bind(workspace_id)
.bind(doc_id)
@@ -74,7 +86,7 @@ pub(super) async fn project_document(
let role: String = row
.try_get("role")
.map_err(|error| RuntimeError::database("decode grant role", error))?;
Ok(doc_role_allows(&role, "Doc.Read").unwrap_or(false).then_some(id))
Ok(readable_role(&role, "grant")?.then_some(id))
})
.collect::<RuntimeResult<Vec<_>>>()?
.into_iter()
@@ -96,7 +108,6 @@ pub(super) async fn project_document(
public_readable: acl_public_readable,
member_default_readable: acl_member_default_readable,
read_user_ids: acl_read_user_ids,
revision: acl_revision,
};
let document_payload = with_acl(
json!({
@@ -113,14 +124,7 @@ pub(super) async fn project_document(
}),
&acl,
);
let document = input(
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}"),
revision,
document_payload,
&acl,
);
let document = input(workspace_id, doc_id, document_payload);
let blocks = projection
.units
.into_iter()
@@ -142,24 +146,21 @@ pub(super) async fn project_document(
}),
&acl,
);
input(
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}/{block_id}"),
revision,
payload,
&acl,
)
input(workspace_id, doc_id, payload)
})
.collect();
Ok(Some((document, blocks)))
}
fn readable_role(role: &str, source: &str) -> RuntimeResult<bool> {
doc_role_allows(role, "Doc.Read")
.map_err(|_| RuntimeError::SearchSourceInvalid(format!("invalid {source} document role")))
}
struct AclFields {
public_readable: bool,
member_default_readable: bool,
read_user_ids: Vec<String>,
revision: i64,
}
fn with_acl(mut payload: Value, acl: &AclFields) -> Value {
@@ -181,27 +182,25 @@ fn with_acl(mut payload: Value, acl: &AclFields) -> Value {
tokens.push("public".to_string());
}
object.insert("acl_read_tokens".to_string(), json!(tokens));
object.insert("acl_revision".to_string(), json!(acl.revision));
payload
}
fn input(
workspace_id: &str,
doc_id: &str,
external_id: &str,
revision: i64,
payload: Value,
acl: &AclFields,
) -> ProjectionInput {
fn input(workspace_id: &str, doc_id: &str, payload: Value) -> ProjectionInput {
ProjectionInput {
external_id: external_id.to_string(),
workspace_id: workspace_id.to_string(),
doc_id: doc_id.to_string(),
revision,
payload,
acl_public_readable: acl.public_readable,
acl_member_default_readable: acl.member_default_readable,
acl_read_user_ids: acl.read_user_ids.clone(),
acl_revision: acl.revision,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn invalid_acl_role_is_a_permanent_search_source_error() {
let error = readable_role("future-role", "grant").unwrap_err();
assert!(matches!(error, RuntimeError::SearchSourceInvalid(_)));
assert!(error.is_permanent_search_source());
}
}

View File

@@ -1,326 +1,690 @@
use serde_json::{Value, json};
use std::time::Duration;
use crate::runtime::{RuntimeError, RuntimeResult};
use reqwest::{Client, redirect::Policy};
use serde_json::{Map, Value, json};
use sha2::{Digest, Sha256};
pub(super) fn prepare_manticore_payload(
payload: &mut Value,
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<()> {
let object = payload.as_object_mut().expect("search payload is an object");
object.remove("acl_read_user_ids");
if let Some(Value::Array(tokens)) = object.remove("acl_read_tokens") {
object.insert(
"acl_read_token_ids".to_string(),
Value::Array(
tokens
.iter()
.filter_map(Value::as_str)
.map(|token| {
token_ids
.get(token)
.copied()
.map(Value::from)
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))
})
.collect::<RuntimeResult<Vec<_>>>()?,
),
);
}
if let Some(Value::Array(tokens)) = object.get("ref_doc_id").cloned() {
object.insert(
"ref_doc_token_ids".to_string(),
Value::Array(
tokens
.iter()
.filter_map(Value::as_str)
.map(|token| {
token_ids
.get(token)
.copied()
.map(Value::from)
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))
})
.collect::<RuntimeResult<Vec<_>>>()?,
),
);
}
for field in ["created_at", "updated_at"] {
if let Some(value) = object.get_mut(field)
&& let Some(milliseconds) = value.as_i64()
{
*value = json!(milliseconds / 1_000);
}
}
for (field, value) in object.iter_mut() {
if let Value::Array(values) = value {
*value = if matches!(field.as_str(), "acl_read_token_ids" | "ref_doc_token_ids") {
continue;
} else if field == "content" {
Value::String(values.iter().filter_map(Value::as_str).collect::<Vec<_>>().join(" "))
} else {
Value::String(
serde_json::to_string(values).map_err(|error| RuntimeError::json("encode manticore array", error))?,
)
};
} else if value.is_null() {
*value = Value::String(String::new());
}
}
Ok(())
use super::{SearchChange, SearchTable, provider_write_error, webpki_tls_config};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
pub(in crate::runtime::backend_runtime::search) struct ManticoreSearchProvider {
client: Client,
endpoint: String,
api_key: String,
username: String,
password: String,
}
pub(super) fn prepare_manticore_search(
dsl: &mut Value,
cursor: Option<Value>,
size: u64,
initial_offset: u64,
requested_fields: &[String],
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<u64> {
normalize_manticore_terms(dsl, token_ids)?;
let object = dsl.as_object_mut().expect("search DSL is an object");
let mut source = object
.get("_source")
.and_then(Value::as_array)
impl ManticoreSearchProvider {
pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
let endpoint = config.endpoint.trim_end_matches('/');
let url = url::Url::parse(endpoint).map_err(|_| RuntimeError::config("invalid search provider endpoint"))?;
if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return Err(RuntimeError::config("invalid search provider endpoint"));
}
let client = Client::builder()
.tls_backend_preconfigured(
webpki_tls_config()
.map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?,
)
.redirect(Policy::none())
.timeout(Duration::from_secs(30))
.build()
.map_err(|error| RuntimeError::invalid_state(format!("search HTTP client failed: {error}")))?;
Ok(Self {
client,
endpoint: endpoint.to_string(),
api_key: config.api_key.clone(),
username: config.username.clone(),
password: config.password.clone(),
})
}
pub(super) async fn search(&self, physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
let request = translate_search_request(physical_table, dsl)?;
let response = self
.request(reqwest::Method::POST, "search")
.json(&request)
.send()
.await;
let response = response.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(if status.as_u16() == 400 {
RuntimeError::SearchUnsupportedQuery
} else {
RuntimeError::SearchProviderUnavailable
});
}
let value: Value =
serde_json::from_slice(&body).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize(value, request.get("offset").and_then(Value::as_u64).unwrap_or_default())
}
pub(super) async fn aggregate(&self, _physical_table: &str, _dsl: Value) -> RuntimeResult<Value> {
Err(RuntimeError::SearchUnsupportedQuery)
}
pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
let response = self
.request(reqwest::Method::POST, "sql?mode=raw")
.header("content-type", "application/x-www-form-urlencoded")
.body(create_table_sql(physical_table, table))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid search provider schema response", error))?;
if value.get("error").is_some() {
return Err(RuntimeError::invalid_state("provider_schema_failed"));
}
Ok(())
}
pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
if changes.is_empty() {
return Ok(());
}
let mut body = String::new();
for change in changes {
let payload = change.upsert_payload()?;
let mut document = manticore_document(payload);
document.insert("external_id".to_string(), json!(change.external_id));
body.push_str(
&serde_json::to_string(&json!({
"replace": {
"table": physical_table,
"id": document_id(&change.external_id),
"doc": document,
}
}))
.map_err(|error| RuntimeError::json("encode manticore document", error))?,
);
body.push('\n');
}
if body.is_empty() {
return Ok(());
}
let response = self
.request(reqwest::Method::POST, "bulk")
.header("content-type", "application/x-ndjson")
.body(body)
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(provider_write_error(status.as_u16()));
}
let value: Value =
serde_json::from_slice(&body).map_err(|error| RuntimeError::json("invalid manticore bulk response", error))?;
validate_bulk_response(&value, changes.len())
}
pub(super) async fn gc_document_history(
&self,
physical_table: &str,
workspace_id: &str,
doc_id: &str,
source_version: i64,
permission_version: i64,
limit: usize,
) -> RuntimeResult<()> {
let query = json!({
"bool": {
"must": [
{"equals":{"workspace_id":workspace_id}},
{"equals":{"doc_id":doc_id}},
],
"must_not": [{"bool":{"must":[
{"equals":{"source_version":source_version}},
{"equals":{"permission_version":permission_version}}
]}}]
}
});
self
.delete_bounded(physical_table, query, limit.max(1))
.await
.map(|_| ())
}
pub(super) async fn gc_workspace(
&self,
physical_table: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> RuntimeResult<bool> {
let limit = limit.max(1);
let query = json!({"bool":{"must":[
{"equals":{"workspace_id":workspace_id}},
{"range":{"source_version":{"lte":source_version_high_water}}}
]}});
let deleted = self.delete_bounded(physical_table, query, limit).await?;
Ok(deleted == limit)
}
async fn delete_bounded(&self, physical_table: &str, query: Value, limit: usize) -> RuntimeResult<usize> {
let response = self
.request(reqwest::Method::POST, "search")
.json(&json!({
"table":physical_table,
"query":query,
"limit":limit,
"sort":[{"id":"asc"}]
}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid manticore GC search response", error))?;
let ids = value
.pointer("/hits/hits")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid manticore GC search response"))?
.iter()
.map(|hit| {
let id = hit
.get("_id")
.or_else(|| hit.get("id"))
.and_then(|id| id.as_u64().or_else(|| id.as_str().and_then(|id| id.parse().ok())))
.ok_or_else(|| RuntimeError::invalid_state("invalid manticore GC search response"))?;
Ok(json!(id))
})
.collect::<RuntimeResult<Vec<_>>>()?;
if ids.is_empty() {
return Ok(0);
}
let response = self
.request(reqwest::Method::POST, "delete")
.json(&json!({"table":physical_table,"id":ids}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value =
serde_json::from_slice(&body).map_err(|error| RuntimeError::json("invalid manticore delete response", error))?;
if value.get("error").is_some() {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(ids.len())
}
fn request(&self, method: reqwest::Method, path: &str) -> reqwest::RequestBuilder {
let mut request = self.client.request(method, format!("{}/{path}", self.endpoint));
if !self.api_key.is_empty() {
request = request.header("Authorization", format!("ApiKey {}", self.api_key));
} else if !self.username.is_empty() {
request = request.basic_auth(&self.username, Some(&self.password));
}
request
}
}
fn create_table_sql(physical_table: &str, table: SearchTable) -> String {
let text_field = table.text_field();
let mut columns = vec![format!("{text_field} text")];
for field in [
"external_id",
"workspace_id",
"workspace_token",
"generation_id",
"doc_id",
"doc_token",
"created_by_user_id",
"updated_by_user_id",
] {
columns.push(format!("{field} string"));
}
for field in ["source_version", "permission_version", "created_at", "updated_at"] {
columns.push(format!("{field} bigint"));
}
columns.push("acl_public_readable integer".to_string());
columns.push("acl_member_default_readable integer".to_string());
columns.push("acl_read_tokens text".to_string());
let extra_fields = if table == SearchTable::Doc {
vec!["summary string", "journal string"]
} else {
vec![
"block_id string",
"block_token string",
"unit_id string",
"projection_version integer",
"source_hash string",
"visibility string",
"element_id string",
"frame_id string",
"source_block_id string",
"flavour string",
"blob string",
"ref_doc_id string",
"ref string",
"parent_flavour string",
"parent_block_id string",
"additional string",
"markdown_preview string",
]
};
columns.extend(extra_fields.into_iter().map(str::to_string));
format!("CREATE TABLE IF NOT EXISTS {physical_table} ({})", columns.join(", "))
}
fn manticore_document(payload: &Value) -> Map<String, Value> {
payload
.as_object()
.into_iter()
.flatten()
.filter_map(Value::as_str)
.map(str::to_string)
.collect::<std::collections::BTreeSet<_>>();
source.extend(requested_fields.iter().cloned());
object.insert("_source".to_string(), json!(source));
object.remove("fields");
if let Some(highlight) = object.get_mut("highlight")
&& let Some(options) = highlight
.get("fields")
.filter(|(_, value)| !value.is_null())
.map(|(key, value)| {
let value = match value {
Value::Bool(value) => json!(i32::from(*value)),
Value::Array(values) => json!(
values
.iter()
.filter(|value| !value.is_null())
.map(|value| value.as_str().map(str::to_string).unwrap_or_else(|| value.to_string()))
.collect::<Vec<_>>()
.join(" ")
),
value => value.clone(),
};
(key.clone(), value)
})
.collect()
}
fn document_id(external_id: &str) -> u64 {
let digest = Sha256::digest(external_id.as_bytes());
let mut bytes = [0; 8];
bytes.copy_from_slice(&digest[..8]);
u64::from_be_bytes(bytes) | 1
}
fn validate_bulk_response(value: &Value, expected_changes: usize) -> RuntimeResult<()> {
let items = value
.get("items")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid manticore bulk response"))?;
let mut affected = 0_u64;
for item in items {
let result = item
.as_object()
.and_then(|item| item.values().next())
.and_then(Value::as_object)
.and_then(|fields| fields.values().next())
.cloned()
{
*highlight = options;
}
let offset = if let Some(Value::String(cursor)) = cursor {
let offset = serde_json::from_str::<Value>(&cursor)
.map_err(|error| RuntimeError::json("invalid search cursor", error))?
.get("offset")
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
let status = result
.get("status")
.and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
if offset.saturating_add(size) > 10_000 {
return Err(RuntimeError::invalid_input("search cursor exceeds 10000"));
.and_then(|status| u16::try_from(status).ok())
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
if result.get("error").is_some() || !(200..300).contains(&status) {
return Err(provider_write_error(status));
}
object.insert("from".to_string(), json!(offset));
offset
} else if cursor.is_some() {
return Err(RuntimeError::invalid_input("invalid search cursor"));
} else {
initial_offset
};
Ok(offset)
}
pub(super) fn manticore_fields(source: Option<&Value>, requested_fields: &[String]) -> Value {
let source = source.and_then(Value::as_object);
Value::Object(
requested_fields
.iter()
.filter_map(|field| {
let mut value = source?.get(field)?.clone();
if matches!(field.as_str(), "created_at" | "updated_at")
&& let Some(seconds) = value.as_i64()
{
value = json!(seconds * 1_000);
} else if let Some(encoded) = value.as_str()
&& encoded.starts_with('[')
&& let Ok(decoded) = serde_json::from_str(encoded)
{
value = decoded;
}
if !value.is_array() {
value = Value::Array(vec![value]);
}
Some((field.clone(), value))
})
.collect(),
)
}
fn normalize_manticore_terms(
value: &mut Value,
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<()> {
if let Some(term) = manticore_term(value, token_ids)? {
*value = term;
return Ok(());
let numeric = ["created", "updated", "deleted"]
.into_iter()
.filter_map(|field| result.get(field).and_then(Value::as_u64))
.sum::<u64>();
affected += if numeric > 0 {
numeric
} else if result.get("created").and_then(Value::as_bool).is_some()
|| result
.get("result")
.and_then(Value::as_str)
.is_some_and(|result| matches!(result, "created" | "updated"))
{
1
} else {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
};
}
match value {
Value::Object(object) => {
if let Some(Value::Object(boolean)) = object.get_mut("bool")
&& boolean.get("boost").and_then(Value::as_f64) == Some(1.0)
{
boolean.remove("boost");
}
if let Some(Value::Object(terms)) = object.get_mut("terms") {
terms.entry("order").or_insert_with(|| json!({"_count":"desc"}));
}
for child in object.values_mut() {
normalize_manticore_terms(child, token_ids)?;
}
}
Value::Array(array) => {
for child in array {
normalize_manticore_terms(child, token_ids)?;
}
}
_ => {}
if value.get("errors").and_then(Value::as_bool) != Some(false) || affected != expected_changes as u64 {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(())
}
fn manticore_term(value: &Value, token_ids: &std::collections::HashMap<String, i64>) -> RuntimeResult<Option<Value>> {
let Some(term) = value.get("term").and_then(Value::as_object) else {
return Ok(None);
fn translate_search_request(physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
let object = dsl
.as_object()
.ok_or_else(|| RuntimeError::invalid_input("invalid search request"))?;
let size = object.get("size").and_then(Value::as_u64).unwrap_or(10);
let offset = if let Some(cursor) = object.get("cursor") {
let cursor = cursor
.as_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
serde_json::from_str::<Value>(cursor)
.ok()
.and_then(|cursor| cursor.get("offset").and_then(Value::as_u64))
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?
} else {
object.get("from").and_then(Value::as_u64).unwrap_or_default()
};
if term.len() != 1 {
return Ok(None);
let query = object.get("query").cloned().unwrap_or_else(|| json!({"match_all":{}}));
let query = translate_query(query)?;
let mut request = json!({
"table":physical_table,
"query":query,
"limit":size,
"offset":offset,
});
if let Some(sort) = object.get("sort") {
request["sort"] = translate_sort(sort)?;
}
let Some((field, clause)) = term.iter().next() else {
return Ok(None);
};
let value = clause.get("value").unwrap_or(clause);
Ok(match value {
Value::String(value) => {
if matches!(field.as_str(), "acl_read_tokens" | "ref_doc_id") {
let token_id = token_ids
.get(value)
.copied()
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))?;
let field = if field == "acl_read_tokens" {
"acl_read_token_ids"
} else {
"ref_doc_token_ids"
let mut source = Vec::new();
if let Some(fields) = object.get("_source").and_then(Value::as_array) {
source.extend(fields.iter().filter_map(Value::as_str).map(str::to_string));
}
if let Some(fields) = object.get("fields").and_then(Value::as_array) {
source.extend(fields.iter().filter_map(Value::as_str).map(str::to_string));
}
source.push("external_id".to_string());
source.push("doc_id".to_string());
source.push("source_version".to_string());
source.push("permission_version".to_string());
source.sort();
source.dedup();
request["_source"] = json!(source);
if let Some(highlight) = object.get("highlight") {
request["highlight"] = translate_highlight(highlight)?;
}
Ok(request)
}
fn translate_highlight(highlight: &Value) -> RuntimeResult<Value> {
let fields = highlight
.get("fields")
.and_then(Value::as_object)
.ok_or(RuntimeError::SearchUnsupportedQuery)?;
let mut request = json!({"fields":fields.keys().collect::<Vec<_>>()});
let tags = fields.values().filter_map(Value::as_object).next();
if let Some(pre_tag) = tags
.and_then(|options| options.get("pre_tags"))
.and_then(Value::as_array)
.and_then(|tags| tags.first())
.and_then(Value::as_str)
{
request["pre_tags"] = json!(pre_tag);
}
if let Some(post_tag) = tags
.and_then(|options| options.get("post_tags"))
.and_then(Value::as_array)
.and_then(|tags| tags.first())
.and_then(Value::as_str)
{
request["post_tags"] = json!(post_tag);
}
Ok(request)
}
fn translate_sort(sort: &Value) -> RuntimeResult<Value> {
let values = sort
.as_array()
.ok_or_else(|| RuntimeError::invalid_input("invalid search sort"))?
.iter()
.filter_map(|value| match value {
Value::String(field) => match field.as_str() {
"_score" => None,
"_id" => Some(json!({"id":"asc"})),
field => Some(json!({field:"asc"})),
},
Value::Object(object) => {
let (field, direction) = object.iter().next()?;
let field = match field.as_str() {
"_id" => "id",
"_score" => return None,
field => field,
};
return Ok(Some(json!({"equals":{field:token_id}})));
}
let (field, value) = match field.as_str() {
"workspace_id" => ("workspace_token", super::super::exact_token(value)),
"doc_id" => ("doc_token", super::super::exact_token(value)),
"block_id" => ("block_token", super::super::exact_token(value)),
_ => (field.as_str(), value.clone()),
};
if let Some(boost) = clause.get("boost").and_then(Value::as_f64) {
Some(json!({"match":{field:{"query":value,"boost":boost}}}))
} else {
Some(json!({"equals":{field:value}}))
Some(json!({field:direction}))
}
_ => None,
})
.collect::<Vec<_>>();
Ok(json!(values))
}
fn translate_query(query: Value) -> RuntimeResult<Value> {
let Some(object) = query.as_object() else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
if object.contains_key("match_all") {
return Ok(json!({"match_all":{}}));
}
if let Some(match_query) = object.get("match") {
let Some((field, value)) = match_query.as_object().and_then(|object| object.iter().next()) else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
let value = value
.get("query")
.cloned()
.or_else(|| value.as_str().map(|value| json!(value)))
.ok_or(RuntimeError::SearchUnsupportedQuery)?;
return Ok(json!({"match":{field:value}}));
}
if let Some(term_query) = object.get("term") {
let Some((field, value)) = term_query.as_object().and_then(|object| object.iter().next()) else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
let value = value
.get("value")
.cloned()
.ok_or(RuntimeError::SearchUnsupportedQuery)?;
if field == "acl_read_tokens" {
let value = value.as_str().ok_or(RuntimeError::SearchUnsupportedQuery)?;
return Ok(json!({"match":{field:value}}));
}
Value::Bool(value) => Some(json!({"equals":{field:u8::from(*value)}})),
Value::Number(value) => Some(json!({"equals":{field:value}})),
_ => None,
return Ok(json!({"equals":{field:manticore_scalar(value)?}}));
}
if let Some(bool_query) = object.get("bool") {
let Some(bool_query) = bool_query.as_object() else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
let mut translated = Map::new();
for occurrence in ["must", "should", "must_not"] {
let Some(clauses) = bool_query.get(occurrence) else {
continue;
};
let clauses = if let Some(array) = clauses.as_array() {
array
.iter()
.map(|clause| translate_query(clause.clone()))
.collect::<RuntimeResult<Vec<_>>>()?
} else {
vec![translate_query(clauses.clone())?]
};
translated.insert(occurrence.to_string(), json!(clauses));
}
return Ok(json!({"bool":translated}));
}
if object.contains_key("boost") {
return translate_query(object.get("boost").cloned().unwrap_or_default());
}
Err(RuntimeError::SearchUnsupportedQuery)
}
fn manticore_scalar(value: Value) -> RuntimeResult<Value> {
Ok(match value {
Value::Bool(value) => json!(i32::from(value)),
Value::String(value) => json!(value),
Value::Number(value) => Value::Number(value),
_ => return Err(RuntimeError::SearchUnsupportedQuery),
})
}
pub(super) fn manticore_exact_tokens(value: &Value) -> Vec<String> {
let mut tokens = Vec::new();
collect_exact_tokens(value, &mut tokens);
tokens
async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>> {
if response
.content_length()
.is_some_and(|length| length > MAX_RESPONSE_BYTES as u64)
{
return Err(RuntimeError::invalid_state("provider_response_too_large"));
}
let mut bytes = Vec::new();
while let Some(chunk) = response
.chunk()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?
{
if bytes.len() + chunk.len() > MAX_RESPONSE_BYTES {
return Err(RuntimeError::invalid_state("provider_response_too_large"));
}
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
fn collect_exact_tokens(value: &Value, tokens: &mut Vec<String>) {
match value {
Value::Object(object) => {
if let Some(token) = object
.get("term")
.and_then(|term| term.get("acl_read_tokens").or_else(|| term.get("ref_doc_id")))
.and_then(|clause| clause.get("value").unwrap_or(clause).as_str())
{
tokens.push(token.to_string());
}
for field in ["acl_read_tokens", "ref_doc_id"] {
if let Some(values) = object.get(field).and_then(Value::as_array) {
tokens.extend(values.iter().filter_map(Value::as_str).map(str::to_string));
}
}
for child in object.values() {
collect_exact_tokens(child, tokens);
}
}
Value::Array(values) => {
for child in values {
collect_exact_tokens(child, tokens);
}
}
_ => {}
}
fn normalize(value: Value, offset: u64) -> RuntimeResult<Value> {
let hits = value
.pointer("/hits/hits")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider response"))?;
let total = value
.pointer("/hits/total")
.and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_state("inexact provider total"))?;
let nodes = hits
.iter()
.map(|hit| {
let source = hit.get("_source").cloned().unwrap_or_else(|| json!({}));
json!({
"id":source.get("external_id").and_then(Value::as_str).unwrap_or_default(),
"score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(),
"fields":fields_from_source(&source),
"highlights":hit.get("highlight").cloned().unwrap_or_else(||json!({})),
"_source":source,
})
})
.collect::<Vec<_>>();
let next_cursor = (offset.saturating_add(nodes.len() as u64) < total)
.then(|| json!({"offset":offset.saturating_add(nodes.len() as u64)}).to_string());
Ok(json!({"total":total,"nodes":nodes,"nextCursor":next_cursor}))
}
fn fields_from_source(source: &Value) -> Value {
let Some(source) = source.as_object() else {
return json!({});
};
source
.iter()
.map(|(field, value)| (field.clone(), json!([value])))
.collect::<Map<_, _>>()
.into()
}
#[cfg(test)]
mod tests {
use super::{super::super::exact_token, *};
use serde_json::json;
use super::{
SearchTable, document_id, manticore_document, translate_query, translate_search_request, translate_sort,
validate_bulk_response,
};
use crate::runtime::RuntimeError;
#[test]
fn payload_and_fields_preserve_terminal_types() {
let mut payload = json!({
"content":["hello","world"],
"ref_doc_id":["doc-a","doc-b"],
"summary":null,
"created_at":2_000,
"updated_at":3_000,
"acl_read_tokens":["member"]
});
prepare_manticore_payload(
&mut payload,
&[
("member".to_string(), 7),
("doc-a".to_string(), 8),
("doc-b".to_string(), 9),
]
.into_iter()
.collect(),
)
.unwrap();
assert_eq!(payload["content"], "hello world");
assert_eq!(payload["ref_doc_id"], "[\"doc-a\",\"doc-b\"]");
assert_eq!(payload["summary"], "");
assert_eq!(payload["created_at"], 2);
assert_eq!(payload["acl_read_token_ids"], json!([7]));
assert_eq!(payload["ref_doc_token_ids"], json!([8, 9]));
let fields = manticore_fields(
Some(&payload),
&[
"ref_doc_id".to_string(),
"summary".to_string(),
"updated_at".to_string(),
],
fn translates_the_shared_basic_query_subset() {
assert_eq!(
translate_query(json!({
"bool":{"must":[
{"term":{"workspace_id":{"value":"workspace"}}},
{"match":{"content":{"query":"hello","boost":1.0}}}
]}
}))
.unwrap(),
json!({
"bool":{"must":[
{"equals":{"workspace_id":"workspace"}},
{"match":{"content":"hello"}}
]}
})
);
assert_eq!(
translate_search_request(
"blocks",
json!({
"query":{"match":{"content":{"query":"hello"}}},
"highlight":{"fields":{"content":{"pre_tags":["<b>"],"post_tags":["</b>"]}}}
})
)
.unwrap()["highlight"],
json!({"fields":["content"],"pre_tags":"<b>","post_tags":"</b>"})
);
assert_eq!(fields["ref_doc_id"], json!(["doc-a", "doc-b"]));
assert_eq!(fields["summary"], json!([""]));
assert_eq!(fields["updated_at"], json!([3_000]));
}
#[test]
fn nested_terms_use_exact_identity_and_acl_tokens() {
let mut dsl = json!({"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":"workspace","boost":2.0}}},
{"bool":{"must_not":[{"term":{"doc_id":{"value":"doc"}}}]}},
{"term":{"acl_read_tokens":{"value":"member"}}},
{"term":{"ref_doc_id":{"value":"ref-doc"}}}
],"boost":1.0}}});
normalize_manticore_terms(
&mut dsl,
&[("member".to_string(), 9), ("ref-doc".to_string(), 10)]
.into_iter()
.collect(),
)
.unwrap();
fn converts_projection_arrays_and_booleans_to_basic_attributes() {
assert_eq!(
dsl,
json!({"query":{"bool":{"must":[
{"match":{"workspace_token":{"query":exact_token("workspace"),"boost":2.0}}},
{"bool":{"must_not":[{"equals":{"doc_token":exact_token("doc")}}]}},
{"equals":{"acl_read_token_ids":9}},
{"equals":{"ref_doc_token_ids":10}}
]}}})
manticore_document(&json!({"acl_read_tokens":["user","member"],"acl_public_readable":false})),
json!({"acl_read_tokens":"user member","acl_public_readable":0})
.as_object()
.unwrap()
.clone()
);
assert!(
validate_bulk_response(
&json!({"errors":false,"items":[{"bulk":{"status":201,"created":2,"updated":0,"deleted":0}}]}),
2,
)
.is_ok()
);
assert!(
validate_bulk_response(
&json!({"errors":false,"items":[
{"replace":{"status":201,"created":true}},
{"replace":{"status":200,"created":false,"result":"updated"}}
]}),
2,
)
.is_ok()
);
assert!(
validate_bulk_response(
&json!({"errors":true,"items":[{"replace":{"status":400,"error":{"reason":"invalid"}}}]}),
1,
)
.is_err()
);
assert!(matches!(
validate_bulk_response(
&json!({"errors":true,"items":[{"replace":{"status":400,"error":{"reason":"invalid"}}}]}),
1
),
Err(RuntimeError::SearchSourceInvalid(_))
));
}
#[test]
fn uses_stable_nonzero_document_ids() {
assert_ne!(document_id("workspace/doc/block"), 0);
assert_eq!(document_id("workspace/doc/block"), document_id("workspace/doc/block"));
assert_ne!(
document_id("generation/workspace/doc/1/1/block"),
document_id("generation/workspace/doc/2/1/block")
);
}
#[test]
fn drops_score_sort_but_keeps_basic_attribute_order() {
assert_eq!(
translate_sort(&json!(["_score",{"updated_at":"desc"},{"_id":"asc"}])).unwrap(),
json!([{ "updated_at":"desc" }, { "id":"asc" }])
);
assert_eq!(SearchTable::Doc.text_field(), "title");
}
}

View File

@@ -1,40 +1,188 @@
mod manticore;
mod remote;
pub(super) use remote::RemoteProvider;
use serde_json::{Value, json};
use super::types::SearchTable;
use self::{manticore::ManticoreSearchProvider, remote::RemoteProvider};
use super::{SearchTable, webpki_tls_config};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
pub(super) fn mapping(table: SearchTable, provider: &str) -> Value {
fn provider_write_error(status: u16) -> RuntimeError {
match status {
408 | 429 | 500..=599 => RuntimeError::SearchProviderUnavailable,
400..=499 => RuntimeError::SearchSourceInvalid("search provider rejected projection".to_string()),
_ => RuntimeError::invalid_state("provider_apply_failed"),
}
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct SearchChange {
pub(super) table: SearchTable,
pub(super) external_id: String,
pub(super) workspace_id: String,
pub(super) doc_id: String,
pub(super) source_version: i64,
pub(super) permission_version: i64,
pub(super) payload: Value,
}
impl SearchChange {
fn upsert_payload(&self) -> RuntimeResult<&Value> {
let payload = &self.payload;
if payload.get("workspace_id").and_then(Value::as_str) != Some(&self.workspace_id)
|| payload.get("doc_id").and_then(Value::as_str) != Some(&self.doc_id)
|| payload.get("source_version").and_then(Value::as_i64) != Some(self.source_version)
|| payload.get("permission_version").and_then(Value::as_i64) != Some(self.permission_version)
{
return Err(RuntimeError::invalid_input(
"provider projection tuple does not match payload",
));
}
let generation_id = payload
.get("generation_id")
.and_then(Value::as_str)
.ok_or_else(|| RuntimeError::invalid_input("provider projection generation is required"))?;
if projection_external_id(
self.table,
generation_id,
&self.workspace_id,
&self.doc_id,
payload.get("block_id").and_then(Value::as_str),
self.source_version,
self.permission_version,
)? != self.external_id
{
return Err(RuntimeError::invalid_input(
"provider external id does not match projection tuple",
));
}
Ok(payload)
}
}
pub(super) enum SearchProvider {
/// Elasticsearch-compatible provider with the RFC6 shared projection
/// contract.
Elasticsearch(RemoteProvider),
/// Manticore Search provides candidate retrieval. Canonical permission facts
/// filter every ACL-scoped result before it leaves the runtime.
ManticoreSearch(ManticoreSearchProvider),
}
impl SearchProvider {
pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
match config.provider.as_str() {
"elasticsearch" => RemoteProvider::new(config).map(Self::Elasticsearch),
"manticoresearch" => ManticoreSearchProvider::new(config).map(Self::ManticoreSearch),
_ => Err(crate::runtime::RuntimeError::config("unsupported search provider")),
}
}
pub(super) async fn search(&self, physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
match self {
Self::Elasticsearch(provider) => provider.search(physical_table, dsl).await,
Self::ManticoreSearch(provider) => provider.search(physical_table, dsl).await,
}
}
pub(super) async fn aggregate(&self, physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
match self {
Self::Elasticsearch(provider) => provider.aggregate(physical_table, dsl).await,
Self::ManticoreSearch(provider) => provider.aggregate(physical_table, dsl).await,
}
}
pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
match self {
Self::Elasticsearch(provider) => provider.provision(physical_table, table).await,
Self::ManticoreSearch(provider) => provider.provision(physical_table, table).await,
}
}
pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
match self {
Self::Elasticsearch(provider) => provider.apply(physical_table, changes).await,
Self::ManticoreSearch(provider) => provider.apply(physical_table, changes).await,
}
}
pub(super) async fn gc_document_history(
&self,
physical_table: &str,
workspace_id: &str,
doc_id: &str,
source_version: i64,
permission_version: i64,
limit: usize,
) -> RuntimeResult<()> {
match self {
Self::Elasticsearch(provider) => {
provider
.gc_document_history(
physical_table,
workspace_id,
doc_id,
source_version,
permission_version,
limit,
)
.await
}
Self::ManticoreSearch(provider) => {
provider
.gc_document_history(
physical_table,
workspace_id,
doc_id,
source_version,
permission_version,
limit,
)
.await
}
}
}
pub(super) async fn gc_workspace(
&self,
physical_table: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> RuntimeResult<bool> {
match self {
Self::Elasticsearch(provider) => {
provider
.gc_workspace(physical_table, workspace_id, source_version_high_water, limit)
.await
}
Self::ManticoreSearch(provider) => {
provider
.gc_workspace(physical_table, workspace_id, source_version_high_water, limit)
.await
}
}
}
}
fn mapping(table: SearchTable) -> Value {
let text_field = table.text_field();
let mut properties = serde_json::Map::from_iter([
("workspace_id".into(), json!({"type":"keyword"})),
("workspace_token".into(), json!({"type":"keyword"})),
("generation_id".into(), json!({"type":"keyword"})),
("source_version".into(), json!({"type":"long"})),
("permission_version".into(), json!({"type":"long"})),
("doc_id".into(), json!({"type":"keyword"})),
("doc_token".into(), json!({"type":"keyword"})),
(text_field.into(), json!({"type":"text"})),
(
"created_at".into(),
json!({"type":if provider == "manticoresearch" { "long" } else { "date" }}),
),
(
"updated_at".into(),
json!({"type":if provider == "manticoresearch" { "long" } else { "date" }}),
),
("created_at".into(), json!({"type":"date"})),
("updated_at".into(), json!({"type":"date"})),
("created_by_user_id".into(), json!({"type":"keyword"})),
("updated_by_user_id".into(), json!({"type":"keyword"})),
("acl_public_readable".into(), json!({"type":"boolean"})),
("acl_member_default_readable".into(), json!({"type":"boolean"})),
(
"acl_read_tokens".into(),
if provider == "manticoresearch" {
json!({"type":"keyword","mva":true})
} else {
json!({"type":"keyword"})
},
),
("acl_revision".into(), json!({"type":"long"})),
("acl_read_tokens".into(), json!({"type":"keyword"})),
]);
if table == SearchTable::Block {
for field in [
@@ -65,61 +213,76 @@ pub(super) fn mapping(table: SearchTable, provider: &str) -> Value {
json!({"mappings":{"properties":properties}})
}
pub(super) fn manticore_schema(table: SearchTable, physical_table: &str) -> String {
let common = r#"
workspace_id string attribute indexed,
workspace_token string attribute indexed,
doc_id string attribute indexed,
doc_token string attribute indexed,"#;
let fields = match table {
SearchTable::Doc => format!(
r#"{common}
title text,
summary string stored,
journal string stored,
created_by_user_id string attribute indexed,
updated_by_user_id string attribute indexed,
created_at timestamp,
updated_at timestamp,
acl_public_readable bool,
acl_member_default_readable bool,
acl_read_token_ids multi64,
acl_revision bigint"#,
),
SearchTable::Block => format!(
r#"{common}
block_id string attribute indexed,
block_token string attribute indexed,
unit_id string attribute indexed,
projection_version bigint,
source_hash string attribute indexed,
visibility string attribute indexed,
element_id string attribute indexed,
frame_id string attribute indexed,
source_block_id string attribute indexed,
content text,
flavour string attribute indexed,
blob string attribute indexed,
ref_doc_id string attribute indexed,
ref_doc_token_ids multi64,
ref string stored,
parent_flavour string attribute indexed,
parent_block_id string attribute indexed,
additional string stored,
markdown_preview string stored,
created_by_user_id string attribute indexed,
updated_by_user_id string attribute indexed,
created_at timestamp,
updated_at timestamp,
acl_public_readable bool,
acl_member_default_readable bool,
acl_read_token_ids multi64,
acl_revision bigint"#,
),
};
format!(
"CREATE TABLE IF NOT EXISTS {physical_table} ({fields}) charset_table='non_cjk, chinese' ngram_len='1' \
ngram_chars='U+1100..U+11FF, U+3130..U+318F, U+A960..U+A97F, U+AC00..U+D7AF, U+D7B0..U+D7FF, U+3040..U+30FF, \
U+0E00..U+0E7F' index_field_lengths='1'"
)
pub(super) fn projection_external_id(
table: SearchTable,
generation_id: &str,
workspace_id: &str,
doc_id: &str,
block_id: Option<&str>,
source_version: i64,
permission_version: i64,
) -> RuntimeResult<String> {
let mut id = format!("{generation_id}/{workspace_id}/{doc_id}/{source_version}/{permission_version}");
if table == SearchTable::Block {
let block_id = block_id
.filter(|block_id| !block_id.is_empty())
.ok_or_else(|| RuntimeError::invalid_input("block projection id requires block_id"))?;
id.push('/');
id.push_str(block_id);
}
Ok(id)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{SearchChange, SearchTable, projection_external_id};
#[test]
fn projection_ids_are_stable_and_keep_block_identity_last() {
let doc = projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 7, 3).unwrap();
assert_eq!(doc, "generation/workspace/doc/7/3");
assert_eq!(
projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 7, 3).unwrap(),
doc
);
assert_eq!(
projection_external_id(
SearchTable::Block,
"generation",
"workspace",
"doc",
Some("block"),
7,
3,
)
.unwrap(),
"generation/workspace/doc/7/3/block"
);
assert_ne!(
projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 8, 3).unwrap(),
doc
);
}
#[test]
fn change_rejects_an_external_id_outside_its_tuple() {
let change = SearchChange {
table: SearchTable::Doc,
external_id: "generation/workspace/doc/6/3".into(),
workspace_id: "workspace".into(),
doc_id: "doc".into(),
source_version: 7,
permission_version: 3,
payload: json!({
"generation_id":"generation",
"workspace_id":"workspace",
"doc_id":"doc",
"source_version":7,
"permission_version":3
}),
};
assert!(change.upsert_payload().is_err());
}
}

View File

@@ -2,12 +2,8 @@ use std::time::Duration;
use reqwest::{Client, redirect::Policy};
use serde_json::{Value, json};
use sqlx::PgPool;
use super::{
super::{store::SearchChange, types::SearchTable},
manticore::{manticore_exact_tokens, manticore_fields, prepare_manticore_payload, prepare_manticore_search},
};
use super::{SearchChange, SearchTable, provider_write_error, webpki_tls_config};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
@@ -15,64 +11,42 @@ const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
pub(in crate::runtime::backend_runtime::search) struct RemoteProvider {
client: Client,
endpoint: String,
provider: String,
api_key: String,
username: String,
password: String,
pool: PgPool,
}
impl RemoteProvider {
pub(in crate::runtime::backend_runtime::search) fn new(
config: &SearchRuntimeConfig,
pool: PgPool,
) -> RuntimeResult<Self> {
pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
let endpoint = config.endpoint.trim_end_matches('/');
let url = url::Url::parse(endpoint).map_err(|_| RuntimeError::config("invalid search provider endpoint"))?;
if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return Err(RuntimeError::config("invalid search provider endpoint"));
}
let mut client = Client::builder()
let client = Client::builder()
.tls_backend_preconfigured(
webpki_tls_config()
.map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?,
)
.redirect(Policy::none())
.timeout(Duration::from_secs(30));
if config.provider == "manticoresearch" {
client = client.pool_max_idle_per_host(0);
}
let client = client
.build()
.map_err(|error| RuntimeError::invalid_state(format!("search HTTP client failed: {error}")))?;
Ok(Self {
client,
endpoint: endpoint.to_string(),
provider: config.provider.clone(),
api_key: config.api_key.clone(),
username: config.username.clone(),
password: config.password.clone(),
pool,
})
}
pub(in crate::runtime::backend_runtime::search) async fn search(
&self,
physical_table: &str,
mut dsl: Value,
) -> RuntimeResult<Value> {
let requested_fields = dsl
.get("fields")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_str)
.map(str::to_string)
.collect::<Vec<_>>();
pub(super) async fn search(&self, physical_table: &str, mut dsl: Value) -> RuntimeResult<Value> {
ensure_projection_source(&mut dsl);
dsl["track_total_hits"] = json!(true);
let size = dsl.get("size").and_then(Value::as_u64).unwrap_or(10);
let mut offset = dsl.get("from").and_then(Value::as_u64).unwrap_or(0);
let cursor = dsl.as_object_mut().and_then(|object| object.remove("cursor"));
if self.provider == "manticoresearch" {
let token_ids = self.resolve_manticore_tokens(manticore_exact_tokens(&dsl)).await?;
offset = prepare_manticore_search(&mut dsl, cursor, size, offset, &requested_fields, &token_ids)?;
} else if let Some(cursor) = cursor {
if let Some(cursor) = cursor {
let cursor = cursor
.as_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
@@ -103,22 +77,20 @@ impl RemoteProvider {
}
let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize(
value,
self.provider == "manticoresearch",
offset,
size,
&requested_fields,
)
normalize(value)
}
pub(in crate::runtime::backend_runtime::search) async fn aggregate(
&self,
physical_table: &str,
mut dsl: Value,
) -> RuntimeResult<Value> {
if self.provider == "manticoresearch" {
return Err(RuntimeError::SearchUnsupportedQuery);
pub(super) async fn aggregate(&self, physical_table: &str, mut dsl: Value) -> RuntimeResult<Value> {
if let Some(top_hits) = dsl.pointer_mut("/aggs/result/aggs/result/top_hits") {
ensure_projection_source(top_hits);
}
let limit = dsl
.pointer("/aggs/result/terms/size")
.and_then(Value::as_u64)
.unwrap_or(10);
let skip = dsl.get("from").and_then(Value::as_u64).unwrap_or_default();
if let Some(size) = dsl.pointer_mut("/aggs/result/terms/size") {
*size = json!(skip.saturating_add(limit).saturating_add(1));
}
dsl["track_total_hits"] = json!(true);
let response = self
@@ -130,32 +102,18 @@ impl RemoteProvider {
let status = response.status();
let bytes = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchUnsupportedQuery);
return Err(if status.as_u16() == 400 {
RuntimeError::SearchUnsupportedQuery
} else {
RuntimeError::SearchProviderUnavailable
});
}
let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize_aggregate(value)
normalize_aggregate(value, skip as usize, limit as usize)
}
pub(in crate::runtime::backend_runtime::search) async fn provision(
&self,
physical_table: &str,
table: SearchTable,
) -> RuntimeResult<()> {
if self.provider == "manticoresearch" {
let response = self
.request(reqwest::Method::POST, "cli")
.header("content-type", "text/plain")
.body(super::manticore_schema(table, physical_table))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
return response
.status()
.is_success()
.then_some(())
.ok_or(RuntimeError::SearchProviderUnavailable);
}
pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
if self
.request(reqwest::Method::HEAD, physical_table)
.send()
@@ -166,7 +124,7 @@ impl RemoteProvider {
}
let response = self
.request(reqwest::Method::PUT, physical_table)
.json(&super::mapping(table, &self.provider))
.json(&super::mapping(table))
.send()
.await;
match response {
@@ -175,73 +133,124 @@ impl RemoteProvider {
}
}
pub(in crate::runtime::backend_runtime::search) async fn apply(
&self,
physical_table: &str,
changes: &[SearchChange],
) -> RuntimeResult<()> {
pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
if changes.is_empty() {
return Ok(());
}
let token_ids = if self.provider == "manticoresearch" {
self
.resolve_manticore_tokens(
changes
.iter()
.filter_map(|change| change.payload.as_ref())
.flat_map(manticore_exact_tokens),
)
.await?
} else {
Default::default()
};
self.apply_elasticsearch(physical_table, changes).await
}
async fn apply_elasticsearch(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
let mut body = String::new();
for change in changes {
if change.operation == "delete" {
body.push_str(
&serde_json::to_string(&json!({"delete":{"_index":physical_table,"_id":change.external_id}}))
.map_err(|error| RuntimeError::json("encode provider delete", error))?,
);
body.push('\n');
} else if let Some(payload) = &change.payload {
body.push_str(
&serde_json::to_string(&json!({"index":{"_index":physical_table,"_id":change.external_id}}))
.map_err(|error| RuntimeError::json("encode provider upsert", error))?,
);
body.push('\n');
let mut payload = super::super::provider_payload(payload);
if self.provider == "manticoresearch" {
prepare_manticore_payload(&mut payload, &token_ids)?;
}
body.push_str(
&serde_json::to_string(&payload).map_err(|error| RuntimeError::json("encode provider document", error))?,
);
body.push('\n');
}
let payload = change.upsert_payload()?;
body.push_str(
&serde_json::to_string(&json!({
"index": {
"_index": physical_table,
"_id": change.external_id
}
}))
.map_err(|error| RuntimeError::json("encode immutable provider upsert", error))?,
);
body.push('\n');
body.push_str(
&serde_json::to_string(payload)
.map_err(|error| RuntimeError::json("encode immutable provider document", error))?,
);
body.push('\n');
}
if body.is_empty() {
return Ok(());
}
let path = if self.provider == "elasticsearch" {
"_bulk?refresh=wait_for"
} else {
"_bulk"
};
let response = self
.request(reqwest::Method::POST, path)
.request(reqwest::Method::POST, "_bulk?refresh=wait_for")
.header("content-type", "application/x-ndjson")
.body(body)
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
if !response.status().is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
return Err(provider_write_error(response.status().as_u16()));
}
let value: Value = response
.json()
.await
.map_err(|error| RuntimeError::invalid_state(format!("invalid provider bulk response: {error}")))?;
if value.get("errors").and_then(Value::as_bool) == Some(true) {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
validate_bulk_response(&value, changes.len())
}
pub(super) async fn gc_document_history(
&self,
physical_table: &str,
workspace_id: &str,
doc_id: &str,
source_version: i64,
permission_version: i64,
limit: usize,
) -> RuntimeResult<()> {
let query = json!({"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"term":{"doc_id":{"value":doc_id}}}
],"must_not":[{"bool":{"must":[
{"term":{"source_version":{"value":source_version}}},
{"term":{"permission_version":{"value":permission_version}}}
]}}]}});
let response = self
.request(
reqwest::Method::POST,
&format!("{physical_table}/_delete_by_query?conflicts=proceed&refresh=true"),
)
.json(&json!({"query":query,"max_docs":limit.max(1)}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
Ok(())
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid provider history GC response", error))?;
validate_delete_response(&value)
}
pub(super) async fn gc_workspace(
&self,
physical_table: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> RuntimeResult<bool> {
let limit = limit.max(1);
let response = self
.request(
reqwest::Method::POST,
&format!("{physical_table}/_delete_by_query?conflicts=proceed&refresh=true"),
)
.json(&json!({
"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"range":{"source_version":{"lte":source_version_high_water}}}
]}},
"max_docs":limit
}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid provider workspace GC response", error))?;
validate_delete_response(&value)?;
let deleted = value
.get("deleted")
.and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider workspace GC response"))?;
Ok(deleted >= limit as u64)
}
fn request(&self, method: reqwest::Method, path: &str) -> reqwest::RequestBuilder {
@@ -253,48 +262,61 @@ impl RemoteProvider {
}
request
}
}
async fn resolve_manticore_tokens(
&self,
tokens: impl IntoIterator<Item = String>,
) -> RuntimeResult<std::collections::HashMap<String, i64>> {
let tokens = tokens.into_iter().collect::<std::collections::BTreeSet<_>>();
if tokens.is_empty() {
return Ok(Default::default());
}
let tokens = tokens.into_iter().collect::<Vec<_>>();
let mut transaction = self
.pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin Manticore exact token resolution", error))?;
sqlx::query(
r#"INSERT INTO search_runtime_acl_tokens(token)
SELECT candidate.token FROM unnest($1::text[]) candidate(token)
LEFT JOIN search_runtime_acl_tokens existing USING(token)
WHERE existing.token IS NULL ON CONFLICT DO NOTHING"#,
)
.bind(&tokens)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("allocate Manticore exact token IDs", error))?;
let rows: Vec<(String, i64)> =
sqlx::query_as("SELECT token,token_id FROM search_runtime_acl_tokens WHERE token=ANY($1)")
.bind(&tokens)
.fetch_all(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load Manticore exact token IDs", error))?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit Manticore exact token resolution", error))?;
if rows.len() != tokens.len() {
return Err(RuntimeError::invalid_state(
"Manticore exact token mapping is incomplete",
));
}
Ok(rows.into_iter().collect())
fn validate_delete_response(value: &Value) -> RuntimeResult<()> {
if value.get("timed_out").and_then(Value::as_bool) == Some(true)
|| value
.get("failures")
.and_then(Value::as_array)
.is_some_and(|failures| !failures.is_empty())
{
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(())
}
fn validate_bulk_response(value: &Value, expected_items: usize) -> RuntimeResult<()> {
let items = value
.get("items")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider bulk response"))?;
if items.len() != expected_items {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
for item in items {
let result = item
.as_object()
.and_then(|item| item.values().next())
.and_then(Value::as_object)
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
let status = result
.get("status")
.and_then(Value::as_u64)
.and_then(|status| u16::try_from(status).ok())
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
if result.get("error").is_some() || !(200..300).contains(&status) {
return Err(provider_write_error(status));
}
}
if value.get("errors").and_then(Value::as_bool) != Some(false) {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(())
}
fn ensure_projection_source(dsl: &mut Value) {
let mut source = dsl
.get("_source")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
for field in ["doc_id", "source_version", "permission_version"] {
if !source.iter().any(|value| value.as_str() == Some(field)) {
source.push(json!(field));
}
}
dsl["_source"] = json!(source);
}
async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>> {
@@ -318,13 +340,7 @@ async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>
Ok(bytes)
}
fn normalize(
value: Value,
manticore: bool,
_offset: u64,
_size: u64,
requested_fields: &[String],
) -> RuntimeResult<Value> {
fn normalize(value: Value) -> RuntimeResult<Value> {
let hits = value
.pointer("/hits/hits")
.and_then(Value::as_array)
@@ -337,11 +353,7 @@ fn normalize(
let nodes = hits
.iter()
.map(|hit| {
let fields = if manticore {
manticore_fields(hit.get("_source"), requested_fields)
} else {
hit.get("fields").cloned().unwrap_or_else(|| json!({}))
};
let fields = hit.get("fields").cloned().unwrap_or_else(|| json!({}));
json!({
"id":hit.get("_id").and_then(Value::as_str).unwrap_or_default(),
"score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(),
@@ -351,23 +363,16 @@ fn normalize(
})
})
.collect::<Vec<_>>();
let next_cursor = if manticore {
(!hits.is_empty())
.then(|| serde_json::to_string(&json!({"offset":_offset + _size})))
.transpose()
.map_err(|error| RuntimeError::json("encode provider cursor", error))?
} else {
hits
.last()
.and_then(|hit| hit.get("sort"))
.map(serde_json::to_string)
.transpose()
.map_err(|error| RuntimeError::json("encode provider cursor", error))?
};
let next_cursor = hits
.last()
.and_then(|hit| hit.get("sort"))
.map(serde_json::to_string)
.transpose()
.map_err(|error| RuntimeError::json("encode provider cursor", error))?;
Ok(json!({"total":total,"nodes":nodes,"nextCursor":next_cursor}))
}
fn normalize_aggregate(value: Value) -> RuntimeResult<Value> {
fn normalize_aggregate(value: Value, skip: usize, limit: usize) -> RuntimeResult<Value> {
let buckets = value
.pointer("/aggregations/result/buckets")
.and_then(Value::as_array)
@@ -382,12 +387,14 @@ fn normalize_aggregate(value: Value) -> RuntimeResult<Value> {
Ok(json!({
"key":bucket.get("key").cloned().unwrap_or(Value::Null),
"count":bucket.get("doc_count").cloned().unwrap_or(json!(0)),
"hits":{"total":bucket.pointer("/result/hits/total/value").or_else(||bucket.pointer("/result/hits/total")).cloned().unwrap_or(json!(0)),
"nodes":hits.iter().map(normalize_hit).collect::<Vec<_>>()}
"hits":hits.iter().map(normalize_hit).collect::<Vec<_>>()
}))
})
.collect::<RuntimeResult<Vec<_>>>()?;
Ok(json!({"total":nodes.len(),"hasMore":false,"buckets":nodes}))
let total = nodes.len();
let has_more = total > skip.saturating_add(limit);
let nodes = nodes.into_iter().skip(skip).take(limit).collect::<Vec<_>>();
Ok(json!({"total":total,"hasMore":has_more,"buckets":nodes}))
}
fn normalize_hit(hit: &Value) -> Value {
@@ -399,3 +406,76 @@ fn normalize_hit(hit: &Value) -> Value {
"_source":hit.get("_source").cloned().unwrap_or_else(||json!({})),
})
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{normalize_aggregate, validate_bulk_response, validate_delete_response};
use crate::runtime::RuntimeError;
#[test]
fn delete_response_rejects_partial_provider_failures() {
assert!(validate_delete_response(&json!({"failures":[]})).is_ok());
assert!(validate_delete_response(&json!({"timed_out":true})).is_err());
assert!(validate_delete_response(&json!({"failures":[{"reason":"stale"}]})).is_err());
}
#[test]
fn bulk_response_requires_every_item_to_succeed() {
for (response, expected_items, valid) in [
(
json!({"errors":false,"items":[{"index":{"status":201}},{"index":{"status":200}}]}),
2,
true,
),
(json!({"errors":false,"items":[{"index":{"status":404}}]}), 1, false),
(
json!({"errors":true,"items":[{"index":{"status":201}},{"index":{"status":400,"error":{}}}]}),
2,
false,
),
(json!({"errors":false}), 1, false),
] {
assert_eq!(validate_bulk_response(&response, expected_items).is_ok(), valid);
}
assert!(matches!(
validate_bulk_response(&json!({"errors":true,"items":[{"index":{"status":400,"error":{}}}]}), 1),
Err(RuntimeError::SearchSourceInvalid(_))
));
assert!(matches!(
validate_bulk_response(&json!({"errors":true,"items":[{"index":{"status":429,"error":{}}}]}), 1),
Err(RuntimeError::SearchProviderUnavailable)
));
}
#[test]
fn aggregate_normalization_reports_extra_buckets() {
let value = json!({
"aggregations":{"result":{"buckets":[
{"key":"one","doc_count":2,"result":{"hits":{"hits":[]}}},
{"key":"two","doc_count":1,"result":{"hits":{"hits":[]}}}
]}}
});
let result = normalize_aggregate(value, 0, 1).unwrap();
assert_eq!(result["total"], 2);
assert_eq!(result["hasMore"], true);
assert_eq!(result["buckets"].as_array().unwrap().len(), 1);
assert!(result["buckets"][0]["hits"].is_array());
}
#[test]
fn aggregate_normalization_applies_bucket_skip_before_limit() {
let value = json!({
"aggregations":{"result":{"buckets":[
{"key":"one","doc_count":2,"result":{"hits":{"hits":[]}}},
{"key":"two","doc_count":1,"result":{"hits":{"hits":[]}}},
{"key":"three","doc_count":1,"result":{"hits":{"hits":[]}}}
]}}
});
let result = normalize_aggregate(value, 1, 1).unwrap();
assert_eq!(result["total"], 3);
assert_eq!(result["hasMore"], true);
assert_eq!(result["buckets"][0]["key"], "two");
}
}

View File

@@ -1,10 +1,105 @@
use serde::Deserialize;
use serde_json::{Value, json};
use super::types::{AggregateRequest, SearchOptions, SearchQuery, SearchRequest, SearchTable};
use crate::runtime::{
RuntimeError, RuntimeResult,
backend_runtime::permission::{AuthorizedSearchScope, DocReadScope},
use super::{
AggregateOptions, AuthorizedSearchScope, DocReadScope, RuntimeAggregateRequest, RuntimeSearchQuery,
RuntimeSearchRequest, SearchOptions, SearchTable,
};
use crate::runtime::{RuntimeError, RuntimeResult};
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct AggregateRequest {
pub(super) table: SearchTable,
query: SearchQuery,
field: String,
options: AggregateOptions,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchRequest {
pub(super) table: SearchTable,
query: SearchQuery,
options: SearchOptions,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchQuery {
#[serde(rename = "type")]
query_type: String,
field: Option<String>,
#[serde(rename = "match")]
match_value: Option<String>,
query: Option<Box<SearchQuery>>,
queries: Option<Vec<SearchQuery>>,
occur: Option<String>,
boost: Option<f64>,
}
impl RuntimeSearchRequest {
pub(super) fn into_search_request(self) -> RuntimeResult<SearchRequest> {
let mut decoded_nodes = 0;
Ok(SearchRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
options: self.options,
})
}
}
impl RuntimeAggregateRequest {
pub(super) fn into_aggregate_request(self) -> RuntimeResult<AggregateRequest> {
let mut decoded_nodes = 0;
Ok(AggregateRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
field: self.field,
options: self.options,
})
}
}
const MAX_QUERY_GRAPH_NODES: usize = 100;
const MAX_QUERY_DEPTH: usize = 100;
const MAX_DECODED_QUERY_NODES: usize = 1_000;
fn decode_query(
nodes: &[RuntimeSearchQuery],
index: u32,
depth: usize,
decoded_nodes: &mut usize,
) -> RuntimeResult<SearchQuery> {
if nodes.len() > MAX_QUERY_GRAPH_NODES || depth > MAX_QUERY_DEPTH || *decoded_nodes >= MAX_DECODED_QUERY_NODES {
return Err(RuntimeError::invalid_input("search query is too complex"));
}
*decoded_nodes += 1;
let node = nodes
.get(index as usize)
.ok_or_else(|| RuntimeError::invalid_input("invalid search query node"))?;
Ok(SearchQuery {
query_type: node.query_type.clone(),
field: node.field.clone(),
match_value: node.match_value.clone(),
query: node
.query
.map(|index| decode_query(nodes, index, depth + 1, decoded_nodes).map(Box::new))
.transpose()?,
queries: node
.queries
.as_ref()
.map(|indices| {
indices
.iter()
.map(|index| decode_query(nodes, *index, depth + 1, decoded_nodes))
.collect()
})
.transpose()?,
occur: node.occur.clone(),
boost: node.boost,
})
}
pub(super) fn compile(request: &SearchRequest, scope: &AuthorizedSearchScope) -> RuntimeResult<Value> {
let query = compile_query(request.table, &request.query)?;
@@ -26,7 +121,7 @@ pub(super) fn compile(request: &SearchRequest, scope: &AuthorizedSearchScope) ->
.map(|field| validate_field(request.table, field).map(str::to_string))
.collect::<RuntimeResult<Vec<_>>>()?;
let mut dsl = json!({
"_source":["workspace_id","doc_id"],
"_source":["workspace_id","doc_id","source_version","permission_version"],
"fields":fields,
"query":{"bool":{"must":must}},
"sort": stable_sort(request.table),
@@ -73,15 +168,19 @@ pub(super) fn compile_aggregate(request: &AggregateRequest, scope: &AuthorizedSe
let hit_dsl = compile(&search, scope)?;
let field = validate_field(request.table, &request.field)?;
let limit = request.options.pagination.limit.unwrap_or(10);
if limit > 10_000 {
return Err(RuntimeError::invalid_input("aggregate limit exceeds 10000"));
let skip = request.options.pagination.skip.unwrap_or(0);
if skip.saturating_add(limit) > 10_000 {
return Err(RuntimeError::invalid_input("aggregate pagination exceeds 10000"));
}
if request.options.pagination.cursor.is_some() {
return Err(RuntimeError::invalid_input("aggregate cursor is unsupported"));
}
Ok(json!({
"query":hit_dsl["query"],
"from":request.options.pagination.skip.unwrap_or(0),
"from":skip,
"size":0,
"aggs":{"result":{"terms":{"field":field,"size":limit},"aggs":{"result":{"top_hits":{
"size":hit_dsl["size"],"_source":hit_dsl["_source"],"fields":hit_dsl["fields"],
"size":hit_dsl["size"],"from":hit_dsl.get("from").cloned().unwrap_or_else(||json!(0)),"_source":hit_dsl["_source"],"fields":hit_dsl["fields"],
"sort":hit_dsl["sort"],"highlight":hit_dsl.get("highlight").cloned().unwrap_or_else(||json!({}))
}}}}}
}))
@@ -234,7 +333,6 @@ mod tests {
fn scope() -> AuthorizedSearchScope {
AuthorizedSearchScope {
workspace_id: "workspace".to_string(),
permission_revision: 1,
docs: DocReadScope::All,
}
}
@@ -325,6 +423,10 @@ mod tests {
dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["highlight"]["fields"]["content"],
json!({"pre_tags":["<b>"],"post_tags":["</b>"]})
);
let mut hit_skip = aggregate.clone();
hit_skip.options.hits.pagination.skip = Some(1);
let hit_skip_dsl = compile_aggregate(&hit_skip, &scope()).unwrap();
assert_eq!(hit_skip_dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["from"], 1);
let mut invalid = aggregate;
invalid.field = "aclReadTokens".to_string();
@@ -332,5 +434,43 @@ mod tests {
invalid.field = "docId".to_string();
invalid.options.pagination.limit = Some(10_001);
assert!(compile_aggregate(&invalid, &scope()).is_err());
invalid.options.pagination.limit = Some(10_000);
invalid.options.pagination.skip = Some(1);
assert!(compile_aggregate(&invalid, &scope()).is_err());
invalid.options.pagination.limit = Some(10);
invalid.options.pagination.skip = None;
invalid.options.pagination.cursor = Some("candidate".to_string());
assert!(compile_aggregate(&invalid, &scope()).is_err());
}
fn runtime_node(query_type: &str) -> RuntimeSearchQuery {
RuntimeSearchQuery {
query_type: query_type.to_string(),
field: None,
match_value: None,
query: None,
queries: None,
occur: None,
boost: None,
}
}
#[test]
fn rejects_invalid_or_overly_complex_query_graphs() {
assert!(decode_query(&[runtime_node("all")], 1, 0, &mut 0).is_err());
let mut oversized = (0..101).map(|_| runtime_node("all")).collect::<Vec<_>>();
oversized[0].query = Some(1);
assert!(decode_query(&oversized, 0, 0, &mut 0).is_err());
let mut recursive = vec![runtime_node("boost")];
recursive[0].query = Some(0);
assert!(decode_query(&recursive, 0, 0, &mut 0).is_err());
let mut shared_child = (0..100).map(|_| runtime_node("boolean")).collect::<Vec<_>>();
for (index, node) in shared_child.iter_mut().enumerate().take(99) {
node.queries = Some(vec![(index + 1) as u32, (index + 1) as u32]);
}
assert!(decode_query(&shared_child, 0, 0, &mut 0).is_err());
}
}

View File

@@ -0,0 +1,376 @@
use std::{
collections::{BTreeMap, BTreeSet},
sync::atomic::Ordering,
};
use sqlx::Row;
use tokio::time::Instant;
use super::{
ActiveGeneration, DocReadScope, RuntimeAggregateRequest, RuntimeSearchRequest, SearchActor, SearchRuntime,
SearchTable, candidates, compile, compile_aggregate, retain_visible_nodes,
};
use crate::runtime::{RuntimeError, RuntimeResult};
const SEARCH_CANDIDATE_PAGE_BUDGET: usize = 3;
impl SearchRuntime {
pub(in crate::runtime::backend_runtime) async fn search_authorized(
&self,
actor_user_id: &str,
workspace_id: &str,
request: RuntimeSearchRequest,
) -> RuntimeResult<serde_json::Value> {
let request = request.into_search_request()?;
for attempt in 0..=1 {
let scope = self
.authorizer
.authorize_search(
&SearchActor::User {
user_id: actor_user_id.to_string(),
},
workspace_id,
)
.await?;
let (generation, required_permission_version) = self.query_gate(workspace_id).await?;
let dsl = compile(&request, &scope)?;
let result = self
.execute_visible_search(
&generation,
workspace_id,
actor_user_id,
&scope.docs,
request.table,
dsl,
)
.await?;
if self
.query_snapshot_is_current(workspace_id, generation.id, required_permission_version)
.await?
{
return Ok(result);
}
if attempt == 1 {
return Err(RuntimeError::SearchPermissionSyncing);
}
}
unreachable!()
}
pub(super) async fn execute_visible_search(
&self,
generation: &ActiveGeneration,
workspace_id: &str,
actor_user_id: &str,
scope: &DocReadScope,
table: SearchTable,
mut dsl: serde_json::Value,
) -> RuntimeResult<serde_json::Value> {
let limit = dsl.get("size").and_then(serde_json::Value::as_u64).unwrap_or(10) as usize;
let mut result = serde_json::json!({"total":0,"nodes":[],"nextCursor":null});
for _ in 0..SEARCH_CANDIDATE_PAGE_BUDGET {
let mut page = self.execute_search(generation, table, dsl.clone()).await?;
self
.retain_canonically_visible(generation, workspace_id, actor_user_id, scope, &mut page)
.await?;
let next_cursor = page.get("nextCursor").cloned().unwrap_or(serde_json::Value::Null);
let nodes = page
.get_mut("nodes")
.and_then(serde_json::Value::as_array_mut)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider response"))?;
result["nodes"]
.as_array_mut()
.expect("search result nodes are initialized")
.append(nodes);
let returned = result["nodes"]
.as_array()
.expect("search result nodes are initialized")
.len();
result["total"] = serde_json::json!(returned);
result["nextCursor"] = next_cursor.clone();
let Some(cursor) = next_cursor.as_str() else {
break;
};
if returned >= limit || limit == 0 {
break;
}
dsl
.as_object_mut()
.expect("compiled search DSL is an object")
.remove("from");
dsl["cursor"] = serde_json::json!(cursor);
dsl["size"] = serde_json::json!(limit - returned);
}
Ok(result)
}
pub(super) async fn retain_canonically_visible(
&self,
generation: &ActiveGeneration,
workspace_id: &str,
actor_user_id: &str,
scope: &DocReadScope,
result: &mut serde_json::Value,
) -> RuntimeResult<()> {
let candidate_tuples = candidates(result)?;
let doc_ids = candidate_tuples
.iter()
.map(|candidate| candidate.doc_id.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
if doc_ids.is_empty() {
return retain_visible_nodes(result, &BTreeSet::new());
}
let rows = sqlx::query(
r#"SELECT state.doc_id,state.published_source_version,state.published_permission_version
FROM search_projection.document_states state
JOIN snapshots source
ON source.workspace_id=state.workspace_id AND source.guid=state.doc_id
WHERE state.generation_id=$1 AND state.workspace_id=$2
AND state.doc_id=ANY($3) AND state.published_source_exists"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(&doc_ids)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load published search document states", error))?;
let published = rows
.into_iter()
.map(|row| {
let doc_id: String = row
.try_get("doc_id")
.map_err(|error| RuntimeError::database("decode published search document", error))?;
let source_version: i64 = row
.try_get("published_source_version")
.map_err(|error| RuntimeError::database("decode published search source version", error))?;
let permission_version: i64 = row
.try_get("published_permission_version")
.map_err(|error| RuntimeError::database("decode published search permission version", error))?;
Ok((doc_id, (source_version, permission_version)))
})
.collect::<RuntimeResult<BTreeMap<_, _>>>()?;
let current_doc_ids = candidate_tuples
.iter()
.filter(|candidate| {
published.get(&candidate.doc_id) == Some(&(candidate.source_version, candidate.permission_version))
})
.map(|candidate| candidate.doc_id.clone())
.collect::<BTreeSet<_>>();
let missing_published = candidate_tuples
.iter()
.filter(|candidate| !published.contains_key(&candidate.doc_id))
.count() as u64;
let projection_mismatch = candidate_tuples
.iter()
.filter(|candidate| {
published
.get(&candidate.doc_id)
.is_some_and(|tuple| *tuple != (candidate.source_version, candidate.permission_version))
})
.count() as u64;
let readable = if matches!(scope, DocReadScope::ProjectedAcl(_)) {
self
.authorizer
.filter_readable_docs(workspace_id, actor_user_id, current_doc_ids.iter().cloned().collect())
.await?
} else {
current_doc_ids.clone()
};
self
.observability
.missing_published
.fetch_add(missing_published, Ordering::Relaxed);
self
.observability
.projection_mismatch
.fetch_add(projection_mismatch, Ordering::Relaxed);
self.observability.canonical_permission.fetch_add(
current_doc_ids.len().saturating_sub(readable.len()) as u64,
Ordering::Relaxed,
);
let visible = candidate_tuples
.into_iter()
.filter(|candidate| {
readable.contains(&candidate.doc_id)
&& published.get(&candidate.doc_id) == Some(&(candidate.source_version, candidate.permission_version))
})
.collect::<BTreeSet<_>>();
retain_visible_nodes(result, &visible)
}
pub(in crate::runtime::backend_runtime) async fn aggregate_authorized(
&self,
actor_user_id: &str,
workspace_id: &str,
request: RuntimeAggregateRequest,
) -> RuntimeResult<serde_json::Value> {
let request = request.into_aggregate_request()?;
for attempt in 0..=1 {
let scope = self
.authorizer
.authorize_search(
&SearchActor::User {
user_id: actor_user_id.to_string(),
},
workspace_id,
)
.await?;
let (generation, required_permission_version) = self.query_gate(workspace_id).await?;
let dsl = compile_aggregate(&request, &scope)?;
let mut result = self.execute_aggregate(&generation, request.table, dsl).await?;
self
.retain_canonically_visible(&generation, workspace_id, actor_user_id, &scope.docs, &mut result)
.await?;
if self
.query_snapshot_is_current(workspace_id, generation.id, required_permission_version)
.await?
{
return Ok(result);
}
if attempt == 1 {
return Err(RuntimeError::SearchPermissionSyncing);
}
}
unreachable!()
}
async fn execute_search(
&self,
generation: &ActiveGeneration,
table: SearchTable,
dsl: serde_json::Value,
) -> RuntimeResult<serde_json::Value> {
let started = Instant::now();
let result = if let Some(remote) = &self.remote {
remote.search(generation.physical_table(table)?, dsl).await
} else if !self.embedded.has_generation(generation.id).await {
Err(RuntimeError::SearchIndexNotReady)
} else {
match serde_json::to_string(&dsl).map_err(|error| RuntimeError::json("encode embedded search", error)) {
Ok(dsl) => match self
.embedded
.search_for_generation(generation.id, table.as_str().to_string(), dsl)
.await
{
Ok(result) => {
serde_json::from_str(&result).map_err(|error| RuntimeError::json("decode embedded search", error))
}
Err(error) => Err(error.into()),
},
Err(error) => Err(error),
}
};
self.observe_provider_request(started);
result
}
async fn execute_aggregate(
&self,
generation: &ActiveGeneration,
table: SearchTable,
dsl: serde_json::Value,
) -> RuntimeResult<serde_json::Value> {
let started = Instant::now();
let result = if let Some(remote) = &self.remote {
remote.aggregate(generation.physical_table(table)?, dsl).await
} else if !self.embedded.has_generation(generation.id).await {
Err(RuntimeError::SearchIndexNotReady)
} else {
match serde_json::to_string(&dsl).map_err(|error| RuntimeError::json("encode embedded aggregate", error)) {
Ok(dsl) => match self
.embedded
.aggregate_for_generation(generation.id, table.as_str().to_string(), dsl)
.await
{
Ok(result) => {
serde_json::from_str(&result).map_err(|error| RuntimeError::json("decode embedded aggregate", error))
}
Err(error) => Err(error.into()),
},
Err(error) => Err(error),
}
};
self.observe_provider_request(started);
let mut value: serde_json::Value = result?;
if let Some(buckets) = value.get_mut("buckets").and_then(serde_json::Value::as_array_mut) {
for bucket in buckets {
let hits = bucket
.as_object_mut()
.and_then(|bucket| bucket.remove("hits"))
.unwrap_or_else(|| serde_json::json!([]));
bucket["hits"] = serde_json::json!({"nodes":hits});
}
}
Ok(value)
}
fn observe_provider_request(&self, started: Instant) {
self.observability.provider_requests.fetch_add(1, Ordering::Relaxed);
self.observability.provider_latency_micros.fetch_add(
started.elapsed().as_micros().min(u64::MAX as u128) as u64,
Ordering::Relaxed,
);
}
async fn query_gate(&self, workspace_id: &str) -> RuntimeResult<(ActiveGeneration, i64)> {
let generation = self.active_generation().await?;
let row = sqlx::query(
r#"SELECT covered,required_permission_version,applied_permission_version,last_error
FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation.id)
.bind(workspace_id)
.fetch_optional(&self.pool)
.await
.map_err(|error| RuntimeError::database("load search workspace gate", error))?
.ok_or(RuntimeError::SearchIndexNotReady)?;
let covered: bool = row
.try_get("covered")
.map_err(|error| RuntimeError::database("decode search coverage", error))?;
let required: i64 = row
.try_get("required_permission_version")
.map_err(|error| RuntimeError::database("decode search required permission", error))?;
let applied: i64 = row
.try_get("applied_permission_version")
.map_err(|error| RuntimeError::database("decode search applied permission", error))?;
let last_error: Option<String> = row
.try_get("last_error")
.map_err(|error| RuntimeError::database("decode search workspace error", error))?;
if last_error.is_some() {
return Err(RuntimeError::SearchIndexFailed(
"search_workspace_reconcile_failed".to_string(),
));
}
if !covered {
return Err(RuntimeError::SearchIndexNotReady);
}
if required > applied {
return Err(RuntimeError::SearchPermissionSyncing);
}
Ok((generation, required))
}
pub(super) async fn query_snapshot_is_current(
&self,
workspace_id: &str,
generation_id: uuid::Uuid,
permission_version: i64,
) -> RuntimeResult<bool> {
sqlx::query_scalar(
r#"SELECT EXISTS(
SELECT 1 FROM search_projection.generations generation
JOIN search_projection.workspace_states state ON state.generation_id=generation.id
WHERE generation.id=$1 AND state.workspace_id=$2 AND generation.state='active'
AND state.required_permission_version=$3
)"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(permission_version)
.fetch_one(&self.pool)
.await
.map_err(|error| RuntimeError::database("recheck search query snapshot", error))
}
}

View File

@@ -0,0 +1,152 @@
use std::collections::BTreeSet;
use serde_json::{Value, json};
use crate::runtime::{RuntimeError, RuntimeResult};
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub(super) struct CandidateTuple {
pub(super) doc_id: String,
pub(super) source_version: i64,
pub(super) permission_version: i64,
}
pub(super) fn candidates(result: &Value) -> RuntimeResult<Vec<CandidateTuple>> {
let mut candidates = Vec::new();
for node in result_nodes(result)? {
candidates.push(candidate(node)?);
}
candidates.sort();
candidates.dedup();
Ok(candidates)
}
pub(super) fn retain_visible_nodes(result: &mut Value, visible: &BTreeSet<CandidateTuple>) -> RuntimeResult<()> {
if let Some(nodes) = result.get_mut("nodes").and_then(Value::as_array_mut) {
retain_nodes(nodes, visible)?;
result["total"] = json!(nodes.len());
return Ok(());
}
let buckets = result
.get_mut("buckets")
.and_then(Value::as_array_mut)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider response"))?;
for bucket in buckets.iter_mut() {
let count = {
let nodes = bucket
.pointer_mut("/hits/nodes")
.and_then(Value::as_array_mut)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider aggregate response"))?;
retain_nodes(nodes, visible)?;
nodes.len()
};
bucket["count"] = json!(count);
bucket["hits"]["total"] = json!(count);
}
buckets.retain(|bucket| {
bucket
.pointer("/hits/nodes")
.and_then(Value::as_array)
.is_some_and(|nodes| !nodes.is_empty())
});
result["total"] = json!(buckets.len());
Ok(())
}
fn result_nodes(result: &Value) -> RuntimeResult<Vec<&Value>> {
if let Some(nodes) = result.get("nodes").and_then(Value::as_array) {
return Ok(nodes.iter().collect());
}
let buckets = result
.get("buckets")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider response"))?;
buckets
.iter()
.map(|bucket| {
bucket
.pointer("/hits/nodes")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider aggregate response"))
})
.collect::<RuntimeResult<Vec<_>>>()
.map(|buckets| buckets.into_iter().flatten().collect())
}
fn retain_nodes(nodes: &mut Vec<Value>, visible: &BTreeSet<CandidateTuple>) -> RuntimeResult<()> {
let mut decoded = Vec::with_capacity(nodes.len());
for node in nodes.iter() {
decoded.push(candidate(node)?);
}
let mut index = 0;
nodes.retain(|_| {
let keep = visible.contains(&decoded[index]);
index += 1;
keep
});
Ok(())
}
fn candidate(node: &Value) -> RuntimeResult<CandidateTuple> {
Ok(CandidateTuple {
doc_id: string_field(node, "doc_id")?.to_string(),
source_version: integer_field(node, "source_version")?,
permission_version: integer_field(node, "permission_version")?,
})
}
fn string_field<'a>(node: &'a Value, field: &str) -> RuntimeResult<&'a str> {
node
.pointer(&format!("/fields/{field}/0"))
.and_then(Value::as_str)
.or_else(|| node.pointer(&format!("/_source/{field}")).and_then(Value::as_str))
.ok_or_else(|| RuntimeError::invalid_state(format!("provider result is missing {field}")))
}
fn integer_field(node: &Value, field: &str) -> RuntimeResult<i64> {
node
.pointer(&format!("/fields/{field}/0"))
.and_then(Value::as_i64)
.or_else(|| node.pointer(&format!("/_source/{field}")).and_then(Value::as_i64))
.filter(|version| *version >= 0)
.ok_or_else(|| RuntimeError::invalid_state(format!("provider result has invalid {field}")))
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{CandidateTuple, candidates, retain_visible_nodes};
#[test]
fn canonical_filter_requires_the_published_projection_tuple() {
let mut result = json!({
"total": 8,
"nextCursor": "{\"offset\":2}",
"nodes": [
{"fields":{"doc_id":["readable"],"source_version":[2],"permission_version":[3]}},
{"_source":{"doc_id":"readable","source_version":1,"permission_version":3}}
]
});
assert_eq!(candidates(&result).unwrap().len(), 2);
retain_visible_nodes(
&mut result,
&[CandidateTuple {
doc_id: "readable".to_string(),
source_version: 2,
permission_version: 3,
}]
.into_iter()
.collect(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(result["nodes"].as_array().unwrap().len(), 1);
assert_eq!(result["nextCursor"], "{\"offset\":2}");
}
#[test]
fn malformed_provider_tuple_fails_closed() {
assert!(candidates(&json!({"nodes":[{"_source":{"doc_id":"doc"}}]})).is_err());
}
}

View File

@@ -1,6 +0,0 @@
mod projection;
pub(super) mod stream;
mod types;
pub(super) use projection::SearchStore;
pub(super) use types::{ProjectionInput, SearchChange, SearchSnapshot, SearchTable};

View File

@@ -1,450 +0,0 @@
use std::collections::{HashMap, HashSet};
use serde_json::Value;
use sqlx::{PgPool, Postgres, Row, Transaction};
use super::{ProjectionInput, SearchChange, SearchSnapshot, SearchTable, stream::allocate};
use crate::runtime::{RuntimeError, RuntimeResult};
pub(in crate::runtime::backend_runtime::search) struct SearchStore {
pool: PgPool,
}
impl SearchStore {
pub(in crate::runtime::backend_runtime::search) fn new(pool: PgPool) -> Self {
Self { pool }
}
pub(in crate::runtime::backend_runtime::search) async fn replace_document(
&self,
mut document: ProjectionInput,
mut blocks: Vec<ProjectionInput>,
) -> RuntimeResult<()> {
if blocks
.iter()
.any(|block| block.workspace_id != document.workspace_id || block.doc_id != document.doc_id)
{
return Err(RuntimeError::invalid_input("block identity does not match document"));
}
let mut transaction = self
.pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search projection transaction", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(format!("{}/{}", document.workspace_id, document.doc_id))
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("lock search document projection", error))?;
let current_document = load_rows(
&mut transaction,
SearchTable::Doc,
&document.workspace_id,
&document.doc_id,
)
.await?;
if let Some(existing) = current_document.get(&document.external_id)
&& existing != &document
&& existing.acl_revision < document.acl_revision
{
document.revision = existing.revision + 1;
for block in &mut blocks {
block.revision = document.revision;
}
}
if let Some(existing) = current_document.get(&document.external_id) {
if existing.revision > document.revision {
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit stale search projection", error))?;
return Ok(());
}
if existing.revision == document.revision {
let current_blocks = load_rows(
&mut transaction,
SearchTable::Block,
&document.workspace_id,
&document.doc_id,
)
.await?;
let incoming_blocks = blocks
.iter()
.map(|block| (block.external_id.clone(), block))
.collect::<HashMap<_, _>>();
let blocks_match = current_blocks.len() == incoming_blocks.len()
&& current_blocks
.iter()
.all(|(id, block)| incoming_blocks.get(id).is_some_and(|incoming| block == *incoming));
if existing != &document || !blocks_match {
return Err(RuntimeError::invalid_state("conflicting search projection revision"));
}
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit duplicate search projection", error))?;
return Ok(());
}
}
self
.replace_table(
&mut transaction,
SearchTable::Doc,
vec![document.clone()],
Some(document.revision),
)
.await?;
self
.replace_table(&mut transaction, SearchTable::Block, blocks, Some(document.revision))
.await?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit search projection transaction", error))
}
pub(in crate::runtime::backend_runtime::search) async fn delete_document(
&self,
workspace_id: &str,
doc_id: &str,
revision: i64,
) -> RuntimeResult<()> {
let mut transaction = self
.pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search delete transaction", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
.bind(format!("{workspace_id}/{doc_id}"))
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("lock search document deletion", error))?;
for table in SearchTable::ORDERED {
let rows = load_rows(&mut transaction, table, workspace_id, doc_id).await?;
let deletions = rows
.into_values()
.filter(|row| row.revision <= revision)
.collect::<Vec<_>>();
self
.apply(&mut transaction, table, Vec::new(), deletions, revision)
.await?;
}
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit search delete transaction", error))
}
async fn replace_table(
&self,
transaction: &mut Transaction<'_, Postgres>,
table: SearchTable,
inputs: Vec<ProjectionInput>,
delete_revision: Option<i64>,
) -> RuntimeResult<()> {
let Some(first) = inputs.first() else {
return Ok(());
};
let current = load_rows(transaction, table, &first.workspace_id, &first.doc_id).await?;
let input_ids = inputs
.iter()
.map(|input| input.external_id.clone())
.collect::<HashSet<_>>();
let mut upserts = Vec::new();
for input in inputs {
match current.get(&input.external_id) {
Some(existing) if existing.revision > input.revision => continue,
Some(existing) if existing.revision == input.revision => {
if existing != &input {
return Err(RuntimeError::invalid_state("conflicting search projection revision"));
}
}
Some(existing) if existing == &input => {}
_ => upserts.push(input),
}
}
let deletions = current
.into_values()
.filter(|row| {
!input_ids.contains(row.external_id.as_str())
&& delete_revision.is_some_and(|revision| row.revision <= revision)
})
.collect();
self
.apply(
transaction,
table,
upserts,
deletions,
delete_revision.unwrap_or_default(),
)
.await
}
async fn apply(
&self,
transaction: &mut Transaction<'_, Postgres>,
table: SearchTable,
upserts: Vec<ProjectionInput>,
deletions: Vec<ProjectionInput>,
delete_revision: i64,
) -> RuntimeResult<()> {
let first_sequence = allocate(transaction, table, upserts.len() + deletions.len()).await?;
let mut sequence = first_sequence;
for input in upserts {
insert_change(
transaction,
table,
sequence,
"upsert",
&input,
Some(&input.payload),
input.revision,
)
.await?;
sqlx::query(
r#"INSERT INTO search_runtime_projections
(table_key, external_id, workspace_id, doc_id, revision, payload,
acl_public_readable, acl_member_default_readable, acl_read_user_ids, acl_revision)
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)
ON CONFLICT (table_key, external_id) DO UPDATE SET
workspace_id=EXCLUDED.workspace_id, doc_id=EXCLUDED.doc_id,
revision=EXCLUDED.revision, payload=EXCLUDED.payload,
acl_public_readable=EXCLUDED.acl_public_readable,
acl_member_default_readable=EXCLUDED.acl_member_default_readable,
acl_read_user_ids=EXCLUDED.acl_read_user_ids,
acl_revision=EXCLUDED.acl_revision, updated_at=now()
WHERE search_runtime_projections.revision < EXCLUDED.revision"#,
)
.bind(table.as_str())
.bind(&input.external_id)
.bind(&input.workspace_id)
.bind(&input.doc_id)
.bind(input.revision)
.bind(&input.payload)
.bind(input.acl_public_readable)
.bind(input.acl_member_default_readable)
.bind(&input.acl_read_user_ids)
.bind(input.acl_revision)
.execute(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("upsert search projection", error))?;
sequence += 1;
}
for input in deletions {
insert_change(
transaction,
table,
sequence,
"delete",
&input,
Some(&input.payload),
delete_revision,
)
.await?;
sqlx::query("DELETE FROM search_runtime_projections WHERE table_key=$1 AND external_id=$2 AND revision <= $3")
.bind(table.as_str())
.bind(&input.external_id)
.bind(delete_revision)
.execute(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("delete search projection", error))?;
sequence += 1;
}
Ok(())
}
pub(in crate::runtime::backend_runtime::search) async fn snapshot(
&self,
table: SearchTable,
) -> RuntimeResult<SearchSnapshot> {
let mut transaction = self
.pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search snapshot", error))?;
sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("configure search snapshot", error))?;
let head = sqlx::query_scalar("SELECT head FROM search_runtime_streams WHERE table_key=$1")
.bind(table.as_str())
.fetch_one(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("read search snapshot head", error))?;
let rows = sqlx::query(
r#"SELECT external_id, workspace_id, doc_id, revision, payload,
acl_public_readable, acl_member_default_readable, acl_read_user_ids, acl_revision
FROM search_runtime_projections WHERE table_key=$1 ORDER BY external_id"#,
)
.bind(table.as_str())
.fetch_all(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("read search projection snapshot", error))?;
let projections = rows.iter().map(decode_projection).collect::<RuntimeResult<Vec<_>>>()?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit search snapshot", error))?;
Ok(SearchSnapshot { head, projections })
}
pub(in crate::runtime::backend_runtime::search) async fn changes(
&self,
table: SearchTable,
after: i64,
limit: i64,
) -> RuntimeResult<(i64, Vec<SearchChange>)> {
if after < 0 || limit <= 0 {
return Err(RuntimeError::invalid_input("invalid search replay cursor"));
}
let state = sqlx::query("SELECT head, retained_from FROM search_runtime_streams WHERE table_key=$1")
.bind(table.as_str())
.fetch_one(&self.pool)
.await
.map_err(|error| RuntimeError::database("read search stream state", error))?;
let head: i64 = state
.try_get("head")
.map_err(|error| RuntimeError::database("decode stream head", error))?;
let retained_from: i64 = state
.try_get("retained_from")
.map_err(|error| RuntimeError::database("decode retained cursor", error))?;
if after < retained_from {
return Err(RuntimeError::SearchReplayGap);
}
let rows = sqlx::query(
r#"SELECT stream_sequence, external_id, workspace_id, doc_id, revision, operation, payload
FROM search_runtime_changes WHERE table_key=$1 AND stream_sequence>$2
ORDER BY stream_sequence LIMIT $3"#,
)
.bind(table.as_str())
.bind(after)
.bind(limit)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("read search stream changes", error))?;
if after < head
&& rows
.first()
.and_then(|row| row.try_get::<i64, _>("stream_sequence").ok())
!= Some(after + 1)
{
return Err(RuntimeError::SearchReplayGap);
}
let changes = rows
.into_iter()
.map(|row| {
Ok(SearchChange {
sequence: row
.try_get("stream_sequence")
.map_err(|error| RuntimeError::database("decode change sequence", error))?,
external_id: row
.try_get("external_id")
.map_err(|error| RuntimeError::database("decode change external id", error))?,
workspace_id: row
.try_get("workspace_id")
.map_err(|error| RuntimeError::database("decode change workspace", error))?,
doc_id: row
.try_get("doc_id")
.map_err(|error| RuntimeError::database("decode change doc", error))?,
revision: row
.try_get("revision")
.map_err(|error| RuntimeError::database("decode change revision", error))?,
operation: row
.try_get("operation")
.map_err(|error| RuntimeError::database("decode change operation", error))?,
payload: row
.try_get("payload")
.map_err(|error| RuntimeError::database("decode change payload", error))?,
})
})
.collect::<RuntimeResult<Vec<_>>>()?;
Ok((head, changes))
}
}
async fn load_rows(
transaction: &mut Transaction<'_, Postgres>,
table: SearchTable,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<HashMap<String, ProjectionInput>> {
let rows = sqlx::query(
r#"SELECT external_id, workspace_id, doc_id, revision, payload,
acl_public_readable, acl_member_default_readable, acl_read_user_ids, acl_revision
FROM search_runtime_projections
WHERE table_key=$1 AND workspace_id=$2 AND doc_id=$3 FOR UPDATE"#,
)
.bind(table.as_str())
.bind(workspace_id)
.bind(doc_id)
.fetch_all(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("load search projections", error))?;
rows
.iter()
.map(|row| decode_projection(row).map(|projection| (projection.external_id.clone(), projection)))
.collect()
}
fn decode_projection(row: &sqlx::postgres::PgRow) -> RuntimeResult<ProjectionInput> {
Ok(ProjectionInput {
external_id: row
.try_get("external_id")
.map_err(|error| RuntimeError::database("decode projection id", error))?,
workspace_id: row
.try_get("workspace_id")
.map_err(|error| RuntimeError::database("decode projection workspace", error))?,
doc_id: row
.try_get("doc_id")
.map_err(|error| RuntimeError::database("decode projection doc", error))?,
revision: row
.try_get("revision")
.map_err(|error| RuntimeError::database("decode projection revision", error))?,
payload: row
.try_get("payload")
.map_err(|error| RuntimeError::database("decode projection payload", error))?,
acl_public_readable: row
.try_get("acl_public_readable")
.map_err(|error| RuntimeError::database("decode projection public ACL", error))?,
acl_member_default_readable: row
.try_get("acl_member_default_readable")
.map_err(|error| RuntimeError::database("decode projection member ACL", error))?,
acl_read_user_ids: row
.try_get("acl_read_user_ids")
.map_err(|error| RuntimeError::database("decode projection ACL users", error))?,
acl_revision: row
.try_get("acl_revision")
.map_err(|error| RuntimeError::database("decode projection ACL revision", error))?,
})
}
async fn insert_change(
transaction: &mut Transaction<'_, Postgres>,
table: SearchTable,
sequence: i64,
operation: &str,
input: &ProjectionInput,
payload: Option<&Value>,
revision: i64,
) -> RuntimeResult<()> {
sqlx::query(
r#"INSERT INTO search_runtime_changes
(table_key,stream_sequence,external_id,workspace_id,doc_id,revision,operation,payload)
VALUES ($1,$2,$3,$4,$5,$6,$7,$8)"#,
)
.bind(table.as_str())
.bind(sequence)
.bind(&input.external_id)
.bind(&input.workspace_id)
.bind(&input.doc_id)
.bind(revision)
.bind(operation)
.bind(payload)
.execute(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("insert search stream change", error))?;
Ok(())
}

View File

@@ -1,26 +0,0 @@
use sqlx::{Postgres, Row, Transaction};
use super::SearchTable;
use crate::runtime::{RuntimeError, RuntimeResult};
pub(in crate::runtime::backend_runtime::search) async fn allocate(
transaction: &mut Transaction<'_, Postgres>,
table: SearchTable,
count: usize,
) -> RuntimeResult<i64> {
if count == 0 {
return Ok(0);
}
let row = sqlx::query(
"UPDATE search_runtime_streams SET head = head + $2, updated_at = now() WHERE table_key = $1 RETURNING head",
)
.bind(table.as_str())
.bind(count as i64)
.fetch_one(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("allocate search stream sequence", error))?;
let head: i64 = row
.try_get("head")
.map_err(|error| RuntimeError::database("decode search stream head", error))?;
Ok(head - count as i64 + 1)
}

View File

@@ -1,54 +0,0 @@
use serde_json::Value;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime::search) enum SearchTable {
Doc,
Block,
}
impl SearchTable {
pub(in crate::runtime::backend_runtime::search) const ORDERED: [Self; 2] = [Self::Doc, Self::Block];
pub(in crate::runtime::backend_runtime::search) fn as_str(self) -> &'static str {
match self {
Self::Doc => "doc",
Self::Block => "block",
}
}
pub(in crate::runtime::backend_runtime::search) fn cursor_index(self) -> usize {
match self {
Self::Doc => 0,
Self::Block => 1,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub(in crate::runtime::backend_runtime::search) struct ProjectionInput {
pub(in crate::runtime::backend_runtime::search) external_id: String,
pub(in crate::runtime::backend_runtime::search) workspace_id: String,
pub(in crate::runtime::backend_runtime::search) doc_id: String,
pub(in crate::runtime::backend_runtime::search) revision: i64,
pub(in crate::runtime::backend_runtime::search) payload: Value,
pub(in crate::runtime::backend_runtime::search) acl_public_readable: bool,
pub(in crate::runtime::backend_runtime::search) acl_member_default_readable: bool,
pub(in crate::runtime::backend_runtime::search) acl_read_user_ids: Vec<String>,
pub(in crate::runtime::backend_runtime::search) acl_revision: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub(in crate::runtime::backend_runtime::search) struct SearchChange {
pub(in crate::runtime::backend_runtime::search) sequence: i64,
pub(in crate::runtime::backend_runtime::search) external_id: String,
pub(in crate::runtime::backend_runtime::search) workspace_id: String,
pub(in crate::runtime::backend_runtime::search) doc_id: Option<String>,
pub(in crate::runtime::backend_runtime::search) revision: i64,
pub(in crate::runtime::backend_runtime::search) operation: String,
pub(in crate::runtime::backend_runtime::search) payload: Option<Value>,
}
pub(in crate::runtime::backend_runtime::search) struct SearchSnapshot {
pub(in crate::runtime::backend_runtime::search) head: i64,
pub(in crate::runtime::backend_runtime::search) projections: Vec<ProjectionInput>,
}

View File

@@ -1,635 +0,0 @@
use serde_json::json;
use sqlx::PgPool;
use super::{
SearchRuntime, generation,
projection::project_document,
provider::RemoteProvider,
query,
store::{ProjectionInput, SearchChange, SearchStore, SearchTable, stream::allocate},
types::SearchRequest,
};
use crate::runtime::{
SearchRuntimeConfig,
backend_runtime::permission::{AclPredicate, AuthorizedSearchScope, DocReadScope},
migrations::migrate_search_tables,
};
static SEARCH_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
#[test]
fn canonical_query_injects_workspace_and_projected_acl() {
let request = SearchRequest::parse(serde_json::json!({
"table":"block",
"query":{"type":"match","field":"content","match":"hello"},
"options":{"fields":["docId","content"],"pagination":{"limit":20}}
}))
.unwrap();
let scope = AuthorizedSearchScope {
workspace_id: "workspace".to_string(),
permission_revision: 7,
docs: DocReadScope::ProjectedAcl(AclPredicate {
actor_user_id: "user".to_string(),
active_member: true,
sharing_enabled: false,
}),
};
let dsl = query::compile(&request, &scope).unwrap();
assert_eq!(dsl["size"], 20);
assert_eq!(
dsl["query"]["bool"]["must"][0]["term"]["workspace_id"]["value"],
"workspace"
);
let acl = &dsl["query"]["bool"]["must"][2]["bool"]["should"];
assert_eq!(acl.as_array().unwrap().len(), 2);
assert!(dsl.to_string().contains("acl_read_tokens"));
}
fn projection(table: SearchTable, id: &str, revision: i64) -> ProjectionInput {
let block_id = (table == SearchTable::Block).then_some(id);
ProjectionInput {
external_id: id.to_string(),
workspace_id: "search-runtime-test-workspace".to_string(),
doc_id: "search-runtime-test-doc".to_string(),
revision,
payload: json!({
"workspace_id": "search-runtime-test-workspace",
"doc_id": "search-runtime-test-doc",
"block_id": block_id,
"revision": revision,
}),
acl_public_readable: false,
acl_member_default_readable: true,
acl_read_user_ids: vec!["search-runtime-test-user".to_string()],
acl_revision: revision,
}
}
async fn pool() -> Option<PgPool> {
let database_url = std::env::var("DATABASE_URL").ok()?;
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::raw_sql(
"DELETE FROM search_runtime_changes; DELETE FROM search_runtime_projections; UPDATE search_runtime_streams SET \
head=0, retained_from=0",
)
.execute(&pool)
.await
.unwrap();
Some(pool)
}
#[tokio::test]
async fn projection_replace_replay_delete_and_stale_revision_are_monotonic() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Some(pool) = pool().await else { return };
let store = SearchStore::new(pool);
store
.replace_document(
projection(SearchTable::Doc, "doc", 2),
vec![projection(SearchTable::Block, "block-a", 2)],
)
.await
.unwrap();
store
.replace_document(
projection(SearchTable::Doc, "doc", 1),
vec![projection(SearchTable::Block, "block-stale", 1)],
)
.await
.unwrap();
store
.replace_document(
projection(SearchTable::Doc, "doc", 2),
vec![projection(SearchTable::Block, "block-a", 2)],
)
.await
.unwrap();
let doc = store.snapshot(SearchTable::Doc).await.unwrap();
let block = store.snapshot(SearchTable::Block).await.unwrap();
assert_eq!((doc.head, doc.projections.len()), (1, 1));
assert_eq!((block.head, block.projections.len()), (1, 1));
let (_, changes) = store.changes(SearchTable::Block, 0, 10).await.unwrap();
assert_eq!(changes.len(), 1);
assert_eq!(changes[0].external_id, "block-a");
store
.delete_document("search-runtime-test-workspace", "search-runtime-test-doc", 3)
.await
.unwrap();
assert!(store.snapshot(SearchTable::Doc).await.unwrap().projections.is_empty());
assert!(store.snapshot(SearchTable::Block).await.unwrap().projections.is_empty());
let (_, changes) = store.changes(SearchTable::Block, 1, 10).await.unwrap();
assert_eq!(changes[0].operation, "delete");
assert_eq!(changes[0].revision, 3);
}
#[tokio::test]
async fn document_projection_loads_snapshot_metadata_and_search_units() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Some(pool) = pool().await else { return };
let suffix = uuid::Uuid::new_v4().simple().to_string();
let workspace_id = format!("search-projection-workspace-{suffix}");
let user_id = format!("search-projection-user-{suffix}");
let doc_id = format!("search-projection-doc-{suffix}");
sqlx::query(
"INSERT INTO users(id,name,email,registered,email_verified,disabled) VALUES($1,'Search Projection \
User',$2,true,now(),false)",
)
.bind(&user_id)
.bind(format!("search-projection-{suffix}@example.com"))
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO workspace_access_policies(workspace_id) VALUES($1) ON CONFLICT DO NOTHING")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO workspace_members(workspace_id,user_id,role,state) VALUES($1,$2,'owner','active')")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
let blob = affine_doc_loader::build_full_doc(
"Projection title",
"Projection body\n\n![Asset](blob://projection-blob)",
&doc_id,
)
.unwrap();
sqlx::query(
"INSERT INTO snapshots(workspace_id,guid,blob,created_by,updated_by,updated_at) \
VALUES($1,$2,$3,$4,$4,clock_timestamp())",
)
.bind(&workspace_id)
.bind(&doc_id)
.bind(blob)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
let (document, blocks) = project_document(&pool, &workspace_id, &doc_id).await.unwrap().unwrap();
assert_eq!(document.payload["title"], "Projection title");
assert_eq!(document.payload["created_by_user_id"], user_id);
assert!(
document.payload["summary"]
.as_str()
.unwrap()
.contains("Projection body")
);
assert!(document.acl_revision > 0);
assert!(blocks.iter().any(|block| block.payload["content"] == "Projection body"));
assert!(blocks.iter().any(|block| block.payload["blob"] == "projection-blob"));
assert!(
blocks
.iter()
.all(|block| block.payload["acl_revision"] == document.payload["acl_revision"])
);
}
#[tokio::test]
async fn stream_sequence_follows_commit_order_and_rollback_has_no_gap() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Some(pool) = pool().await else { return };
let mut first = pool.begin().await.unwrap();
assert_eq!(allocate(&mut first, SearchTable::Doc, 1).await.unwrap(), 1);
let second_pool = pool.clone();
let second = tokio::spawn(async move {
let mut transaction = second_pool.begin().await.unwrap();
let sequence = allocate(&mut transaction, SearchTable::Doc, 1).await.unwrap();
transaction.commit().await.unwrap();
sequence
});
tokio::task::yield_now().await;
let visible_head: i64 = sqlx::query_scalar("SELECT head FROM search_runtime_streams WHERE table_key='doc'")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(visible_head, 0);
first.commit().await.unwrap();
assert_eq!(second.await.unwrap(), 2);
sqlx::query("UPDATE search_runtime_streams SET head=0 WHERE table_key='doc'")
.execute(&pool)
.await
.unwrap();
let mut rolled_back = pool.begin().await.unwrap();
assert_eq!(allocate(&mut rolled_back, SearchTable::Doc, 1).await.unwrap(), 1);
rolled_back.rollback().await.unwrap();
let mut committed = pool.begin().await.unwrap();
assert_eq!(allocate(&mut committed, SearchTable::Doc, 1).await.unwrap(), 1);
committed.commit().await.unwrap();
}
#[tokio::test]
async fn replay_rejects_retention_gap() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Some(pool) = pool().await else { return };
sqlx::query("UPDATE search_runtime_streams SET head=5, retained_from=3 WHERE table_key='doc'")
.execute(&pool)
.await
.unwrap();
let error = SearchStore::new(pool)
.changes(SearchTable::Doc, 2, 10)
.await
.unwrap_err();
assert!(matches!(error, crate::runtime::RuntimeError::SearchReplayGap));
}
#[tokio::test]
async fn concurrent_generation_prepare_reuses_pending_and_restart_preserves_active() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Some(pool) = pool().await else { return };
sqlx::query("DELETE FROM search_runtime_generations")
.execute(&pool)
.await
.unwrap();
let config = SearchRuntimeConfig::default();
let (first, second) = tokio::join!(
generation::prepare(&pool, &config, None),
generation::prepare(&pool, &config, None)
);
let first = first.unwrap();
let second = second.unwrap();
assert_eq!(first.id, second.id);
generation::activate(&pool, &first).await.unwrap();
let restarted = generation::prepare(&pool, &config, None).await.unwrap();
assert_eq!(restarted.id, first.id);
generation::activate(&pool, &restarted).await.unwrap();
let active: i64 = sqlx::query_scalar("SELECT count(*) FROM search_runtime_generations WHERE state='active'")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(active, 1);
let unavailable = SearchRuntimeConfig {
provider: "elasticsearch".into(),
endpoint: "http://127.0.0.1:1".into(),
..Default::default()
};
let remote = RemoteProvider::new(&unavailable, pool.clone()).unwrap();
assert!(generation::prepare(&pool, &unavailable, Some(&remote)).await.is_err());
let states: (i64, i64, i64) = sqlx::query_as(
"SELECT count(*) FILTER (WHERE state='active'), count(*) FILTER (WHERE state='pending'), count(*) FILTER (WHERE \
state='failed') FROM search_runtime_generations",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(states, (1, 0, 1));
}
#[tokio::test]
async fn embedded_replicas_share_one_checkpoint_and_rebuild_a_corrupt_snapshot() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Some(pool) = pool().await else { return };
sqlx::raw_sql(
"DELETE FROM search_runtime_generations; DELETE FROM search_runtime_checkpoints; DELETE FROM \
search_runtime_changes; DELETE FROM search_runtime_projections; UPDATE search_runtime_streams SET head=0, \
retained_from=0",
)
.execute(&pool)
.await
.unwrap();
SearchStore::new(pool.clone())
.replace_document(
ProjectionInput {
payload: json!({
"workspace_id":"replica-workspace","doc_id":"replica-doc","title":"replica search",
"summary":"","created_by_user_id":"user","updated_by_user_id":"user",
"created_at":1,"updated_at":1,"acl_public_readable":false,
"acl_member_default_readable":true,"acl_read_tokens":["member"],
"acl_revision":1
}),
external_id: "replica-workspace/replica-doc".into(),
workspace_id: "replica-workspace".into(),
doc_id: "replica-doc".into(),
revision: 1,
acl_public_readable: false,
acl_member_default_readable: true,
acl_read_user_ids: vec![],
acl_revision: 1,
},
vec![],
)
.await
.unwrap();
let config = SearchRuntimeConfig::default();
let first = SearchRuntime::new(pool.clone(), config.clone()).unwrap();
let second = SearchRuntime::new(pool.clone(), config.clone()).unwrap();
let third = SearchRuntime::new(pool.clone(), config.clone()).unwrap();
let (first_result, second_result, third_result) =
tokio::join!(first.initialize(), second.initialize(), third.initialize());
first_result.unwrap();
second_result.unwrap();
third_result.unwrap();
SearchStore::new(pool.clone())
.replace_document(
ProjectionInput {
payload: json!({
"workspace_id":"replica-workspace","doc_id":"replica-doc-2","title":"replica search second",
"summary":"","created_by_user_id":"user","updated_by_user_id":"user",
"created_at":2,"updated_at":2,"acl_public_readable":false,
"acl_member_default_readable":true,"acl_read_tokens":["member"],
"acl_revision":1
}),
external_id: "replica-workspace/replica-doc-2".into(),
workspace_id: "replica-workspace".into(),
doc_id: "replica-doc-2".into(),
revision: 2,
acl_public_readable: false,
acl_member_default_readable: true,
acl_read_user_ids: vec![],
acl_revision: 1,
},
vec![],
)
.await
.unwrap();
first.sync().await.unwrap();
sqlx::query("UPDATE search_runtime_streams SET retained_from=head WHERE table_key='doc'")
.execute(&pool)
.await
.unwrap();
second.sync().await.unwrap();
let second_replica_result: serde_json::Value = serde_json::from_str(
&second
.embedded
.search(
"doc".into(),
json!({"query":{"match_all":{}},"fields":["doc_id"],"sort":["doc_id"],"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(second_replica_result["total"], 2);
let checkpoints: i64 = sqlx::query_scalar("SELECT count(*) FROM search_runtime_checkpoints")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(checkpoints, 1);
sqlx::query("UPDATE search_runtime_checkpoints SET checkpoint_blob='\\x010203' WHERE table_key='doc'")
.execute(&pool)
.await
.unwrap();
let recovered = SearchRuntime::new(pool, config).unwrap();
recovered.initialize().await.unwrap();
let result: serde_json::Value = serde_json::from_str(
&recovered
.embedded
.search(
"doc".into(),
json!({"query":{"match_all":{}},"fields":["doc_id"],"sort":["doc_id"],"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 2);
}
#[tokio::test]
async fn remote_providers_apply_search_and_delete_the_same_contract() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let require_remote = std::env::var("SEARCH_REQUIRE_REMOTE_TESTS").as_deref() == Ok("1");
let pool = match pool().await {
Some(pool) => pool,
None if require_remote => panic!("DATABASE_URL is required"),
None => return,
};
let mut tested_providers = 0;
for (provider, variable) in [("elasticsearch", "SEARCH_ES_URL"), ("manticoresearch", "SEARCH_MS_URL")] {
let Ok(endpoint) = std::env::var(variable) else {
continue;
};
tested_providers += 1;
let remote = RemoteProvider::new(
&SearchRuntimeConfig {
provider: provider.to_string(),
endpoint: endpoint.clone(),
..Default::default()
},
pool.clone(),
)
.unwrap();
let table = format!("affine_search_contract_{}", uuid::Uuid::new_v4().simple());
let block_table = format!("affine_search_contract_block_{}", uuid::Uuid::new_v4().simple());
let cleanup_tables = [table.clone(), block_table.clone()];
let contract = tokio::spawn(async move {
remote.provision(&table, super::types::SearchTable::Doc).await.unwrap();
let upsert = SearchChange {
sequence: 1,
external_id: "workspace/doc".into(),
workspace_id: "workspace".into(),
doc_id: Some("doc".into()),
revision: 1,
operation: "upsert".into(),
payload: Some(json!({
"workspace_id":"workspace","workspace_token":super::exact_token("workspace"),
"doc_id":"doc","doc_token":super::exact_token("doc"),"title":"search contract",
"summary":"","created_by_user_id":"user","updated_by_user_id":"user",
"created_at":1,"updated_at":1,"acl_public_readable":false,
"acl_member_default_readable":true,"acl_read_tokens":["member"],"acl_revision":1
})),
};
let mut second = upsert.clone();
second.sequence = 2;
second.external_id = "workspace/doc-2".into();
second.doc_id = Some("doc-2".into());
second.payload.as_mut().unwrap()["doc_id"] = json!("doc-2");
second.payload.as_mut().unwrap()["doc_token"] = json!(super::exact_token("doc-2"));
remote.apply(&table, &[upsert.clone(), second.clone()]).await.unwrap();
let dsl = json!({
"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":"workspace"}}},
{"match":{"title":{"query":"contract"}}},
{"bool":{"should":[{"term":{"acl_read_tokens":{"value":"member"}}}]}}
],"boost":1.0}},
"fields":["doc_id","title"],"_source":["workspace_id","doc_id"],
"highlight":{"fields":{"title":{"pre_tags":["<b>"],"post_tags":["</b>"]}}},
"sort":["doc_id"],"size":1
});
let result = remote.search(&table, dsl.clone()).await.unwrap();
assert_eq!(result["total"], 2, "provider {provider}");
assert!(result["nodes"][0]["fields"]["doc_id"].is_array(), "provider {provider}");
assert!(
result["nodes"][0]["highlights"]["title"].is_array(),
"provider {provider}"
);
let first_doc = result["nodes"][0]["fields"]["doc_id"][0].clone();
let mut next_dsl = dsl.clone();
let first_cursor = result["nextCursor"].clone();
next_dsl["cursor"] = first_cursor.clone();
let next = remote.search(&table, next_dsl).await.unwrap();
assert_ne!(
first_doc, next["nodes"][0]["fields"]["doc_id"][0],
"provider {provider}"
);
assert_ne!(first_cursor, next["nextCursor"], "provider {provider}");
let aggregate_dsl = json!({
"query":{"term":{"workspace_id":{"value":"workspace"}}},
"size":0,
"aggs":{"result":{"terms":{"field":"doc_id","size":10},"aggs":{"result":{"top_hits":{
"size":1,"_source":["workspace_id","doc_id"],"fields":["doc_id","title"],"sort":["doc_id"]
}}}}}
});
if provider == "elasticsearch" {
let aggregate = remote.aggregate(&table, aggregate_dsl).await.unwrap();
assert_eq!(aggregate["total"], 2);
assert_eq!(aggregate["buckets"].as_array().unwrap().len(), 2);
assert!(aggregate["buckets"][0]["hits"]["nodes"][0]["fields"]["doc_id"].is_array());
} else {
assert!(matches!(
remote.aggregate(&table, aggregate_dsl).await,
Err(crate::runtime::RuntimeError::SearchUnsupportedQuery)
));
}
remote
.provision(&block_table, super::types::SearchTable::Block)
.await
.unwrap();
let block = SearchChange {
sequence: 1,
external_id: "workspace/doc/block".into(),
workspace_id: "workspace".into(),
doc_id: Some("doc".into()),
revision: 1,
operation: "upsert".into(),
payload: Some(json!({
"workspace_id":"workspace","workspace_token":super::exact_token("workspace"),
"doc_id":"doc","doc_token":super::exact_token("doc"),
"block_id":"block","block_token":super::exact_token("block"),
"content":"笔记应用 다람쥐 いろはにほへと https://linear.app/affine-design/issue/AF-1379/slash-commands",
"flavour":"affine:paragraph",
"ref_doc_id":["ref-a","ref-b"],"blob":["blob-a","blob-b"],
"created_by_user_id":"user","updated_by_user_id":"user",
"created_at":2_000,"updated_at":3_000,"acl_public_readable":false,
"acl_member_default_readable":true,"acl_read_tokens":["member"],"acl_revision":1
})),
};
remote.apply(&block_table, std::slice::from_ref(&block)).await.unwrap();
let exists = remote
.search(
&block_table,
json!({
"query":{"exists":{"field":"ref_doc_id"}},
"fields":["block_id","ref_doc_id"],"_source":["workspace_id","doc_id"],
"sort":["block_id"],"size":10
}),
)
.await
.unwrap();
assert_eq!(exists["total"], 1, "provider {provider}");
let exact_ref = remote
.search(
&block_table,
json!({
"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":"workspace"}}},
{"term":{"ref_doc_id":{"value":"ref-a"}}},
{"bool":{"must_not":[{"term":{"doc_id":{"value":"other-doc"}}}]}}
]}},
"fields":["block_id","ref_doc_id"],"_source":["workspace_id","doc_id"],
"sort":["block_id"],"size":10
}),
)
.await
.unwrap();
assert_eq!(exact_ref["total"], 1, "provider {provider}");
let terms = if provider == "elasticsearch" {
["记", "https://linear.app"].as_slice()
} else {
["쥐", "へ", "https://linear.app"].as_slice()
};
for term in terms {
let language = remote
.search(
&block_table,
json!({
"query":{"match":{"content":{"query":term}}},
"fields":["block_id","ref_doc_id","blob","created_at","updated_at"],
"_source":["workspace_id","doc_id"],
"highlight":{"fields":{"content":{"pre_tags":["<b>"],"post_tags":["</b>"]}}},
"sort":["block_id"],"size":10
}),
)
.await
.unwrap();
assert_eq!(language["total"], 1, "provider {provider}, term {term}");
assert_eq!(language["nodes"][0]["fields"]["ref_doc_id"], json!(["ref-a", "ref-b"]));
assert!(language["nodes"][0]["fields"]["created_at"].is_array());
assert!(language["nodes"][0]["highlights"]["content"].is_array());
}
let mut revoked = upsert.clone();
revoked.sequence = 3;
revoked.revision = 3;
revoked.payload.as_mut().unwrap()["acl_read_tokens"] = json!([]);
remote.apply(&table, &[revoked]).await.unwrap();
let revoked_result = remote.search(&table, dsl.clone()).await.unwrap();
assert_eq!(revoked_result["total"], 1, "provider {provider}");
assert_eq!(revoked_result["nodes"][0]["fields"]["doc_id"][0], "doc-2");
let deletion = SearchChange {
operation: "delete".into(),
payload: None,
sequence: 4,
revision: 4,
..upsert.clone()
};
let second_deletion = SearchChange {
operation: "delete".into(),
payload: None,
sequence: 5,
revision: 5,
..second
};
remote.apply(&table, &[deletion, second_deletion]).await.unwrap();
let result = remote
.search(
&table,
json!({"query":{"match_all":{}},"fields":["doc_id"],"sort":["doc_id"],"size":10}),
)
.await
.unwrap();
assert_eq!(result["total"], 0, "provider {provider}");
})
.await;
let client = reqwest::Client::new();
if provider == "elasticsearch" {
for physical_table in &cleanup_tables {
let response = client
.delete(format!("{endpoint}/{physical_table}"))
.send()
.await
.unwrap();
assert!(response.status().is_success() || response.status() == reqwest::StatusCode::NOT_FOUND);
}
} else {
for physical_table in &cleanup_tables {
client
.post(format!("{endpoint}/cli"))
.header("content-type", "text/plain")
.body(format!("DROP TABLE IF EXISTS {physical_table}"))
.send()
.await
.unwrap()
.error_for_status()
.unwrap();
}
}
contract.unwrap();
}
if require_remote {
assert_eq!(tested_providers, 2, "SEARCH_ES_URL and SEARCH_MS_URL are required");
}
}

View File

@@ -1,7 +1,5 @@
use serde::{Deserialize, Serialize};
use crate::runtime::{RuntimeError, RuntimeResult};
#[napi_derive::napi(string_enum = "snake_case")]
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
@@ -26,20 +24,6 @@ impl SearchTable {
}
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchQuery {
#[serde(rename = "type")]
pub(super) query_type: String,
pub(super) field: Option<String>,
#[serde(rename = "match")]
pub(super) match_value: Option<String>,
pub(super) query: Option<Box<SearchQuery>>,
pub(super) queries: Option<Vec<SearchQuery>>,
pub(super) occur: Option<String>,
pub(super) boost: Option<f64>,
}
#[napi_derive::napi(object)]
#[derive(Clone, Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -89,23 +73,6 @@ pub struct AggregateOptions {
pub pagination: SearchPagination,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct AggregateRequest {
pub(super) table: SearchTable,
pub(super) query: SearchQuery,
pub(super) field: String,
pub(super) options: AggregateOptions,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchRequest {
pub(super) table: SearchTable,
pub(super) query: SearchQuery,
pub(super) options: SearchOptions,
}
#[napi_derive::napi(object)]
pub struct RuntimeSearchQuery {
pub query_type: String,
@@ -133,109 +100,3 @@ pub struct RuntimeAggregateRequest {
pub field: String,
pub options: AggregateOptions,
}
#[cfg(test)]
impl SearchRequest {
pub(super) fn parse(value: serde_json::Value) -> RuntimeResult<Self> {
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid search request", error))
}
}
impl RuntimeSearchRequest {
pub(super) fn into_search_request(self) -> RuntimeResult<SearchRequest> {
let mut decoded_nodes = 0;
Ok(SearchRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
options: self.options,
})
}
}
impl RuntimeAggregateRequest {
pub(super) fn into_aggregate_request(self) -> RuntimeResult<AggregateRequest> {
let mut decoded_nodes = 0;
Ok(AggregateRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
field: self.field,
options: self.options,
})
}
}
const MAX_QUERY_GRAPH_NODES: usize = 100;
const MAX_QUERY_DEPTH: usize = 100;
const MAX_DECODED_QUERY_NODES: usize = 1_000;
fn decode_query(
nodes: &[RuntimeSearchQuery],
index: u32,
depth: usize,
decoded_nodes: &mut usize,
) -> RuntimeResult<SearchQuery> {
if nodes.len() > MAX_QUERY_GRAPH_NODES || depth > MAX_QUERY_DEPTH || *decoded_nodes >= MAX_DECODED_QUERY_NODES {
return Err(RuntimeError::invalid_input("search query is too complex"));
}
*decoded_nodes += 1;
let node = nodes
.get(index as usize)
.ok_or_else(|| RuntimeError::invalid_input("invalid search query node"))?;
Ok(SearchQuery {
query_type: node.query_type.clone(),
field: node.field.clone(),
match_value: node.match_value.clone(),
query: node
.query
.map(|index| decode_query(nodes, index, depth + 1, decoded_nodes).map(Box::new))
.transpose()?,
queries: node
.queries
.as_ref()
.map(|indices| {
indices
.iter()
.map(|index| decode_query(nodes, *index, depth + 1, decoded_nodes))
.collect()
})
.transpose()?,
occur: node.occur.clone(),
boost: node.boost,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn node(query_type: &str) -> RuntimeSearchQuery {
RuntimeSearchQuery {
query_type: query_type.to_string(),
field: None,
match_value: None,
query: None,
queries: None,
occur: None,
boost: None,
}
}
#[test]
fn rejects_invalid_or_overly_complex_query_graphs() {
assert!(decode_query(&[node("all")], 1, 0, &mut 0).is_err());
let mut oversized = (0..101).map(|_| node("all")).collect::<Vec<_>>();
oversized[0].query = Some(1);
assert!(decode_query(&oversized, 0, 0, &mut 0).is_err());
let mut recursive = vec![node("boost")];
recursive[0].query = Some(0);
assert!(decode_query(&recursive, 0, 0, &mut 0).is_err());
let mut shared_child = (0..100).map(|_| node("boolean")).collect::<Vec<_>>();
for (index, node) in shared_child.iter_mut().enumerate().take(99) {
node.queries = Some(vec![(index + 1) as u32, (index + 1) as u32]);
}
assert!(decode_query(&shared_child, 0, 0, &mut 0).is_err());
}
}

View File

@@ -1,272 +0,0 @@
use sqlx::PgPool;
use tokio::sync::RwLock;
use super::{
generation::ActiveGeneration,
provider::RemoteProvider,
store::{SearchChange, SearchStore, SearchTable},
};
use crate::{
runtime::{RuntimeError, RuntimeResult},
search_index::EmbeddedSearchIndex,
};
pub(super) async fn rebuild(
pool: &PgPool,
store: &SearchStore,
embedded: &EmbeddedSearchIndex,
remote: Option<&RemoteProvider>,
generation: &ActiveGeneration,
embedded_cursors: &RwLock<[i64; 2]>,
restore_checkpoint: bool,
) -> RuntimeResult<()> {
for table in SearchTable::ORDERED {
if restore_checkpoint
&& remote.is_none()
&& let Some(cursor) = super::checkpoint::restore(pool, embedded, table).await?
{
set_cursor(pool, remote, generation, embedded_cursors, table, cursor).await?;
continue;
}
let snapshot = store.snapshot(table).await?;
if let Some(remote) = remote {
let changes = snapshot
.projections
.into_iter()
.enumerate()
.map(|(offset, projection)| SearchChange {
sequence: offset as i64 + 1,
external_id: projection.external_id,
workspace_id: projection.workspace_id,
doc_id: Some(projection.doc_id),
revision: projection.revision,
operation: "upsert".into(),
payload: Some(projection.payload),
})
.collect::<Vec<_>>();
for batch in changes.chunks(1000) {
remote
.apply(generation.physical_table(runtime_table(table))?, batch)
.await?;
}
} else {
embedded.reset(table.as_str().to_string()).await?;
for documents in snapshot.projections.chunks(1000) {
embedded
.write(
table.as_str().to_string(),
serde_json::to_string(
&documents
.iter()
.map(|projection| super::provider_payload(&projection.payload))
.collect::<Vec<_>>(),
)
.map_err(|error| RuntimeError::json("encode embedded snapshot", error))?,
)
.await?;
}
}
set_cursor(pool, remote, generation, embedded_cursors, table, snapshot.head).await?;
}
sync(pool, store, embedded, remote, generation, embedded_cursors).await?;
Ok(())
}
pub(super) async fn sync(
pool: &PgPool,
store: &SearchStore,
embedded: &EmbeddedSearchIndex,
remote: Option<&RemoteProvider>,
generation: &ActiveGeneration,
embedded_cursors: &RwLock<[i64; 2]>,
) -> RuntimeResult<()> {
for table in SearchTable::ORDERED {
loop {
let cursor = cursor(pool, remote, generation, embedded_cursors, table).await?;
let (head, changes) = store.changes(table, cursor, 1000).await?;
if changes.is_empty() {
if cursor < head {
return Err(RuntimeError::invalid_state(
"search provider cursor did not reach stream head",
));
}
break;
}
if let Some(remote) = remote {
remote
.apply(generation.physical_table(runtime_table(table))?, &changes)
.await?;
} else {
apply_embedded(embedded, table, &changes).await?;
}
set_cursor(
pool,
remote,
generation,
embedded_cursors,
table,
changes.last().expect("non-empty changes").sequence,
)
.await?;
}
}
if remote.is_none() {
super::checkpoint::persist(pool, embedded, *embedded_cursors.read().await).await?;
} else {
super::checkpoint::gc(pool).await?;
}
Ok(())
}
async fn apply_embedded(
embedded: &EmbeddedSearchIndex,
table: SearchTable,
changes: &[SearchChange],
) -> RuntimeResult<()> {
let mut upserts = Vec::new();
for change in changes {
if change.operation == "delete" {
if !upserts.is_empty() {
embedded
.write(
table.as_str().to_string(),
serde_json::to_string(&upserts).map_err(|error| RuntimeError::json("encode embedded changes", error))?,
)
.await?;
upserts.clear();
}
embedded
.delete(table.as_str().to_string(), change.external_id.clone())
.await?;
} else if let Some(payload) = &change.payload {
upserts.push(super::provider_payload(payload));
}
}
if !upserts.is_empty() {
embedded
.write(
table.as_str().to_string(),
serde_json::to_string(&upserts).map_err(|error| RuntimeError::json("encode embedded changes", error))?,
)
.await?;
}
Ok(())
}
async fn set_cursor(
pool: &PgPool,
remote: Option<&RemoteProvider>,
generation: &ActiveGeneration,
embedded_cursors: &RwLock<[i64; 2]>,
table: SearchTable,
cursor: i64,
) -> RuntimeResult<()> {
if remote.is_none() {
let mut cursors = embedded_cursors.write().await;
cursors[table.cursor_index()] = cursors[table.cursor_index()].max(cursor);
return Ok(());
}
sqlx::query(
"UPDATE search_runtime_provider_cursors SET source_cursor=GREATEST(source_cursor,$3), updated_at=now() WHERE \
generation_id=$1 AND table_key=$2",
)
.bind(generation.id)
.bind(table.as_str())
.bind(cursor)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("advance search provider cursor", error))?;
Ok(())
}
async fn cursor(
pool: &PgPool,
remote: Option<&RemoteProvider>,
generation: &ActiveGeneration,
embedded_cursors: &RwLock<[i64; 2]>,
table: SearchTable,
) -> RuntimeResult<i64> {
if remote.is_none() {
return Ok(embedded_cursors.read().await[table.cursor_index()]);
}
sqlx::query_scalar(
"SELECT source_cursor FROM search_runtime_provider_cursors WHERE generation_id=$1 AND table_key=$2",
)
.bind(generation.id)
.bind(table.as_str())
.fetch_one(pool)
.await
.map_err(|error| RuntimeError::database("load search provider cursor", error))
}
fn runtime_table(table: SearchTable) -> super::types::SearchTable {
match table {
SearchTable::Doc => super::types::SearchTable::Doc,
SearchTable::Block => super::types::SearchTable::Block,
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[tokio::test]
async fn embedded_changes_follow_stream_order() {
let embedded = EmbeddedSearchIndex::new();
let changes = vec![
SearchChange {
sequence: 1,
external_id: "workspace/doc".into(),
workspace_id: "workspace".into(),
doc_id: Some("doc".into()),
revision: 1,
operation: "delete".into(),
payload: Some(json!({
"workspace_id": "workspace",
"doc_id": "doc",
"title": "deleted",
"created_at": 1,
"updated_at": 1
})),
},
SearchChange {
sequence: 2,
external_id: "workspace/doc".into(),
workspace_id: "workspace".into(),
doc_id: Some("doc".into()),
revision: 2,
operation: "upsert".into(),
payload: Some(json!({
"workspace_id": "workspace",
"doc_id": "doc",
"title": "restored",
"created_at": 1,
"updated_at": 2
})),
},
];
apply_embedded(&embedded, SearchTable::Doc, &changes).await.unwrap();
let result: serde_json::Value = serde_json::from_str(
&embedded
.search(
"doc".into(),
json!({
"query": {"match_all": {}},
"fields": ["doc_id", "title"],
"sort": ["doc_id"],
"size": 10
})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(result["nodes"][0]["fields"]["title"], json!(["restored"]));
}
}

View File

@@ -0,0 +1,603 @@
use std::collections::{HashMap, HashSet};
use serde_json::{Value, json};
use sqlx::{PgPool, Row};
use super::{
ActiveGeneration, RECONCILE_BATCH, SearchProvider, SearchTable, WorkspacePhase, WorkspaceReconcileContext,
WorkspaceStep, mark_workspace_failed, project_document, renew_workspace_lease, upsert_document,
};
use crate::{
runtime::{RuntimeError, RuntimeResult},
search_index::EmbeddedSearchIndex,
};
const WORKSPACE_GC_BATCH: usize = 100;
const GENERATION_GC_BATCH: usize = 20;
pub(in crate::runtime::backend_runtime::search) async fn sweep_generation_orphans(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
) -> RuntimeResult<()> {
let (table, cursor): (String, Option<String>) =
sqlx::query_as("SELECT gc_table,gc_cursor FROM search_projection.generations WHERE id=$1")
.bind(generation.id)
.fetch_one(pool)
.await
.map_err(|error| RuntimeError::database("load search generation GC progress", error))?;
let table = match table.as_str() {
"doc" => SearchTable::Doc,
"block" => SearchTable::Block,
_ => return Err(RuntimeError::invalid_state("invalid search generation GC table")),
};
let mut dsl = json!({
"query":{"match_all":{}},
"fields":["workspace_id"],
"size":GENERATION_GC_BATCH,
"sort":if remote.is_some() { json!([{"doc_id":"asc"},{"_id":"asc"}]) } else { json!(["doc_id","id"]) }
});
if let Some(cursor) = cursor.as_deref() {
dsl["cursor"] = json!(cursor);
}
let result = search_provider(embedded, remote, generation, table, dsl).await?;
let workspace_ids = result
.get("nodes")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|node| provider_field_string(node, "workspace_id"))
.collect::<HashSet<_>>();
let existing = sqlx::query_scalar::<_, String>("SELECT id FROM workspaces WHERE id=ANY($1)")
.bind(workspace_ids.iter().cloned().collect::<Vec<_>>())
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load canonical workspaces for generation GC", error))?
.into_iter()
.collect::<HashSet<_>>();
for workspace_id in workspace_ids.difference(&existing) {
let Some(source_version_high_water) = capture_orphan_gc_high_water(pool, workspace_id).await? else {
continue;
};
match remote {
Some(provider) => {
provider
.gc_workspace(
generation.physical_table(table)?,
workspace_id,
source_version_high_water,
WORKSPACE_GC_BATCH,
)
.await?;
}
None => {
embedded
.gc_workspace_for_generation(
generation.id,
table.as_str(),
workspace_id,
source_version_high_water,
WORKSPACE_GC_BATCH,
)
.await
.map_err(|error| RuntimeError::invalid_state(format!("embedded generation GC failed: {error}")))?;
}
}
}
let next_cursor = result.get("nextCursor").and_then(Value::as_str);
let (next_table, next_cursor) = match next_cursor {
Some(cursor) => (table.as_str(), Some(cursor)),
None if table == SearchTable::Doc => (SearchTable::Block.as_str(), None),
None => (SearchTable::Doc.as_str(), None),
};
sqlx::query("UPDATE search_projection.generations SET gc_table=$2,gc_cursor=$3 WHERE id=$1")
.bind(generation.id)
.bind(next_table)
.bind(next_cursor)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("persist search generation GC progress", error))?;
Ok(())
}
pub(super) async fn capture_orphan_gc_high_water(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Option<i64>> {
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin generation orphan GC fence", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("lock generation orphan GC fence", error))?;
let workspace_exists: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM workspaces WHERE id=$1)")
.bind(workspace_id)
.fetch_one(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("confirm generation orphan workspace", error))?;
let high_water = if workspace_exists {
None
} else {
Some(
sqlx::query_scalar("SELECT nextval('search_projection.source_mutation_version')")
.fetch_one(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("capture generation orphan GC high water", error))?,
)
};
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit generation orphan GC fence", error))?;
Ok(high_water)
}
pub(super) async fn sweep_deleted_workspace(
context: &WorkspaceReconcileContext<'_>,
table: SearchTable,
quiet: bool,
) -> RuntimeResult<WorkspaceStep> {
if !renew_workspace_lease(context.pool, context.generation.id, context.workspace_id, context.fence).await? {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Deleted { table, quiet }));
}
let Some(source_version_high_water) = capture_orphan_gc_high_water(context.pool, context.workspace_id).await? else {
return Ok(WorkspaceStep::Complete);
};
let may_have_more = match context.remote {
Some(provider) => {
provider
.gc_workspace(
context.generation.physical_table(table)?,
context.workspace_id,
source_version_high_water,
WORKSPACE_GC_BATCH,
)
.await?
}
None => context
.embedded
.gc_workspace_for_generation(
context.generation.id,
table.as_str(),
context.workspace_id,
source_version_high_water,
WORKSPACE_GC_BATCH,
)
.await
.map_err(|error| RuntimeError::invalid_state(format!("embedded search workspace GC failed: {error}")))?,
};
if may_have_more {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Deleted { table, quiet }));
}
if table == SearchTable::Doc {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Deleted {
table: SearchTable::Block,
quiet,
}));
}
if !quiet {
return Ok(WorkspaceStep::Quiet(WorkspacePhase::Deleted {
table: SearchTable::Doc,
quiet: true,
}));
}
Ok(WorkspaceStep::Complete)
}
pub(super) const CANONICAL_SNAPSHOT_BATCH_SQL: &str = r#"SELECT snapshot.guid,state.target_source_version,state.target_permission_version
FROM snapshots snapshot
LEFT JOIN search_projection.document_states state
ON state.generation_id=$2 AND state.workspace_id=snapshot.workspace_id AND state.doc_id=snapshot.guid
WHERE snapshot.workspace_id=$1 AND snapshot.guid <> $1 AND ($3::text IS NULL OR snapshot.guid > $3)
ORDER BY snapshot.guid LIMIT $4"#;
pub(super) async fn reconcile_stale_provider_rows(
context: &WorkspaceReconcileContext<'_>,
table: SearchTable,
cursor: Option<String>,
) -> RuntimeResult<WorkspaceStep> {
let pool = context.pool;
let embedded = context.embedded;
let remote = context.remote;
let generation = context.generation;
let workspace_id = context.workspace_id;
let workspace_fence = context.fence;
let (doc_ids, next_cursor) =
provider_document_page(embedded, remote, generation, workspace_id, table, cursor.as_deref()).await?;
let doc_ids = doc_ids.into_iter().collect::<HashSet<_>>();
let rows = sqlx::query(
r#"SELECT state.doc_id,state.published_source_version,state.published_permission_version,
state.published_source_exists
FROM search_projection.document_states state
WHERE state.generation_id=$1 AND state.workspace_id=$2 AND state.doc_id=ANY($3)
AND state.target_source_version=state.published_source_version
AND state.target_source_exists=state.published_source_exists
AND state.target_permission_version=state.published_permission_version"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(doc_ids.iter().cloned().collect::<Vec<_>>())
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load published tuples for search history GC", error))?;
let published = rows
.into_iter()
.map(|row| {
Ok((
row
.try_get::<String, _>("doc_id")
.map_err(|error| RuntimeError::database("decode search history document", error))?,
(
row
.try_get::<i64, _>("published_source_version")
.map_err(|error| RuntimeError::database("decode search history source version", error))?,
row
.try_get::<i64, _>("published_permission_version")
.map_err(|error| RuntimeError::database("decode search history permission version", error))?,
row
.try_get::<bool, _>("published_source_exists")
.map_err(|error| RuntimeError::database("decode search history source existence", error))?,
),
))
})
.collect::<RuntimeResult<HashMap<_, _>>>()?;
for doc_id in doc_ids {
if !renew_workspace_lease(pool, generation.id, workspace_id, workspace_fence).await? {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Stale { table, cursor }));
}
if let Some((source_version, permission_version, true)) = published.get(&doc_id).copied() {
let _ = gc_provider_document_history(
embedded,
remote,
generation,
table,
ProjectionExpectation {
workspace_id,
doc_id: &doc_id,
source_version,
permission_version,
},
)
.await;
} else if let Err(error) = upsert_document(pool, embedded, remote, generation, workspace_id, &doc_id).await {
if error.is_permanent_search_source() {
mark_workspace_failed(pool, generation.id, workspace_id, workspace_fence).await?;
return Ok(WorkspaceStep::Failed);
}
return Err(error);
}
}
if let Some(next_cursor) = next_cursor {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Stale {
table,
cursor: Some(next_cursor),
}));
}
if table == SearchTable::Doc {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Stale {
table: SearchTable::Block,
cursor: None,
}));
}
Ok(WorkspaceStep::Complete)
}
async fn gc_provider_document_history(
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
table: SearchTable,
published: ProjectionExpectation<'_>,
) -> RuntimeResult<()> {
let ProjectionExpectation {
workspace_id,
doc_id,
source_version,
permission_version,
} = published;
match remote {
Some(provider) => {
provider
.gc_document_history(
generation.physical_table(table)?,
workspace_id,
doc_id,
source_version,
permission_version,
WORKSPACE_GC_BATCH,
)
.await
}
None => embedded
.gc_document_history_for_generation(
generation.id,
table.as_str(),
workspace_id,
doc_id,
(source_version, permission_version),
WORKSPACE_GC_BATCH,
)
.await
.map_err(|error| RuntimeError::invalid_state(format!("embedded search history GC failed: {error}"))),
}
}
async fn provider_document_page(
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
workspace_id: &str,
table: SearchTable,
cursor: Option<&str>,
) -> RuntimeResult<(Vec<String>, Option<String>)> {
let mut dsl = json!({
"query":{"term":{"workspace_id":{"value":workspace_id}}},
"fields":["doc_id"],
"size":RECONCILE_BATCH,
"sort":if remote.is_some() { json!([{"doc_id":"asc"},{"_id":"asc"}]) } else { json!(["doc_id","id"]) }
});
if let Some(cursor) = cursor {
dsl["cursor"] = json!(cursor);
}
let result = search_provider(embedded, remote, generation, table, dsl).await?;
let doc_ids = result
.get("nodes")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|node| provider_field_string(node, "doc_id"))
.collect();
let next_cursor = result.get("nextCursor").and_then(Value::as_str).map(str::to_string);
Ok((doc_ids, next_cursor))
}
fn provider_field_string(node: &Value, field: &str) -> Option<String> {
node
.pointer(&format!("/fields/{field}/0"))
.and_then(Value::as_str)
.or_else(|| node.pointer(&format!("/_source/{field}")).and_then(Value::as_str))
.map(str::to_string)
}
async fn search_provider(
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
table: SearchTable,
dsl: Value,
) -> RuntimeResult<Value> {
match remote {
Some(remote) => remote.search(generation.physical_table(table)?, dsl).await,
None => {
let result = embedded
.search_for_generation(generation.id, table.as_str().to_string(), dsl.to_string())
.await
.map_err(|error| RuntimeError::invalid_state(format!("embedded search lookup failed: {error}")))?;
serde_json::from_str(&result).map_err(|error| RuntimeError::json("decode embedded search", error))
}
}
}
pub(super) struct ProjectionExpectation<'a> {
pub(super) workspace_id: &'a str,
pub(super) doc_id: &'a str,
pub(super) source_version: i64,
pub(super) permission_version: i64,
}
pub(super) async fn provider_projection_matches(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
expected: ProjectionExpectation<'_>,
) -> RuntimeResult<bool> {
let ProjectionExpectation {
workspace_id,
doc_id,
source_version,
permission_version,
} = expected;
let dsl = json!({
"_source":["source_version","permission_version"],
"fields":["source_version","permission_version"],
"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"term":{"doc_id":{"value":doc_id}}},
{"term":{"source_version":{"value":source_version}}},
{"term":{"permission_version":{"value":permission_version}}}
]}},
"size":1,
});
let result = search_provider(embedded, remote, generation, SearchTable::Doc, dsl).await?;
let Some(node) = result
.get("nodes")
.and_then(Value::as_array)
.and_then(|nodes| nodes.first())
else {
return Ok(false);
};
let Some(provider_source_version) = provider_field_i64(node, "source_version") else {
return Ok(false);
};
let Some(provider_permission_version) = provider_field_i64(node, "permission_version") else {
return Ok(false);
};
if provider_source_version != source_version || provider_permission_version != permission_version {
return Ok(false);
}
let Some((_, blocks)) = project_document(pool, workspace_id, doc_id).await? else {
return Ok(false);
};
let expected_block_ids = blocks
.iter()
.filter_map(|block| block.payload.get("block_id").and_then(Value::as_str))
.map(str::to_string)
.collect::<HashSet<_>>();
let block_result = search_provider(
embedded,
remote,
generation,
SearchTable::Block,
json!({
"_source":["block_id","source_version","permission_version"],
"fields":["block_id","source_version","permission_version"],
"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"term":{"doc_id":{"value":doc_id}}},
{"term":{"source_version":{"value":source_version}}},
{"term":{"permission_version":{"value":permission_version}}}
]}},
"size":expected_block_ids.len().saturating_add(1),
}),
)
.await?;
let nodes = block_result
.get("nodes")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
Ok(provider_block_projection_matches(
&nodes,
&expected_block_ids,
source_version,
permission_version,
))
}
pub(super) fn provider_block_projection_matches(
nodes: &[Value],
expected_block_ids: &HashSet<String>,
source_version: i64,
permission_version: i64,
) -> bool {
if nodes.len() != expected_block_ids.len() {
return false;
}
let actual_block_ids = nodes
.iter()
.filter_map(|node| {
let block_id = provider_field_string(node, "block_id")?;
let version = provider_field_i64(node, "source_version")?;
let provider_permission_version = provider_field_i64(node, "permission_version")?;
(version == source_version && provider_permission_version == permission_version).then_some(block_id)
})
.collect::<HashSet<_>>();
actual_block_ids == *expected_block_ids
}
fn provider_field_i64(node: &Value, field: &str) -> Option<i64> {
node
.pointer(&format!("/fields/{field}/0"))
.and_then(Value::as_i64)
.or_else(|| node.pointer(&format!("/_source/{field}")).and_then(Value::as_i64))
}
pub(super) async fn reconcile_source_documents(
context: &WorkspaceReconcileContext<'_>,
permission_version: i64,
after_doc_id: Option<String>,
) -> RuntimeResult<WorkspaceStep> {
let pool = context.pool;
let embedded = context.embedded;
let remote = context.remote;
let generation = context.generation;
let workspace_id = context.workspace_id;
let workspace_fence = context.fence;
let rows = sqlx::query(CANONICAL_SNAPSHOT_BATCH_SQL)
.bind(workspace_id)
.bind(generation.id)
.bind(&after_doc_id)
.bind(RECONCILE_BATCH + 1)
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load anti-entropy search source batch", error))?;
let complete = rows.len() <= RECONCILE_BATCH as usize;
let mut after_doc_id = after_doc_id;
for row in rows.into_iter().take(RECONCILE_BATCH as usize) {
let doc_id: String = row
.try_get("guid")
.map_err(|error| RuntimeError::database("decode anti-entropy search source document", error))?;
let source_version: Option<i64> = row
.try_get("target_source_version")
.map_err(|error| RuntimeError::database("decode anti-entropy search source version", error))?;
let target_permission_version: Option<i64> = row
.try_get("target_permission_version")
.map_err(|error| RuntimeError::database("decode anti-entropy search permission version", error))?;
if !renew_workspace_lease(pool, generation.id, workspace_id, workspace_fence).await? {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Source { after_doc_id }));
}
let projection_result = match (source_version, target_permission_version) {
(Some(source_version), Some(target_permission_version)) => {
provider_projection_matches(
pool,
embedded,
remote,
generation,
ProjectionExpectation {
workspace_id,
doc_id: &doc_id,
source_version,
permission_version: target_permission_version.max(permission_version),
},
)
.await
}
_ => Ok(false),
};
let projection_matches = match projection_result {
Ok(matches) => matches,
Err(error) if error.is_permanent_search_source() => {
mark_workspace_failed(pool, generation.id, workspace_id, workspace_fence).await?;
return Ok(WorkspaceStep::Failed);
}
Err(error) => return Err(error),
};
if !projection_matches {
if !renew_workspace_lease(pool, generation.id, workspace_id, workspace_fence).await? {
return Ok(WorkspaceStep::Continue(WorkspacePhase::Source { after_doc_id }));
}
if let Err(error) = upsert_document(pool, embedded, remote, generation, workspace_id, &doc_id).await {
if error.is_permanent_search_source() {
mark_workspace_failed(pool, generation.id, workspace_id, workspace_fence).await?;
return Ok(WorkspaceStep::Failed);
}
return Err(error);
}
}
after_doc_id = Some(doc_id);
}
if complete {
Ok(WorkspaceStep::Complete)
} else {
Ok(WorkspaceStep::Continue(WorkspacePhase::Source { after_doc_id }))
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn block_projection_check_rejects_missing_or_stale_rows() {
let expected = HashSet::from(["one".to_string(), "two".to_string()]);
let complete = vec![
json!({"_source":{"block_id":"one","source_version":7,"permission_version":3}}),
json!({"_source":{"block_id":"two","source_version":7,"permission_version":3}}),
];
assert!(provider_block_projection_matches(&complete, &expected, 7, 3));
assert!(!provider_block_projection_matches(&complete[..1], &expected, 7, 3));
assert!(!provider_block_projection_matches(
&[
json!({"_source":{"block_id":"one","source_version":7,"permission_version":3}}),
json!({"_source":{"block_id":"two","source_version":6,"permission_version":3}}),
],
&expected,
7,
3,
));
}
}

View File

@@ -0,0 +1,696 @@
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{
ActiveGeneration, DOCUMENT_LEASE_SECONDS, ProjectionInput, SearchChange, SearchProvider,
SearchTable as ProviderTable, project_document, projection_external_id, provider_payload,
};
use crate::{
runtime::{RuntimeError, RuntimeResult},
search_index::EmbeddedSearchIndex,
};
struct DocumentClaim {
fence: i64,
target_source_version: i64,
target_source_exists: bool,
target_permission_version: i64,
}
struct ProjectionTuple {
source_version: i64,
source_exists: bool,
permission_version: i64,
}
const HISTORY_GC_BATCH: usize = 100;
pub(super) async fn upsert_document(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<()> {
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search document enqueue", error))?;
sqlx::query("SELECT pg_advisory_xact_lock_shared(hashtextextended('search-projection-generation', 0))")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("lock search document enqueue", error))?;
let updated = sqlx::query(
r#"UPDATE search_projection.document_states state
SET target_permission_version=GREATEST(state.target_permission_version,workspace.required_permission_version),
last_error=NULL, available_at=now(), updated_at=now()
FROM search_projection.workspace_states workspace
WHERE state.generation_id=$1 AND state.workspace_id=$2 AND state.doc_id=$3
AND workspace.generation_id=state.generation_id AND workspace.workspace_id=state.workspace_id"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(doc_id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("refresh search document target", error))?
.rows_affected();
if updated == 0 {
sqlx::query(
r#"INSERT INTO search_projection.document_states
(generation_id,workspace_id,doc_id,target_source_version,target_source_exists,target_permission_version)
SELECT $1,$2,$3,nextval('search_projection.source_mutation_version'),
EXISTS(SELECT 1 FROM snapshots WHERE workspace_id=$2 AND guid=$3),
required_permission_version
FROM search_projection.workspace_states
WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(doc_id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("enqueue search document", error))?;
}
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit search document enqueue", error))?;
let Some(claim) = claim_document(pool, generation.id, workspace_id, doc_id).await? else {
let published = sqlx::query_as::<_, (i64, bool, i64)>(
r#"SELECT target_source_version,target_source_exists,target_permission_version
FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3
AND target_source_version=published_source_version
AND target_source_exists=published_source_exists
AND target_permission_version=published_permission_version
AND (lease_expires_at IS NULL OR lease_expires_at <= now())"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load published search repair tuple", error))?;
if let Some((source_version, source_exists, permission_version)) = published {
let changes = projection_changes(
pool,
generation,
workspace_id,
doc_id,
ProjectionTuple {
source_version,
source_exists,
permission_version,
},
)
.await?;
apply_changes(embedded, remote, generation, changes).await?;
gc_document_history(pool, embedded, remote, generation, workspace_id, doc_id).await;
}
return Ok(());
};
let changes = projection_changes(
pool,
generation,
workspace_id,
doc_id,
ProjectionTuple {
source_version: claim.target_source_version,
source_exists: claim.target_source_exists,
permission_version: claim.target_permission_version,
},
)
.await?;
renew_document_claim(pool, generation.id, workspace_id, doc_id, claim.fence).await?;
apply_changes(embedded, remote, generation, changes).await?;
complete_document(pool, generation.id, workspace_id, doc_id, &claim).await?;
gc_document_history(pool, embedded, remote, generation, workspace_id, doc_id).await;
Ok(())
}
async fn gc_document_history(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
workspace_id: &str,
doc_id: &str,
) {
let Ok(Some((source_version, permission_version))) = sqlx::query_as::<_, (i64, i64)>(
r#"SELECT published_source_version,published_permission_version
FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3
AND target_source_version=published_source_version
AND target_source_exists=published_source_exists
AND target_permission_version=published_permission_version"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(pool)
.await
else {
return;
};
for table in [ProviderTable::Doc, ProviderTable::Block] {
let result = match remote {
Some(provider) => {
let Ok(physical_table) = generation.physical_table(table) else {
return;
};
provider
.gc_document_history(
physical_table,
workspace_id,
doc_id,
source_version,
permission_version,
HISTORY_GC_BATCH,
)
.await
}
None => embedded
.gc_document_history_for_generation(
generation.id,
table.as_str(),
workspace_id,
doc_id,
(source_version, permission_version),
HISTORY_GC_BATCH,
)
.await
.map_err(|error| RuntimeError::invalid_state(format!("embedded search history GC failed: {error}"))),
};
if result.is_err() {
return;
}
}
}
async fn projection_changes(
pool: &PgPool,
generation: &ActiveGeneration,
workspace_id: &str,
doc_id: &str,
projection_tuple: ProjectionTuple,
) -> RuntimeResult<Vec<SearchChange>> {
if !projection_tuple.source_exists {
return Ok(Vec::new());
}
let Some((document, blocks)) = project_document(pool, workspace_id, doc_id).await? else {
return Ok(Vec::new());
};
let mut changes = Vec::with_capacity(blocks.len() + 1);
changes.push(change(
document,
generation,
projection_tuple.source_version,
projection_tuple.permission_version,
ProviderTable::Doc,
)?);
changes.extend(
blocks
.into_iter()
.map(|block| {
change(
block,
generation,
projection_tuple.source_version,
projection_tuple.permission_version,
ProviderTable::Block,
)
})
.collect::<RuntimeResult<Vec<_>>>()?,
);
Ok(changes)
}
async fn renew_document_claim(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
doc_id: &str,
fence: i64,
) -> RuntimeResult<()> {
let updated = sqlx::query(
r#"UPDATE search_projection.document_states
SET lease_expires_at=now()+make_interval(secs=>$5), updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3 AND claim_fence=$4
AND lease_expires_at > now()"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(doc_id)
.bind(fence)
.bind(DOCUMENT_LEASE_SECONDS)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("renew search document claim", error))?
.rows_affected();
if updated != 1 {
return Err(RuntimeError::invalid_state("search document claim lost"));
}
Ok(())
}
async fn claim_document(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<Option<DocumentClaim>> {
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search document claim", error))?;
let row = sqlx::query(
r#"SELECT generation_id,workspace_id,doc_id,target_source_version,target_source_exists,target_permission_version
FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3 AND available_at <= now()
AND (lease_expires_at IS NULL OR lease_expires_at <= now())
AND (target_source_version <> published_source_version
OR target_source_exists <> published_source_exists
OR target_permission_version <> published_permission_version)
FOR UPDATE SKIP LOCKED"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut *tx)
.await
.map_err(|error| RuntimeError::database("claim search document publication", error))?;
let Some(row) = row else {
tx.rollback()
.await
.map_err(|error| RuntimeError::database("rollback empty search document claim", error))?;
return Ok(None);
};
let fence: i64 = sqlx::query_scalar("SELECT nextval('search_projection.claim_fence')")
.fetch_one(&mut *tx)
.await
.map_err(|error| RuntimeError::database("allocate search claim fence", error))?;
let target_source_version: i64 = row
.try_get("target_source_version")
.map_err(|error| RuntimeError::database("decode search state source version", error))?;
let target_source_exists: bool = row
.try_get("target_source_exists")
.map_err(|error| RuntimeError::database("decode search state source existence", error))?;
let permission_version: i64 = row
.try_get("target_permission_version")
.map_err(|error| RuntimeError::database("decode search publication permission version", error))?;
let workspace_id: String = row
.try_get("workspace_id")
.map_err(|error| RuntimeError::database("decode search publication workspace", error))?;
let doc_id: String = row
.try_get("doc_id")
.map_err(|error| RuntimeError::database("decode search publication document", error))?;
sqlx::query(
r#"UPDATE search_projection.document_states
SET claim_fence=$4, lease_owner=$5, lease_expires_at=now()+make_interval(secs=>$6),
attempt_count=attempt_count+1, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3"#,
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.bind(fence)
.bind(format!("native-search-{}", std::process::id()))
.bind(DOCUMENT_LEASE_SECONDS)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("write search document claim", error))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("commit search document claim", error))?;
Ok(Some(DocumentClaim {
fence,
target_source_version,
target_source_exists,
target_permission_version: permission_version,
}))
}
async fn complete_document(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
doc_id: &str,
claim: &DocumentClaim,
) -> RuntimeResult<()> {
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search document completion", error))?;
let completed = sqlx::query(
r#"UPDATE search_projection.document_states
SET published_source_version=$5, published_source_exists=$6,
published_permission_version=$7, claim_fence=NULL, lease_owner=NULL,
lease_expires_at=NULL, last_error=NULL, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3 AND claim_fence=$4
AND target_source_version=$5 AND target_source_exists=$6
AND target_permission_version=$7"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(doc_id)
.bind(claim.fence)
.bind(claim.target_source_version)
.bind(claim.target_source_exists)
.bind(claim.target_permission_version)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("publish search document state", error))?
.rows_affected();
if completed == 0 {
sqlx::query(
r#"UPDATE search_projection.document_states
SET available_at=now(), lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3 AND claim_fence=$4"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(doc_id)
.bind(claim.fence)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("reschedule changed search document state", error))?;
}
if completed == 1 {
sqlx::query(
r#"UPDATE search_projection.workspace_states state
SET applied_permission_version=GREATEST(state.applied_permission_version, state.required_permission_version),
updated_at=now()
WHERE state.generation_id=$1 AND state.workspace_id=$2
AND state.required_permission_version <= $3
AND state.pending_scope='none'
AND NOT EXISTS(
SELECT 1 FROM search_projection.document_states document
WHERE document.generation_id=state.generation_id AND document.workspace_id=state.workspace_id
AND document.target_permission_version <> document.published_permission_version
)"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(claim.target_permission_version)
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("advance search permission state", error))?;
}
tx.commit()
.await
.map_err(|error| RuntimeError::database("commit search document completion", error))
}
fn change(
input: ProjectionInput,
generation: &ActiveGeneration,
source_version: i64,
permission_version: i64,
table: ProviderTable,
) -> RuntimeResult<SearchChange> {
let mut payload = provider_payload(&input.payload);
payload["generation_id"] = serde_json::json!(generation.id.to_string());
payload["source_version"] = serde_json::json!(source_version);
payload["permission_version"] = serde_json::json!(permission_version);
let block_id = payload.get("block_id").and_then(serde_json::Value::as_str);
let external_id = projection_external_id(
table,
&generation.id.to_string(),
&input.workspace_id,
&input.doc_id,
block_id,
source_version,
permission_version,
)?;
Ok(SearchChange {
table,
external_id,
workspace_id: input.workspace_id,
doc_id: input.doc_id,
source_version,
permission_version,
payload,
})
}
async fn apply_changes(
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
changes: Vec<SearchChange>,
) -> RuntimeResult<()> {
let mut by_table = [Vec::new(), Vec::new()];
for change in changes {
let index = match change.table {
ProviderTable::Doc => 0,
ProviderTable::Block => 1,
};
by_table[index].push(change);
}
if let Some(remote) = remote {
for (index, changes) in by_table.into_iter().enumerate() {
if !changes.is_empty() {
let table = if index == 0 {
ProviderTable::Doc
} else {
ProviderTable::Block
};
remote.apply(generation.physical_table(table)?, &changes).await?;
}
}
return Ok(());
}
for (index, changes) in by_table.into_iter().enumerate() {
let table = if index == 0 { "doc" } else { "block" };
let mut upserts = Vec::new();
for change in changes {
upserts.push(change.payload);
}
if !upserts.is_empty() {
embedded
.write_for_generation(
generation.id,
table.to_string(),
serde_json::to_string(&upserts).map_err(|error| RuntimeError::json("encode embedded changes", error))?,
)
.await?;
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use sqlx::PgPool;
use uuid::Uuid;
use super::{DocumentClaim, claim_document, complete_document};
use crate::runtime::{backend_runtime::search::SEARCH_TEST_LOCK, migrations::migrate_search_tables};
#[tokio::test]
async fn document_completion_only_advances_document_scoped_permission_work() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
let generation_id = Uuid::new_v4();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','failed',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
for (workspace_id, pending_scope) in [("workspace-scope", "workspace"), ("document-scope", "none")] {
sqlx::query(
r#"INSERT INTO search_projection.workspace_states(
generation_id,workspace_id,covered,required_permission_version,
applied_permission_version,pending_scope
) VALUES($1,$2,true,5,4,$3)"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(pending_scope)
.execute(&pool)
.await
.unwrap();
for (doc_id, fence) in [("first", 101_i64), ("second", 102_i64)] {
sqlx::query(
r#"INSERT INTO search_projection.document_states(
generation_id,workspace_id,doc_id,target_source_version,target_source_exists,
target_permission_version,claim_fence
) VALUES($1,$2,$3,1,true,5,$4)"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(doc_id)
.bind(fence)
.execute(&pool)
.await
.unwrap();
}
complete_document(
&pool,
generation_id,
workspace_id,
"first",
&DocumentClaim {
fence: 101,
target_source_version: 1,
target_source_exists: true,
target_permission_version: 5,
},
)
.await
.unwrap();
let applied: i64 = sqlx::query_scalar(
"SELECT applied_permission_version FROM search_projection.workspace_states WHERE generation_id=$1 AND \
workspace_id=$2",
)
.bind(generation_id)
.bind(workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(applied, 4);
complete_document(
&pool,
generation_id,
workspace_id,
"second",
&DocumentClaim {
fence: 102,
target_source_version: 1,
target_source_exists: true,
target_permission_version: 5,
},
)
.await
.unwrap();
let applied: i64 = sqlx::query_scalar(
"SELECT applied_permission_version FROM search_projection.workspace_states WHERE generation_id=$1 AND \
workspace_id=$2",
)
.bind(generation_id)
.bind(workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(applied, if pending_scope == "none" { 5 } else { 4 });
}
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
}
#[tokio::test]
async fn document_claim_skips_locked_rows_and_global_fence_prevents_aba_completion() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
let generation_id = Uuid::new_v4();
let workspace_id = format!("claim-document-{}", Uuid::new_v4().simple());
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','failed',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO search_projection.workspace_states(generation_id,workspace_id,pending_scope) VALUES($1,$2,'none')",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO \
search_projection.document_states(generation_id,workspace_id,doc_id,target_source_version,\
target_source_exists) VALUES($1,$2,'doc',1,true)",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let mut blocker = pool.begin().await.unwrap();
sqlx::query(
"SELECT 1 FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 AND doc_id='doc' \
FOR UPDATE",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&mut *blocker)
.await
.unwrap();
assert!(
claim_document(&pool, generation_id, &workspace_id, "doc")
.await
.unwrap()
.is_none()
);
blocker.rollback().await.unwrap();
let first_claim = claim_document(&pool, generation_id, &workspace_id, "doc")
.await
.unwrap()
.unwrap();
complete_document(&pool, generation_id, &workspace_id, "doc", &first_claim)
.await
.unwrap();
sqlx::query(
"UPDATE search_projection.document_states SET target_source_version=2, claim_fence=NULL, lease_owner=NULL, \
lease_expires_at=NULL WHERE generation_id=$1 AND workspace_id=$2 AND doc_id='doc'",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let second_claim = claim_document(&pool, generation_id, &workspace_id, "doc")
.await
.unwrap()
.unwrap();
assert!(second_claim.fence > first_claim.fence);
complete_document(&pool, generation_id, &workspace_id, "doc", &first_claim)
.await
.unwrap();
let remaining_fence: i64 = sqlx::query_scalar(
"SELECT claim_fence FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 AND \
doc_id='doc'",
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(remaining_fence, second_claim.fence);
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
}
}

View File

@@ -0,0 +1,41 @@
mod anti_entropy;
mod document;
mod workspace;
mod workspace_state;
pub(super) const DOCUMENT_LEASE_SECONDS: i64 = 300;
pub(super) const LEASE_SECONDS: i64 = 300;
pub(super) const RECONCILE_BATCH: i64 = 100;
struct WorkspaceReconcileContext<'a> {
pool: &'a sqlx::PgPool,
embedded: &'a crate::search_index::EmbeddedSearchIndex,
remote: Option<&'a SearchProvider>,
generation: &'a ActiveGeneration,
workspace_id: &'a str,
fence: i64,
}
enum WorkspaceStep {
Continue(WorkspacePhase),
Quiet(WorkspacePhase),
Complete,
Failed,
}
pub(super) use anti_entropy::sweep_generation_orphans;
use anti_entropy::{
CANONICAL_SNAPSHOT_BATCH_SQL, ProjectionExpectation, provider_projection_matches, reconcile_source_documents,
reconcile_stale_provider_rows, sweep_deleted_workspace,
};
use document::upsert_document;
pub(super) use workspace::reconcile_workspace;
use workspace_state::{
WorkspacePhase, checkpoint_workspace, checkpoint_workspace_after, claim_workspace, complete_workspace,
delete_workspace_state, mark_workspace_failed, renew_workspace_lease,
};
use super::{
ActiveGeneration, ProjectionInput, SearchChange, SearchProvider, SearchTable, project_document,
projection_external_id, provider_payload,
};

View File

@@ -0,0 +1,877 @@
use sqlx::{PgPool, Row};
use super::{
ActiveGeneration, CANONICAL_SNAPSHOT_BATCH_SQL, DOCUMENT_LEASE_SECONDS, ProjectionExpectation, RECONCILE_BATCH,
SearchProvider, SearchTable, WorkspacePhase, WorkspaceReconcileContext, WorkspaceStep, checkpoint_workspace,
checkpoint_workspace_after, claim_workspace, complete_workspace, delete_workspace_state, mark_workspace_failed,
provider_projection_matches, reconcile_source_documents, reconcile_stale_provider_rows, renew_workspace_lease,
sweep_deleted_workspace, upsert_document,
};
use crate::{
runtime::{RuntimeError, RuntimeResult, storage_runtime::load_current_doc},
search_index::EmbeddedSearchIndex,
};
const DUE_DOCUMENT_PUBLICATION_BATCH_SQL: &str = r#"SELECT doc_id
FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND available_at <= now()
AND (lease_expires_at IS NULL OR lease_expires_at <= now())
AND (target_source_version <> published_source_version
OR target_source_exists <> published_source_exists
OR target_permission_version <> published_permission_version)
AND ($3::text IS NULL OR doc_id > $3)
ORDER BY doc_id LIMIT $4"#;
pub(in crate::runtime::backend_runtime::search) async fn reconcile_workspace(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
workspace_id: &str,
) -> RuntimeResult<bool> {
sqlx::query(
r#"INSERT INTO search_projection.workspace_states(generation_id,workspace_id)
VALUES ($1,$2) ON CONFLICT DO NOTHING"#,
)
.bind(generation.id)
.bind(workspace_id)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("ensure search workspace state", error))?;
let Some(claim) = claim_workspace(pool, generation.id, workspace_id).await? else {
return Ok(false);
};
let context = WorkspaceReconcileContext {
pool,
embedded,
remote,
generation,
workspace_id,
fence: claim.fence,
};
if let WorkspacePhase::Deleted { table, quiet } = claim.progress.phase.clone() {
match sweep_deleted_workspace(&context, table, quiet).await? {
WorkspaceStep::Continue(progress) => {
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim.progress.with_phase(progress).value(),
)
.await?;
}
WorkspaceStep::Quiet(progress) => {
checkpoint_workspace_after(
pool,
generation.id,
workspace_id,
claim.fence,
claim.progress.with_phase(progress).value(),
DOCUMENT_LEASE_SECONDS,
)
.await?;
}
WorkspaceStep::Complete => {
delete_workspace_state(pool, generation.id, workspace_id, claim.fence).await?;
}
WorkspaceStep::Failed => unreachable!("workspace GC does not classify source failures"),
}
return Ok(true);
}
if let Err(error) = validate_workspace_root(pool, workspace_id).await {
if error.is_permanent_search_source() {
mark_workspace_failed(pool, generation.id, workspace_id, claim.fence).await?;
return Ok(true);
}
return Err(error);
}
if let WorkspacePhase::Stale { table, cursor } = claim.progress.phase.clone() {
match reconcile_stale_provider_rows(&context, table, cursor).await? {
WorkspaceStep::Continue(progress) => {
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim.progress.with_phase(progress).value(),
)
.await?;
return Ok(true);
}
WorkspaceStep::Failed => return Ok(true),
WorkspaceStep::Quiet(_) => unreachable!("stale reconciliation does not schedule a quiet period"),
WorkspaceStep::Complete => {}
}
complete_workspace(
pool,
generation,
workspace_id,
claim.progress.captured_root_revision,
claim.progress.captured_permission_version,
claim.fence,
)
.await?;
return Ok(true);
}
if let WorkspacePhase::Source { after_doc_id } = claim.progress.phase.clone() {
match reconcile_source_documents(&context, claim.progress.captured_permission_version, after_doc_id).await? {
WorkspaceStep::Continue(progress) => {
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim.progress.with_phase(progress).value(),
)
.await?;
return Ok(true);
}
WorkspaceStep::Failed => return Ok(true),
WorkspaceStep::Quiet(_) => unreachable!("source reconciliation does not schedule a quiet period"),
WorkspaceStep::Complete => {}
}
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim
.progress
.with_phase(WorkspacePhase::Stale {
table: SearchTable::Doc,
cursor: None,
})
.value(),
)
.await?;
return Ok(true);
}
let (mut after_publication_doc_id, mut after_doc_id, scan_workspace) = match claim.progress.phase.clone() {
WorkspacePhase::Publications {
after_publication_doc_id,
resume_after_doc_id,
scan_workspace,
} => (after_publication_doc_id, resume_after_doc_id, scan_workspace),
WorkspacePhase::Documents { after_doc_id } => (None, after_doc_id, true),
WorkspacePhase::Source { .. } => unreachable!("source progress is handled before document reconcile"),
WorkspacePhase::Stale { .. } => unreachable!("stale progress is handled before document reconcile"),
WorkspacePhase::Deleted { .. } => unreachable!("deleted progress is handled before document reconcile"),
};
let publication_rows = sqlx::query(DUE_DOCUMENT_PUBLICATION_BATCH_SQL)
.bind(generation.id)
.bind(workspace_id)
.bind(&after_publication_doc_id)
.bind(RECONCILE_BATCH + 1)
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load due search document publications", error))?;
let publications_complete = publication_rows.len() <= RECONCILE_BATCH as usize;
for row in publication_rows.into_iter().take(RECONCILE_BATCH as usize) {
let doc_id: String = row
.try_get("doc_id")
.map_err(|error| RuntimeError::database("decode due search document publication", error))?;
if !renew_workspace_lease(pool, generation.id, workspace_id, claim.fence).await? {
return Ok(false);
}
if !process_document(pool, embedded, remote, generation, workspace_id, &doc_id, claim.fence).await? {
return Ok(true);
}
after_publication_doc_id = Some(doc_id);
}
if !publications_complete {
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim
.progress
.with_phase(WorkspacePhase::Publications {
after_publication_doc_id,
resume_after_doc_id: after_doc_id,
scan_workspace,
})
.value(),
)
.await?;
return Ok(true);
}
let has_due_publications: bool = sqlx::query_scalar(
r#"SELECT EXISTS(
SELECT 1 FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND available_at <= now()
AND (lease_expires_at IS NULL OR lease_expires_at <= now())
AND (target_source_version <> published_source_version
OR target_source_exists <> published_source_exists
OR target_permission_version <> published_permission_version)
)"#,
)
.bind(generation.id)
.bind(workspace_id)
.fetch_one(pool)
.await
.map_err(|error| RuntimeError::database("check due search document publications", error))?;
if has_due_publications {
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim
.progress
.with_phase(WorkspacePhase::Publications {
after_publication_doc_id: None,
resume_after_doc_id: after_doc_id,
scan_workspace,
})
.value(),
)
.await?;
return Ok(true);
}
if !scan_workspace {
complete_workspace(
pool,
generation,
workspace_id,
claim.progress.captured_root_revision,
claim.progress.captured_permission_version,
claim.fence,
)
.await?;
return Ok(true);
}
let rows = sqlx::query(CANONICAL_SNAPSHOT_BATCH_SQL)
.bind(workspace_id)
.bind(generation.id)
.bind(&after_doc_id)
.bind(RECONCILE_BATCH + 1)
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load canonical search workspace batch", error))?;
let complete = rows.len() <= RECONCILE_BATCH as usize;
for row in rows.into_iter().take(RECONCILE_BATCH as usize) {
let doc_id: String = row
.try_get("guid")
.map_err(|error| RuntimeError::database("decode canonical search workspace document", error))?;
let source_version: Option<i64> = row
.try_get("target_source_version")
.map_err(|error| RuntimeError::database("decode canonical search document version", error))?;
let target_permission_version: Option<i64> = row
.try_get("target_permission_version")
.map_err(|error| RuntimeError::database("decode canonical search document permission version", error))?;
if !renew_workspace_lease(pool, generation.id, workspace_id, claim.fence).await? {
return Ok(false);
}
let projection_result = match (source_version, target_permission_version) {
(Some(source_version), Some(target_permission_version)) => {
provider_projection_matches(
pool,
embedded,
remote,
generation,
ProjectionExpectation {
workspace_id,
doc_id: &doc_id,
source_version,
permission_version: target_permission_version.max(claim.progress.captured_permission_version),
},
)
.await
}
_ => Ok(false),
};
let projection_matches = match projection_result {
Ok(matches) => matches,
Err(error) if error.is_permanent_search_source() => {
mark_workspace_failed(pool, generation.id, workspace_id, claim.fence).await?;
return Ok(true);
}
Err(error) => return Err(error),
};
if !projection_matches {
if !renew_workspace_lease(pool, generation.id, workspace_id, claim.fence).await? {
return Ok(false);
}
if !process_document(pool, embedded, remote, generation, workspace_id, &doc_id, claim.fence).await? {
return Ok(true);
}
}
after_doc_id = Some(doc_id);
}
if !complete {
checkpoint_workspace(
pool,
generation.id,
workspace_id,
claim.fence,
claim
.progress
.with_phase(WorkspacePhase::Documents { after_doc_id })
.value(),
)
.await?;
return Ok(true);
}
complete_workspace(
pool,
generation,
workspace_id,
claim.progress.captured_root_revision,
claim.progress.captured_permission_version,
claim.fence,
)
.await?;
Ok(true)
}
async fn validate_workspace_root(pool: &PgPool, workspace_id: &str) -> RuntimeResult<()> {
let root = load_current_doc(pool, workspace_id, workspace_id)
.await
.map_err(|error| match error {
RuntimeError::InvalidState(message) => RuntimeError::SearchSourceInvalid(message),
error => error,
})?
.ok_or_else(|| RuntimeError::SearchSourceInvalid("workspace root doc is missing".to_string()))?;
let projection = affine_doc_loader::project_workspace_root(root.blob, true)
.map_err(|error| RuntimeError::SearchSourceInvalid(format!("workspace root projection failed: {error}")))?;
if !projection.complete {
return Err(RuntimeError::SearchSourceInvalid(
"workspace root projection is incomplete".to_string(),
));
}
Ok(())
}
async fn process_document(
pool: &PgPool,
embedded: &EmbeddedSearchIndex,
remote: Option<&SearchProvider>,
generation: &ActiveGeneration,
workspace_id: &str,
doc_id: &str,
workspace_fence: i64,
) -> RuntimeResult<bool> {
match upsert_document(pool, embedded, remote, generation, workspace_id, doc_id).await {
Ok(()) => Ok(true),
Err(error) if error.is_permanent_search_source() => {
mark_workspace_failed(pool, generation.id, workspace_id, workspace_fence).await?;
Ok(false)
}
Err(error) => Err(error),
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use serde_json::{Value, json};
use sqlx::PgPool;
use uuid::Uuid;
use super::{super::workspace_state::WorkspaceProgress, *};
use crate::runtime::{backend_runtime::search::SEARCH_TEST_LOCK, migrations::migrate_search_tables};
#[tokio::test]
async fn deleted_workspace_waits_for_quiet_sweep_before_dropping_state() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
let suffix = Uuid::new_v4().simple().to_string();
let workspace_id = format!("deleted-workspace-{suffix}");
let doc_id = format!("doc-{suffix}");
let generation_id = Uuid::new_v4();
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version,manifest)
VALUES($1,'embedded','failed',decode(repeat('00',32),'hex'),1,'{}')"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO search_projection.workspace_states(generation_id,workspace_id) VALUES($1,$2)")
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let embedded = Arc::new(EmbeddedSearchIndex::default());
embedded.prepare_generation(generation_id).await;
let projection = json!({
"generation_id":generation_id.to_string(),"workspace_id":workspace_id,"doc_id":doc_id,
"source_version":1,"permission_version":1,"title":"late"
});
embedded
.write_for_generation(generation_id, "doc".to_string(), json!([projection]).to_string())
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let generation = ActiveGeneration {
id: generation_id,
manifest: json!({}),
};
assert!(
reconcile_workspace(&pool, &embedded, None, &generation, &workspace_id)
.await
.unwrap()
);
assert!(
reconcile_workspace(&pool, &embedded, None, &generation, &workspace_id)
.await
.unwrap()
);
let quiet: bool = sqlx::query_scalar(
"SELECT progress->>'quiet'='true' AND available_at > now() FROM search_projection.workspace_states WHERE \
generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(quiet);
let late_projection = json!({
"generation_id":generation_id.to_string(),"workspace_id":workspace_id,"doc_id":doc_id,
"source_version":1,"permission_version":1,"title":"late"
});
embedded
.write_for_generation(generation_id, "doc".to_string(), json!([late_projection]).to_string())
.await
.unwrap();
sqlx::query(
"UPDATE search_projection.workspace_states SET available_at=now() WHERE generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
assert!(
reconcile_workspace(&pool, &embedded, None, &generation, &workspace_id)
.await
.unwrap()
);
assert!(
reconcile_workspace(&pool, &embedded, None, &generation, &workspace_id)
.await
.unwrap()
);
let state_exists: bool = sqlx::query_scalar(
"SELECT EXISTS(SELECT 1 FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2)",
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(!state_exists);
let result: Value = serde_json::from_str(
&embedded
.search_for_generation(
generation_id,
"doc".to_string(),
json!({"query":{"term":{"workspace_id":{"value":workspace_id}}},"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 0);
let race_workspace_id = format!("recreated-during-sweep-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&race_workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO search_projection.workspace_states(generation_id,workspace_id) VALUES($1,$2)")
.bind(generation_id)
.bind(&race_workspace_id)
.execute(&pool)
.await
.unwrap();
embedded
.write_for_generation(
generation_id,
"doc".to_string(),
json!([{
"generation_id":generation_id.to_string(),"workspace_id":race_workspace_id,"doc_id":"old",
"source_version":1,"permission_version":1,"title":"old incarnation"
}])
.to_string(),
)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&race_workspace_id)
.execute(&pool)
.await
.unwrap();
let mut blocker = pool.begin().await.unwrap();
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *blocker)
.await
.unwrap();
let sweep_pool = pool.clone();
let sweep_embedded = Arc::clone(&embedded);
let sweep_generation = generation.clone();
let sweep_workspace_id = race_workspace_id.clone();
let sweep = tokio::spawn(async move {
reconcile_workspace(
&sweep_pool,
&sweep_embedded,
None,
&sweep_generation,
&sweep_workspace_id,
)
.await
.unwrap()
});
for _ in 0..100 {
let claimed: bool = sqlx::query_scalar(
"SELECT COALESCE(lease_expires_at > now(),false) FROM search_projection.workspace_states WHERE \
generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&race_workspace_id)
.fetch_one(&pool)
.await
.unwrap();
if claimed {
break;
}
tokio::task::yield_now().await;
}
let claimed: bool = sqlx::query_scalar(
"SELECT COALESCE(lease_expires_at > now(),false) FROM search_projection.workspace_states WHERE generation_id=$1 \
AND workspace_id=$2",
)
.bind(generation_id)
.bind(&race_workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(claimed);
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&race_workspace_id)
.execute(&pool)
.await
.unwrap();
blocker.commit().await.unwrap();
assert!(sweep.await.unwrap());
let result: Value = serde_json::from_str(
&embedded
.search_for_generation(
generation_id,
"doc".to_string(),
json!({"query":{"term":{"workspace_id":{"value":race_workspace_id}}},"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&race_workspace_id)
.execute(&pool)
.await
.unwrap();
let source_version_high_water = super::super::anti_entropy::capture_orphan_gc_high_water(&pool, &workspace_id)
.await
.unwrap()
.unwrap();
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let mut recreated_writer = pool.begin().await.unwrap();
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *recreated_writer)
.await
.unwrap();
let recreated_source_version: i64 =
sqlx::query_scalar("SELECT nextval('search_projection.source_mutation_version')")
.fetch_one(&mut *recreated_writer)
.await
.unwrap();
recreated_writer.commit().await.unwrap();
assert!(recreated_source_version > source_version_high_water);
let old_projection = json!({
"generation_id":generation_id.to_string(),"workspace_id":workspace_id,"doc_id":format!("{doc_id}-old"),
"source_version":source_version_high_water,"permission_version":1,"title":"old incarnation"
});
let recreated_projection = json!({
"generation_id":generation_id.to_string(),"workspace_id":workspace_id,"doc_id":doc_id,
"source_version":recreated_source_version,"permission_version":1,"title":"recreated"
});
embedded
.write_for_generation(
generation_id,
"doc".to_string(),
json!([old_projection, recreated_projection]).to_string(),
)
.await
.unwrap();
embedded
.gc_workspace_for_generation(generation_id, "doc", &workspace_id, source_version_high_water, 100)
.await
.unwrap();
let result: Value = serde_json::from_str(
&embedded
.search_for_generation(
generation_id,
"doc".to_string(),
json!({"query":{"term":{"workspace_id":{"value":workspace_id}}},"fields":["source_version"],"size":10})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(
result["nodes"][0]["fields"]["source_version"],
json!([recreated_source_version])
);
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let post_finalize_projections = (0..125)
.map(|index| {
json!({
"generation_id":generation_id.to_string(),"workspace_id":workspace_id,
"doc_id":format!("{doc_id}-{index:03}"),
"source_version":1,"permission_version":1,"title":"post-finalize"
})
})
.collect::<Vec<_>>();
embedded
.write_for_generation(
generation_id,
"doc".to_string(),
json!(post_finalize_projections).to_string(),
)
.await
.unwrap();
for _ in 0..5 {
super::super::sweep_generation_orphans(&pool, &embedded, None, &generation)
.await
.unwrap();
}
let result: Value = serde_json::from_str(
&embedded
.search_for_generation(
generation_id,
"doc".to_string(),
json!({"query":{"term":{"workspace_id":{"value":workspace_id}}},"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 0);
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let recreated_projection = json!({
"generation_id":generation_id.to_string(),"workspace_id":workspace_id,"doc_id":doc_id,
"source_version":2,"permission_version":1,"title":"recreated"
});
embedded
.write_for_generation(
generation_id,
"doc".to_string(),
json!([recreated_projection]).to_string(),
)
.await
.unwrap();
super::super::sweep_generation_orphans(&pool, &embedded, None, &generation)
.await
.unwrap();
super::super::sweep_generation_orphans(&pool, &embedded, None, &generation)
.await
.unwrap();
let result: Value = serde_json::from_str(
&embedded
.search_for_generation(
generation_id,
"doc".to_string(),
json!({"query":{"term":{"workspace_id":{"value":workspace_id}}},"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
}
#[tokio::test]
async fn permanent_source_failures_are_terminal_in_every_workspace_phase() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
for (phase_name, phase, corrupt_update) in [
(
"publications",
WorkspacePhase::Publications {
after_publication_doc_id: None,
resume_after_doc_id: None,
scan_workspace: false,
},
false,
),
("documents", WorkspacePhase::Documents { after_doc_id: None }, true),
("source", WorkspacePhase::Source { after_doc_id: None }, false),
] {
let suffix = Uuid::new_v4().simple().to_string();
let generation_id = Uuid::new_v4();
let workspace_id = format!("source-failure-{phase_name}-{suffix}");
let doc_id = format!("doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let doc_blob = if corrupt_update {
affine_doc_loader::build_full_doc("title", "body", &doc_id).unwrap()
} else {
vec![0]
};
let root_blob = affine_doc_loader::add_doc_to_root_doc(Vec::new(), &doc_id, None).unwrap();
sqlx::query(
r#"INSERT INTO snapshots(workspace_id,guid,blob,updated_at)
VALUES($1,$1,$3,now()),($1,$2,$4,now())"#,
)
.bind(&workspace_id)
.bind(&doc_id)
.bind(root_blob)
.bind(doc_blob)
.execute(&pool)
.await
.unwrap();
if corrupt_update {
sqlx::query("INSERT INTO updates(workspace_id,guid,blob,created_at) VALUES($1,$2,decode('00','hex'),now())")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
}
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version,manifest)
VALUES($1,'embedded','failed',decode(repeat('00',32),'hex'),1,$2)"#,
)
.bind(generation_id)
.bind(json!({}))
.execute(&pool)
.await
.unwrap();
let progress = WorkspaceProgress::new(0, 0, phase).value();
sqlx::query(
r#"INSERT INTO search_projection.workspace_states(
generation_id,workspace_id,pending_scope,progress
) VALUES($1,$2,'workspace',$3)"#,
)
.bind(generation_id)
.bind(&workspace_id)
.bind(progress)
.execute(&pool)
.await
.unwrap();
if phase_name == "publications" {
sqlx::query(
r#"INSERT INTO search_projection.document_states(
generation_id,workspace_id,doc_id,target_source_version,target_source_exists
) VALUES($1,$2,$3,1,true)"#,
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
}
let embedded = EmbeddedSearchIndex::default();
embedded.prepare_generation(generation_id).await;
let generation = ActiveGeneration {
id: generation_id,
manifest: json!({}),
};
assert!(
reconcile_workspace(&pool, &embedded, None, &generation, &workspace_id)
.await
.unwrap()
);
let failed: (bool, Option<String>) = sqlx::query_as(
"SELECT covered,last_error FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
failed,
(false, Some("search_workspace_reconcile_failed".to_string())),
"{phase_name}"
);
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
}
}
}

View File

@@ -0,0 +1,588 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{ActiveGeneration, LEASE_SECONDS, SearchTable};
use crate::runtime::{RuntimeError, RuntimeResult};
const ANTI_ENTROPY_INTERVAL_SECONDS: i64 = 300;
pub(super) struct WorkspaceClaim {
pub(super) fence: i64,
pub(super) progress: WorkspaceProgress,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(super) struct WorkspaceProgress {
version: u8,
pub(super) captured_root_revision: i64,
pub(super) captured_permission_version: i64,
#[serde(flatten)]
pub(super) phase: WorkspacePhase,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub(super) enum WorkspacePhase {
Publications {
after_publication_doc_id: Option<String>,
resume_after_doc_id: Option<String>,
scan_workspace: bool,
},
Documents {
after_doc_id: Option<String>,
},
Source {
after_doc_id: Option<String>,
},
Stale {
table: SearchTable,
cursor: Option<String>,
},
Deleted {
table: SearchTable,
quiet: bool,
},
}
impl WorkspaceProgress {
pub(super) fn new(captured_root_revision: i64, captured_permission_version: i64, phase: WorkspacePhase) -> Self {
Self {
version: 1,
captured_root_revision,
captured_permission_version,
phase,
}
}
fn from_value(progress: Value) -> RuntimeResult<Self> {
let progress: Self =
serde_json::from_value(progress).map_err(|_| RuntimeError::invalid_state("invalid search workspace progress"))?;
if progress.version != 1 {
return Err(RuntimeError::invalid_state(
"unsupported search workspace progress version",
));
}
Ok(progress)
}
pub(super) fn with_phase(&self, phase: WorkspacePhase) -> Self {
Self::new(self.captured_root_revision, self.captured_permission_version, phase)
}
pub(super) fn value(&self) -> Value {
serde_json::to_value(self).expect("workspace progress is serializable")
}
}
pub(super) async fn claim_workspace(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
) -> RuntimeResult<Option<WorkspaceClaim>> {
let mut tx = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search workspace claim", error))?;
let row = sqlx::query(
r#"SELECT covered,pending_scope,required_permission_version,progress,
target_root_revision AS current_root_revision,
EXISTS(
SELECT 1 FROM search_projection.document_states document
WHERE document.generation_id=state.generation_id
AND document.workspace_id=state.workspace_id
AND document.available_at <= now()
AND (document.target_source_version <> document.published_source_version
OR document.target_source_exists <> document.published_source_exists
OR document.target_permission_version <> document.published_permission_version)
) AS has_due_publications
FROM search_projection.workspace_states state
WHERE generation_id=$1 AND workspace_id=$2
AND (available_at <= now() OR EXISTS(
SELECT 1 FROM search_projection.document_states document
WHERE document.generation_id=state.generation_id
AND document.workspace_id=state.workspace_id
AND document.available_at <= now()
AND (document.target_source_version <> document.published_source_version
OR document.target_source_exists <> document.published_source_exists
OR document.target_permission_version <> document.published_permission_version)
))
AND (lease_expires_at IS NULL OR lease_expires_at <= now())
FOR UPDATE SKIP LOCKED"#,
)
.bind(generation_id)
.bind(workspace_id)
.fetch_optional(&mut *tx)
.await
.map_err(|error| RuntimeError::database("claim search workspace", error))?;
let Some(row) = row else {
tx.rollback()
.await
.map_err(|error| RuntimeError::database("rollback empty search workspace claim", error))?;
return Ok(None);
};
let fence: i64 = sqlx::query_scalar("SELECT nextval('search_projection.claim_fence')")
.fetch_one(&mut *tx)
.await
.map_err(|error| RuntimeError::database("allocate search workspace fence", error))?;
let permission_version: i64 = row
.try_get("required_permission_version")
.map_err(|error| RuntimeError::database("decode search workspace permission version", error))?;
let root_revision: i64 = row
.try_get("current_root_revision")
.map_err(|error| RuntimeError::database("decode search workspace root revision", error))?;
let covered: bool = row
.try_get("covered")
.map_err(|error| RuntimeError::database("decode search workspace coverage", error))?;
let pending_scope: String = row
.try_get("pending_scope")
.map_err(|error| RuntimeError::database("decode search workspace pending scope", error))?;
let has_due_publications: bool = row
.try_get("has_due_publications")
.map_err(|error| RuntimeError::database("decode due search document publications", error))?;
let progress: Option<Value> = row
.try_get("progress")
.map_err(|error| RuntimeError::database("decode search workspace progress", error))?;
let progress = match progress {
Some(progress) => WorkspaceProgress::from_value(progress)?,
None if covered && pending_scope == "none" => WorkspaceProgress::new(
root_revision,
permission_version,
if has_due_publications {
WorkspacePhase::Publications {
after_publication_doc_id: None,
resume_after_doc_id: None,
scan_workspace: false,
}
} else {
WorkspacePhase::Source { after_doc_id: None }
},
),
None => WorkspaceProgress::new(
root_revision,
permission_version,
WorkspacePhase::Documents { after_doc_id: None },
),
};
sqlx::query(
r#"UPDATE search_projection.workspace_states
SET claim_fence=$3, lease_owner=$4, lease_expires_at=now()+make_interval(secs=>$5),
progress=$6, attempt_count=attempt_count+1, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.bind(format!("native-search-{}", std::process::id()))
.bind(LEASE_SECONDS)
.bind(progress.value())
.execute(&mut *tx)
.await
.map_err(|error| RuntimeError::database("write search workspace claim", error))?;
tx.commit()
.await
.map_err(|error| RuntimeError::database("commit search workspace claim", error))?;
Ok(Some(WorkspaceClaim { fence, progress }))
}
pub(super) async fn checkpoint_workspace(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
fence: i64,
progress: Value,
) -> RuntimeResult<()> {
sqlx::query(
r#"UPDATE search_projection.workspace_states
SET progress=$4, available_at=now(), lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND claim_fence=$3"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.bind(progress)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("checkpoint search workspace reconcile", error))?;
Ok(())
}
pub(super) async fn checkpoint_workspace_after(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
fence: i64,
progress: Value,
delay_seconds: i64,
) -> RuntimeResult<()> {
sqlx::query(
r#"UPDATE search_projection.workspace_states
SET progress=$4, available_at=now()+make_interval(secs=>$5),
lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND claim_fence=$3"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.bind(progress)
.bind(delay_seconds)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("checkpoint deleted search workspace quiet period", error))?;
Ok(())
}
pub(super) async fn delete_workspace_state(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
fence: i64,
) -> RuntimeResult<()> {
sqlx::query(
r#"DELETE FROM search_projection.workspace_states state
WHERE state.generation_id=$1 AND state.workspace_id=$2 AND state.claim_fence=$3
AND NOT EXISTS (SELECT 1 FROM workspaces WHERE id=$2)"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("delete completed search workspace state", error))?;
Ok(())
}
pub(super) async fn renew_workspace_lease(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
fence: i64,
) -> RuntimeResult<bool> {
let renewed = sqlx::query(
r#"UPDATE search_projection.workspace_states
SET lease_expires_at=now()+make_interval(secs=>$4), updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND claim_fence=$3
AND lease_expires_at > now()"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.bind(LEASE_SECONDS)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("renew search workspace lease", error))?;
Ok(renewed.rows_affected() == 1)
}
pub(super) async fn complete_workspace(
pool: &PgPool,
generation: &ActiveGeneration,
workspace_id: &str,
root_revision: i64,
permission_version: i64,
fence: i64,
) -> RuntimeResult<()> {
sqlx::query(
r#"UPDATE search_projection.workspace_states state
SET target_root_revision=GREATEST(state.target_root_revision,$3),
applied_root_revision=CASE
WHEN state.target_root_revision <= $3 AND state.required_permission_version <= $4
THEN GREATEST(state.applied_root_revision,$3) ELSE state.applied_root_revision END,
covered=CASE
WHEN state.target_root_revision <= $3 AND state.required_permission_version <= $4
THEN true ELSE state.covered END,
applied_permission_version=CASE
WHEN state.required_permission_version <= $4
AND NOT EXISTS (
SELECT 1 FROM search_projection.document_states document
WHERE document.generation_id=state.generation_id
AND document.workspace_id=state.workspace_id
AND document.target_permission_version <> document.published_permission_version
)
THEN GREATEST(state.applied_permission_version,state.required_permission_version)
ELSE state.applied_permission_version END,
pending_scope=CASE
WHEN state.required_permission_version > $4 THEN 'permission'
WHEN state.target_root_revision > $3 THEN 'workspace'
WHEN EXISTS (
SELECT 1 FROM search_projection.document_states document
WHERE document.generation_id=state.generation_id AND document.workspace_id=state.workspace_id
AND (document.target_source_version <> document.published_source_version
OR document.target_source_exists <> document.published_source_exists
OR document.target_permission_version <> document.published_permission_version)
) THEN 'workspace'
ELSE 'none' END,
progress=NULL,
available_at=CASE
WHEN state.required_permission_version <= $4
AND state.target_root_revision <= $3
AND NOT EXISTS (
SELECT 1 FROM search_projection.document_states document
WHERE document.generation_id=state.generation_id AND document.workspace_id=state.workspace_id
AND (document.target_source_version <> document.published_source_version
OR document.target_source_exists <> document.published_source_exists
OR document.target_permission_version <> document.published_permission_version)
)
THEN now()+make_interval(secs=>$5) ELSE now() END,
lease_owner=NULL, lease_expires_at=NULL, last_error=NULL, updated_at=now()
WHERE state.generation_id=$1 AND state.workspace_id=$2 AND state.claim_fence=$6"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(root_revision)
.bind(permission_version)
.bind(ANTI_ENTROPY_INTERVAL_SECONDS)
.bind(fence)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("complete search workspace reconcile", error))?;
Ok(())
}
pub(super) async fn mark_workspace_failed(
pool: &PgPool,
generation_id: Uuid,
workspace_id: &str,
fence: i64,
) -> RuntimeResult<()> {
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin failed search workspace update", error))?;
let generation_state: Option<String> = sqlx::query_scalar(
r#"SELECT generation.state
FROM search_projection.workspace_states state
JOIN search_projection.generations generation ON generation.id=state.generation_id
WHERE state.generation_id=$1 AND state.workspace_id=$2 AND state.claim_fence=$3
FOR UPDATE"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load failed search workspace claim", error))?;
let Some(generation_state) = generation_state else {
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit empty search workspace update", error))?;
return Ok(());
};
sqlx::query(
r#"UPDATE search_projection.workspace_states
SET covered=false, pending_scope='workspace', last_error='search_workspace_reconcile_failed',
progress=NULL, available_at='infinity'::timestamptz,
lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2 AND claim_fence=$3"#,
)
.bind(generation_id)
.bind(workspace_id)
.bind(fence)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("mark search workspace failed", error))?;
sqlx::query(
r#"UPDATE search_projection.document_states
SET available_at='infinity'::timestamptz, lease_owner=NULL, lease_expires_at=NULL, updated_at=now()
WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation_id)
.bind(workspace_id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("pause failed search document publications", error))?;
if generation_state != "active" {
sqlx::query(
r#"UPDATE search_projection.generations
SET state='failed', last_error='search_workspace_reconcile_failed'
WHERE id=$1 AND state='building'"#,
)
.bind(generation_id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("mark search generation failed", error))?;
}
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit failed search workspace update", error))?;
Ok(())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
use crate::runtime::{backend_runtime::search::SEARCH_TEST_LOCK, migrations::migrate_search_tables};
#[test]
fn progress_round_trips_versioned_context_for_each_phase() {
let phases = [
WorkspacePhase::Publications {
after_publication_doc_id: Some("doc-2".to_string()),
resume_after_doc_id: Some("doc-7".to_string()),
scan_workspace: true,
},
WorkspacePhase::Source {
after_doc_id: Some("doc-7".to_string()),
},
WorkspacePhase::Stale {
table: SearchTable::Block,
cursor: Some("cursor".to_string()),
},
WorkspacePhase::Deleted {
table: SearchTable::Doc,
quiet: true,
},
];
for phase in phases {
let progress = WorkspaceProgress::new(11, 13, phase);
assert_eq!(WorkspaceProgress::from_value(progress.value()).unwrap(), progress);
}
assert!(WorkspaceProgress::from_value(json!({"kind":"unknown"})).is_err());
assert!(
WorkspaceProgress::from_value(json!({
"version":2,
"captured_root_revision":1,
"captured_permission_version":1,
"kind":"documents",
"after_doc_id":null
}))
.is_err()
);
}
#[tokio::test]
async fn claim_preserves_captured_versions_and_failure_is_terminal() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
let generation_id = Uuid::new_v4();
let workspace_id = format!("claim-workspace-{}", Uuid::new_v4().simple());
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$1,decode('00','hex'),'2026-01-01 UTC')",
)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let initial_root: i64 = sqlx::query_scalar(
"SELECT floor(extract(epoch FROM updated_at) * 1000)::bigint FROM snapshots WHERE workspace_id=$1 AND guid=$1",
)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','failed',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.workspace_states(
generation_id,workspace_id,target_root_revision,required_permission_version,pending_scope
) VALUES($1,$2,$3,3,'workspace')"#,
)
.bind(generation_id)
.bind(&workspace_id)
.bind(initial_root)
.execute(&pool)
.await
.unwrap();
let first = claim_workspace(&pool, generation_id, &workspace_id)
.await
.unwrap()
.unwrap();
assert!(
renew_workspace_lease(&pool, generation_id, &workspace_id, first.fence)
.await
.unwrap()
);
sqlx::query(
"UPDATE search_projection.workspace_states SET lease_expires_at=now()-interval '1 second' WHERE \
generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
assert!(
!renew_workspace_lease(&pool, generation_id, &workspace_id, first.fence)
.await
.unwrap()
);
let checkpoint = first
.progress
.with_phase(WorkspacePhase::Documents {
after_doc_id: Some("doc-7".to_string()),
})
.value();
checkpoint_workspace(&pool, generation_id, &workspace_id, first.fence, checkpoint)
.await
.unwrap();
sqlx::query("UPDATE snapshots SET updated_at='2026-01-02 UTC' WHERE workspace_id=$1 AND guid=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"UPDATE search_projection.workspace_states SET required_permission_version=4 WHERE generation_id=$1 AND \
workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let second = claim_workspace(&pool, generation_id, &workspace_id)
.await
.unwrap()
.unwrap();
assert_eq!(second.progress.captured_root_revision, initial_root);
assert_eq!(second.progress.captured_permission_version, 3);
mark_workspace_failed(&pool, generation_id, &workspace_id, second.fence)
.await
.unwrap();
let failed: (bool, Option<String>, Option<Value>, bool) = sqlx::query_as(
r#"SELECT covered,last_error,progress,available_at='infinity'::timestamptz
FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
failed,
(false, Some("search_workspace_reconcile_failed".to_string()), None, true)
);
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
}
}

View File

@@ -1,85 +0,0 @@
WITH targets AS (
SELECT UNNEST($1::varchar[]) AS workspace_id
),
snapshot_stats AS (
SELECT workspace_id,
COUNT(*) AS snapshot_count,
COALESCE(SUM(COALESCE(size, octet_length(blob))), 0) AS snapshot_size
FROM snapshots
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
blob_stats AS (
SELECT workspace_id,
COUNT(*) FILTER (WHERE deleted_at IS NULL AND status = 'completed') AS blob_count,
COALESCE(SUM(size) FILTER (WHERE deleted_at IS NULL AND status = 'completed'), 0) AS blob_size
FROM blobs
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
member_stats AS (
SELECT workspace_id, COUNT(*) AS member_count
FROM workspace_user_permissions
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
public_page_stats AS (
SELECT workspace_id, COUNT(*) AS public_page_count
FROM workspace_pages
WHERE public = TRUE AND workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
feature_stats AS (
SELECT workspace_id,
ARRAY_AGG(DISTINCT name ORDER BY name) FILTER (WHERE activated) AS features
FROM workspace_features
WHERE workspace_id IN (SELECT workspace_id FROM targets)
GROUP BY workspace_id
),
aggregated AS (
SELECT t.workspace_id,
COALESCE(ss.snapshot_count, 0) AS snapshot_count,
COALESCE(ss.snapshot_size, 0) AS snapshot_size,
COALESCE(bs.blob_count, 0) AS blob_count,
COALESCE(bs.blob_size, 0) AS blob_size,
COALESCE(ms.member_count, 0) AS member_count,
COALESCE(pp.public_page_count, 0) AS public_page_count,
COALESCE(fs.features, ARRAY[]::text[]) AS features
FROM targets t
LEFT JOIN snapshot_stats ss ON ss.workspace_id = t.workspace_id
LEFT JOIN blob_stats bs ON bs.workspace_id = t.workspace_id
LEFT JOIN member_stats ms ON ms.workspace_id = t.workspace_id
LEFT JOIN public_page_stats pp ON pp.workspace_id = t.workspace_id
LEFT JOIN feature_stats fs ON fs.workspace_id = t.workspace_id
)
INSERT INTO workspace_admin_stats (
workspace_id,
snapshot_count,
snapshot_size,
blob_count,
blob_size,
member_count,
public_page_count,
features,
updated_at
)
SELECT
workspace_id,
snapshot_count,
snapshot_size,
blob_count,
blob_size,
member_count,
public_page_count,
features,
NOW()
FROM aggregated
ON CONFLICT (workspace_id) DO UPDATE SET
snapshot_count = EXCLUDED.snapshot_count,
snapshot_size = EXCLUDED.snapshot_size,
blob_count = EXCLUDED.blob_count,
blob_size = EXCLUDED.blob_size,
member_count = EXCLUDED.member_count,
public_page_count = EXCLUDED.public_page_count,
features = EXCLUDED.features,
updated_at = EXCLUDED.updated_at

View File

@@ -1,530 +0,0 @@
use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use tokio::time::{Duration as TokioDuration, sleep};
use super::{
BackendRuntime, RuntimeError, RuntimeResult,
constants::{WORKSPACE_STATS_LEASE_KEY, WORKSPACE_STATS_LOCK_NAMESPACE, WORKSPACE_STATS_REFRESH_LOCK_KEY},
napi_error,
types::{
CoordinationLeaseGrant, RuntimeWorkspaceStatsDailyRecalibrationResult, RuntimeWorkspaceStatsRecalibrationResult,
RuntimeWorkspaceStatsRefreshResult, RuntimeWorkspaceStatsSnapshotResult,
},
};
const UPSERT_WORKSPACE_ADMIN_STATS_SQL: &str = include_str!("sql/upsert_workspace_admin_stats.sql");
#[napi_derive::napi]
impl BackendRuntime {
#[napi]
pub async fn refresh_workspace_admin_stats_dirty(
&self,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsRefreshResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats dirty refresh limit must be positive"));
}
let Some(lease) = self
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog: 0,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.refresh_dirty(batch_limit)
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(result?)
}
#[napi]
pub async fn recalibrate_workspace_admin_stats(
&self,
last_sid: i64,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsRecalibrationResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats recalibration limit must be positive"));
}
let Some(lease) = self
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.recalibrate(last_sid, batch_limit)
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(result?)
}
#[napi]
pub async fn write_workspace_admin_stats_daily_snapshot(
&self,
owner: String,
lease_ttl_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsSnapshotResult> {
let Some(lease) = self
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner, lease_ttl_ms)
.await?
else {
return Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted: 0,
skipped: true,
});
};
let result = async {
WorkspaceStatsStore::new(self.pool().await?)
.write_daily_snapshot()
.await
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(result?)
}
#[napi]
pub async fn recalibrate_workspace_admin_stats_daily(
&self,
batch_limit: i64,
owner: String,
lease_ttl_ms: i64,
lock_retry_times: i64,
lock_retry_delay_ms: i64,
) -> napi::Result<RuntimeWorkspaceStatsDailyRecalibrationResult> {
if batch_limit <= 0 {
return Err(napi_error("workspace stats daily recalibration limit must be positive"));
}
if lock_retry_times <= 0 {
return Err(napi_error(
"workspace stats daily recalibration retry times must be positive",
));
}
if lock_retry_delay_ms < 0 {
return Err(napi_error(
"workspace stats daily recalibration retry delay must be non-negative",
));
}
let Some(lease) = acquire_workspace_stats_lease_with_retry(
self,
owner.clone(),
lease_ttl_ms,
lock_retry_times,
lock_retry_delay_ms,
)
.await?
else {
return Ok(RuntimeWorkspaceStatsDailyRecalibrationResult {
processed: 0,
last_sid: 0,
snapshotted: 0,
skipped: true,
});
};
let result: RuntimeResult<RuntimeWorkspaceStatsDailyRecalibrationResult> = async {
let store = WorkspaceStatsStore::new(self.pool().await?);
let mut processed = 0;
let mut last_sid = 0;
loop {
let batch = retry_workspace_stats_operation(lock_retry_times, lock_retry_delay_ms, || {
store.recalibrate(last_sid, batch_limit)
})
.await?;
if batch.skipped {
return Ok(RuntimeWorkspaceStatsDailyRecalibrationResult {
processed,
last_sid,
snapshotted: 0,
skipped: true,
});
}
if batch.processed == 0 {
break;
}
processed += batch.processed;
last_sid = batch.last_sid;
if batch.processed < batch_limit {
break;
}
}
let snapshot =
retry_workspace_stats_operation(lock_retry_times, lock_retry_delay_ms, || store.write_daily_snapshot()).await?;
Ok(RuntimeWorkspaceStatsDailyRecalibrationResult {
processed,
last_sid,
snapshotted: snapshot.snapshotted,
skipped: snapshot.skipped,
})
}
.await;
release_workspace_stats_lease(self, lease).await?;
Ok(result?)
}
}
#[derive(FromRow)]
struct WorkspaceSid {
id: String,
sid: i32,
}
struct WorkspaceStatsStore {
pool: PgPool,
}
impl WorkspaceStatsStore {
fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn refresh_dirty(&self, batch_limit: i64) -> RuntimeResult<RuntimeWorkspaceStatsRefreshResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh transaction failed", err))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog: 0,
skipped: true,
});
}
let backlog = count_dirty(&mut tx).await?;
let dirty = load_dirty(&mut tx, batch_limit).await?;
if dirty.is_empty() {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRefreshResult {
processed: 0,
backlog,
skipped: false,
});
}
upsert_stats(&mut tx, &dirty).await?;
clear_dirty(&mut tx, &dirty).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats dirty refresh commit failed", err))?;
Ok(RuntimeWorkspaceStatsRefreshResult {
processed: dirty.len() as i64,
backlog,
skipped: false,
})
}
async fn recalibrate(
&self,
last_sid: i64,
batch_limit: i64,
) -> RuntimeResult<RuntimeWorkspaceStatsRecalibrationResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration transaction failed", err))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: true,
});
}
let workspaces = fetch_workspace_batch(&mut tx, last_sid, batch_limit).await?;
if workspaces.is_empty() {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration commit failed", err))?;
return Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: 0,
last_sid,
skipped: false,
});
}
let ids = workspaces
.iter()
.map(|workspace| workspace.id.clone())
.collect::<Vec<_>>();
let next_sid = workspaces
.last()
.map(|workspace| workspace.sid as i64)
.unwrap_or(last_sid);
upsert_stats(&mut tx, &ids).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats recalibration commit failed", err))?;
Ok(RuntimeWorkspaceStatsRecalibrationResult {
processed: ids.len() as i64,
last_sid: next_sid,
skipped: false,
})
}
async fn write_daily_snapshot(&self) -> RuntimeResult<RuntimeWorkspaceStatsSnapshotResult> {
let mut tx = self
.pool
.begin()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot transaction failed", err))?;
if !try_transaction_lock(&mut tx).await? {
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot commit failed", err))?;
return Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted: 0,
skipped: true,
});
}
let snapshotted = write_daily_snapshot(&mut tx).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot commit failed", err))?;
Ok(RuntimeWorkspaceStatsSnapshotResult {
snapshotted,
skipped: false,
})
}
}
async fn release_workspace_stats_lease(runtime: &BackendRuntime, lease: CoordinationLeaseGrant) -> RuntimeResult<()> {
let _ = runtime
.release_coordination_lease_inner(lease.key, lease.owner, lease.fencing_token)
.await?;
Ok(())
}
async fn acquire_workspace_stats_lease_with_retry(
runtime: &BackendRuntime,
owner: String,
lease_ttl_ms: i64,
retry_times: i64,
retry_delay_ms: i64,
) -> RuntimeResult<Option<CoordinationLeaseGrant>> {
for attempt in 0..retry_times {
let lease = runtime
.acquire_coordination_lease_inner(WORKSPACE_STATS_LEASE_KEY.to_string(), owner.clone(), lease_ttl_ms)
.await?;
if lease.is_some() {
return Ok(lease);
}
if attempt < retry_times - 1 && retry_delay_ms > 0 {
sleep(TokioDuration::from_millis(retry_delay_ms as u64)).await;
}
}
Ok(None)
}
async fn retry_workspace_stats_operation<T, F, Fut>(
retry_times: i64,
retry_delay_ms: i64,
mut operation: F,
) -> RuntimeResult<T>
where
T: WorkspaceStatsSkippable,
F: FnMut() -> Fut,
Fut: std::future::Future<Output = RuntimeResult<T>>,
{
for attempt in 0..retry_times {
let result = operation().await?;
if !result.skipped() || attempt == retry_times - 1 {
return Ok(result);
}
if retry_delay_ms > 0 {
sleep(TokioDuration::from_millis(retry_delay_ms as u64)).await;
}
}
unreachable!("workspace stats retry loop validates retry_times > 0")
}
trait WorkspaceStatsSkippable {
fn skipped(&self) -> bool;
}
impl WorkspaceStatsSkippable for RuntimeWorkspaceStatsRecalibrationResult {
fn skipped(&self) -> bool {
self.skipped
}
}
impl WorkspaceStatsSkippable for RuntimeWorkspaceStatsSnapshotResult {
fn skipped(&self) -> bool {
self.skipped
}
}
async fn try_transaction_lock(tx: &mut Transaction<'_, Postgres>) -> RuntimeResult<bool> {
let row = sqlx::query(
r#"
SELECT pg_try_advisory_xact_lock(($1::bigint << 32) + $2::bigint) AS locked
"#,
)
.bind(WORKSPACE_STATS_LOCK_NAMESPACE)
.bind(WORKSPACE_STATS_REFRESH_LOCK_KEY)
.fetch_one(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats transaction lock failed", err))?;
Ok(row.get::<bool, _>("locked"))
}
async fn load_dirty(tx: &mut Transaction<'_, Postgres>, limit: i64) -> RuntimeResult<Vec<String>> {
let rows = sqlx::query(
r#"
SELECT workspace_id
FROM workspace_admin_stats_dirty
ORDER BY updated_at ASC
LIMIT $1
FOR UPDATE SKIP LOCKED
"#,
)
.bind(limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats load dirty workspaces failed", err))?;
Ok(rows.into_iter().map(|row| row.get("workspace_id")).collect())
}
async fn count_dirty(tx: &mut Transaction<'_, Postgres>) -> RuntimeResult<i64> {
let row = sqlx::query("SELECT COUNT(*) AS total FROM workspace_admin_stats_dirty")
.fetch_one(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats count dirty workspaces failed", err))?;
Ok(row.get::<i64, _>("total"))
}
async fn clear_dirty(tx: &mut Transaction<'_, Postgres>, workspace_ids: &[String]) -> RuntimeResult<()> {
sqlx::query(
r#"
DELETE FROM workspace_admin_stats_dirty
WHERE workspace_id = ANY($1::varchar[])
"#,
)
.bind(workspace_ids)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats clear dirty workspaces failed", err))?;
Ok(())
}
async fn upsert_stats(tx: &mut Transaction<'_, Postgres>, workspace_ids: &[String]) -> RuntimeResult<()> {
if workspace_ids.is_empty() {
return Ok(());
}
sqlx::query(UPSERT_WORKSPACE_ADMIN_STATS_SQL)
.bind(workspace_ids)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats upsert stats failed", err))?;
Ok(())
}
async fn fetch_workspace_batch(
tx: &mut Transaction<'_, Postgres>,
last_sid: i64,
limit: i64,
) -> RuntimeResult<Vec<WorkspaceSid>> {
sqlx::query_as::<_, WorkspaceSid>(
r#"
SELECT id, sid
FROM workspaces
WHERE sid > $1
ORDER BY sid
LIMIT $2
"#,
)
.bind(last_sid)
.bind(limit)
.fetch_all(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats fetch workspace batch failed", err))
}
async fn write_daily_snapshot(tx: &mut Transaction<'_, Postgres>) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
INSERT INTO workspace_admin_stats_daily (
workspace_id,
date,
snapshot_size,
blob_size,
member_count,
updated_at
)
SELECT
workspace_id,
CURRENT_DATE,
snapshot_size,
blob_size,
member_count,
NOW()
FROM workspace_admin_stats
ON CONFLICT (workspace_id, date)
DO UPDATE SET
snapshot_size = EXCLUDED.snapshot_size,
blob_size = EXCLUDED.blob_size,
member_count = EXCLUDED.member_count,
updated_at = EXCLUDED.updated_at
"#,
)
.execute(&mut **tx)
.await
.map_err(|err| RuntimeError::database("WorkspaceStats daily snapshot failed", err))?;
Ok(result.rows_affected() as i64)
}

View File

@@ -319,6 +319,7 @@ impl TryFrom<CopilotManagedProfileConfigFile> for CopilotManagedProfileConfig {
#[derive(Clone, Debug)]
pub(crate) struct InviteQuotaConfig {
pub(crate) new_account_action_delay_seconds: i64,
pub(crate) high_risk_target_domains: Vec<String>,
pub(crate) subject_hash_salt: String,
pub(crate) mail_class_mapping: BTreeMap<String, String>,
@@ -327,6 +328,7 @@ pub(crate) struct InviteQuotaConfig {
impl Default for InviteQuotaConfig {
fn default() -> Self {
Self {
new_account_action_delay_seconds: 24 * 60 * 60,
high_risk_target_domains: [
"qq.com",
"proton.me",
@@ -479,12 +481,19 @@ fn deployment_from_env() -> Deployment {
#[derive(Default, Deserialize)]
struct AppConfigFile {
auth: Option<AuthConfigFile>,
db: Option<DbConfigFile>,
crypto: Option<CryptoConfigFile>,
copilot: Option<CopilotRuntimeConfigFile>,
indexer: Option<SearchRuntimeConfigFile>,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct AuthConfigFile {
new_account_action_delay: Option<i64>,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase", default)]
struct SearchRuntimeConfigFile {
@@ -542,7 +551,11 @@ impl AppConfigFile {
}
fn invite_quota_config(&self) -> InviteQuotaConfig {
InviteQuotaConfig::default()
let mut config = InviteQuotaConfig::default();
if let Some(delay) = self.auth.as_ref().and_then(|auth| auth.new_account_action_delay) {
config.new_account_action_delay_seconds = delay.max(0);
}
config
}
}
@@ -907,6 +920,16 @@ mod tests {
let enabled_without_provider: SearchRuntimeConfig = enabled_without_provider.indexer.unwrap().into();
assert!(enabled_without_provider.enabled);
assert_eq!(enabled_without_provider.provider, "embedded");
let manticore = app_config_from_flat_overrides([
("indexer.enabled", serde_json::json!(true)),
("indexer.provider.type", serde_json::json!("manticoresearch")),
("indexer.provider.endpoint", serde_json::json!("http://localhost:9308")),
])
.unwrap();
let manticore: SearchRuntimeConfig = manticore.indexer.unwrap().into();
assert!(manticore.enabled);
assert_eq!(manticore.provider, "manticoresearch");
}
#[test]
@@ -985,14 +1008,16 @@ mod tests {
}
#[test]
fn invite_quota_policy_is_internal_not_app_configurable() {
fn invite_abuse_policy_is_internal_while_action_delay_is_configurable() {
let app_config = app_config_from_flat_overrides([
("auth.newAccountActionDelay", serde_json::json!(123)),
("auth.untrustedPolicyOverride", serde_json::json!("runtime-salt-v2")),
("auth.untrustedDomainList", serde_json::json!(["Example.COM."])),
])
.unwrap();
let config = app_config.invite_quota_config();
assert_eq!(config.new_account_action_delay_seconds, 123);
assert!(!config.high_risk_target_domains.contains(&"example.com".to_string()));
assert_ne!(config.subject_hash_salt, "runtime-salt-v2");
assert_eq!(

View File

@@ -21,15 +21,24 @@ pub(crate) enum RuntimeError {
#[error("search permission state unavailable")]
SearchPermissionUnavailable,
#[error("search index is not ready")]
SearchIndexNotReady,
#[error("search permission projection is syncing")]
SearchPermissionSyncing,
#[error("search index failed: {0}")]
SearchIndexFailed(String),
#[error("search source is invalid: {0}")]
SearchSourceInvalid(String),
#[error("search provider unavailable")]
SearchProviderUnavailable,
#[error("search query is not supported by the active provider")]
SearchUnsupportedQuery,
#[error("search stream replay gap")]
SearchReplayGap,
#[error("{context}: {source}")]
Database {
context: String,
@@ -111,9 +120,11 @@ impl RuntimeError {
}
Self::SearchWorkspaceDenied
| Self::SearchPermissionUnavailable
| Self::SearchIndexNotReady
| Self::SearchPermissionSyncing
| Self::SearchProviderUnavailable
| Self::SearchUnsupportedQuery
| Self::SearchReplayGap => false,
| Self::SearchUnsupportedQuery => false,
Self::SearchIndexFailed(_) | Self::SearchSourceInvalid(_) => false,
_ => false,
}
}
@@ -127,6 +138,10 @@ impl RuntimeError {
} if source.code().as_deref() == Some("40001")
)
}
pub(crate) fn is_permanent_search_source(&self) -> bool {
matches!(self, Self::SearchSourceInvalid(_))
}
}
pub(crate) fn to_napi_error(error: RuntimeError) -> Error {

View File

@@ -0,0 +1,13 @@
use rustls::{ClientConfig, RootCertStore};
pub(in crate::runtime) fn webpki_tls_config() -> Result<ClientConfig, rustls::Error> {
let roots = RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
};
Ok(
ClientConfig::builder_with_provider(rustls::crypto::aws_lc_rs::default_provider().into())
.with_safe_default_protocol_versions()?
.with_root_certificates(roots)
.with_no_client_auth(),
)
}

View File

@@ -5,8 +5,7 @@ use super::{RuntimeError, RuntimeResult, types::EmbeddingHealth};
pub(crate) const RUNTIME_MIGRATIONS: &str = include_str!("sql/runtime_migrations.sql");
const EMBEDDING_MIGRATION: &str = include_str!("sql/embedding.sql");
const SEARCH_MIGRATION: &str = include_str!("sql/search.sql");
const SEARCH_ACL_TOKEN_MIGRATION: &str = include_str!("sql/search_acl_tokens.sql");
const SEARCH_PROJECTION_MIGRATION: &str = include_str!("sql/search_projection.sql");
const EMBEDDING_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0046;
const SEARCH_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0053;
#[cfg(test)]
@@ -70,15 +69,11 @@ pub(crate) async fn embedding_schema_health(pool: &PgPool) -> RuntimeResult<Embe
pub(crate) async fn migrate_search_tables(pool: &PgPool) -> RuntimeResult<()> {
migrate_component(
pool,
"search",
"search_projection",
SEARCH_ADVISORY_LOCK,
&[(1, &[SEARCH_MIGRATION]), (2, &[SEARCH_ACL_TOKEN_MIGRATION])],
&[(1, &[SEARCH_PROJECTION_MIGRATION])],
)
.await?;
sqlx::query("INSERT INTO search_runtime_streams(table_key) VALUES ('doc'), ('block') ON CONFLICT DO NOTHING")
.execute(pool)
.await
.map_err(|error| RuntimeError::database("Repair search runtime streams", error))?;
Ok(())
}
@@ -118,6 +113,10 @@ async fn migrate_component(
.begin()
.await
.map_err(|error| RuntimeError::database("Native migration transaction failed", error))?;
transaction
.execute("SET LOCAL lock_timeout = '5s'")
.await
.map_err(|error| RuntimeError::database("Native migration lock timeout failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(advisory_lock)
.execute(&mut *transaction)
@@ -269,11 +268,31 @@ mod tests {
}
#[test]
fn search_schema_uses_transactional_stream_heads() {
assert!(SEARCH_MIGRATION.contains("CREATE TABLE search_runtime_streams"));
assert!(SEARCH_MIGRATION.contains("PRIMARY KEY (table_key, stream_sequence)"));
assert!(!SEARCH_MIGRATION.contains("BIGSERIAL"));
assert!(!SEARCH_MIGRATION.contains("CREATE SEQUENCE"));
fn search_schema_uses_terminal_control_plane() {
assert!(SEARCH_PROJECTION_MIGRATION.contains("DROP SCHEMA IF EXISTS search_projection CASCADE"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SCHEMA search_projection"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.generations"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.workspace_states"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.document_states"));
assert!(
SEARCH_PROJECTION_MIGRATION.contains("CREATE SEQUENCE IF NOT EXISTS search_projection.source_mutation_version")
);
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SEQUENCE IF NOT EXISTS search_projection.permission_version"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SEQUENCE IF NOT EXISTS search_projection.claim_fence"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TRIGGER search_projection_snapshot_capture"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TRIGGER search_projection_membership_capture"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("ON CONFLICT (generation_id, workspace_id, doc_id) DO UPDATE"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("WHERE state IN ('building', 'active')"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("target_source_version <> published_source_version"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("gc_table TEXT NOT NULL DEFAULT 'doc'"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("gc_cursor TEXT"));
assert!(
SEARCH_PROJECTION_MIGRATION
.contains("NEW.state = 'active' AND (TG_OP = 'INSERT' OR OLD.state IS DISTINCT FROM NEW.state)")
);
assert!(!SEARCH_PROJECTION_MIGRATION.contains("clock_timestamp()"));
assert!(!SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_runtime_projections"));
assert!(!SEARCH_PROJECTION_MIGRATION.contains("payload JSONB NOT NULL"));
}
#[tokio::test]
@@ -285,11 +304,367 @@ mod tests {
let (first, second) = tokio::join!(migrate_search_tables(&pool), migrate_search_tables(&pool));
first.unwrap();
second.unwrap();
let versions: Vec<i32> =
sqlx::query_scalar("SELECT version FROM native_schema_migrations WHERE component='search' ORDER BY version")
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(versions, vec![1, 2]);
let versions: Vec<i32> = sqlx::query_scalar(
"SELECT version FROM native_schema_migrations WHERE component='search_projection' ORDER BY version",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(versions, vec![1]);
}
#[tokio::test]
async fn permission_triggers_route_and_coalesce_document_and_workspace_scopes() {
let _guard = crate::runtime::backend_runtime::SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await
.unwrap();
let suffix = uuid::Uuid::new_v4().simple().to_string();
let generation_id = uuid::Uuid::new_v4();
let workspace_id = format!("trigger-workspace-{suffix}");
let doc_id = format!("trigger-doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO search_projection.workspace_states(generation_id,workspace_id,covered,pending_scope) \
VALUES($1,$2,true,'none')",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('00','hex'),now())")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO doc_access_policies(workspace_id,doc_id,visibility,member_default_role) \
VALUES($1,$2,'private','none')",
)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("UPDATE doc_access_policies SET member_default_role='reader' WHERE workspace_id=$1 AND doc_id=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let document_scope: (String, i64) = sqlx::query_as(
r#"SELECT state.pending_scope,count(task.doc_id)
FROM search_projection.workspace_states state
LEFT JOIN search_projection.document_states task
ON task.generation_id=state.generation_id AND task.workspace_id=state.workspace_id
WHERE state.generation_id=$1 AND state.workspace_id=$2
GROUP BY state.pending_scope"#,
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(document_scope, ("none".to_string(), 1));
let moved_workspace_id = format!("trigger-moved-workspace-{suffix}");
let moved_doc_id = format!("trigger-moved-doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&moved_workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('00','hex'),now())")
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.execute(&pool)
.await
.unwrap();
let old_version: i64 = sqlx::query_scalar(
"SELECT target_permission_version FROM search_projection.document_states WHERE generation_id=$1 AND \
workspace_id=$2 AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
let moved_version: i64 = sqlx::query_scalar(
"SELECT target_permission_version FROM search_projection.document_states WHERE generation_id=$1 AND \
workspace_id=$2 AND doc_id=$3",
)
.bind(generation_id)
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.fetch_one(&pool)
.await
.unwrap();
sqlx::query("UPDATE doc_access_policies SET workspace_id=$3,doc_id=$4 WHERE workspace_id=$1 AND doc_id=$2")
.bind(&workspace_id)
.bind(&doc_id)
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.execute(&pool)
.await
.unwrap();
let versions: (i64, i64) = sqlx::query_as(
r#"SELECT
(SELECT target_permission_version FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3),
(SELECT target_permission_version FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$4 AND doc_id=$5)"#,
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(versions.0 > old_version);
assert!(versions.1 > moved_version);
sqlx::query("INSERT INTO workspace_access_policies(workspace_id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let pending_scope: String = sqlx::query_scalar(
"SELECT pending_scope FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(pending_scope, "workspace");
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&moved_workspace_id)
.execute(&pool)
.await
.unwrap();
}
#[tokio::test]
async fn source_mutations_use_unique_versions_and_rollback_only_leaves_a_sequence_gap() {
let _guard = crate::runtime::backend_runtime::SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await
.unwrap();
let suffix = uuid::Uuid::new_v4().simple().to_string();
let generation_id = uuid::Uuid::new_v4();
let workspace_id = format!("source-version-workspace-{suffix}");
let doc_id = format!("source-version-doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO search_projection.workspace_states(generation_id,workspace_id,pending_scope) VALUES($1,$2,'none')",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('00','hex'),'2026-01-01 UTC')",
)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let inserted: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
sqlx::query("UPDATE snapshots SET blob=decode('01','hex') WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let updated: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(updated > inserted);
sqlx::query("DELETE FROM snapshots WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let (deleted, target_exists, published_exists): (i64, bool, bool) = sqlx::query_as(
"SELECT target_source_version,target_source_exists,published_source_exists FROM \
search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(deleted > updated);
assert!(!target_exists);
assert!(!published_exists);
sqlx::query(
"INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('02','hex'),'2026-01-01 UTC')",
)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let recreated: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(recreated > deleted);
let mut transaction = pool.begin().await.unwrap();
sqlx::query("UPDATE snapshots SET blob=decode('03','hex') WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&mut *transaction)
.await
.unwrap();
let rolled_back: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&mut *transaction)
.await
.unwrap();
transaction.rollback().await.unwrap();
let after_rollback: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(after_rollback, recreated);
assert!(rolled_back > recreated);
sqlx::query("UPDATE snapshots SET blob=decode('04','hex') WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let after_gap: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(after_gap > rolled_back);
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let (remaining_states, delete_gc_scheduled): (i64, bool) = sqlx::query_as(
r#"SELECT count(*),bool_and(progress->>'kind'='deleted')
FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(remaining_states, 1);
assert!(delete_gc_scheduled);
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
}
}

View File

@@ -4,6 +4,7 @@ pub mod storage_runtime;
pub(crate) mod config;
mod config_descriptor;
pub(crate) mod error;
mod http;
pub(crate) mod migrations;
pub(crate) mod object_storage;
pub(crate) mod types;
@@ -15,3 +16,4 @@ pub(crate) use config::{
use config::{SUPPORTED_BYOK_PROVIDERS, validate_copilot_config};
pub use config_descriptor::{AppConfigDescriptor, app_config_descriptors, validate_app_config_value};
pub(crate) use error::{RuntimeError, RuntimeResult, napi_error, to_napi_error};
pub(in crate::runtime) use http::webpki_tls_config;

View File

@@ -8,7 +8,6 @@ use reqwest::{
Client as ReqwestClient, Method, StatusCode,
header::{CONTENT_LENGTH, CONTENT_TYPE, ETAG, HeaderMap, HeaderName, HeaderValue, LAST_MODIFIED},
};
use rustls::RootCertStore;
use rusty_s3::{
Bucket, Credentials,
actions::{
@@ -27,6 +26,7 @@ use super::{
ObjectListPage, ObjectMetadata, ObjectPrefix, ObjectPutMetadata, PresignedObjectRequest, completed_multipart_parts,
trim_etag,
},
webpki_tls_config,
};
const DEFAULT_REQUEST_TIMEOUT_MS: u64 = 30_000;
@@ -60,7 +60,10 @@ struct ReqwestStorageHttpClient {
impl ReqwestStorageHttpClient {
fn new(request_timeout_ms: Option<u64>) -> ObjectStorageResult<Self> {
let builder = ReqwestClient::builder()
.tls_backend_preconfigured(Self::webpki_tls_config()?)
.tls_backend_preconfigured(
webpki_tls_config()
.map_err(|err| ObjectStorageError::Config(format!("ObjectStorage TLS config failed: {err}")))?,
)
.timeout(Duration::from_millis(
request_timeout_ms.unwrap_or(DEFAULT_REQUEST_TIMEOUT_MS),
));
@@ -69,19 +72,6 @@ impl ReqwestStorageHttpClient {
})
}
fn webpki_tls_config() -> ObjectStorageResult<rustls::ClientConfig> {
let roots = RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
};
Ok(
rustls::ClientConfig::builder_with_provider(rustls::crypto::aws_lc_rs::default_provider().into())
.with_safe_default_protocol_versions()
.map_err(|err| ObjectStorageError::Config(format!("ObjectStorage TLS config failed: {err}")))?
.with_root_certificates(roots)
.with_no_client_auth(),
)
}
async fn execute(&self, request: StorageHttpRequest) -> ObjectStorageResult<StorageHttpResponse> {
let mut builder = self.client.request(request.method, request.url);
for (key, value) in request.headers {

View File

@@ -14,4 +14,6 @@ pub(in crate::runtime) use backend::{FsStorageConfig, StorageBackendConfig};
pub(in crate::runtime) use config::ObjectStorageConfig;
pub(crate) use service::ObjectStorageService;
pub(super) use super::webpki_tls_config;
pub(in crate::runtime) const MAX_BLOB_SIZE: i64 = i32::MAX as i64;

View File

@@ -1,252 +0,0 @@
CREATE TABLE search_runtime_streams (
table_key TEXT PRIMARY KEY CHECK (table_key IN ('doc', 'block')),
head BIGINT NOT NULL DEFAULT 0 CHECK (head >= 0),
retained_from BIGINT NOT NULL DEFAULT 0 CHECK (retained_from >= 0),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CHECK (retained_from <= head)
);
INSERT INTO search_runtime_streams(table_key)
VALUES ('doc'), ('block')
ON CONFLICT DO NOTHING;
CREATE TABLE search_runtime_projections (
table_key TEXT NOT NULL CHECK (table_key IN ('doc', 'block')),
external_id TEXT NOT NULL,
workspace_id VARCHAR NOT NULL,
doc_id VARCHAR NOT NULL,
revision BIGINT NOT NULL CHECK (revision >= 0),
payload JSONB NOT NULL,
acl_public_readable BOOLEAN NOT NULL DEFAULT false,
acl_member_default_readable BOOLEAN NOT NULL DEFAULT false,
acl_read_user_ids TEXT[] NOT NULL DEFAULT '{}',
acl_revision BIGINT NOT NULL DEFAULT 0 CHECK (acl_revision >= 0),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (table_key, external_id)
);
CREATE INDEX search_runtime_projections_workspace_doc
ON search_runtime_projections(workspace_id, doc_id, table_key);
CREATE TABLE search_runtime_changes (
table_key TEXT NOT NULL CHECK (table_key IN ('doc', 'block')),
stream_sequence BIGINT NOT NULL CHECK (stream_sequence > 0),
external_id TEXT NOT NULL,
workspace_id VARCHAR NOT NULL,
doc_id VARCHAR,
revision BIGINT NOT NULL CHECK (revision >= 0),
operation TEXT NOT NULL CHECK (operation IN ('upsert', 'delete', 'invalidate')),
payload JSONB,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (table_key, stream_sequence)
);
CREATE INDEX search_runtime_changes_workspace
ON search_runtime_changes(workspace_id, table_key, stream_sequence);
CREATE TABLE search_runtime_generations (
generation_id UUID PRIMARY KEY,
provider TEXT NOT NULL CHECK (provider IN ('embedded', 'elasticsearch', 'manticoresearch')),
state TEXT NOT NULL CHECK (state IN ('pending', 'active', 'draining', 'failed')),
config_fingerprint TEXT NOT NULL,
schema_fingerprint TEXT NOT NULL,
manifest JSONB NOT NULL DEFAULT '{}',
applied_permission_revision BIGINT NOT NULL DEFAULT 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
activated_at TIMESTAMPTZ
);
CREATE INDEX search_runtime_generations_fingerprint
ON search_runtime_generations(provider, config_fingerprint, schema_fingerprint);
CREATE UNIQUE INDEX search_runtime_single_active_generation
ON search_runtime_generations ((state)) WHERE state = 'active';
CREATE UNIQUE INDEX search_runtime_single_pending_generation
ON search_runtime_generations ((state)) WHERE state = 'pending';
CREATE TABLE search_runtime_provider_cursors (
generation_id UUID NOT NULL REFERENCES search_runtime_generations(generation_id) ON DELETE CASCADE,
table_key TEXT NOT NULL CHECK (table_key IN ('doc', 'block')),
source_cursor BIGINT NOT NULL DEFAULT 0 CHECK (source_cursor >= 0),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (generation_id, table_key)
);
CREATE TABLE search_runtime_permission_cursors (
generation_id UUID NOT NULL REFERENCES search_runtime_generations(generation_id) ON DELETE CASCADE,
workspace_id VARCHAR NOT NULL,
permission_revision BIGINT NOT NULL DEFAULT 0 CHECK (permission_revision >= 0),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (generation_id, workspace_id)
);
CREATE TABLE search_runtime_checkpoints (
table_key TEXT PRIMARY KEY CHECK (table_key IN ('doc', 'block')),
schema_fingerprint TEXT NOT NULL,
source_cursor BIGINT NOT NULL CHECK (source_cursor >= 0),
checkpoint_sequence BIGINT NOT NULL CHECK (checkpoint_sequence >= 0),
checkpoint_blob BYTEA NOT NULL,
checksum TEXT NOT NULL,
blob_size BIGINT NOT NULL CHECK (blob_size >= 0),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE search_runtime_checkpoints ALTER COLUMN checkpoint_blob SET STORAGE EXTERNAL;
CREATE TABLE workspace_permission_revisions (
workspace_id VARCHAR PRIMARY KEY REFERENCES workspaces(id) ON DELETE CASCADE ON UPDATE CASCADE,
revision BIGINT NOT NULL DEFAULT 0,
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE TABLE workspace_permission_changes (
workspace_id VARCHAR NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE ON UPDATE CASCADE,
revision BIGINT NOT NULL,
doc_id VARCHAR,
scope TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (workspace_id, revision)
);
CREATE INDEX workspace_permission_changes_created_at
ON workspace_permission_changes(created_at);
INSERT INTO workspace_permission_revisions(workspace_id, revision)
SELECT id, 0 FROM workspaces;
CREATE FUNCTION record_workspace_permission_change(
target_workspace_id VARCHAR,
target_doc_id VARCHAR,
target_scope TEXT
)
RETURNS void
LANGUAGE plpgsql
AS $$
DECLARE
next_revision BIGINT;
BEGIN
IF target_workspace_id IS NULL OR
NOT EXISTS (SELECT 1 FROM workspaces WHERE id = target_workspace_id) THEN
RETURN;
END IF;
INSERT INTO workspace_permission_revisions(workspace_id, revision)
VALUES (target_workspace_id, 1)
ON CONFLICT (workspace_id) DO UPDATE
SET revision = workspace_permission_revisions.revision + 1,
updated_at = now()
RETURNING revision INTO next_revision;
INSERT INTO workspace_permission_changes(workspace_id, revision, doc_id, scope)
VALUES (target_workspace_id, next_revision, target_doc_id, target_scope);
END;
$$;
CREATE FUNCTION initialize_workspace_permission_revision()
RETURNS trigger
LANGUAGE plpgsql
AS $$
BEGIN
INSERT INTO workspace_permission_revisions(workspace_id, revision) VALUES (NEW.id, 0);
RETURN NEW;
END;
$$;
CREATE FUNCTION bump_workspace_permission_revision()
RETURNS trigger
LANGUAGE plpgsql
AS $$
DECLARE
old_workspace_id VARCHAR;
new_workspace_id VARCHAR;
old_doc_id VARCHAR;
new_doc_id VARCHAR;
target_scope TEXT;
BEGIN
IF TG_TABLE_NAME = 'entitlements' THEN
old_workspace_id := CASE WHEN TG_OP <> 'INSERT' AND OLD.target_type = 'workspace' THEN OLD.target_id END;
new_workspace_id := CASE WHEN TG_OP <> 'DELETE' AND NEW.target_type = 'workspace' THEN NEW.target_id END;
target_scope := 'capability';
ELSIF TG_TABLE_NAME = 'workspace_members' THEN
old_workspace_id := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace_id := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
target_scope := 'membership';
ELSIF TG_TABLE_NAME = 'workspace_access_policies' THEN
old_workspace_id := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace_id := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
target_scope := 'workspace_policy';
ELSIF TG_TABLE_NAME = 'doc_access_policies' THEN
old_workspace_id := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace_id := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
old_doc_id := CASE WHEN TG_OP <> 'INSERT' THEN OLD.doc_id END;
new_doc_id := CASE WHEN TG_OP <> 'DELETE' THEN NEW.doc_id END;
target_scope := 'doc_policy';
ELSIF TG_TABLE_NAME = 'doc_grants' THEN
old_workspace_id := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace_id := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
old_doc_id := CASE WHEN TG_OP <> 'INSERT' THEN OLD.doc_id END;
new_doc_id := CASE WHEN TG_OP <> 'DELETE' THEN NEW.doc_id END;
target_scope := 'doc_grant';
END IF;
IF old_workspace_id IS NOT NULL AND old_workspace_id IS DISTINCT FROM new_workspace_id THEN
PERFORM record_workspace_permission_change(old_workspace_id, old_doc_id, target_scope);
END IF;
IF new_workspace_id IS NOT NULL THEN
PERFORM record_workspace_permission_change(new_workspace_id, new_doc_id, target_scope);
ELSIF old_workspace_id IS NOT NULL THEN
PERFORM record_workspace_permission_change(old_workspace_id, old_doc_id, target_scope);
END IF;
RETURN CASE WHEN TG_OP = 'DELETE' THEN OLD ELSE NEW END;
END;
$$;
CREATE TRIGGER workspaces_initialize_permission_revision
AFTER INSERT ON workspaces
FOR EACH ROW EXECUTE FUNCTION initialize_workspace_permission_revision();
CREATE TRIGGER workspace_members_permission_revision_mutation
AFTER INSERT OR DELETE ON workspace_members
FOR EACH ROW EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER workspace_members_permission_revision_update
AFTER UPDATE OF workspace_id, user_id, role, state, source ON workspace_members
FOR EACH ROW WHEN (ROW(OLD.workspace_id, OLD.user_id, OLD.role, OLD.state, OLD.source)
IS DISTINCT FROM ROW(NEW.workspace_id, NEW.user_id, NEW.role, NEW.state, NEW.source))
EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER workspace_access_policies_permission_revision_mutation
AFTER INSERT OR DELETE ON workspace_access_policies
FOR EACH ROW EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER workspace_access_policies_permission_revision_update
AFTER UPDATE OF workspace_id, visibility, sharing_enabled, member_default_doc_role ON workspace_access_policies
FOR EACH ROW WHEN (ROW(OLD.workspace_id, OLD.visibility, OLD.sharing_enabled, OLD.member_default_doc_role)
IS DISTINCT FROM ROW(NEW.workspace_id, NEW.visibility, NEW.sharing_enabled, NEW.member_default_doc_role))
EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER doc_access_policies_permission_revision_mutation
AFTER INSERT OR DELETE ON doc_access_policies
FOR EACH ROW EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER doc_access_policies_permission_revision_update
AFTER UPDATE OF workspace_id, doc_id, visibility, public_role, member_default_role ON doc_access_policies
FOR EACH ROW WHEN (ROW(OLD.workspace_id, OLD.doc_id, OLD.visibility, OLD.public_role, OLD.member_default_role)
IS DISTINCT FROM ROW(NEW.workspace_id, NEW.doc_id, NEW.visibility, NEW.public_role, NEW.member_default_role))
EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER doc_grants_permission_revision_mutation
AFTER INSERT OR DELETE ON doc_grants
FOR EACH ROW EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER doc_grants_permission_revision_update
AFTER UPDATE OF workspace_id, doc_id, principal_type, principal_id, role ON doc_grants
FOR EACH ROW WHEN (ROW(OLD.workspace_id, OLD.doc_id, OLD.principal_type, OLD.principal_id, OLD.role)
IS DISTINCT FROM ROW(NEW.workspace_id, NEW.doc_id, NEW.principal_type, NEW.principal_id, NEW.role))
EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER entitlements_permission_revision_mutation
AFTER INSERT OR DELETE ON entitlements
FOR EACH ROW EXECUTE FUNCTION bump_workspace_permission_revision();
CREATE TRIGGER entitlements_permission_revision_update
AFTER UPDATE OF target_type, target_id, source, plan, status, signed_payload, validated_at, expires_at, grace_until ON entitlements
FOR EACH ROW WHEN (ROW(OLD.target_type, OLD.target_id, OLD.source, OLD.plan, OLD.status, OLD.signed_payload, OLD.validated_at, OLD.expires_at, OLD.grace_until)
IS DISTINCT FROM ROW(NEW.target_type, NEW.target_id, NEW.source, NEW.plan, NEW.status, NEW.signed_payload, NEW.validated_at, NEW.expires_at, NEW.grace_until))
EXECUTE FUNCTION bump_workspace_permission_revision();

View File

@@ -1,5 +0,0 @@
CREATE TABLE search_runtime_acl_tokens (
token TEXT PRIMARY KEY,
token_id BIGINT GENERATED ALWAYS AS IDENTITY UNIQUE CHECK (token_id > 0),
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);

View File

@@ -0,0 +1,346 @@
SET LOCAL lock_timeout = '5s';
DROP TRIGGER IF EXISTS workspaces_initialize_permission_revision ON workspaces;
DROP TRIGGER IF EXISTS workspace_members_permission_revision_mutation ON workspace_members;
DROP TRIGGER IF EXISTS workspace_members_permission_revision_update ON workspace_members;
DROP TRIGGER IF EXISTS workspace_access_policies_permission_revision_mutation ON workspace_access_policies;
DROP TRIGGER IF EXISTS workspace_access_policies_permission_revision_update ON workspace_access_policies;
DROP TRIGGER IF EXISTS doc_access_policies_permission_revision_mutation ON doc_access_policies;
DROP TRIGGER IF EXISTS doc_access_policies_permission_revision_update ON doc_access_policies;
DROP TRIGGER IF EXISTS doc_grants_permission_revision_mutation ON doc_grants;
DROP TRIGGER IF EXISTS doc_grants_permission_revision_update ON doc_grants;
DROP TRIGGER IF EXISTS entitlements_permission_revision_mutation ON entitlements;
DROP TRIGGER IF EXISTS entitlements_permission_revision_update ON entitlements;
DROP TRIGGER IF EXISTS search_projection_workspace_delete_capture ON workspaces;
DROP FUNCTION IF EXISTS initialize_workspace_permission_revision();
DROP FUNCTION IF EXISTS bump_workspace_permission_revision();
DROP FUNCTION IF EXISTS record_workspace_permission_change(VARCHAR, VARCHAR, TEXT);
DROP TABLE IF EXISTS search_runtime_permission_cursors;
DROP TABLE IF EXISTS search_runtime_provider_cursors;
DROP TABLE IF EXISTS search_runtime_checkpoints;
DROP TABLE IF EXISTS search_runtime_changes;
DROP TABLE IF EXISTS search_runtime_projections;
DROP TABLE IF EXISTS search_runtime_streams;
DROP TABLE IF EXISTS search_runtime_generations;
DROP TABLE IF EXISTS search_runtime_acl_tokens;
DROP TABLE IF EXISTS workspace_permission_changes;
DROP TABLE IF EXISTS workspace_permission_revisions;
DROP SCHEMA IF EXISTS search_projection CASCADE;
CREATE SCHEMA search_projection;
CREATE SEQUENCE IF NOT EXISTS search_projection.source_mutation_version;
CREATE SEQUENCE IF NOT EXISTS search_projection.permission_version;
CREATE SEQUENCE IF NOT EXISTS search_projection.claim_fence;
CREATE TABLE search_projection.generations (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
provider TEXT NOT NULL CHECK (provider IN ('embedded', 'elasticsearch', 'manticoresearch')),
state TEXT NOT NULL CHECK (state IN ('building', 'active', 'draining', 'failed')),
config_hash BYTEA NOT NULL CHECK (octet_length(config_hash) = 32),
schema_version INTEGER NOT NULL CHECK (schema_version > 0),
manifest JSONB NOT NULL DEFAULT '{}',
scan_high_water_sid INTEGER,
scan_cursor_sid INTEGER,
gc_table TEXT NOT NULL DEFAULT 'doc' CHECK (gc_table IN ('doc', 'block')),
gc_cursor TEXT,
last_error TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
activated_at TIMESTAMPTZ,
drained_at TIMESTAMPTZ,
CHECK (scan_cursor_sid IS NULL OR scan_high_water_sid IS NULL OR scan_cursor_sid <= scan_high_water_sid)
);
CREATE UNIQUE INDEX generations_one_active
ON search_projection.generations ((state)) WHERE state = 'active';
CREATE UNIQUE INDEX generations_one_candidate
ON search_projection.generations ((1)) WHERE state = 'building';
CREATE TABLE search_projection.workspace_states (
generation_id UUID NOT NULL REFERENCES search_projection.generations(id) ON DELETE CASCADE,
workspace_id TEXT NOT NULL,
covered BOOLEAN NOT NULL DEFAULT false,
target_root_revision BIGINT NOT NULL DEFAULT 0,
applied_root_revision BIGINT NOT NULL DEFAULT 0,
required_permission_version BIGINT NOT NULL DEFAULT 0,
applied_permission_version BIGINT NOT NULL DEFAULT 0,
pending_scope TEXT NOT NULL DEFAULT 'workspace' CHECK (pending_scope IN ('none', 'permission', 'workspace')),
progress JSONB,
available_at TIMESTAMPTZ NOT NULL DEFAULT now(),
claim_fence BIGINT,
lease_owner TEXT,
lease_expires_at TIMESTAMPTZ,
attempt_count INTEGER NOT NULL DEFAULT 0 CHECK (attempt_count >= 0),
last_error TEXT,
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (generation_id, workspace_id),
CHECK (applied_root_revision <= target_root_revision),
CHECK (applied_permission_version <= required_permission_version),
CHECK (claim_fence IS NULL OR claim_fence > 0)
);
-- Keep pending work and time-based anti-entropy in one bounded schedule index;
-- PostgreSQL partial predicates cannot depend on now().
CREATE INDEX workspace_states_schedule
ON search_projection.workspace_states (generation_id, available_at, workspace_id);
CREATE TABLE search_projection.document_states (
generation_id UUID NOT NULL REFERENCES search_projection.generations(id) ON DELETE CASCADE,
workspace_id TEXT NOT NULL,
doc_id TEXT NOT NULL,
target_source_version BIGINT NOT NULL,
target_source_exists BOOLEAN NOT NULL,
target_permission_version BIGINT NOT NULL DEFAULT 0,
published_source_version BIGINT NOT NULL DEFAULT 0,
published_source_exists BOOLEAN NOT NULL DEFAULT false,
published_permission_version BIGINT NOT NULL DEFAULT 0,
available_at TIMESTAMPTZ NOT NULL DEFAULT now(),
claim_fence BIGINT,
lease_owner TEXT,
lease_expires_at TIMESTAMPTZ,
attempt_count INTEGER NOT NULL DEFAULT 0 CHECK (attempt_count >= 0),
last_error TEXT,
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (generation_id, workspace_id, doc_id),
FOREIGN KEY (generation_id, workspace_id)
REFERENCES search_projection.workspace_states(generation_id, workspace_id) ON DELETE CASCADE,
CHECK (target_source_version >= 0),
CHECK (published_source_version >= 0 AND published_source_version <= target_source_version),
CHECK (published_permission_version >= 0 AND published_permission_version <= target_permission_version),
CHECK (NOT target_source_exists OR target_source_version > 0),
CHECK (NOT published_source_exists OR published_source_version > 0),
CHECK (claim_fence IS NULL OR claim_fence > 0)
);
CREATE INDEX document_states_pending_schedule
ON search_projection.document_states (generation_id, available_at, workspace_id, doc_id)
WHERE target_source_version <> published_source_version
OR target_source_exists <> published_source_exists
OR target_permission_version <> published_permission_version;
CREATE OR REPLACE FUNCTION search_projection.ensure_workspace_state(target_generation UUID, target_workspace TEXT)
RETURNS void LANGUAGE SQL AS $$
INSERT INTO search_projection.workspace_states(generation_id, workspace_id)
VALUES (target_generation, target_workspace)
ON CONFLICT (generation_id, workspace_id) DO NOTHING
$$;
CREATE OR REPLACE FUNCTION search_projection.capture_snapshot_mutation()
RETURNS trigger LANGUAGE plpgsql AS $$
DECLARE
target_workspace TEXT := COALESCE(NEW.workspace_id, OLD.workspace_id);
target_doc TEXT := COALESCE(NEW.guid, OLD.guid);
target_version BIGINT;
candidate RECORD;
BEGIN
PERFORM pg_advisory_xact_lock_shared(hashtextextended('search-projection-generation', 0));
target_version := nextval('search_projection.source_mutation_version');
FOR candidate IN
SELECT id, state FROM search_projection.generations
WHERE state IN ('building', 'active')
LOOP
PERFORM search_projection.ensure_workspace_state(candidate.id, target_workspace);
IF target_doc = target_workspace THEN
UPDATE search_projection.workspace_states
SET target_root_revision = GREATEST(target_root_revision, target_version),
pending_scope = 'workspace', progress = NULL,
claim_fence = nextval('search_projection.claim_fence'),
lease_owner = NULL, lease_expires_at = NULL,
last_error = NULL, available_at = now(), updated_at = now()
WHERE generation_id = candidate.id AND workspace_id = target_workspace;
ELSE
UPDATE search_projection.workspace_states
SET progress = NULL, claim_fence = nextval('search_projection.claim_fence'),
lease_owner = NULL, lease_expires_at = NULL,
last_error = NULL, available_at = now(), updated_at = now()
WHERE generation_id = candidate.id AND workspace_id = target_workspace;
INSERT INTO search_projection.document_states(
generation_id, workspace_id, doc_id, target_source_version,
target_source_exists, target_permission_version
)
SELECT candidate.id, target_workspace, target_doc, target_version, TG_OP <> 'DELETE',
state.required_permission_version
FROM search_projection.workspace_states state
WHERE state.generation_id = candidate.id AND state.workspace_id = target_workspace
ON CONFLICT (generation_id, workspace_id, doc_id) DO UPDATE
SET target_source_version = EXCLUDED.target_source_version,
target_source_exists = EXCLUDED.target_source_exists,
target_permission_version = GREATEST(search_projection.document_states.target_permission_version, EXCLUDED.target_permission_version),
published_source_exists = CASE WHEN EXCLUDED.target_source_exists THEN search_projection.document_states.published_source_exists ELSE false END,
claim_fence = NULL, lease_owner = NULL, lease_expires_at = NULL,
last_error = NULL, available_at = now(), updated_at = now();
END IF;
END LOOP;
RETURN COALESCE(NEW, OLD);
END;
$$;
CREATE OR REPLACE FUNCTION search_projection.capture_permission_mutation()
RETURNS trigger LANGUAGE plpgsql AS $$
DECLARE
old_workspace TEXT;
new_workspace TEXT;
old_doc TEXT;
new_doc TEXT;
target_scope TEXT;
version BIGINT;
candidate RECORD;
target RECORD;
BEGIN
PERFORM pg_advisory_xact_lock_shared(hashtextextended('search-projection-generation', 0));
version := nextval('search_projection.permission_version');
IF TG_TABLE_NAME = 'entitlements' THEN
old_workspace := CASE WHEN TG_OP <> 'INSERT' AND OLD.target_type = 'workspace' THEN OLD.target_id END;
new_workspace := CASE WHEN TG_OP <> 'DELETE' AND NEW.target_type = 'workspace' THEN NEW.target_id END;
target_scope := 'capability';
ELSIF TG_TABLE_NAME = 'workspace_members' THEN
old_workspace := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
target_scope := 'membership';
ELSIF TG_TABLE_NAME = 'workspace_access_policies' THEN
old_workspace := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
target_scope := 'workspace_policy';
ELSIF TG_TABLE_NAME = 'doc_access_policies' THEN
old_workspace := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
old_doc := CASE WHEN TG_OP <> 'INSERT' THEN OLD.doc_id END;
new_doc := CASE WHEN TG_OP <> 'DELETE' THEN NEW.doc_id END;
target_scope := 'doc_policy';
ELSIF TG_TABLE_NAME = 'doc_grants' THEN
old_workspace := CASE WHEN TG_OP <> 'INSERT' THEN OLD.workspace_id END;
new_workspace := CASE WHEN TG_OP <> 'DELETE' THEN NEW.workspace_id END;
old_doc := CASE WHEN TG_OP <> 'INSERT' THEN OLD.doc_id END;
new_doc := CASE WHEN TG_OP <> 'DELETE' THEN NEW.doc_id END;
target_scope := 'doc_grant';
END IF;
FOR candidate IN
SELECT id, state FROM search_projection.generations
WHERE state IN ('building', 'active')
LOOP
FOR target IN
SELECT DISTINCT mutation.workspace_id, mutation.doc_id
FROM (VALUES (old_workspace, old_doc), (new_workspace, new_doc)) mutation(workspace_id, doc_id)
WHERE mutation.workspace_id IS NOT NULL
LOOP
PERFORM search_projection.ensure_workspace_state(candidate.id, target.workspace_id);
UPDATE search_projection.workspace_states
SET required_permission_version = GREATEST(required_permission_version, version),
pending_scope = CASE
WHEN target_scope IN ('membership', 'capability') THEN 'permission'
WHEN target_scope = 'workspace_policy' THEN 'workspace'
ELSE pending_scope
END,
last_error = NULL, available_at = now(), updated_at = now()
WHERE generation_id = candidate.id AND workspace_id = target.workspace_id;
IF target_scope IN ('doc_policy', 'doc_grant') THEN
UPDATE search_projection.document_states
SET target_permission_version = GREATEST(target_permission_version, version),
claim_fence = NULL, lease_owner = NULL, lease_expires_at = NULL,
last_error = NULL, available_at = now(), updated_at = now()
WHERE generation_id = candidate.id
AND workspace_id = target.workspace_id
AND doc_id = target.doc_id;
END IF;
END LOOP;
END LOOP;
RETURN COALESCE(NEW, OLD);
END;
$$;
CREATE OR REPLACE FUNCTION search_projection.generation_build_complete(target_generation UUID)
RETURNS boolean LANGUAGE SQL STABLE AS $$
SELECT COALESCE((
SELECT scan_high_water_sid IS NOT NULL
AND scan_cursor_sid IS NOT NULL
AND scan_cursor_sid >= scan_high_water_sid
FROM search_projection.generations
WHERE id = target_generation
), false)
AND NOT EXISTS (
SELECT 1 FROM search_projection.workspace_states
WHERE generation_id = target_generation
AND (NOT covered OR pending_scope <> 'none'
OR required_permission_version > applied_permission_version
OR last_error IS NOT NULL)
)
AND NOT EXISTS (
SELECT 1 FROM search_projection.document_states
WHERE generation_id = target_generation
AND (target_source_version <> published_source_version
OR target_source_exists <> published_source_exists
OR target_permission_version <> published_permission_version)
)
$$;
CREATE OR REPLACE FUNCTION search_projection.capture_workspace_delete()
RETURNS trigger LANGUAGE plpgsql AS $$
BEGIN
UPDATE search_projection.workspace_states
SET covered = false,
pending_scope = 'workspace',
progress = jsonb_build_object(
'version', 1,
'captured_root_revision', target_root_revision,
'captured_permission_version', required_permission_version,
'kind', 'deleted',
'table', 'doc',
'quiet', false
),
claim_fence = nextval('search_projection.claim_fence'),
lease_owner = NULL,
lease_expires_at = NULL,
last_error = NULL,
available_at = now(),
updated_at = now()
WHERE workspace_id = OLD.id;
RETURN OLD;
END;
$$;
CREATE OR REPLACE FUNCTION search_projection.guard_generation_state_transition()
RETURNS trigger LANGUAGE plpgsql AS $$
BEGIN
IF NEW.state = 'active' AND (TG_OP = 'INSERT' OR OLD.state IS DISTINCT FROM NEW.state) THEN
IF TG_OP = 'INSERT' OR OLD.state <> 'building'
OR NOT search_projection.generation_build_complete(NEW.id) THEN
RAISE EXCEPTION 'search generation is not ready for activation';
END IF;
END IF;
RETURN NEW;
END;
$$;
CREATE TRIGGER search_projection_generation_state_guard
BEFORE INSERT OR UPDATE ON search_projection.generations
FOR EACH ROW EXECUTE FUNCTION search_projection.guard_generation_state_transition();
CREATE TRIGGER search_projection_snapshot_capture
AFTER INSERT OR UPDATE OR DELETE ON snapshots
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_snapshot_mutation();
CREATE TRIGGER search_projection_workspace_delete_capture
AFTER DELETE ON workspaces
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_workspace_delete();
CREATE TRIGGER search_projection_membership_capture
AFTER INSERT OR UPDATE OR DELETE ON workspace_members
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_permission_mutation();
CREATE TRIGGER search_projection_workspace_policy_capture
AFTER INSERT OR UPDATE OR DELETE ON workspace_access_policies
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_permission_mutation();
CREATE TRIGGER search_projection_doc_policy_capture
AFTER INSERT OR UPDATE OR DELETE ON doc_access_policies
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_permission_mutation();
CREATE TRIGGER search_projection_doc_grant_capture
AFTER INSERT OR UPDATE OR DELETE ON doc_grants
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_permission_mutation();
CREATE TRIGGER search_projection_entitlement_capture
AFTER INSERT OR UPDATE OR DELETE ON entitlements
FOR EACH ROW EXECUTE FUNCTION search_projection.capture_permission_mutation();

View File

@@ -6,9 +6,9 @@ use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use uuid::Uuid;
use super::{
CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupAckResult, RuntimeDocumentCleanupEffect,
RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult,
StorageRuntime, load_workspace_live_doc_ids, merge_current_doc, napi_error,
CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult,
RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult, StorageRuntime, load_workspace_live_doc_ids,
merge_current_doc, napi_error,
};
#[derive(FromRow)]
@@ -416,7 +416,6 @@ async fn delete_doc_rows(tx: &mut Transaction<'_, Postgres>, candidate: &Candida
"cleanupVersion": cleanup_version,
"commentAttachmentKeys": attachment_keys,
"commentObjectsDone": false,
"searchDone": false,
}))
.execute(&mut **tx)
.await
@@ -592,51 +591,43 @@ async fn execute_one(
Ok(Some((candidate, deleted_rows)))
}
fn payload_effect(effect: PendingEffect) -> RuntimeResult<RuntimeDocumentCleanupEffect> {
fn payload_effect(effect: &PendingEffect) -> RuntimeResult<RuntimeDocumentCleanupEffect> {
let cleanup_version = effect
.cleanup_payload
.get("cleanupVersion")
.and_then(Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("Document cleanup effect payload has no cleanupVersion"))?;
Ok(RuntimeDocumentCleanupEffect {
workspace_id: effect.workspace_id,
doc_id: effect.doc_id,
workspace_id: effect.workspace_id.clone(),
doc_id: effect.doc_id.clone(),
cleanup_version: cleanup_version.to_string(),
comment_objects_done: effect
.cleanup_payload
.get("commentObjectsDone")
.and_then(Value::as_bool)
.unwrap_or(false),
search_done: effect
.cleanup_payload
.get("searchDone")
.and_then(Value::as_bool)
.unwrap_or(false),
})
}
async fn complete_effect(
async fn complete_comment_objects(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
cleanup_version: &str,
path: &str,
) -> RuntimeResult<bool> {
let completed = sqlx::query_scalar::<_, bool>(
r#"
UPDATE document_cleanup_candidates
SET cleanup_payload = jsonb_set(cleanup_payload, ARRAY[$4], 'true'),
SET cleanup_payload = jsonb_set(cleanup_payload, '{commentObjectsDone}', 'true'),
error = NULL, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending'
AND cleanup_payload->>'cleanupVersion' = $3
RETURNING COALESCE((cleanup_payload->>'commentObjectsDone')::boolean, false)
AND COALESCE((cleanup_payload->>'searchDone')::boolean, false)
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(cleanup_version)
.bind(path)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup effect completion failed", err))?;
@@ -658,7 +649,7 @@ async fn complete_effect(
Ok(completed)
}
async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffect) -> RuntimeResult<()> {
async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffect) -> RuntimeResult<bool> {
let cleanup_version = effect
.cleanup_payload
.get("cleanupVersion")
@@ -695,18 +686,11 @@ async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffec
.begin()
.await
.map_err(|err| RuntimeError::database("Document cleanup object effect transaction failed", err))?;
complete_effect(
&mut tx,
&effect.workspace_id,
&effect.doc_id,
cleanup_version,
"commentObjectsDone",
)
.await?;
let completed = complete_comment_objects(&mut tx, &effect.workspace_id, &effect.doc_id, cleanup_version).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("Document cleanup object effect commit failed", err))?;
Ok(())
Ok(completed)
}
async fn load_pending_effects(
@@ -789,75 +773,33 @@ impl StorageRuntime {
}
let effects = load_pending_effects(&pool, workspace_id.as_deref(), limit).await?;
for effect in effects {
if let Err(err) = process_comment_objects(self, &effect).await {
result.failed += 1;
sqlx::query(
r#"
UPDATE document_cleanup_candidates
SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2
"#,
)
.bind(&effect.workspace_id)
.bind(&effect.doc_id)
.bind(err.to_string())
.execute(&pool)
.await
.map_err(|db_err| RuntimeError::database("Document cleanup effect failure write failed", db_err))?;
let mut result_effect = payload_effect(&effect)?;
match process_comment_objects(self, &effect).await {
Ok(comment_objects_done) => {
result_effect.comment_objects_done = comment_objects_done;
result.effects.push(result_effect);
}
Err(err) => {
result.failed += 1;
result.effects.push(result_effect);
sqlx::query(
r#"
UPDATE document_cleanup_candidates
SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2
"#,
)
.bind(&effect.workspace_id)
.bind(&effect.doc_id)
.bind(err.to_string())
.execute(&pool)
.await
.map_err(|db_err| RuntimeError::database("Document cleanup effect failure write failed", db_err))?;
}
}
}
result.effects = load_pending_effects(&pool, workspace_id.as_deref(), limit)
.await?
.into_iter()
.map(payload_effect)
.collect::<RuntimeResult<_>>()?;
Ok(result)
}
#[napi]
pub async fn ack_document_cleanup_effect(
&self,
workspace_id: String,
doc_id: String,
cleanup_version: String,
effect: String,
) -> napi::Result<RuntimeDocumentCleanupAckResult> {
let path = match effect.as_str() {
"search" => "searchDone",
_ => return Err(napi_error("document cleanup effect must be search")),
};
let pool = self.pool().await?;
let mut tx = pool
.begin()
.await
.map_err(|err| RuntimeError::database("Document cleanup ack transaction failed", err))?;
let current_version = sqlx::query_scalar::<_, String>(
r#"
SELECT cleanup_payload->>'cleanupVersion'
FROM document_cleanup_candidates
WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending'
"#,
)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_optional(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup ack candidate load failed", err))?;
let Some(current_version) = current_version else {
tx.rollback()
.await
.map_err(|err| RuntimeError::database("Document cleanup duplicate ack rollback failed", err))?;
return Ok(RuntimeDocumentCleanupAckResult { completed: true });
};
if current_version != cleanup_version {
return Err(napi_error("document cleanup effect candidate version mismatch"));
}
let completed = complete_effect(&mut tx, &workspace_id, &doc_id, &cleanup_version, path).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("Document cleanup ack commit failed", err))?;
Ok(RuntimeDocumentCleanupAckResult { completed })
}
}
#[cfg(test)]
@@ -870,7 +812,11 @@ mod tests {
use tokio::sync::Mutex;
use super::*;
use crate::runtime::{migrations::migrate_runtime_tables, storage_runtime::StorageRuntimeConfig};
use crate::runtime::{
backend_runtime::SEARCH_TEST_LOCK,
migrations::{migrate_runtime_tables, migrate_search_tables},
storage_runtime::StorageRuntimeConfig,
};
async fn runtime_from_database_url() -> AnyResult<Option<(StorageRuntime, PgPool)>> {
let Ok(database_url) = std::env::var("DATABASE_URL") else {
@@ -1233,10 +1179,17 @@ mod tests {
#[tokio::test]
async fn document_cleanup_execute_postgres_semantics() -> AnyResult<()> {
let _search_guard = SEARCH_TEST_LOCK.lock().await;
let Some((runtime, pool)) = runtime_from_database_url().await? else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return Ok(());
};
migrate_search_tables(&pool)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await?;
let suffix = Uuid::new_v4().to_string();
let (user_id, workspace_id) = insert_user_workspace(&pool, &suffix).await?;
let object_root = tempfile::tempdir()?;
@@ -1462,6 +1415,14 @@ mod tests {
.bind(doc_id)
.execute(&pool)
.await?;
let generation_id = Uuid::new_v4();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await?;
let executed = runtime
.execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10)
.await
@@ -1470,7 +1431,16 @@ mod tests {
assert_eq!(executed.failed, 0);
assert_eq!(executed.effects.len(), 1);
assert!(executed.effects[0].comment_objects_done);
assert!(!executed.effects[0].search_done);
let search_delete_tasks = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM search_projection.document_states WHERE generation_id = $1 AND workspace_id = $2 AND \
doc_id = $3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(doc_id)
.fetch_one(&pool)
.await?;
assert_eq!(search_delete_tasks, 1);
assert!(
runtime
.head_object("blob".to_string(), attachment_object_key)
@@ -1518,29 +1488,6 @@ mod tests {
assert_eq!(count, 0, "{table}.{column} should be nulled");
}
let effect = &executed.effects[0];
let search = runtime
.ack_document_cleanup_effect(
workspace_id.clone(),
doc_id.to_string(),
effect.cleanup_version.clone(),
"search".to_string(),
)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert!(search.completed);
// the copilot effect was removed; acknowledging it must be rejected
assert!(
runtime
.ack_document_cleanup_effect(
workspace_id.clone(),
doc_id.to_string(),
effect.cleanup_version.clone(),
"copilot".to_string(),
)
.await
.is_err()
);
let candidate_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
)
@@ -1588,8 +1535,7 @@ mod tests {
assert_eq!(failed_object_delete.failed, 1);
assert!(!failed_object_delete.effects[0].comment_objects_done);
let retained = sqlx::query(
"SELECT status, attempt_count, error, cleanup_payload->>'cleanupVersion' AS cleanup_version FROM \
document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
"SELECT status, attempt_count, error FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
)
.bind(&workspace_id)
.bind(retry_doc_id)
@@ -1598,18 +1544,6 @@ mod tests {
assert_eq!(retained.get::<String, _>("status"), "effects_pending");
assert_eq!(retained.get::<i32, _>("attempt_count"), 1);
assert!(retained.get::<Option<String>, _>("error").is_some());
let retry_cleanup_version = retained.get::<String, _>("cleanup_version");
let ack = runtime
.ack_document_cleanup_effect(
workspace_id.clone(),
retry_doc_id.to_string(),
retry_cleanup_version.clone(),
"search".to_string(),
)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert!(!ack.completed);
sqlx::query(
"UPDATE document_cleanup_candidates SET cleanup_payload = jsonb_set(cleanup_payload, '{commentAttachmentKeys}', \
@@ -1625,7 +1559,8 @@ mod tests {
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(retried.executed, 0);
assert_eq!(retried.failed, 0);
assert!(retried.effects.is_empty());
assert_eq!(retried.effects.len(), 1);
assert!(retried.effects[0].comment_objects_done);
let retry_candidate_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
)
@@ -1635,6 +1570,10 @@ mod tests {
.await?;
assert_eq!(retry_candidate_count, 0);
sqlx::query("DELETE FROM search_projection.generations WHERE id = $1")
.bind(generation_id)
.execute(&pool)
.await?;
cleanup_workspace_fixture(&pool, &user_id, &workspace_id).await?;
Ok(())
}

View File

@@ -32,10 +32,10 @@ pub(super) use super::{
napi_error, to_napi_error,
types::{
RuntimeBlobCleanupExecuteResult, RuntimeBlobCleanupPlanResult, RuntimeBlobCleanupResult, RuntimeBlobCompleteResult,
RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupAckResult,
RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult,
RuntimeMultipartUploadInit, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry,
RuntimeObjectMetadata, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupEffect,
RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeMultipartUploadInit,
RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry, RuntimeObjectMetadata,
RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
},
};

View File

@@ -263,6 +263,13 @@ pub struct RuntimeWorkspaceInviteQuotaUsage {
pub target_domains: Vec<RuntimeQuotaTargetDomainInput>,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceActionDecision {
pub allowed: bool,
pub retry_after_seconds: Option<i32>,
pub reason: Option<String>,
}
#[napi_derive::napi(object)]
pub struct RuntimeInviteAbuseActionRequired {
pub action: String,
@@ -469,7 +476,6 @@ pub struct RuntimeDocumentCleanupEffect {
pub doc_id: String,
pub cleanup_version: String,
pub comment_objects_done: bool,
pub search_done: bool,
}
#[napi_derive::napi(object)]
@@ -484,11 +490,6 @@ pub struct RuntimeDocumentCleanupExecuteResult {
pub effects: Vec<RuntimeDocumentCleanupEffect>,
}
#[napi_derive::napi(object)]
pub struct RuntimeDocumentCleanupAckResult {
pub completed: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeBlobCleanupPlanResult {
pub run_id: Option<String>,
@@ -520,34 +521,6 @@ pub struct RuntimeDocCompactionResult {
pub history_created: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsRefreshResult {
pub processed: i64,
pub backlog: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsRecalibrationResult {
pub processed: i64,
pub last_sid: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsSnapshotResult {
pub snapshotted: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsDailyRecalibrationResult {
pub processed: i64,
pub last_sid: i64,
pub snapshotted: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeEmbeddingProgress {
pub total: i64,

View File

@@ -3,12 +3,13 @@ mod query;
mod result;
mod schema;
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use memory_indexer::{MemoryIndex, Mutation, TermsAggregation};
use napi::{Status, bindgen_prelude::Buffer};
use serde_json::Value as JsonValue;
use napi::Status;
use serde_json::{Value as JsonValue, json};
use tokio::sync::RwLock;
use uuid::Uuid;
use self::{
document::compile_document,
@@ -49,12 +50,12 @@ impl TableIndex {
}
}
struct IndexManager {
struct GenerationIndex {
doc: TableIndex,
block: TableIndex,
}
impl IndexManager {
impl GenerationIndex {
fn new() -> Self {
Self {
doc: TableIndex::new(TableSchema::doc()),
@@ -71,9 +72,32 @@ impl IndexManager {
}
}
pub(crate) struct EmbeddedIndexCheckpoint {
pub sequence: i64,
pub data: Buffer,
struct IndexManager {
generations: RwLock<HashMap<Uuid, Arc<GenerationIndex>>>,
}
impl IndexManager {
fn new() -> Self {
Self {
generations: RwLock::new(HashMap::new()),
}
}
async fn generation(&self, generation_id: Uuid, create: bool) -> Option<Arc<GenerationIndex>> {
if let Some(generation) = self.generations.read().await.get(&generation_id).cloned() {
return Some(generation);
}
if !create {
return None;
}
let mut generations = self.generations.write().await;
Some(
generations
.entry(generation_id)
.or_insert_with(|| Arc::new(GenerationIndex::new()))
.clone(),
)
}
}
pub(crate) struct EmbeddedSearchIndex {
@@ -87,26 +111,52 @@ impl EmbeddedSearchIndex {
}
}
pub(crate) async fn restore(&self, table: String, checkpoint: Buffer) -> napi::Result<()> {
let table = self.manager.table(&table)?;
let index =
MemoryIndex::from_checkpoint(table.schema.schema.clone(), checkpoint.as_ref()).map_err(IndexError::from)?;
*table.index.write().await = index;
Ok(())
pub(crate) async fn prepare_generation(&self, generation_id: Uuid) {
self.manager.generation(generation_id, true).await;
}
pub(crate) async fn reset(&self, table: String) -> napi::Result<()> {
let table = self.manager.table(&table)?;
*table.index.write().await = MemoryIndex::new(table.schema.schema.clone());
Ok(())
pub(crate) async fn has_generation(&self, generation_id: Uuid) -> bool {
self.manager.generation(generation_id, false).await.is_some()
}
pub(crate) async fn retain_generation(&self, generation_id: Uuid) {
self
.manager
.generations
.write()
.await
.retain(|id, _| *id == generation_id);
}
#[cfg(test)]
pub(crate) async fn write(&self, table: String, documents_json: String) -> napi::Result<()> {
let table = self.manager.table(&table)?;
self.write_for_generation(Uuid::nil(), table, documents_json).await
}
pub(crate) async fn write_for_generation(
&self,
generation_id: Uuid,
table: String,
documents_json: String,
) -> napi::Result<()> {
let generation = self
.manager
.generation(generation_id, true)
.await
.expect("created embedded search generation");
let table = generation.table(&table)?;
let documents: Vec<JsonValue> = serde_json::from_str(&documents_json)?;
let generation_id = generation_id.to_string();
let documents = documents
.into_iter()
.map(|document| compile_document(&table.schema, document))
.map(|document| {
if document.get("generation_id").and_then(JsonValue::as_str) != Some(generation_id.as_str()) {
return Err(IndexError::InvalidInput(
"embedded projection generation does not match target index".into(),
));
}
compile_document(&table.schema, document)
})
.collect::<Result<Vec<_>>>()?;
table
.index
@@ -117,14 +167,148 @@ impl EmbeddedSearchIndex {
Ok(())
}
pub(crate) async fn delete(&self, table: String, id: String) -> napi::Result<()> {
self.manager.table(&table)?.index.write().await.delete(&id);
pub(crate) async fn gc_document_history_for_generation(
&self,
generation_id: Uuid,
table_name: &str,
workspace_id: &str,
doc_id: &str,
published_tuple: (i64, i64),
limit: usize,
) -> napi::Result<()> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(table_name)?;
let source_version_field = table.schema.field("source_version")?;
let permission_version_field = table.schema.field("permission_version")?;
let query = compile_query(
&table.schema,
&json!({
"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"term":{"doc_id":{"value":doc_id}}}
]}
}),
)?;
let mut index = table.index.write().await;
let matches = index
.search(
&query,
memory_indexer::SearchOptions {
limit: limit.max(1),
offset: 0,
after: None,
sort: vec![memory_indexer::Sort::DocumentId],
stored_fields: vec![source_version_field, permission_version_field],
highlight_fields: Vec::new(),
},
)
.map_err(IndexError::from)?;
for hit in matches.hits {
let mut hit_source_version = None;
let mut hit_permission_version = None;
for (field, values) in hit.fields {
let value = values.first().and_then(|value| match value {
memory_indexer::Value::I64(value) => Some(*value),
_ => None,
});
if field == source_version_field {
hit_source_version = value;
} else if field == permission_version_field {
hit_permission_version = value;
}
}
if (hit_source_version, hit_permission_version) != (Some(published_tuple.0), Some(published_tuple.1)) {
index.delete(&hit.id);
}
}
Ok(())
}
pub(crate) async fn gc_workspace_for_generation(
&self,
generation_id: Uuid,
table_name: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> napi::Result<bool> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(table_name)?;
let source_version_field = table.schema.field("source_version")?;
let query = compile_query(&table.schema, &json!({"term":{"workspace_id":{"value":workspace_id}}}))?;
let mut index = table.index.write().await;
let matches = index
.search(
&query,
memory_indexer::SearchOptions {
limit: limit.max(1),
offset: 0,
after: None,
sort: vec![
memory_indexer::Sort::Field {
field: source_version_field,
order: memory_indexer::SortOrder::Asc,
},
memory_indexer::Sort::DocumentId,
],
stored_fields: vec![source_version_field],
highlight_fields: Vec::new(),
},
)
.map_err(IndexError::from)?;
let may_have_more = matches.hits.len() == limit.max(1)
&& matches.hits.iter().any(|hit| {
hit.fields.iter().any(|(field, values)| {
*field == source_version_field
&& values.first().is_none_or(
|value| !matches!(value, memory_indexer::Value::I64(version) if *version > source_version_high_water),
)
})
});
for hit in matches.hits {
let source_version = hit
.fields
.iter()
.find(|(field, _)| *field == source_version_field)
.and_then(|(_, values)| values.first())
.and_then(|value| match value {
memory_indexer::Value::I64(value) => Some(*value),
_ => None,
});
if source_version.is_none_or(|source_version| source_version <= source_version_high_water) {
index.delete(&hit.id);
}
}
Ok(may_have_more)
}
#[cfg(test)]
pub(crate) async fn search(&self, table: String, dsl_json: String) -> napi::Result<String> {
let table = self.manager.table(&table)?;
let dsl: JsonValue = serde_json::from_str(&dsl_json)?;
self.search_for_generation(Uuid::nil(), table, dsl_json).await
}
pub(crate) async fn search_for_generation(
&self,
generation_id: Uuid,
table: String,
dsl_json: String,
) -> napi::Result<String> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(&table)?;
let mut dsl: JsonValue = serde_json::from_str(&dsl_json)?;
ensure_projection_fields(&mut dsl)?;
let query = compile_query(
&table.schema,
dsl
@@ -144,9 +328,27 @@ impl EmbeddedSearchIndex {
))?)
}
#[cfg(test)]
pub(crate) async fn aggregate(&self, table: String, dsl_json: String) -> napi::Result<String> {
let table = self.manager.table(&table)?;
let dsl: JsonValue = serde_json::from_str(&dsl_json)?;
self.aggregate_for_generation(Uuid::nil(), table, dsl_json).await
}
pub(crate) async fn aggregate_for_generation(
&self,
generation_id: Uuid,
table: String,
dsl_json: String,
) -> napi::Result<String> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(&table)?;
let mut dsl: JsonValue = serde_json::from_str(&dsl_json)?;
if let Some(top_hits) = dsl.pointer_mut("/aggs/result/aggs/result/top_hits") {
ensure_projection_fields(top_hits)?;
}
let query = compile_query(
&table.schema,
dsl
@@ -187,38 +389,23 @@ impl EmbeddedSearchIndex {
&highlight_tags(dsl.pointer("/aggs/result/aggs/result/top_hits").unwrap_or(&dsl)),
))?)
}
}
pub(crate) async fn checkpoint(&self, table: String) -> napi::Result<EmbeddedIndexCheckpoint> {
let checkpoint = self
.manager
.table(&table)?
.index
.read()
.await
.checkpoint()
.map_err(IndexError::from)?;
Ok(EmbeddedIndexCheckpoint {
sequence: checkpoint.sequence as i64,
data: checkpoint.bytes.into(),
})
}
pub(crate) async fn optimize(&self, table: String) -> napi::Result<()> {
self.manager.table(&table)?.index.write().await.optimize();
Ok(())
}
pub(crate) async fn mark_checkpoint_persisted(&self, table: String, sequence: i64) -> napi::Result<()> {
self
.manager
.table(&table)?
.index
.write()
.await
.mark_checkpoint_persisted(sequence as u64)
.map_err(IndexError::from)?;
Ok(())
fn ensure_projection_fields(dsl: &mut JsonValue) -> Result<()> {
let fields = dsl
.as_object_mut()
.ok_or_else(|| IndexError::InvalidInput("search DSL must be an object".into()))?
.entry("fields")
.or_insert_with(|| json!([]));
let fields = fields
.as_array_mut()
.ok_or_else(|| IndexError::InvalidInput("search fields must be an array".into()))?;
for field in ["doc_id", "source_version", "permission_version"] {
if !fields.iter().any(|value| value.as_str() == Some(field)) {
fields.push(json!(field));
}
}
Ok(())
}
fn highlight_tags(dsl: &JsonValue) -> HighlightTags {
@@ -248,13 +435,17 @@ impl Default for EmbeddedSearchIndex {
#[cfg(test)]
mod tests {
use serde_json::{Value, json};
use uuid::Uuid;
use super::EmbeddedSearchIndex;
fn doc(workspace: &str, id: &str, title: &str, updated_at: i64) -> Value {
json!({
"generation_id": Uuid::nil().to_string(),
"workspace_id": workspace,
"doc_id": id,
"source_version": 1,
"permission_version": 1,
"title": title,
"summary": title,
"created_by_user_id": "user",
@@ -277,7 +468,7 @@ mod tests {
}
#[tokio::test]
async fn exact_search_cursor_and_checkpoint_roundtrip() {
async fn exact_search_cursor_and_aggregate() {
let index = EmbeddedSearchIndex::new();
index
.write(
@@ -296,7 +487,7 @@ mod tests {
let first: Value =
serde_json::from_str(&index.search("doc".into(), search(query.clone(), None)).await.unwrap()).unwrap();
assert_eq!(first["total"], 2);
assert_eq!(first["nodes"][0]["id"], "workspace-1/two");
assert_eq!(first["nodes"][0]["id"], format!("{}/workspace-1/two/1/1", Uuid::nil()));
assert_eq!(first["nodes"][0]["fields"]["doc_id"], json!(["two"]));
let cursor = first["nextCursor"].as_str().unwrap();
@@ -307,17 +498,7 @@ mod tests {
.unwrap(),
)
.unwrap();
assert_eq!(second["nodes"][0]["id"], "workspace-1/one");
let checkpoint = index.checkpoint("doc".into()).await.unwrap();
index.reset("doc".into()).await.unwrap();
let empty: Value =
serde_json::from_str(&index.search("doc".into(), search(query.clone(), None)).await.unwrap()).unwrap();
assert_eq!(empty["total"], 0);
index.restore("doc".into(), checkpoint.data).await.unwrap();
let restored: Value =
serde_json::from_str(&index.search("doc".into(), search(query, None)).await.unwrap()).unwrap();
assert_eq!(restored["total"], 2);
assert_eq!(second["nodes"][0]["id"], format!("{}/workspace-1/one/1/1", Uuid::nil()));
let aggregate: Value = serde_json::from_str(
&index
@@ -343,12 +524,16 @@ mod tests {
}
#[tokio::test]
async fn write_is_atomic_and_corrupt_checkpoint_is_rejected() {
async fn write_is_atomic() {
let index = EmbeddedSearchIndex::new();
index
.write(
"doc".into(),
json!([{ "workspace_id": "workspace", "doc_id": "null-values", "summary": [null] }]).to_string(),
json!([{
"generation_id":Uuid::nil().to_string(),"workspace_id":"workspace","doc_id":"null-values",
"source_version":1,"permission_version":1,"summary":[null]
}])
.to_string(),
)
.await
.unwrap();
@@ -357,7 +542,6 @@ mod tests {
serde_json::from_str(&index.search("doc".into(), search(all.clone(), None)).await.unwrap()).unwrap();
assert_eq!(result["total"], 1);
index.reset("doc".into()).await.unwrap();
let documents = json!([
doc("workspace", "valid", "hello", 1),
{ "workspace_id": "workspace", "doc_id": "invalid", "unknown": true }
@@ -365,7 +549,124 @@ mod tests {
assert!(index.write("doc".into(), documents.to_string()).await.is_err());
let result: Value = serde_json::from_str(&index.search("doc".into(), search(all, None)).await.unwrap()).unwrap();
assert_eq!(result["total"], 0);
assert!(index.restore("doc".into(), vec![1, 2, 3].into()).await.is_err());
assert_eq!(result["total"], 1);
}
#[tokio::test]
async fn immutable_tuples_do_not_overwrite_each_other() {
let index = EmbeddedSearchIndex::new();
let mut old = doc("workspace", "doc", "old", 1);
old["source_version"] = json!(1);
let mut current = doc("workspace", "doc", "current", 1);
current["source_version"] = json!(2);
index
.write("doc".into(), json!([old, current]).to_string())
.await
.unwrap();
let result: Value = serde_json::from_str(
&index
.search(
"doc".into(),
json!({"query":{"term":{"doc_id":{"value":"doc"}}},"fields":["title","source_version"],"size":10})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 2);
assert_ne!(result["nodes"][0]["id"], result["nodes"][1]["id"]);
let retired = Uuid::new_v4();
index.prepare_generation(retired).await;
index.retain_generation(Uuid::nil()).await;
assert!(index.has_generation(Uuid::nil()).await);
assert!(!index.has_generation(retired).await);
}
#[tokio::test]
async fn history_gc_keeps_the_published_tuple() {
let index = EmbeddedSearchIndex::new();
let mut stale = doc("workspace", "doc", "old", 1);
stale["source_version"] = json!(1);
let mut current = doc("workspace", "doc", "current", 1);
current["source_version"] = json!(2);
index
.write("doc".into(), json!([stale, current]).to_string())
.await
.unwrap();
index
.gc_document_history_for_generation(Uuid::nil(), "doc", "workspace", "doc", (2, 1), 10)
.await
.unwrap();
let result: Value = serde_json::from_str(
&index
.search(
"doc".into(),
json!({"query":{"term":{"doc_id":{"value":"doc"}}},"fields":["title"],"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(result["nodes"][0]["fields"]["title"], json!(["current"]));
index
.write(
"doc".into(),
json!([
doc("deleted-workspace", "one", "one", 1),
doc("deleted-workspace", "two", "two", 1)
])
.to_string(),
)
.await
.unwrap();
assert!(
index
.gc_workspace_for_generation(Uuid::nil(), "doc", "deleted-workspace", i64::MAX, 1)
.await
.unwrap()
);
assert!(
index
.gc_workspace_for_generation(Uuid::nil(), "doc", "deleted-workspace", i64::MAX, 1)
.await
.unwrap()
);
assert!(
!index
.gc_workspace_for_generation(Uuid::nil(), "doc", "deleted-workspace", i64::MAX, 1)
.await
.unwrap()
);
let old = doc("recreated-workspace", "old", "old", 1);
let mut recreated = doc("recreated-workspace", "new", "new", 1);
recreated["source_version"] = json!(3);
index
.write("doc".into(), json!([old, recreated]).to_string())
.await
.unwrap();
assert!(
!index
.gc_workspace_for_generation(Uuid::nil(), "doc", "recreated-workspace", 2, 10)
.await
.unwrap()
);
let result: Value = serde_json::from_str(
&index
.search(
"doc".into(),
json!({"query":{"term":{"workspace_id":{"value":"recreated-workspace"}}},"fields":["doc_id"],"size":10})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(result["nodes"][0]["fields"]["doc_id"], json!(["new"]));
}
}

View File

@@ -16,6 +16,9 @@ impl TableSchema {
let mut fields = HashMap::new();
keyword(&mut builder, &mut fields, "workspace_id", true, false);
keyword(&mut builder, &mut fields, "workspace_token", true, false);
keyword(&mut builder, &mut fields, "generation_id", true, false);
integer(&mut builder, &mut fields, "source_version", true, true);
integer(&mut builder, &mut fields, "permission_version", true, false);
keyword(&mut builder, &mut fields, "doc_id", true, true);
keyword(&mut builder, &mut fields, "doc_token", true, false);
fields.insert(
@@ -29,10 +32,19 @@ impl TableSchema {
keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_member_default_readable");
integer(&mut builder, &mut fields, "acl_revision", true, false);
integer(&mut builder, &mut fields, "created_at", true, true);
integer(&mut builder, &mut fields, "updated_at", true, true);
Self::finish(builder, fields, &["workspace_id", "doc_id"])
Self::finish(
builder,
fields,
&[
"generation_id",
"workspace_id",
"doc_id",
"source_version",
"permission_version",
],
)
}
pub fn block() -> Self {
@@ -58,6 +70,9 @@ impl TableSchema {
}
keyword(&mut builder, &mut fields, "doc_id", true, true);
keyword(&mut builder, &mut fields, "workspace_token", true, false);
keyword(&mut builder, &mut fields, "generation_id", true, false);
integer(&mut builder, &mut fields, "source_version", true, true);
integer(&mut builder, &mut fields, "permission_version", true, false);
keyword(&mut builder, &mut fields, "doc_token", true, false);
keyword(&mut builder, &mut fields, "block_id", true, true);
keyword(&mut builder, &mut fields, "block_token", true, false);
@@ -74,8 +89,18 @@ impl TableSchema {
keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_member_default_readable");
integer(&mut builder, &mut fields, "acl_revision", true, false);
Self::finish(builder, fields, &["workspace_id", "doc_id", "block_id"])
Self::finish(
builder,
fields,
&[
"generation_id",
"workspace_id",
"doc_id",
"source_version",
"permission_version",
"block_id",
],
)
}
fn finish(
@@ -95,10 +120,14 @@ impl TableSchema {
.id_fields
.iter()
.map(|field| {
document
let value = document
.get(*field)
.and_then(serde_json::Value::as_str)
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} is required")))
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} is required")))?;
value
.as_str()
.map(str::to_string)
.or_else(|| value.as_i64().map(|value| value.to_string()))
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} has invalid identity type")))
})
.collect::<Result<Vec<_>>>()
.map(|parts| parts.join("/"))

View File

@@ -11,5 +11,3 @@
# MAILER_USER="noreply@toeverything.info"
# MAILER_PASSWORD="affine"
# MAILER_SECURE=false
# AFFINE_INDEXER_ENABLED=true

View File

@@ -0,0 +1,45 @@
# Manticore Search provider E2E
This fixture validates the shared immutable projection contract against a real
Manticore Search RT table. Aggregate queries remain unsupported and are validated
through the typed fallback boundary instead.
Start the external dependency:
```bash
docker compose -f packages/backend/server/e2e/manticore-provider/compose.yml up -d --wait
```
Use a disposable PostgreSQL database and run the provider-gated E2E:
```bash
DATABASE_URL=postgresql://ds:ds@localhost:55433/affine_rfc6_manticore_e2e \
yarn workspace @affine/server prisma migrate deploy
```
Configure `packages/backend/server/config.json` with:
```json
{
"indexer": {
"enabled": true,
"provider": {
"type": "manticoresearch",
"endpoint": "http://127.0.0.1:9308"
}
}
}
```
Then run:
```bash
DATABASE_URL=postgresql://ds:ds@localhost:55433/affine_rfc6_manticore_e2e \
yarn af server e2e src/__tests__/e2e/indexer/manticore-provider.spec.ts
```
Stop only this disposable dependency when finished:
```bash
docker compose -f packages/backend/server/e2e/manticore-provider/compose.yml down -v
```

View File

@@ -0,0 +1,33 @@
services:
postgres:
image: pgvector/pgvector:pg16
environment:
POSTGRES_DB: affine_rfc6_manticore_e2e
POSTGRES_USER: ds
POSTGRES_PASSWORD: ds
ports:
- '55433:5432'
healthcheck:
test: ['CMD-SHELL', 'pg_isready -U ds -d affine_rfc6_manticore_e2e']
interval: 2s
timeout: 2s
retries: 30
start_period: 5s
manticore:
image: manticoresearch/manticore:29.0.2
environment:
searchd.listen: 'http:9308'
searchd.network_timeout: '30'
ports:
- '9308:9308'
healthcheck:
test:
[
'CMD-SHELL',
"wget -qO- --post-data='SHOW TABLES' 'http://127.0.0.1:9308/sql?mode=raw' >/dev/null",
]
interval: 2s
timeout: 2s
retries: 30
start_period: 5s

View File

@@ -0,0 +1,39 @@
# Remote provider E2E
These tests exercise the Elasticsearch-compatible provider over HTTP against a
local OpenSearch instance. Toxiproxy is included so the lease test can inject
provider latency without changing production code.
```bash
docker compose -f packages/backend/server/e2e/remote-provider/compose.yml up -d --wait
DATABASE_URL=postgresql://ds:ds@localhost:55432/affine_rfc6_remote_e2e \
yarn workspace @affine/server prisma migrate deploy
```
Configure `packages/backend/server/config.json` with:
```json
{
"indexer": {
"enabled": true,
"provider": {
"type": "elasticsearch",
"endpoint": "http://127.0.0.1:8666"
}
}
}
```
Then run:
```bash
DATABASE_URL=postgresql://ds:ds@localhost:55432/affine_rfc6_remote_e2e \
yarn af server e2e src/__tests__/e2e/indexer/remote-provider.spec.ts
docker compose -f packages/backend/server/e2e/remote-provider/compose.yml down -v
```
The suite uses a disposable PostgreSQL database and creates generation-specific
indices. It directly deletes or mutates provider rows only as a fault-injection
fixture; production code still repairs them through the normal reconcile path.

View File

@@ -0,0 +1,73 @@
services:
postgres:
image: pgvector/pgvector:pg16
environment:
POSTGRES_DB: affine_rfc6_remote_e2e
POSTGRES_USER: ds
POSTGRES_PASSWORD: ds
ports:
- '55432:5432'
healthcheck:
test: ['CMD-SHELL', 'pg_isready -U ds -d affine_rfc6_remote_e2e']
interval: 2s
timeout: 2s
retries: 30
start_period: 5s
opensearch:
image: opensearchproject/opensearch:2.19.1
environment:
discovery.type: single-node
DISABLE_INSTALL_DEMO_CONFIG: 'true'
DISABLE_SECURITY_PLUGIN: 'true'
bootstrap.memory_lock: 'true'
OPENSEARCH_JAVA_OPTS: '-Xms512m -Xmx512m'
ulimits:
memlock:
soft: -1
hard: -1
ports:
- '9200:9200'
healthcheck:
test:
[
'CMD-SHELL',
'curl -fsS http://localhost:9200/_cluster/health >/dev/null',
]
interval: 5s
timeout: 5s
retries: 30
start_period: 10s
toxiproxy:
image: ghcr.io/shopify/toxiproxy:2.12.0
ports:
- '8474:8474'
- '8666:8666'
proxy-init:
image: curlimages/curl:8.12.1
depends_on:
opensearch:
condition: service_healthy
toxiproxy:
condition: service_started
entrypoint: ['/bin/sh', '-ec']
command:
- >-
until curl -fsS http://toxiproxy:8474/version >/dev/null; do sleep 1; done;
curl -fsS http://toxiproxy:8474/proxies/opensearch >/dev/null 2>&1 ||
curl -fsS -X POST http://toxiproxy:8474/proxies
-H 'content-type: application/json'
-d '{"name":"opensearch","listen":"0.0.0.0:8666","upstream":"opensearch:9200"}';
tail -f /dev/null
healthcheck:
test:
[
'CMD-SHELL',
'curl -fsS http://toxiproxy:8474/proxies/opensearch >/dev/null',
]
interval: 2s
timeout: 2s
retries: 30
start_period: 2s

View File

@@ -4,26 +4,32 @@
DO $$
BEGIN
IF to_regclass('public.ai_contexts') IS NOT NULL AND EXISTS (
WITH referenced_blobs AS (
SELECT DISTINCT
session.workspace_id,
value #>> '{}' AS blob_key
FROM ai_contexts context
JOIN ai_sessions_metadata session ON session.id = context.session_id
CROSS JOIN LATERAL jsonb_path_query(
context.config::jsonb,
'$.** ? (@.type() == "string")'
) AS referenced_value(value)
)
SELECT 1
FROM ai_contexts context
JOIN ai_sessions_metadata session ON session.id = context.session_id
FROM referenced_blobs referenced
JOIN blobs blob
ON blob.workspace_id = session.workspace_id
ON blob.workspace_id = referenced.workspace_id
AND blob.key = referenced.blob_key
AND blob.deleted_at IS NULL
AND blob.status = 'completed'
WHERE jsonb_path_exists(
context.config::jsonb,
'$.** ? (@ == $blobKey)',
jsonb_build_object('blobKey', to_jsonb(blob.key::text))
)
AND NOT EXISTS (
SELECT 1
FROM workspace_artifacts artifact
WHERE artifact.workspace_id = session.workspace_id
AND artifact.status = 'ready'
AND artifact.storage_scope = 'blob'
AND artifact.storage_key = concat(session.workspace_id, '/', blob.key)
)
WHERE NOT EXISTS (
SELECT 1
FROM workspace_artifacts artifact
WHERE artifact.workspace_id = referenced.workspace_id
AND artifact.status = 'ready'
AND artifact.storage_scope = 'blob'
AND artifact.storage_key = concat(referenced.workspace_id, '/', blob.key)
)
) THEN
RAISE EXCEPTION
'legacy context blob artifact admission is incomplete; run the data migration before cleanup';

View File

@@ -10,7 +10,7 @@ import {
type Config,
type EventBus,
type JobQueue,
SearchProviderNotFound,
SearchProviderUnavailable,
} from '../../base';
import { ServerFeature, type ServerService } from '../../core';
import type { DocReader } from '../../core/doc';
@@ -450,7 +450,7 @@ test('document tools enforce the user-selected hard scope', async t => {
readableAc,
{
searchDocsByKeyword: async () => {
throw new SearchProviderNotFound();
throw new SearchProviderUnavailable();
},
} as unknown as IndexerService,
vectorSearch,

View File

@@ -1,8 +1,8 @@
import { getCurrentUserQuery } from '@affine/graphql';
import { JobExecutor } from '../../../base/job/queue/executor';
import { JobHandlerScanner } from '../../../base/job/queue/scanner';
import { DatabaseDocReader, DocReader } from '../../../core/doc';
import { RealtimeGateway } from '../../../core/realtime/gateway';
import { createApp } from '../create-app';
import { e2e } from '../test';
@@ -42,7 +42,7 @@ e2e('should init worker service', async t => {
await withFlavor('worker', async app => {
const res = await app.GET('/info').expect(200);
t.is(res.body.flavor, 'worker');
t.truthy(app.get(JobHandlerScanner).getHandler('indexer.indexDoc'));
t.throws(() => app.get(RealtimeGateway));
await t.throwsAsync(app.gql({ query: getCurrentUserQuery }));
await app.PUT('/api/storage/upload').expect(404);
@@ -53,7 +53,6 @@ e2e('should init allinone service with worker handlers', async t => {
await withFlavor('allinone', async app => {
const res = await app.GET('/info').expect(200);
t.is(res.body.flavor, 'allinone');
t.truthy(app.get(JobHandlerScanner).getHandler('indexer.indexDoc'));
});
});
@@ -65,6 +64,7 @@ e2e('should init graphql service', async t => {
const user = await app.gql({ query: getCurrentUserQuery });
t.is(user.currentUser, null);
t.truthy(app.get(RealtimeGateway));
});
});

View File

@@ -1,16 +1,27 @@
import { serverConfigQuery, ServerFeature } from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { IndexerService } from '../../../plugins/indexer';
import { app, e2e } from '../test';
e2e('should indexer feature enabled by default', async t => {
const { serverConfig } = await app.gql({ query: serverConfigQuery });
e2e(
'should expose the indexer feature when its projection is ready',
async t => {
const enabled = app.get(Config).indexer.enabled;
if (enabled) {
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
await app.get(IndexerService).onApplicationBootstrap();
}
const { serverConfig } = await app.gql({ query: serverConfigQuery });
t.is(
serverConfig.features.includes(ServerFeature.Indexer),
true,
JSON.stringify(serverConfig, null, 2)
);
});
t.is(
serverConfig.features.includes(ServerFeature.Indexer),
enabled,
JSON.stringify(serverConfig, null, 2)
);
}
);
e2e('should comment feature enabled by default', async t => {
const { serverConfig } = await app.gql({ query: serverConfigQuery });

View File

@@ -362,7 +362,9 @@ export async function createApp(
await app.init();
await app.get(BackendRuntimeProvider, { strict: false }).runMigrations();
await app.get(StorageRuntimeProvider, { strict: false }).runMigrations();
await app.get(IndexerService, { strict: false }).onApplicationBootstrap();
if (globalThis.env.isApi || globalThis.env.isFrontend) {
await app.get(IndexerService, { strict: false }).onApplicationBootstrap();
}
} catch (error) {
await app.close();
throw error;

View File

@@ -4,14 +4,20 @@ import {
SearchTable,
} from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { createDocWithMarkdown } from '../../../native';
import { IndexerService } from '../../../plugins/indexer/service';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
e2e('should aggregate by docId', async t => {
const indexerE2e = app.get(Config).indexer.enabled ? e2e : e2e.skip;
indexerE2e('should aggregate by docId', async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner });
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
for (const [docId, markdown] of [
['doc-0', 'hello world\n\nhello again'],
['doc-1', 'hello world'],
@@ -23,7 +29,7 @@ e2e('should aggregate by docId', async t => {
user: owner,
blob: createDocWithMarkdown(docId, markdown, docId),
});
await app.get(IndexerService).indexDoc(workspace.id, docId);
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
}
const result = await app.gql({

View File

@@ -0,0 +1,186 @@
import {
indexerSearchDocsQuery,
indexerSearchQuery,
SearchQueryType,
SearchTable,
} from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { DocRole } from '../../../models';
import { createDocWithMarkdown } from '../../../native';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
const manticoreSearchEnabled =
app.get(Config).indexer.enabled &&
app.get(Config).indexer.provider.type === 'manticoresearch';
const manticoreSearchE2e = manticoreSearchEnabled ? e2e : e2e.skip;
async function indexDocument(
workspaceId: string,
user: { id: string },
docId: string,
markdown: string,
defaultRole = DocRole.Manager
) {
await app.create(Mockers.DocMeta, { workspaceId, docId, defaultRole });
await app.create(Mockers.DocSnapshot, {
workspaceId,
docId,
user,
blob: createDocWithMarkdown(docId, markdown, docId),
});
await reconcileSearch();
}
async function reconcileSearch() {
const runtime = app.get(BackendRuntimeProvider);
for (let attempt = 0; attempt < 50; attempt++) {
await runtime.reconcileSearchProjection(1000);
if ((await runtime.searchStatus()).ready) return;
}
throw new Error('search projection did not become ready');
}
async function searchPage(
workspaceId: string,
match: string,
pagination: { limit: number; cursor?: string }
) {
const result = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspaceId,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match,
},
options: {
fields: ['docId', 'blockId'],
pagination,
},
},
},
});
return result.workspace.search;
}
async function searchCount(workspaceId: string, match: string) {
return (await searchPage(workspaceId, match, { limit: 20 })).pagination.count;
}
async function searchDocsCount(workspaceId: string, keyword: string) {
const result = await app.gql({
query: indexerSearchDocsQuery,
variables: {
id: workspaceId,
input: { keyword, limit: 20 },
},
});
return result.workspace.searchDocs.length;
}
manticoreSearchE2e(
'indexes and searches through the Manticore Search provider',
async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
const docId = `manticore-basic-${Date.now()}`;
const marker = 'manticorebasicmarker';
await indexDocument(
workspace.id,
owner,
docId,
`${marker} first block\n\nsecond block`
);
await indexDocument(
workspace.id,
owner,
`${docId}-second`,
`${marker} from a second document`
);
t.is(await searchCount(workspace.id, marker), 2);
const basicDocs = await app.gql({
query: indexerSearchDocsQuery,
variables: {
id: workspace.id,
input: { keyword: marker, limit: 2 },
},
});
t.is(basicDocs.workspace.searchDocs.length, 2);
await t.throwsAsync(
app.gql({
query: indexerSearchDocsQuery,
variables: {
id: workspace.id,
input: { keyword: marker, limit: 0 },
},
})
);
const firstPage = await searchPage(workspace.id, marker, { limit: 1 });
const nextCursor = firstPage.pagination.nextCursor;
t.truthy(nextCursor);
const secondPage = await searchPage(workspace.id, marker, {
limit: 1,
cursor: nextCursor ?? undefined,
});
t.is(secondPage.nodes.length, 1);
t.not(
firstPage.nodes[0]?.fields.docId[0],
secondPage.nodes[0]?.fields.docId[0]
);
}
);
manticoreSearchE2e(
'enforces ACL changes through the Manticore Search provider',
async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
await app.create(Mockers.TeamWorkspace, { id: workspace.id });
const marker = `manticore-acl-${Date.now()}`;
await indexDocument(
workspace.id,
owner,
`${marker}-doc`,
marker,
DocRole.None
);
const member = await app.signup();
await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: member.id,
});
await reconcileSearch();
t.is(await searchCount(workspace.id, marker), 0);
t.is(await searchDocsCount(workspace.id, marker), 0);
await app.create(Mockers.DocUser, {
workspaceId: workspace.id,
docId: `${marker}-doc`,
userId: member.id,
type: DocRole.Reader,
});
await reconcileSearch();
t.true((await searchCount(workspace.id, marker)) > 0);
t.true((await searchDocsCount(workspace.id, marker)) > 0);
await app.models.docUser.delete(workspace.id, `${marker}-doc`, member.id);
await reconcileSearch();
t.is(await searchCount(workspace.id, marker), 0);
t.is(await searchDocsCount(workspace.id, marker), 0);
}
);

View File

@@ -0,0 +1,83 @@
import {
indexerAggregateQuery,
SearchQueryType,
SearchTable,
} from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { createDocWithMarkdown } from '../../../native';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
const indexer = app.get(Config).indexer;
const remoteE2e =
indexer.enabled && indexer.provider.type === 'elasticsearch' ? e2e : e2e.skip;
async function indexDocument(
workspaceId: string,
user: { id: string },
docId: string,
markdown: string
) {
await app.create(Mockers.DocMeta, { workspaceId, docId });
await app.create(Mockers.DocSnapshot, {
workspaceId,
docId,
user,
blob: createDocWithMarkdown(docId, markdown, docId),
});
const runtime = app.get(BackendRuntimeProvider);
for (let attempt = 0; attempt < 50; attempt++) {
await runtime.reconcileSearchProjection(1000);
if ((await runtime.searchStatus()).ready) return;
}
throw new Error('search projection did not become ready');
}
remoteE2e(
'exposes remote aggregation through the GraphQL contract',
async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
const suffix = Date.now();
for (const index of [0, 1, 2]) {
await indexDocument(
workspace.id,
owner,
`remote-aggregate-${suffix}-${index}`,
`remote aggregate marker ${index}`
);
}
const result = await app.gql({
query: indexerAggregateQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'remote aggregate marker',
},
field: 'docId',
options: {
pagination: { limit: 2, skip: 0 },
hits: {
pagination: { limit: 1, skip: 0 },
fields: ['docId', 'blockId'],
},
},
},
},
});
t.is(result.workspace.aggregate.buckets.length, 2);
t.is(result.workspace.aggregate.pagination.count, 2);
t.true(result.workspace.aggregate.pagination.hasMore);
}
);

View File

@@ -1,15 +1,19 @@
import { indexerSearchDocsQuery } from '@affine/graphql';
import { ConfigFactory } from '../../../base';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { createDocWithMarkdown } from '../../../native';
import { SearchProviderType } from '../../../plugins/indexer/config';
import { IndexerService } from '../../../plugins/indexer/service';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
e2e('should search docs by keyword', async t => {
const indexerE2e = app.get(Config).indexer.enabled ? e2e : e2e.skip;
indexerE2e('should search docs by keyword', async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner });
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
for (const docId of ['doc-0', 'doc-1', 'doc-2']) {
await app.create(Mockers.DocMeta, { workspaceId: workspace.id, docId });
await app.create(Mockers.DocSnapshot, {
@@ -18,33 +22,26 @@ e2e('should search docs by keyword', async t => {
user: owner,
blob: createDocWithMarkdown(docId, `${docId} hello`, docId),
});
await app.get(IndexerService).indexDoc(workspace.id, docId);
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
}
const search = app.gql({
query: indexerSearchDocsQuery,
variables: { id: workspace.id, input: { keyword: 'hello', limit: 2 } },
});
if (
app.get(ConfigFactory).config.indexer.provider.type ===
SearchProviderType.Manticoresearch
) {
await t.throwsAsync(search, {
message: /Invalid indexer input: unsupported_query/,
});
return;
}
const result = await search;
t.is(result.workspace.searchDocs.length, 2);
t.true(result.workspace.searchDocs.every(doc => doc.highlight.length > 0));
});
e2e(
indexerE2e(
'should search docs by keyword failed when workspace is no permission',
async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner });
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
await app.signup();
await t.throwsAsync(
app.gql({

View File

@@ -4,12 +4,15 @@ import {
SearchTable,
} from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { DocRole } from '../../../models';
import { createDocWithMarkdown } from '../../../native';
import { IndexerService } from '../../../plugins/indexer/service';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
const indexerE2e = app.get(Config).indexer.enabled ? e2e : e2e.skip;
async function indexDoc(
workspaceId: string,
user: { id: string },
@@ -24,12 +27,15 @@ async function indexDoc(
user,
blob: createDocWithMarkdown(docId, markdown, docId),
});
await app.get(IndexerService).indexDoc(workspaceId, docId);
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
}
e2e('should search with query', async t => {
indexerE2e('should search with query', async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner });
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
await indexDoc(
workspace.id,
owner,
@@ -114,100 +120,116 @@ e2e('should search with query', async t => {
);
});
e2e('should filter no read permission docs on team workspace', async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner });
await app.create(Mockers.TeamWorkspace, { id: workspace.id });
await indexDoc(
workspace.id,
owner,
'private-doc',
'team secret searchable',
DocRole.None
);
indexerE2e(
'should filter no read permission docs on team workspace',
async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
await app.create(Mockers.TeamWorkspace, { id: workspace.id });
await indexDoc(
workspace.id,
owner,
'private-doc',
'team secret searchable',
DocRole.None
);
const member = await app.signup();
await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: member.id,
});
await app.get(IndexerService).reconcileWorkspace(workspace.id);
const denied = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'secret',
const member = await app.signup();
await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: member.id,
});
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const denied = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'secret',
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
},
});
t.is(denied.workspace.search.pagination.count, 0);
});
t.is(denied.workspace.search.pagination.count, 0);
await app.create(Mockers.DocUser, {
workspaceId: workspace.id,
docId: 'private-doc',
userId: member.id,
type: DocRole.Reader,
});
await app.get(IndexerService).reconcileWorkspace(workspace.id);
const allowed = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'secret',
await app.create(Mockers.DocUser, {
workspaceId: workspace.id,
docId: 'private-doc',
userId: member.id,
type: DocRole.Reader,
});
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const allowed = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'secret',
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
},
});
t.true(allowed.workspace.search.pagination.count > 0);
});
t.true(allowed.workspace.search.pagination.count > 0);
await app.models.docUser.delete(workspace.id, 'private-doc', member.id);
await app.get(IndexerService).reconcileWorkspace(workspace.id);
const revoked = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'secret',
await app.models.docUser.delete(workspace.id, 'private-doc', member.id);
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const revoked = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.block,
query: {
type: SearchQueryType.match,
field: 'content',
match: 'secret',
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
},
});
t.is(revoked.workspace.search.pagination.count, 0);
});
});
t.is(revoked.workspace.search.pagination.count, 0);
}
);
e2e('should return empty results when search not match any docs', async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner });
await app.get(IndexerService).reconcileWorkspace(workspace.id);
const result = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.doc,
query: { type: SearchQueryType.match, field: 'title', match: 'absent' },
options: { fields: ['docId'], pagination: { limit: 10 } },
indexerE2e(
'should return empty results when search not match any docs',
async t => {
const owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const result = await app.gql({
query: indexerSearchQuery,
variables: {
id: workspace.id,
input: {
table: SearchTable.doc,
query: {
type: SearchQueryType.match,
field: 'title',
match: 'absent',
},
options: { fields: ['docId'], pagination: { limit: 10 } },
},
},
},
});
t.is(result.workspace.search.pagination.count, 0);
t.deepEqual(result.workspace.search.nodes, []);
});
});
t.is(result.workspace.search.pagination.count, 0);
t.deepEqual(result.workspace.search.nodes, []);
}
);

View File

@@ -6,14 +6,18 @@ export const TEST_LOG_LEVEL: LogLevel =
(process.env.TEST_LOG_LEVEL as LogLevel) ?? 'fatal';
async function flushDB(client: PrismaClient) {
const result: { tablename: string }[] =
await client.$queryRaw`SELECT tablename
const result: { schemaname: string; tablename: string }[] =
await client.$queryRaw`SELECT schemaname, tablename
FROM pg_catalog.pg_tables
WHERE schemaname != 'pg_catalog'
AND schemaname != 'information_schema'`;
const query = `TRUNCATE TABLE ${result
.map(({ tablename }) => tablename)
.filter(name => !name.includes('migrations'))
.filter(({ tablename }) => !tablename.includes('migrations'))
.map(({ schemaname, tablename }) =>
[schemaname, tablename]
.map(identifier => `"${identifier.replaceAll('"', '""')}"`)
.join('.')
)
.join(', ')}`;
for (let attempt = 0; attempt < 3; attempt++) {

View File

@@ -45,7 +45,7 @@ import { NotificationModule } from './core/notification';
import { PermissionModule } from './core/permission';
import { QueueDashboardModule } from './core/queue-dashboard';
import { QuotaModule } from './core/quota';
import { RealtimeModule } from './core/realtime';
import { RealtimeGatewayModule, RealtimeModule } from './core/realtime';
import { SelfhostModule } from './core/selfhost';
import { StaticFileModule } from './core/static-files';
import { StorageApiModule, StorageWorkerModule } from './core/storage';
@@ -61,7 +61,7 @@ import { CalendarModule } from './plugins/calendar';
import { CaptchaModule } from './plugins/captcha';
import { CopilotModule } from './plugins/copilot';
import { GCloudModule } from './plugins/gcloud';
import { IndexerModule, IndexerWorkerModule } from './plugins/indexer';
import { IndexerModule } from './plugins/indexer';
import { LicenseModule } from './plugins/license';
import { OAuthModule } from './plugins/oauth';
import { PaymentModule } from './plugins/payment';
@@ -167,10 +167,10 @@ export function buildAppModule(env: Env) {
factor
// basic
.use(...FunctionalityModules)
.useIf(() => !workerOnly, RealtimeGatewayModule)
// online roles publish indexer events; only the worker registers consumers
// Search API and worker runtime are separate from the queue worker application.
.useIf(() => env.isApi || env.isFrontend, IndexerModule)
.useIf(() => env.isWorker, IndexerWorkerModule)
// the worker owns doc consumers and schedulers
.useIf(() => env.isWorker, DocJobsModule)

View File

@@ -14,6 +14,7 @@ export type UserFriendlyErrorBaseType =
| 'no_permission'
| 'quota_exceeded'
| 'authentication_required'
| 'service_unavailable'
| 'internal_server_error';
type ErrorArgType = 'string' | 'number' | 'boolean';
@@ -36,6 +37,7 @@ const BaseTypeToHttpStatusMap: Record<UserFriendlyErrorBaseType, HttpStatus> = {
no_permission: HttpStatus.FORBIDDEN,
quota_exceeded: HttpStatus.PAYMENT_REQUIRED,
authentication_required: HttpStatus.UNAUTHORIZED,
service_unavailable: HttpStatus.SERVICE_UNAVAILABLE,
internal_server_error: HttpStatus.INTERNAL_SERVER_ERROR,
};
@@ -969,6 +971,25 @@ export const USER_FRIENDLY_ERRORS = {
},
// indexer errors
search_index_not_ready: {
type: 'service_unavailable',
args: { spaceId: 'string' },
message: ({ spaceId }) =>
`Search index for Space ${spaceId} is not ready yet.`,
},
search_permission_syncing: {
type: 'service_unavailable',
message: 'Search permissions are still syncing. Please try again shortly.',
},
search_provider_unavailable: {
type: 'service_unavailable',
message: 'Search provider is temporarily unavailable.',
},
search_index_failed: {
type: 'service_unavailable',
args: { diagnosticId: 'string' },
message: 'Search index is temporarily unavailable.',
},
search_provider_not_found: {
type: 'resource_not_found',
message: 'Search provider not found.',

View File

@@ -1123,6 +1123,38 @@ export class InvalidAppConfigInput extends UserFriendlyError {
super('invalid_input', 'invalid_app_config_input', message, args);
}
}
@ObjectType()
class SearchIndexNotReadyDataType {
@Field() spaceId!: string
}
export class SearchIndexNotReady extends UserFriendlyError {
constructor(args: SearchIndexNotReadyDataType, message?: string | ((args: SearchIndexNotReadyDataType) => string)) {
super('service_unavailable', 'search_index_not_ready', message, args);
}
}
export class SearchPermissionSyncing extends UserFriendlyError {
constructor(message?: string) {
super('service_unavailable', 'search_permission_syncing', message);
}
}
export class SearchProviderUnavailable extends UserFriendlyError {
constructor(message?: string) {
super('service_unavailable', 'search_provider_unavailable', message);
}
}
@ObjectType()
class SearchIndexFailedDataType {
@Field() diagnosticId!: string
}
export class SearchIndexFailed extends UserFriendlyError {
constructor(args: SearchIndexFailedDataType, message?: string | ((args: SearchIndexFailedDataType) => string)) {
super('service_unavailable', 'search_index_failed', message, args);
}
}
export class SearchProviderNotFound extends UserFriendlyError {
constructor(message?: string) {
@@ -1320,6 +1352,10 @@ export enum ErrorNames {
MENTION_USER_ONESELF_DENIED,
INVALID_APP_CONFIG,
INVALID_APP_CONFIG_INPUT,
SEARCH_INDEX_NOT_READY,
SEARCH_PERMISSION_SYNCING,
SEARCH_PROVIDER_UNAVAILABLE,
SEARCH_INDEX_FAILED,
SEARCH_PROVIDER_NOT_FOUND,
INVALID_SEARCH_PROVIDER_REQUEST,
INVALID_INDEXER_INPUT,
@@ -1335,5 +1371,5 @@ registerEnumType(ErrorNames, {
export const ErrorDataUnionType = createUnionType({
name: 'ErrorDataUnion',
types: () =>
[GraphqlBadRequestDataType, HttpRequestErrorDataType, SsrfBlockedErrorDataType, ResponseTooLargeErrorDataType, ImageFormatNotSupportedDataType, QueryTooLongDataType, ValidationErrorDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, InvalidOauthCallbackCodeDataType, MissingOauthQueryParameterDataType, InvalidOauthResponseDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocActionDeniedDataType, DocUpdateBlockedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, NoMoreSeatDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CalendarProviderRequestErrorDataType, NoCopilotProviderAvailableDataType, CopilotFailedToGenerateEmbeddingDataType, CopilotDocNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderNotSupportedDataType, CopilotProviderSideErrorDataType, CopilotFailedToAddWorkspaceArtifactDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseToActivateDataType, InvalidLicenseUpdateParamsDataType, UnsupportedClientVersionDataType, UnsupportedServerVersionDataType, MentionUserDocAccessDeniedDataType, InvalidAppConfigDataType, InvalidAppConfigInputDataType, InvalidSearchProviderRequestDataType, InvalidIndexerInputDataType] as const,
[GraphqlBadRequestDataType, HttpRequestErrorDataType, SsrfBlockedErrorDataType, ResponseTooLargeErrorDataType, ImageFormatNotSupportedDataType, QueryTooLongDataType, ValidationErrorDataType, WrongSignInCredentialsDataType, UnknownOauthProviderDataType, InvalidOauthCallbackCodeDataType, MissingOauthQueryParameterDataType, InvalidOauthResponseDataType, InvalidEmailDataType, InvalidPasswordLengthDataType, WorkspacePermissionNotFoundDataType, SpaceNotFoundDataType, MemberNotFoundInSpaceDataType, NotInSpaceDataType, AlreadyInSpaceDataType, SpaceAccessDeniedDataType, SpaceOwnerNotFoundDataType, SpaceShouldHaveOnlyOneOwnerDataType, DocNotFoundDataType, DocActionDeniedDataType, DocUpdateBlockedDataType, VersionRejectedDataType, InvalidHistoryTimestampDataType, DocHistoryNotFoundDataType, BlobNotFoundDataType, ExpectToGrantDocUserRolesDataType, ExpectToRevokeDocUserRolesDataType, ExpectToUpdateDocUserRoleDataType, NoMoreSeatDataType, UnsupportedSubscriptionPlanDataType, SubscriptionAlreadyExistsDataType, SubscriptionNotExistsDataType, SameSubscriptionRecurringDataType, SubscriptionPlanNotFoundDataType, CalendarProviderRequestErrorDataType, NoCopilotProviderAvailableDataType, CopilotFailedToGenerateEmbeddingDataType, CopilotDocNotFoundDataType, CopilotMessageNotFoundDataType, CopilotPromptNotFoundDataType, CopilotProviderNotSupportedDataType, CopilotProviderSideErrorDataType, CopilotFailedToAddWorkspaceArtifactDataType, RuntimeConfigNotFoundDataType, InvalidRuntimeConfigTypeDataType, InvalidLicenseToActivateDataType, InvalidLicenseUpdateParamsDataType, UnsupportedClientVersionDataType, UnsupportedServerVersionDataType, MentionUserDocAccessDeniedDataType, InvalidAppConfigDataType, InvalidAppConfigInputDataType, SearchIndexNotReadyDataType, SearchIndexFailedDataType, InvalidSearchProviderRequestDataType, InvalidIndexerInputDataType] as const,
});

View File

@@ -71,14 +71,6 @@ defineModuleConfig('job', {
schema,
},
'queues.indexer': {
desc: 'The config for indexer job queue',
default: {
concurrency: 1,
},
schema,
},
'queues.notification': {
desc: 'The config for notification job queue',
default: {

View File

@@ -27,7 +27,6 @@ export enum Queue {
NOTIFICATION = 'notification',
DOC = 'doc',
COPILOT = 'copilot',
INDEXER = 'indexer',
CALENDAR = 'calendar',
BACKENDRUNTIME = 'backendRuntime',
INVITE_ABUSE = 'inviteAbuse',

View File

@@ -1,11 +1,7 @@
import { ServerRole } from '../../../env';
import { Queue, QUEUES } from './def';
export const WORKER_QUEUES = [
Queue.DOC,
Queue.INDEXER,
Queue.BACKENDRUNTIME,
] as const;
export const WORKER_QUEUES = [Queue.DOC, Queue.BACKENDRUNTIME] as const;
export function queuesForRole(role: ServerRole | undefined): Queue[] {
switch (role) {

View File

@@ -63,6 +63,7 @@ export type KnownMetricScopes =
| 'storage'
| 'process'
| 'permission'
| 'search'
| 'workspace';
const metricCreators: MetricCreators = {

View File

@@ -18,7 +18,7 @@ export interface AuthConfig {
allowSignupForOauth: boolean;
requireEmailDomainVerification: boolean;
requireEmailVerification: boolean;
newAccountShareActionDelay: number;
newAccountActionDelay: number;
trustedCloudflareHeaders: boolean;
signInRateLimit: ConfigItem<{
ttl: number;
@@ -56,8 +56,8 @@ defineModuleConfig('auth', {
desc: 'Whether require email verification before accessing restricted resources(not implemented).',
default: true,
},
newAccountShareActionDelay: {
desc: 'Minimum account age in seconds before new accounts can invite members or create share links.',
newAccountActionDelay: {
desc: 'Minimum account age in seconds before new accounts can invite members, create invite links, or publish documents. Set to 0 to disable.',
default: 24 * 60 * 60,
shape: z.number().int().min(0),
},

View File

@@ -5,6 +5,7 @@ import {
BackendRuntimeEmbeddingProducer,
BackendRuntimeEmbeddingService,
BackendRuntimeHousekeepingJob,
BackendRuntimeSearchJob,
} from './job';
import {
BACKEND_RUNTIME_CONFIG_PATHS,
@@ -33,7 +34,11 @@ export class BackendRuntimeProducerModule {}
@Module({
imports: [BackendRuntimeModule],
providers: [BackendRuntimeEmbeddingJob, BackendRuntimeHousekeepingJob],
providers: [
BackendRuntimeEmbeddingJob,
BackendRuntimeHousekeepingJob,
BackendRuntimeSearchJob,
],
})
export class BackendRuntimeWorkerModule {}
@@ -42,6 +47,7 @@ export {
BackendRuntimeEmbeddingProducer,
BackendRuntimeEmbeddingService,
BackendRuntimeHousekeepingJob,
BackendRuntimeSearchJob,
} from './job';
export {
BACKEND_RUNTIME_CONFIG_PATHS,

View File

@@ -2,11 +2,13 @@ import { Injectable, Logger } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule';
import {
ConfigFactory,
CopilotSelectedSourcesFailed,
CopilotSelectedSourcesLimitExceeded,
CopilotSelectedSourcesProcessing,
CopilotSelectedSourcesUnavailable,
JobQueue,
metrics,
OnEvent,
OnJob,
} from '../../base';
@@ -30,6 +32,9 @@ declare global {
'backendRuntime.reconcileDocumentEmbeddings': {
workspaceId: string;
};
'backendRuntime.reconcileSearchProjection': {
limit?: number;
};
}
}
@@ -257,3 +262,98 @@ export class BackendRuntimeHousekeepingJob {
return total;
}
}
@Injectable()
export class BackendRuntimeSearchJob {
constructor(
private readonly rt: BackendRuntimeProvider,
private readonly queue: JobQueue,
private readonly config: ConfigFactory
) {}
@Cron(CronExpression.EVERY_30_SECONDS)
async scheduleReconciliation() {
if (!this.config.config.indexer.enabled) return;
await this.queue.add(
'backendRuntime.reconcileSearchProjection',
{ limit: 100 },
{ jobId: 'backend-runtime-search-reconciliation', removeOnFail: true }
);
}
@OnJob('backendRuntime.reconcileSearchProjection')
async reconcileProjection({
limit = 100,
}: Jobs['backendRuntime.reconcileSearchProjection']) {
if (!this.config.config.indexer.enabled) return 0;
const startedAt = performance.now();
try {
const reconciled = await this.rt.reconcileSearchProjection(limit);
const status = (await this.rt.searchStatus()) as {
ready?: boolean;
state?: string;
metrics?: {
scanCursor?: number;
scanHighWater?: number;
pendingPublications?: number;
gcBacklog?: number;
providerRequests?: number;
providerLatencyMicrosAvg?: number;
generationGcFailures?: number;
filterDrops?: {
missingPublished?: number;
projectionMismatch?: number;
canonicalPermission?: number;
};
};
};
metrics.search.counter('reconcile_runs').add(1);
metrics.search.gauge('reconciled_workspaces').record(reconciled);
metrics.search.gauge('generation_ready').record(status.ready ? 1 : 0, {
state: status.state ?? 'unknown',
});
metrics.search
.histogram('reconcile_latency_ms')
.record(performance.now() - startedAt);
const projection = status.metrics;
if (projection) {
metrics.search.gauge('scan_cursor').record(projection.scanCursor ?? 0);
metrics.search
.gauge('scan_high_water')
.record(projection.scanHighWater ?? 0);
metrics.search
.gauge('pending_publications')
.record(projection.pendingPublications ?? 0);
metrics.search.gauge('gc_backlog').record(projection.gcBacklog ?? 0);
metrics.search
.gauge('provider_requests')
.record(projection.providerRequests ?? 0);
metrics.search
.gauge('provider_latency_micros_avg')
.record(projection.providerLatencyMicrosAvg ?? 0);
metrics.search
.gauge('generation_gc_failures')
.record(projection.generationGcFailures ?? 0);
metrics.search
.gauge('filter_drops')
.record(projection.filterDrops?.missingPublished ?? 0, {
reason: 'missing_published',
});
metrics.search
.gauge('filter_drops')
.record(projection.filterDrops?.projectionMismatch ?? 0, {
reason: 'projection_mismatch',
});
metrics.search
.gauge('filter_drops')
.record(projection.filterDrops?.canonicalPermission ?? 0, {
reason: 'canonical_permission',
});
}
return reconciled;
} catch (error) {
metrics.search.counter('reconcile_failures').add(1);
throw error;
}
}
}

Some files were not shown because too many files have changed in this diff Show More