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 -->
This commit is contained in:
DarkSky authored and GitHub committed 2026-08-23 17:54:49 +08:00
1 parent a8eef53966
commit b530198a3b
107 files changed
+8612 -5306

No files matched your search

-15
View File
@@ -12,7 +12,6 @@ services:
environment: environment:
DATABASE_URL: postgresql://affine:affine@db:5432/affine DATABASE_URL: postgresql://affine:affine@db:5432/affine
REDIS_SERVER_HOST: redis REDIS_SERVER_HOST: redis
AFFINE_INDEXER_SEARCH_ENDPOINT: http://indexer:9308
db: db:
image: pgvector/pgvector:pg16 image: pgvector/pgvector:pg16
@@ -26,19 +25,5 @@ services:
redis: redis:
image: 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: volumes:
postgres-data: postgres-data:
manticoresearch_data:
-3
View File
@@ -10,6 +10,3 @@ DB_DATABASE_NAME=affine
# enable for arm64, e.g.: macOS M1+ # enable for arm64, e.g.: macOS M1+
# ELASTIC_VERSION_ARM64=-arm64 # ELASTIC_VERSION_ARM64=-arm64
# ELASTIC_PLATFORM=linux/arm64 # ELASTIC_PLATFORM=linux/arm64
# manticoresearch
MANTICORE_VERSION=10.1.0
-17
View File
@@ -32,22 +32,6 @@ services:
volumes: volumes:
- mailpit_data:/data - 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: # elasticsearch:
# image: docker.elastic.co/elasticsearch/elasticsearch:${ELASTIC_VERSION:-9.0.1}${ELASTIC_VERSION_ARM64} # image: docker.elastic.co/elasticsearch/elasticsearch:${ELASTIC_VERSION:-9.0.1}${ELASTIC_VERSION_ARM64}
# platform: ${ELASTIC_PLATFORM} # platform: ${ELASTIC_PLATFORM}
@@ -94,6 +78,5 @@ networks:
volumes: volumes:
postgres_data: postgres_data:
manticoresearch_data:
mailpit_data: mailpit_data:
elasticsearch_data: elasticsearch_data:
+6 -23
View File
@@ -86,18 +86,6 @@
"concurrency": 1 "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": { "queues.notification": {
"type": "object", "type": "object",
"description": "The config for notification job queue\n@default {\"concurrency\":10}", "description": "The config for notification job queue\n@default {\"concurrency\":10}",
@@ -1488,38 +1476,33 @@
"properties": { "properties": {
"enabled": { "enabled": {
"type": "boolean", "type": "boolean",
"description": "Enable indexer plugin\n@default false\n@environment `AFFINE_INDEXER_ENABLED`", "description": "Enable indexer plugin\n@default false",
"default": false "default": false
}, },
"provider.type": { "provider.type": {
"type": "string", "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" "default": "embedded"
}, },
"provider.endpoint": { "provider.endpoint": {
"type": "string", "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": "" "default": ""
}, },
"provider.apiKey": { "provider.apiKey": {
"type": "string", "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": "" "default": ""
}, },
"provider.username": { "provider.username": {
"type": "string", "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": "" "default": ""
}, },
"provider.password": { "provider.password": {
"type": "string", "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": "" "default": ""
},
"autoIndex.batchSize": {
"type": "number",
"description": "Number of workspaces automatically indexed per batch\n@default 10",
"default": 10
} }
} }
}, },
@@ -33,6 +33,9 @@ runs:
shell: bash shell: bash
run: | run: |
case '${{ inputs.indexer-provider }}' in case '${{ inputs.indexer-provider }}' in
disabled)
default_indexer='{"enabled":false,"provider":{"type":"embedded","endpoint":""}}'
;;
embedded) embedded)
default_indexer='{"enabled":true,"provider":{"type":"embedded","endpoint":""}}' default_indexer='{"enabled":true,"provider":{"type":"embedded","endpoint":""}}'
;; ;;
+87 -30
View File
@@ -933,8 +933,6 @@ jobs:
NODE_ENV: test NODE_ENV: test
DATABASE_URL: postgresql://affine:affine@localhost:5432/affine DATABASE_URL: postgresql://affine:affine@localhost:5432/affine
REDIS_SERVER_HOST: localhost REDIS_SERVER_HOST: localhost
SERVER_CONFIG: >-
{"indexer":{"enabled":true,"provider":{"type":"elasticsearch","endpoint":"http://localhost:9200"}}}
services: services:
postgres: postgres:
image: pgvector/pgvector:pg16 image: pgvector/pgvector:pg16
@@ -1005,8 +1003,74 @@ jobs:
name: affine name: affine
fail_ci_if_error: false fail_ci_if_error: false
server-test-manticore: server-test-embedded:
name: Server Test with Manticore 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 runs-on: ubuntu-latest
needs: needs:
- build-server-native - build-server-native
@@ -1014,9 +1078,6 @@ jobs:
NODE_ENV: test NODE_ENV: test
DATABASE_URL: postgresql://affine:affine@localhost:5432/affine DATABASE_URL: postgresql://affine:affine@localhost:5432/affine
REDIS_SERVER_HOST: localhost REDIS_SERVER_HOST: localhost
SEARCH_MS_URL: http://localhost:9308
SERVER_CONFIG: >-
{"indexer":{"enabled":true,"provider":{"type":"manticoresearch","endpoint":"http://localhost:9308"}}}
services: services:
postgres: postgres:
image: pgvector/pgvector:pg16 image: pgvector/pgvector:pg16
@@ -1038,15 +1099,18 @@ jobs:
ports: ports:
- 1025:1025 - 1025:1025
- 8025:8025 - 8025:8025
indexer: manticore:
image: manticoresearch/manticore:10.1.0 image: manticoresearch/manticore:29.0.2
options: >- env:
--health-cmd "wget -q -O /dev/null http://localhost:9308/" 'searchd.listen': 'http:9308'
--health-interval 5s 'searchd.network_timeout': '30'
--health-timeout 3s
--health-retries 12
ports: ports:
- 9308:9308 - 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: steps:
- uses: actions/checkout@v7 - uses: actions/checkout@v7
@@ -1057,13 +1121,6 @@ jobs:
electron-install: false electron-install: false
full-cache: true 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 - name: Download server-native.node
uses: actions/download-artifact@v4 uses: actions/download-artifact@v4
with: with:
@@ -1071,17 +1128,14 @@ jobs:
path: ./packages/backend/native path: ./packages/backend/native
- name: Prepare Server Test Environment - name: Prepare Server Test Environment
uses: ./.github/actions/server-test-env
with: with:
indexer-provider: manticoresearch indexer-provider: manticoresearch
uses: ./.github/actions/server-test-env
- name: Run server tests with Manticore - name: Run Manticore Search provider E2E
run: yarn affine @affine/server e2e:coverage src/__tests__/e2e/indexer/search-docs.spec.ts --forbid-only run: yarn affine @affine/server e2e:coverage src/__tests__/e2e/indexer/manticore-provider.spec.ts --forbid-only
env:
- name: Run Manticore provider contract CARGO_TARGET_DIR: '${{ github.workspace }}/target'
run: >-
cargo test --manifest-path packages/backend/native/Cargo.toml
remote_providers_apply_search_and_delete_the_same_contract --lib
- name: Upload server test coverage results - name: Upload server test coverage results
uses: codecov/codecov-action@v5 uses: codecov/codecov-action@v5
@@ -1135,6 +1189,8 @@ jobs:
path: ./packages/backend/native path: ./packages/backend/native
- name: Prepare Server Test Environment - name: Prepare Server Test Environment
with:
indexer-provider: disabled
uses: ./.github/actions/server-test-env uses: ./.github/actions/server-test-env
- name: Run server tests - name: Run server tests
@@ -1597,7 +1653,8 @@ jobs:
- native-unit-test - native-unit-test
- server-test - server-test
- server-test-elasticsearch - server-test-elasticsearch
- server-test-manticore - server-test-embedded
- server-test-manticoresearch
- server-e2e-test - server-e2e-test
- rust-test - rust-test
- rust-test-filter - rust-test-filter
+1 -38
View File
@@ -56,10 +56,6 @@ export declare class BackendRuntime {
getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null> getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean> revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean>
cleanupExpiredRuntimeStates(limit: number): Promise<number> 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) constructor(privateKey?: string | undefined | null, configPaths?: Array<string> | undefined | null)
start(): Promise<void> start(): Promise<void>
stop(): Promise<void> stop(): Promise<void>
@@ -68,10 +64,7 @@ export declare class BackendRuntime {
runMigrations(): Promise<void> runMigrations(): Promise<void>
searchAuthorized(actorUserId: string, workspaceId: string, request: RuntimeSearchRequest): Promise<SearchOperationOutput> searchAuthorized(actorUserId: string, workspaceId: string, request: RuntimeSearchRequest): Promise<SearchOperationOutput>
aggregateAuthorized(actorUserId: string, workspaceId: string, request: RuntimeAggregateRequest): Promise<SearchOperationOutput> aggregateAuthorized(actorUserId: string, workspaceId: string, request: RuntimeAggregateRequest): Promise<SearchOperationOutput>
indexSearchDocument(workspaceId: string, docId: string): Promise<void> reconcileSearchProjection(limit?: number | undefined | null): Promise<number>
deleteSearchDocument(workspaceId: string, docId: string): Promise<void>
reconcileSearchWorkspace(workspaceId: string): Promise<void>
deleteSearchWorkspace(workspaceId: string): Promise<void>
filterReadableDocs(actorUserId: string, workspaceId: string, docIds: Array<string>): Promise<Array<string>> filterReadableDocs(actorUserId: string, workspaceId: string, docIds: Array<string>): Promise<Array<string>>
searchStatus(): Promise<any> searchStatus(): Promise<any>
embeddingHealth(): Promise<EmbeddingHealth> embeddingHealth(): Promise<EmbeddingHealth>
@@ -113,7 +106,6 @@ export declare class StorageRuntime {
rebuildWorkspaceDocBlobRefs(workspaceId: string, limit: number): Promise<RuntimeDocBlobRefsResult> rebuildWorkspaceDocBlobRefs(workspaceId: string, limit: number): Promise<RuntimeDocBlobRefsResult>
reconcileWorkspaceDocuments(workspaceId: string): Promise<RuntimeDocumentCleanupReconcileResult> reconcileWorkspaceDocuments(workspaceId: string): Promise<RuntimeDocumentCleanupReconcileResult>
executeDocumentCleanupCandidates(workspaceId: string | undefined | null, gracePeriodDays: number, limit: number): Promise<RuntimeDocumentCleanupExecuteResult> executeDocumentCleanupCandidates(workspaceId: string | undefined | null, gracePeriodDays: number, limit: number): Promise<RuntimeDocumentCleanupExecuteResult>
ackDocumentCleanupEffect(workspaceId: string, docId: string, cleanupVersion: string, effect: string): Promise<RuntimeDocumentCleanupAckResult>
constructor() constructor()
start(): Promise<void> start(): Promise<void>
configure(configJson: string): void configure(configJson: string): void
@@ -1261,16 +1253,11 @@ export interface RuntimeDocHistoryInput {
historyMaxAgeMs: number historyMaxAgeMs: number
} }
export interface RuntimeDocumentCleanupAckResult {
completed: boolean
}
export interface RuntimeDocumentCleanupEffect { export interface RuntimeDocumentCleanupEffect {
workspaceId: string workspaceId: string
docId: string docId: string
cleanupVersion: string cleanupVersion: string
commentObjectsDone: boolean commentObjectsDone: boolean
searchDone: boolean
} }
export interface RuntimeDocumentCleanupExecuteResult { export interface RuntimeDocumentCleanupExecuteResult {
@@ -1551,30 +1538,6 @@ export interface RuntimeWorkspaceInviteQuotaUsage {
targetDomains: Array<RuntimeQuotaTargetDomainInput> 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 declare function safeFetch(request: SafeFetchRequest): Promise<SafeFetchResponse>
export type SafeFetchMethod = 'get'| export type SafeFetchMethod = 'get'|
@@ -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 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_ID_PURPOSE: &str = "workspace_invite_link:id";
pub(super) const WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE: &str = "workspace_invite_link:workspace"; 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;
@@ -16,7 +16,6 @@ mod scope_compiler;
mod search; mod search;
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
mod workspace_stats;
use std::{ use std::{
sync::{Arc, RwLock}, sync::{Arc, RwLock},
time::Duration, time::Duration,
@@ -26,6 +25,8 @@ use byok::LocalLeasePayload;
use copilot::{backend_provider, executable_protocol}; use copilot::{backend_provider, executable_protocol};
use embedding::register_artifact_source; use embedding::register_artifact_source;
use napi::{Result, bindgen_prelude::Buffer}; use napi::{Result, bindgen_prelude::Buffer};
#[cfg(test)]
pub(crate) use search::SEARCH_TEST_LOCK;
use search::SearchRuntime; use search::SearchRuntime;
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions}; use sqlx::{PgPool, Row, postgres::PgPoolOptions};
@@ -66,10 +67,12 @@ fn search_operation_output(result: RuntimeResult<serde_json::Value>) -> SearchOp
error_code: Some( error_code: Some(
match error { match error {
RuntimeError::SearchWorkspaceDenied => "workspace_denied", 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::SearchProviderUnavailable => "provider_unavailable",
RuntimeError::SearchUnsupportedQuery => "unsupported_query", RuntimeError::SearchUnsupportedQuery => "unsupported_query",
RuntimeError::SearchReplayGap => "provider_unavailable",
RuntimeError::InvalidInput(_) | RuntimeError::Json { .. } => "invalid_request", RuntimeError::InvalidInput(_) | RuntimeError::Json { .. } => "invalid_request",
_ => "internal", _ => "internal",
} }
@@ -358,46 +361,14 @@ impl BackendRuntime {
} }
#[napi] #[napi]
pub async fn index_search_document(&self, workspace_id: String, doc_id: String) -> Result<()> { pub async fn reconcile_search_projection(&self, limit: Option<i32>) -> Result<i32> {
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<()> {
self.require_background()?; self.require_background()?;
self self
.search_runtime() .search_runtime()
.await? .await?
.reconcile_workspace(permission::SystemSearchCapability::ReconcileIndex, &workspace_id) .reconcile_pending(limit.unwrap_or(100))
.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)
.await .await
.map(|count| count as i32)
.map_err(to_napi_error) .map_err(to_napi_error)
} }
@@ -29,7 +29,7 @@ impl PermissionAuthorizer {
) -> RuntimeResult<AuthorizedSearchScope> { ) -> RuntimeResult<AuthorizedSearchScope> {
match actor { match actor {
SearchActor::User { user_id } => { 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")) let owner_or_admin = matches!(snapshot.evaluation.workspace.role.as_deref(), Some("owner" | "admin"))
&& snapshot.evaluation.workspace.member_state.as_deref() == Some("active"); && snapshot.evaluation.workspace.member_state.as_deref() == Some("active");
let decision = evaluate_permission(snapshot.evaluation) let decision = evaluate_permission(snapshot.evaluation)
@@ -54,17 +54,12 @@ impl PermissionAuthorizer {
}; };
Ok(AuthorizedSearchScope { Ok(AuthorizedSearchScope {
workspace_id: workspace_id.to_string(), workspace_id: workspace_id.to_string(),
permission_revision: snapshot.revision,
docs, 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( pub(in crate::runtime::backend_runtime) async fn filter_readable_docs(
&self, &self,
workspace_id: &str, workspace_id: &str,
@@ -72,6 +67,9 @@ impl PermissionAuthorizer {
doc_ids: Vec<String>, doc_ids: Vec<String>,
) -> RuntimeResult<BTreeSet<String>> { ) -> RuntimeResult<BTreeSet<String>> {
let snapshot = self.store.permission_snapshot(workspace_id, user_id, &doc_ids).await?; 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)?; let output = evaluate_permission(snapshot.evaluation).map_err(|_| RuntimeError::SearchPermissionUnavailable)?;
Ok( Ok(
output output
@@ -3,9 +3,7 @@ mod store;
mod types; mod types;
pub(super) use authorizer::PermissionAuthorizer; pub(super) use authorizer::PermissionAuthorizer;
#[cfg(test)] pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor};
pub(super) use types::AclPredicate;
pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor, SystemSearchCapability};
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
@@ -20,10 +20,6 @@ impl PermissionStore {
Self { pool } 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( pub(super) async fn permission_snapshot(
&self, &self,
workspace_id: &str, workspace_id: &str,
@@ -40,12 +36,10 @@ impl PermissionStore {
.await .await
.map_err(|error| RuntimeError::database("configure permission snapshot", error))?; .map_err(|error| RuntimeError::database("configure permission snapshot", error))?;
let row = sqlx::query( let row = sqlx::query(
r#"SELECT revision.revision, r#"SELECT policy.visibility, coalesce(policy.sharing_enabled, true) AS sharing_enabled,
policy.visibility, coalesce(policy.sharing_enabled, true) AS sharing_enabled,
coalesce(policy.member_default_doc_role, 'manager') AS member_default_doc_role, coalesce(policy.member_default_doc_role, 'manager') AS member_default_doc_role,
member.role, member.state member.role, member.state
FROM workspaces workspace 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 workspace_access_policies policy ON policy.workspace_id=workspace.id
LEFT JOIN LATERAL ( LEFT JOIN LATERAL (
SELECT role,state FROM workspace_members SELECT role,state FROM workspace_members
@@ -60,10 +54,6 @@ impl PermissionStore {
.await .await
.map_err(|error| RuntimeError::database("load workspace permission facts", error))? .map_err(|error| RuntimeError::database("load workspace permission facts", error))?
.ok_or_else(|| RuntimeError::invalid_input("workspace_not_found"))?; .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 let sharing_enabled: bool = row
.try_get("sharing_enabled") .try_get("sharing_enabled")
.map_err(|error| RuntimeError::database("decode workspace sharing", error))?; .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")); member_state.as_deref() == Some("active") && matches!(role.as_deref(), Some("member" | "admin" | "owner"));
Ok(PermissionSnapshot { Ok(PermissionSnapshot {
revision,
capability, capability,
evaluation: PermissionEvaluationInputV1 { evaluation: PermissionEvaluationInputV1 {
version: 1, version: 1,
@@ -163,15 +152,6 @@ impl PermissionStore {
sharing_enabled, 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( async fn load_doc_acl_capability(
@@ -7,7 +7,14 @@ static PERMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_
async fn setup() -> Option<(PgPool, String, String)> { async fn setup() -> Option<(PgPool, String, String)> {
let database_url = std::env::var("DATABASE_URL").ok()?; let database_url = std::env::var("DATABASE_URL").ok()?;
let pool = PgPool::connect(&database_url).await.unwrap(); 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 suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("search-permission-user-{suffix}"); let user_id = format!("search-permission-user-{suffix}");
let workspace_id = format!("search-permission-workspace-{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_eq!(predicate.actor_user_id, user_id);
assert!(predicate.active_member); 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] #[tokio::test]
@@ -112,13 +127,20 @@ async fn inactive_member_is_denied_and_unknown_capability_fails_closed() {
} }
#[tokio::test] #[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 _guard = PERMISSION_TEST_LOCK.lock().await;
let Some((pool, workspace_id, user_id)) = setup().await else { let Some((pool, workspace_id, user_id)) = setup().await else {
return; return;
}; };
let authorizer = PermissionAuthorizer::new(pool.clone()); 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( sqlx::query(
"INSERT INTO doc_grants(workspace_id,doc_id,principal_type,principal_id,role) VALUES($1,'doc','user',$2,'reader')", "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) .execute(&pool)
.await .await
.unwrap(); .unwrap();
let after = authorizer.revision(&workspace_id).await.unwrap(); let readable = authorizer
assert_eq!(after, before + 1); .filter_readable_docs(&workspace_id, &user_id, vec!["doc".to_string(), "hidden".to_string()])
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)
.await .await
.unwrap(); .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("UPDATE doc_access_policies SET member_default_role='reader' WHERE workspace_id=$1 AND doc_id='hidden'")
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)
.bind(&workspace_id) .bind(&workspace_id)
.execute(&pool) .execute(&pool)
.await .await
.unwrap(); .unwrap();
assert_eq!(authorizer.revision(&workspace_id).await.unwrap(), after + 1); let readable = authorizer
assert_eq!(authorizer.revision(&moved_workspace_id).await.unwrap(), 1); .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()
);
} }
@@ -1,10 +1,5 @@
use crate::permission::PermissionEvaluationInputV1; use crate::permission::PermissionEvaluationInputV1;
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum SystemSearchCapability {
ReconcileIndex,
}
#[derive(Clone, Debug, Eq, PartialEq)] #[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum SearchActor { pub(in crate::runtime::backend_runtime) enum SearchActor {
User { user_id: String }, User { user_id: String },
@@ -33,12 +28,10 @@ pub(in crate::runtime::backend_runtime) enum DocReadScope {
#[derive(Clone, Debug, Eq, PartialEq)] #[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) struct AuthorizedSearchScope { pub(in crate::runtime::backend_runtime) struct AuthorizedSearchScope {
pub(in crate::runtime::backend_runtime) workspace_id: String, 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(in crate::runtime::backend_runtime) docs: DocReadScope,
} }
pub(super) struct PermissionSnapshot { pub(super) struct PermissionSnapshot {
pub(super) revision: i64,
pub(super) capability: DocAclCapability, pub(super) capability: DocAclCapability,
pub(super) evaluation: PermissionEvaluationInputV1, pub(super) evaluation: PermissionEvaluationInputV1,
pub(super) actor_user_id: String, pub(super) actor_user_id: String,
@@ -250,7 +250,15 @@ mod tests {
}; };
let _guard = crate::runtime::migrations::EMBEDDING_TEST_LOCK.lock().await; let _guard = crate::runtime::migrations::EMBEDDING_TEST_LOCK.lock().await;
let pool = PgPool::connect(&database_url).await.unwrap(); 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 suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("scope-user-{suffix}"); let user_id = format!("scope-user-{suffix}");
let collaborator_id = format!("scope-collaborator-{suffix}"); let collaborator_id = format!("scope-collaborator-{suffix}");
@@ -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()
}
@@ -3,7 +3,7 @@ use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row}; use sqlx::{PgPool, Row};
use uuid::Uuid; use uuid::Uuid;
use super::{SCHEMA_FINGERPRINT, provider::RemoteProvider, types::SearchTable}; use super::{SCHEMA_FINGERPRINT, SearchProvider, SearchTable};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig}; use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
#[derive(Clone)] #[derive(Clone)]
@@ -22,82 +22,120 @@ impl ActiveGeneration {
} }
} }
pub(super) async fn prepare( pub(super) async fn ensure(
pool: &PgPool, pool: &PgPool,
config: &SearchRuntimeConfig, config: &SearchRuntimeConfig,
remote: Option<&RemoteProvider>, remote: Option<&SearchProvider>,
rebuild_embedded: bool,
) -> RuntimeResult<ActiveGeneration> { ) -> RuntimeResult<ActiveGeneration> {
let fingerprint = config_fingerprint(config); let expected_config_hash = config_hash(config);
let mut transaction = pool let mut transaction = pool
.begin() .begin()
.await .await
.map_err(|error| RuntimeError::database("begin search generation", error))?; .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) .execute(&mut *transaction)
.await .await
.map_err(|error| RuntimeError::database("lock search generation", error))?; .map_err(|error| RuntimeError::database("lock search generation", error))?;
let existing = sqlx::query( let existing = sqlx::query(
r#"SELECT generation_id,provider,manifest FROM search_runtime_generations r#"SELECT id,provider,config_hash,schema_version,manifest,state
WHERE state IN ('active','pending') AND provider=$1 AND config_fingerprint=$2 AND schema_fingerprint=$3 FROM search_projection.generations
ORDER BY (state='active') DESC LIMIT 1"#, WHERE state='building'
ORDER BY created_at DESC
LIMIT 1"#,
) )
.bind(&config.provider)
.bind(&fingerprint)
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(&mut *transaction) .fetch_optional(&mut *transaction)
.await .await
.map_err(|error| RuntimeError::database("load search generation", error))?; .map_err(|error| RuntimeError::database("load search generation", error))?;
let generation = if let Some(row) = existing { let generation = if let Some(row) = existing {
decode(row)? let provider: String = row
} else { .try_get("provider")
let pending: bool = .map_err(|error| RuntimeError::database("decode search generation provider", error))?;
sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM search_runtime_generations WHERE state='pending')") let existing_config_hash: Vec<u8> = row
.fetch_one(&mut *transaction) .try_get("config_hash")
.await .map_err(|error| RuntimeError::database("decode search generation config hash", error))?;
.map_err(|error| RuntimeError::database("check pending search generation", error))?; let schema_version: i32 = row
if pending { .try_get("schema_version")
return Err(RuntimeError::invalid_state("search_generation_change_in_progress")); .map_err(|error| RuntimeError::database("decode search generation schema version", error))?;
} if provider != config.provider
let generation_id = Uuid::new_v4(); || existing_config_hash != expected_config_hash
let suffix = generation_id.simple().to_string(); || schema_version != SCHEMA_FINGERPRINT
let manifest = if config.provider == "embedded" { {
json!({"doc":"doc","block":"block"}) 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 { } else {
json!({ let state: String = row
"doc":format!("affine_search_doc_{suffix}"), .try_get("state")
"block":format!("affine_search_block_{suffix}"), .map_err(|error| RuntimeError::database("decode search generation state", error))?;
}) if rebuild_embedded {
}; if state == "building" {
sqlx::query( sqlx::query(
r#"INSERT INTO search_runtime_generations "UPDATE search_projection.generations SET state='failed', last_error='embedded generation lost on \
(generation_id,provider,state,config_fingerprint,schema_fingerprint,manifest) restart' WHERE id=$1 AND state='building'",
VALUES ($1,$2,'pending',$3,$4,$5)"#, )
) .bind(
.bind(generation_id) row
.bind(&config.provider) .try_get::<Uuid, _>("id")
.bind(&fingerprint) .map_err(|error| RuntimeError::database("decode search generation id", error))?,
.bind(SCHEMA_FINGERPRINT) )
.bind(&manifest) .execute(&mut *transaction)
.execute(&mut *transaction) .await
.await .map_err(|error| RuntimeError::database("fail stale embedded generation", error))?;
.map_err(|error| RuntimeError::database("create pending search generation", error))?; create_generation(&mut transaction, config, expected_config_hash.clone()).await?
for table in [SearchTable::Doc, SearchTable::Block] { } else {
sqlx::query("INSERT INTO search_runtime_provider_cursors(generation_id,table_key) VALUES ($1,$2)") decode(row)?
.bind(generation_id) }
.bind(table.as_str()) } else {
.execute(&mut *transaction) decode(row)?
.await }
.map_err(|error| RuntimeError::database("initialize search generation cursor", error))?;
} }
ActiveGeneration { } else {
id: generation_id, let active = sqlx::query(
manifest, 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 transaction
.commit() .commit()
.await .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 { if let Some(remote) = remote {
for table in [SearchTable::Doc, SearchTable::Block] { for table in [SearchTable::Doc, SearchTable::Block] {
@@ -114,14 +152,14 @@ pub(super) async fn load_active(
pool: &PgPool, pool: &PgPool,
config: &SearchRuntimeConfig, config: &SearchRuntimeConfig,
) -> RuntimeResult<Option<ActiveGeneration>> { ) -> RuntimeResult<Option<ActiveGeneration>> {
let fingerprint = config_fingerprint(config);
let row = sqlx::query( let row = sqlx::query(
r#"SELECT generation_id,provider,manifest FROM search_runtime_generations r#"SELECT id,provider,config_hash,manifest
WHERE state='active' AND provider=$1 AND config_fingerprint=$2 AND schema_fingerprint=$3 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"#, ORDER BY activated_at DESC NULLS LAST LIMIT 1"#,
) )
.bind(&config.provider) .bind(&config.provider)
.bind(&fingerprint) .bind(config_hash(config))
.bind(SCHEMA_FINGERPRINT) .bind(SCHEMA_FINGERPRINT)
.fetch_optional(pool) .fetch_optional(pool)
.await .await
@@ -129,28 +167,122 @@ pub(super) async fn load_active(
row.map(decode).transpose() row.map(decode).transpose()
} }
pub(super) async fn fail(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> { async fn fail(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
sqlx::query("UPDATE search_runtime_generations SET state='failed' WHERE generation_id=$1 AND state='pending'") sqlx::query(
.bind(generation.id) "UPDATE search_projection.generations SET state='failed', last_error=$2 WHERE id=$1 AND state NOT IN \
.execute(pool) ('active','draining')",
.await )
.map_err(|error| RuntimeError::database("fail pending search generation", error))?; .bind(generation.id)
.bind("search generation build failed")
.execute(pool)
.await
.map_err(|error| RuntimeError::database("fail search generation", error))?;
Ok(()) 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 let mut transaction = pool
.begin() .begin()
.await .await
.map_err(|error| RuntimeError::database("begin search generation activation", error))?; .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") sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.bind(generation.id)
.execute(&mut *transaction) .execute(&mut *transaction)
.await .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( sqlx::query(
"UPDATE search_runtime_generations SET state='active', activated_at=coalesce(activated_at,now()) WHERE \ "UPDATE search_projection.generations SET state='draining', drained_at=now() WHERE state='active' AND id<>$1",
generation_id=$1 AND state IN ('pending','active')", )
.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) .bind(generation.id)
.execute(&mut *transaction) .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)) .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(); let mut hash = Sha256::new();
for value in [ for value in [
&config.provider, &config.provider,
@@ -174,16 +306,108 @@ fn config_fingerprint(config: &SearchRuntimeConfig) -> String {
hash.update(value.as_bytes()); hash.update(value.as_bytes());
hash.update([0]); hash.update([0]);
} }
hash.finalize().iter().map(|byte| format!("{byte:02x}")).collect() hash.finalize().to_vec()
} }
fn decode(row: sqlx::postgres::PgRow) -> RuntimeResult<ActiveGeneration> { fn decode(row: sqlx::postgres::PgRow) -> RuntimeResult<ActiveGeneration> {
Ok(ActiveGeneration { Ok(ActiveGeneration {
id: row id: row
.try_get("generation_id") .try_get("id")
.map_err(|error| RuntimeError::database("decode search generation id", error))?, .map_err(|error| RuntimeError::database("decode search generation id", error))?,
manifest: row manifest: row
.try_get("manifest") .try_get("manifest")
.map_err(|error| RuntimeError::database("decode search generation manifest", error))?, .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();
}
}
@@ -1,19 +1,32 @@
mod checkpoint;
mod generation; mod generation;
mod projection; mod projection;
mod provider; mod provider;
mod query; mod query;
mod query_runtime;
mod result_filter;
mod runtime; mod runtime;
mod store;
mod types; mod types;
mod worker; 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; pub(super) use runtime::SearchRuntime;
use types::{AggregateOptions, RuntimeSearchQuery, SearchOptions, SearchTable};
pub(super) use types::{RuntimeAggregateRequest, RuntimeSearchRequest}; pub(super) use types::{RuntimeAggregateRequest, RuntimeSearchRequest};
use worker::{reconcile_workspace, sweep_generation_orphans};
pub(super) use super::webpki_tls_config; use super::{
permission::{AuthorizedSearchScope, DocReadScope, PermissionAuthorizer, SearchActor},
webpki_tls_config,
};
const SCHEMA_FINGERPRINT: &str = "search-runtime-v5"; 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 { fn exact_token(value: &str) -> String {
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
@@ -32,4 +45,19 @@ fn provider_payload(payload: &serde_json::Value) -> serde_json::Value {
} }
#[cfg(test)] #[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);
}
}
}
@@ -1,32 +1,42 @@
use serde_json::{Value, json}; use serde_json::{Value, json};
use sqlx::{PgPool, Row}; use sqlx::{PgPool, Row};
use super::store::ProjectionInput;
use crate::{ use crate::{
permission::doc_role_allows, permission::doc_role_allows,
runtime::{RuntimeError, RuntimeResult, storage_runtime::load_current_doc}, 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( pub(super) async fn project_document(
pool: &PgPool, pool: &PgPool,
workspace_id: &str, workspace_id: &str,
doc_id: &str, doc_id: &str,
) -> RuntimeResult<Option<(ProjectionInput, Vec<ProjectionInput>)>> { ) -> 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); return Ok(None);
}; };
let revision = current.updated_at.timestamp_millis(); let revision = current.updated_at.timestamp_millis();
let projection = let projection =
affine_doc_loader::project_document_search(current.blob, doc_id.to_string(), revision.to_string()) 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}")))?; .map_err(|error| RuntimeError::SearchSourceInvalid(format!("document projection failed: {error}")))?;
let metadata = sqlx::query( let Some(metadata) = sqlx::query(
r#"SELECT snapshot.created_at,snapshot.updated_at,snapshot.created_by,snapshot.updated_by, 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, coalesce(doc_policy.visibility,'private') AS visibility,
doc_policy.public_role, doc_policy.public_role,
coalesce(doc_policy.member_default_role,workspace_policy.member_default_doc_role,'manager') AS member_default_role coalesce(doc_policy.member_default_role,workspace_policy.member_default_doc_role,'manager') AS member_default_role
FROM snapshots snapshot 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 workspace_access_policies workspace_policy ON workspace_policy.workspace_id=snapshot.workspace_id
LEFT JOIN doc_access_policies doc_policy LEFT JOIN doc_access_policies doc_policy
ON doc_policy.workspace_id=snapshot.workspace_id AND doc_policy.doc_id=snapshot.guid 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) .fetch_optional(pool)
.await .await
.map_err(|error| RuntimeError::database("load search projection metadata", error))? .map_err(|error| RuntimeError::database("load search projection metadata", error))?
.ok_or_else(|| RuntimeError::invalid_state("search snapshot metadata unavailable"))?; else {
let acl_revision = metadata return Ok(None);
.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"))?;
let visibility: String = metadata let visibility: String = metadata
.try_get("visibility") .try_get("visibility")
.map_err(|error| RuntimeError::database("decode search doc visibility", error))?; .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 let member_default_role: String = metadata
.try_get("member_default_role") .try_get("member_default_role")
.map_err(|error| RuntimeError::database("decode search member default role", error))?; .map_err(|error| RuntimeError::database("decode search member default role", error))?;
let acl_public_readable = visibility == "public" let acl_public_readable = if visibility == "public" {
&& public_role match public_role.as_deref() {
.as_deref() Some(role) => readable_role(role, "public")?,
.is_some_and(|role| doc_role_allows(role, "Doc.Read").unwrap_or(false)); None => false,
let acl_member_default_readable = doc_role_allows(&member_default_role, "Doc.Read") }
.map_err(|_| RuntimeError::invalid_state("permission_state_unavailable"))?; } else {
false
};
let acl_member_default_readable = readable_role(&member_default_role, "member default")?;
let grants = sqlx::query( 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(workspace_id)
.bind(doc_id) .bind(doc_id)
@@ -74,7 +86,7 @@ pub(super) async fn project_document(
let role: String = row let role: String = row
.try_get("role") .try_get("role")
.map_err(|error| RuntimeError::database("decode grant role", error))?; .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<_>>>()? .collect::<RuntimeResult<Vec<_>>>()?
.into_iter() .into_iter()
@@ -96,7 +108,6 @@ pub(super) async fn project_document(
public_readable: acl_public_readable, public_readable: acl_public_readable,
member_default_readable: acl_member_default_readable, member_default_readable: acl_member_default_readable,
read_user_ids: acl_read_user_ids, read_user_ids: acl_read_user_ids,
revision: acl_revision,
}; };
let document_payload = with_acl( let document_payload = with_acl(
json!({ json!({
@@ -113,14 +124,7 @@ pub(super) async fn project_document(
}), }),
&acl, &acl,
); );
let document = input( let document = input(workspace_id, doc_id, document_payload);
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}"),
revision,
document_payload,
&acl,
);
let blocks = projection let blocks = projection
.units .units
.into_iter() .into_iter()
@@ -142,24 +146,21 @@ pub(super) async fn project_document(
}), }),
&acl, &acl,
); );
input( input(workspace_id, doc_id, payload)
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}/{block_id}"),
revision,
payload,
&acl,
)
}) })
.collect(); .collect();
Ok(Some((document, blocks))) 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 { struct AclFields {
public_readable: bool, public_readable: bool,
member_default_readable: bool, member_default_readable: bool,
read_user_ids: Vec<String>, read_user_ids: Vec<String>,
revision: i64,
} }
fn with_acl(mut payload: Value, acl: &AclFields) -> Value { 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()); tokens.push("public".to_string());
} }
object.insert("acl_read_tokens".to_string(), json!(tokens)); object.insert("acl_read_tokens".to_string(), json!(tokens));
object.insert("acl_revision".to_string(), json!(acl.revision));
payload payload
} }
fn input( fn input(workspace_id: &str, doc_id: &str, payload: Value) -> ProjectionInput {
workspace_id: &str,
doc_id: &str,
external_id: &str,
revision: i64,
payload: Value,
acl: &AclFields,
) -> ProjectionInput {
ProjectionInput { ProjectionInput {
external_id: external_id.to_string(),
workspace_id: workspace_id.to_string(), workspace_id: workspace_id.to_string(),
doc_id: doc_id.to_string(), doc_id: doc_id.to_string(),
revision,
payload, 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());
} }
} }
@@ -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( use super::{SearchChange, SearchTable, provider_write_error, webpki_tls_config};
payload: &mut Value, use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<()> { const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
let object = payload.as_object_mut().expect("search payload is an object");
object.remove("acl_read_user_ids"); pub(in crate::runtime::backend_runtime::search) struct ManticoreSearchProvider {
if let Some(Value::Array(tokens)) = object.remove("acl_read_tokens") { client: Client,
object.insert( endpoint: String,
"acl_read_token_ids".to_string(), api_key: String,
Value::Array( username: String,
tokens password: String,
.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(())
} }
pub(super) fn prepare_manticore_search( impl ManticoreSearchProvider {
dsl: &mut Value, pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
cursor: Option<Value>, let endpoint = config.endpoint.trim_end_matches('/');
size: u64, let url = url::Url::parse(endpoint).map_err(|_| RuntimeError::config("invalid search provider endpoint"))?;
initial_offset: u64, if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
requested_fields: &[String], return Err(RuntimeError::config("invalid search provider endpoint"));
token_ids: &std::collections::HashMap<String, i64>, }
) -> RuntimeResult<u64> { let client = Client::builder()
normalize_manticore_terms(dsl, token_ids)?; .tls_backend_preconfigured(
let object = dsl.as_object_mut().expect("search DSL is an object"); webpki_tls_config()
let mut source = object .map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?,
.get("_source") )
.and_then(Value::as_array) .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() .into_iter()
.flatten() .flatten()
.filter_map(Value::as_str) .filter(|(_, value)| !value.is_null())
.map(str::to_string) .map(|(key, value)| {
.collect::<std::collections::BTreeSet<_>>(); let value = match value {
source.extend(requested_fields.iter().cloned()); Value::Bool(value) => json!(i32::from(*value)),
object.insert("_source".to_string(), json!(source)); Value::Array(values) => json!(
object.remove("fields"); values
if let Some(highlight) = object.get_mut("highlight") .iter()
&& let Some(options) = highlight .filter(|value| !value.is_null())
.get("fields") .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(Value::as_object)
.and_then(|fields| fields.values().next()) .ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
.cloned() let status = result
{ .get("status")
*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")
.and_then(Value::as_u64) .and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?; .and_then(|status| u16::try_from(status).ok())
if offset.saturating_add(size) > 10_000 { .ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
return Err(RuntimeError::invalid_input("search cursor exceeds 10000")); if result.get("error").is_some() || !(200..300).contains(&status) {
return Err(provider_write_error(status));
} }
object.insert("from".to_string(), json!(offset)); let numeric = ["created", "updated", "deleted"]
offset .into_iter()
} else if cursor.is_some() { .filter_map(|field| result.get(field).and_then(Value::as_u64))
return Err(RuntimeError::invalid_input("invalid search cursor")); .sum::<u64>();
} else { affected += if numeric > 0 {
initial_offset numeric
}; } else if result.get("created").and_then(Value::as_bool).is_some()
Ok(offset) || result
} .get("result")
.and_then(Value::as_str)
pub(super) fn manticore_fields(source: Option<&Value>, requested_fields: &[String]) -> Value { .is_some_and(|result| matches!(result, "created" | "updated"))
let source = source.and_then(Value::as_object); {
Value::Object( 1
requested_fields } else {
.iter() return Err(RuntimeError::invalid_state("provider_apply_failed"));
.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(());
} }
match value { if value.get("errors").and_then(Value::as_bool) != Some(false) || affected != expected_changes as u64 {
Value::Object(object) => { return Err(RuntimeError::invalid_state("provider_apply_failed"));
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)?;
}
}
_ => {}
} }
Ok(()) Ok(())
} }
fn manticore_term(value: &Value, token_ids: &std::collections::HashMap<String, i64>) -> RuntimeResult<Option<Value>> { fn translate_search_request(physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
let Some(term) = value.get("term").and_then(Value::as_object) else { let object = dsl
return Ok(None); .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 { let query = object.get("query").cloned().unwrap_or_else(|| json!({"match_all":{}}));
return Ok(None); 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 { let mut source = Vec::new();
return Ok(None); 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));
let value = clause.get("value").unwrap_or(clause); }
Ok(match value { if let Some(fields) = object.get("fields").and_then(Value::as_array) {
Value::String(value) => { source.extend(fields.iter().filter_map(Value::as_str).map(str::to_string));
if matches!(field.as_str(), "acl_read_tokens" | "ref_doc_id") { }
let token_id = token_ids source.push("external_id".to_string());
.get(value) source.push("doc_id".to_string());
.copied() source.push("source_version".to_string());
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))?; source.push("permission_version".to_string());
let field = if field == "acl_read_tokens" { source.sort();
"acl_read_token_ids" source.dedup();
} else { request["_source"] = json!(source);
"ref_doc_token_ids" 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}}))); Some(json!({field:direction}))
}
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}}))
} }
_ => 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)}})), return Ok(json!({"equals":{field:manticore_scalar(value)?}}));
Value::Number(value) => Some(json!({"equals":{field:value}})), }
_ => None, 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> { async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>> {
let mut tokens = Vec::new(); if response
collect_exact_tokens(value, &mut tokens); .content_length()
tokens .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>) { fn normalize(value: Value, offset: u64) -> RuntimeResult<Value> {
match value { let hits = value
Value::Object(object) => { .pointer("/hits/hits")
if let Some(token) = object .and_then(Value::as_array)
.get("term") .ok_or_else(|| RuntimeError::invalid_state("invalid provider response"))?;
.and_then(|term| term.get("acl_read_tokens").or_else(|| term.get("ref_doc_id"))) let total = value
.and_then(|clause| clause.get("value").unwrap_or(clause).as_str()) .pointer("/hits/total")
{ .and_then(Value::as_u64)
tokens.push(token.to_string()); .ok_or_else(|| RuntimeError::invalid_state("inexact provider total"))?;
} let nodes = hits
for field in ["acl_read_tokens", "ref_doc_id"] { .iter()
if let Some(values) = object.get(field).and_then(Value::as_array) { .map(|hit| {
tokens.extend(values.iter().filter_map(Value::as_str).map(str::to_string)); let source = hit.get("_source").cloned().unwrap_or_else(|| json!({}));
} json!({
} "id":source.get("external_id").and_then(Value::as_str).unwrap_or_default(),
for child in object.values() { "score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(),
collect_exact_tokens(child, tokens); "fields":fields_from_source(&source),
} "highlights":hit.get("highlight").cloned().unwrap_or_else(||json!({})),
} "_source":source,
Value::Array(values) => { })
for child in values { })
collect_exact_tokens(child, tokens); .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)] #[cfg(test)]
mod tests { 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] #[test]
fn payload_and_fields_preserve_terminal_types() { fn translates_the_shared_basic_query_subset() {
let mut payload = json!({ assert_eq!(
"content":["hello","world"], translate_query(json!({
"ref_doc_id":["doc-a","doc-b"], "bool":{"must":[
"summary":null, {"term":{"workspace_id":{"value":"workspace"}}},
"created_at":2_000, {"match":{"content":{"query":"hello","boost":1.0}}}
"updated_at":3_000, ]}
"acl_read_tokens":["member"] }))
}); .unwrap(),
prepare_manticore_payload( json!({
&mut payload, "bool":{"must":[
&[ {"equals":{"workspace_id":"workspace"}},
("member".to_string(), 7), {"match":{"content":"hello"}}
("doc-a".to_string(), 8), ]}
("doc-b".to_string(), 9), })
] );
.into_iter() assert_eq!(
.collect(), translate_search_request(
) "blocks",
.unwrap(); json!({
assert_eq!(payload["content"], "hello world"); "query":{"match":{"content":{"query":"hello"}}},
assert_eq!(payload["ref_doc_id"], "[\"doc-a\",\"doc-b\"]"); "highlight":{"fields":{"content":{"pre_tags":["<b>"],"post_tags":["</b>"]}}}
assert_eq!(payload["summary"], ""); })
assert_eq!(payload["created_at"], 2); )
assert_eq!(payload["acl_read_token_ids"], json!([7])); .unwrap()["highlight"],
assert_eq!(payload["ref_doc_token_ids"], json!([8, 9])); json!({"fields":["content"],"pre_tags":"<b>","post_tags":"</b>"})
let fields = manticore_fields(
Some(&payload),
&[
"ref_doc_id".to_string(),
"summary".to_string(),
"updated_at".to_string(),
],
); );
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] #[test]
fn nested_terms_use_exact_identity_and_acl_tokens() { fn converts_projection_arrays_and_booleans_to_basic_attributes() {
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();
assert_eq!( assert_eq!(
dsl, manticore_document(&json!({"acl_read_tokens":["user","member"],"acl_public_readable":false})),
json!({"query":{"bool":{"must":[ json!({"acl_read_tokens":"user member","acl_public_readable":0})
{"match":{"workspace_token":{"query":exact_token("workspace"),"boost":2.0}}}, .as_object()
{"bool":{"must_not":[{"equals":{"doc_token":exact_token("doc")}}]}}, .unwrap()
{"equals":{"acl_read_token_ids":9}}, .clone()
{"equals":{"ref_doc_token_ids":10}}
]}}})
); );
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");
} }
} }
@@ -1,40 +1,188 @@
mod manticore; mod manticore;
mod remote; mod remote;
pub(super) use remote::RemoteProvider;
use serde_json::{Value, json}; 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 text_field = table.text_field();
let mut properties = serde_json::Map::from_iter([ let mut properties = serde_json::Map::from_iter([
("workspace_id".into(), json!({"type":"keyword"})), ("workspace_id".into(), json!({"type":"keyword"})),
("workspace_token".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_id".into(), json!({"type":"keyword"})),
("doc_token".into(), json!({"type":"keyword"})), ("doc_token".into(), json!({"type":"keyword"})),
(text_field.into(), json!({"type":"text"})), (text_field.into(), json!({"type":"text"})),
( ("created_at".into(), json!({"type":"date"})),
"created_at".into(), ("updated_at".into(), json!({"type":"date"})),
json!({"type":if provider == "manticoresearch" { "long" } else { "date" }}),
),
(
"updated_at".into(),
json!({"type":if provider == "manticoresearch" { "long" } else { "date" }}),
),
("created_by_user_id".into(), json!({"type":"keyword"})), ("created_by_user_id".into(), json!({"type":"keyword"})),
("updated_by_user_id".into(), json!({"type":"keyword"})), ("updated_by_user_id".into(), json!({"type":"keyword"})),
("acl_public_readable".into(), json!({"type":"boolean"})), ("acl_public_readable".into(), json!({"type":"boolean"})),
("acl_member_default_readable".into(), json!({"type":"boolean"})), ("acl_member_default_readable".into(), json!({"type":"boolean"})),
( ("acl_read_tokens".into(), json!({"type":"keyword"})),
"acl_read_tokens".into(),
if provider == "manticoresearch" {
json!({"type":"keyword","mva":true})
} else {
json!({"type":"keyword"})
},
),
("acl_revision".into(), json!({"type":"long"})),
]); ]);
if table == SearchTable::Block { if table == SearchTable::Block {
for field in [ for field in [
@@ -65,61 +213,76 @@ pub(super) fn mapping(table: SearchTable, provider: &str) -> Value {
json!({"mappings":{"properties":properties}}) json!({"mappings":{"properties":properties}})
} }
pub(super) fn manticore_schema(table: SearchTable, physical_table: &str) -> String { pub(super) fn projection_external_id(
let common = r#" table: SearchTable,
workspace_id string attribute indexed, generation_id: &str,
workspace_token string attribute indexed, workspace_id: &str,
doc_id string attribute indexed, doc_id: &str,
doc_token string attribute indexed,"#; block_id: Option<&str>,
let fields = match table { source_version: i64,
SearchTable::Doc => format!( permission_version: i64,
r#"{common} ) -> RuntimeResult<String> {
title text, let mut id = format!("{generation_id}/{workspace_id}/{doc_id}/{source_version}/{permission_version}");
summary string stored, if table == SearchTable::Block {
journal string stored, let block_id = block_id
created_by_user_id string attribute indexed, .filter(|block_id| !block_id.is_empty())
updated_by_user_id string attribute indexed, .ok_or_else(|| RuntimeError::invalid_input("block projection id requires block_id"))?;
created_at timestamp, id.push('/');
updated_at timestamp, id.push_str(block_id);
acl_public_readable bool, }
acl_member_default_readable bool, Ok(id)
acl_read_token_ids multi64, }
acl_revision bigint"#,
), #[cfg(test)]
SearchTable::Block => format!( mod tests {
r#"{common} use serde_json::json;
block_id string attribute indexed,
block_token string attribute indexed, use super::{SearchChange, SearchTable, projection_external_id};
unit_id string attribute indexed,
projection_version bigint, #[test]
source_hash string attribute indexed, fn projection_ids_are_stable_and_keep_block_identity_last() {
visibility string attribute indexed, let doc = projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 7, 3).unwrap();
element_id string attribute indexed, assert_eq!(doc, "generation/workspace/doc/7/3");
frame_id string attribute indexed, assert_eq!(
source_block_id string attribute indexed, projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 7, 3).unwrap(),
content text, doc
flavour string attribute indexed, );
blob string attribute indexed, assert_eq!(
ref_doc_id string attribute indexed, projection_external_id(
ref_doc_token_ids multi64, SearchTable::Block,
ref string stored, "generation",
parent_flavour string attribute indexed, "workspace",
parent_block_id string attribute indexed, "doc",
additional string stored, Some("block"),
markdown_preview string stored, 7,
created_by_user_id string attribute indexed, 3,
updated_by_user_id string attribute indexed, )
created_at timestamp, .unwrap(),
updated_at timestamp, "generation/workspace/doc/7/3/block"
acl_public_readable bool, );
acl_member_default_readable bool, assert_ne!(
acl_read_token_ids multi64, projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 8, 3).unwrap(),
acl_revision bigint"#, doc
), );
}; }
format!(
"CREATE TABLE IF NOT EXISTS {physical_table} ({fields}) charset_table='non_cjk, chinese' ngram_len='1' \ #[test]
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, \ fn change_rejects_an_external_id_outside_its_tuple() {
U+0E00..U+0E7F' index_field_lengths='1'" 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());
}
} }
@@ -2,12 +2,8 @@ use std::time::Duration;
use reqwest::{Client, redirect::Policy}; use reqwest::{Client, redirect::Policy};
use serde_json::{Value, json}; use serde_json::{Value, json};
use sqlx::PgPool;
use super::{ use super::{SearchChange, SearchTable, provider_write_error, webpki_tls_config};
super::{store::SearchChange, types::SearchTable, webpki_tls_config},
manticore::{manticore_exact_tokens, manticore_fields, prepare_manticore_payload, prepare_manticore_search},
};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig}; use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024; const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
@@ -15,68 +11,42 @@ const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
pub(in crate::runtime::backend_runtime::search) struct RemoteProvider { pub(in crate::runtime::backend_runtime::search) struct RemoteProvider {
client: Client, client: Client,
endpoint: String, endpoint: String,
provider: String,
api_key: String, api_key: String,
username: String, username: String,
password: String, password: String,
pool: PgPool,
} }
impl RemoteProvider { impl RemoteProvider {
pub(in crate::runtime::backend_runtime::search) fn new( pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
config: &SearchRuntimeConfig,
pool: PgPool,
) -> RuntimeResult<Self> {
let endpoint = config.endpoint.trim_end_matches('/'); let endpoint = config.endpoint.trim_end_matches('/');
let url = url::Url::parse(endpoint).map_err(|_| RuntimeError::config("invalid search provider endpoint"))?; 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() { if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return Err(RuntimeError::config("invalid search provider endpoint")); return Err(RuntimeError::config("invalid search provider endpoint"));
} }
let mut client = Client::builder() let client = Client::builder()
.tls_backend_preconfigured( .tls_backend_preconfigured(
webpki_tls_config() webpki_tls_config()
.map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?, .map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?,
) )
.redirect(Policy::none()) .redirect(Policy::none())
.timeout(Duration::from_secs(30)); .timeout(Duration::from_secs(30));
if config.provider == "manticoresearch" {
client = client.pool_max_idle_per_host(0);
}
let client = client let client = client
.build() .build()
.map_err(|error| RuntimeError::invalid_state(format!("search HTTP client failed: {error}")))?; .map_err(|error| RuntimeError::invalid_state(format!("search HTTP client failed: {error}")))?;
Ok(Self { Ok(Self {
client, client,
endpoint: endpoint.to_string(), endpoint: endpoint.to_string(),
provider: config.provider.clone(),
api_key: config.api_key.clone(), api_key: config.api_key.clone(),
username: config.username.clone(), username: config.username.clone(),
password: config.password.clone(), password: config.password.clone(),
pool,
}) })
} }
pub(in crate::runtime::backend_runtime::search) async fn search( pub(super) async fn search(&self, physical_table: &str, mut dsl: Value) -> RuntimeResult<Value> {
&self, ensure_projection_source(&mut dsl);
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<_>>();
dsl["track_total_hits"] = json!(true); 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")); let cursor = dsl.as_object_mut().and_then(|object| object.remove("cursor"));
if self.provider == "manticoresearch" { if let Some(cursor) = cursor {
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 {
let cursor = cursor let cursor = cursor
.as_str() .as_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?; .ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
@@ -107,22 +77,20 @@ impl RemoteProvider {
} }
let value: Value = let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?; serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize( normalize(value)
value,
self.provider == "manticoresearch",
offset,
size,
&requested_fields,
)
} }
pub(in crate::runtime::backend_runtime::search) async fn aggregate( pub(super) async fn aggregate(&self, physical_table: &str, mut dsl: Value) -> RuntimeResult<Value> {
&self, if let Some(top_hits) = dsl.pointer_mut("/aggs/result/aggs/result/top_hits") {
physical_table: &str, ensure_projection_source(top_hits);
mut dsl: Value, }
) -> RuntimeResult<Value> { let limit = dsl
if self.provider == "manticoresearch" { .pointer("/aggs/result/terms/size")
return Err(RuntimeError::SearchUnsupportedQuery); .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); dsl["track_total_hits"] = json!(true);
let response = self let response = self
@@ -134,32 +102,18 @@ impl RemoteProvider {
let status = response.status(); let status = response.status();
let bytes = read_response(response).await?; let bytes = read_response(response).await?;
if !status.is_success() { if !status.is_success() {
return Err(RuntimeError::SearchUnsupportedQuery); return Err(if status.as_u16() == 400 {
RuntimeError::SearchUnsupportedQuery
} else {
RuntimeError::SearchProviderUnavailable
});
} }
let value: Value = let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?; 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( pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
&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);
}
if self if self
.request(reqwest::Method::HEAD, physical_table) .request(reqwest::Method::HEAD, physical_table)
.send() .send()
@@ -170,7 +124,7 @@ impl RemoteProvider {
} }
let response = self let response = self
.request(reqwest::Method::PUT, physical_table) .request(reqwest::Method::PUT, physical_table)
.json(&super::mapping(table, &self.provider)) .json(&super::mapping(table))
.send() .send()
.await; .await;
match response { match response {
@@ -179,73 +133,124 @@ impl RemoteProvider {
} }
} }
pub(in crate::runtime::backend_runtime::search) async fn apply( pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
&self,
physical_table: &str,
changes: &[SearchChange],
) -> RuntimeResult<()> {
if changes.is_empty() { if changes.is_empty() {
return Ok(()); return Ok(());
} }
let token_ids = if self.provider == "manticoresearch" { self.apply_elasticsearch(physical_table, changes).await
self }
.resolve_manticore_tokens(
changes async fn apply_elasticsearch(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
.iter()
.filter_map(|change| change.payload.as_ref())
.flat_map(manticore_exact_tokens),
)
.await?
} else {
Default::default()
};
let mut body = String::new(); let mut body = String::new();
for change in changes { for change in changes {
if change.operation == "delete" { let payload = change.upsert_payload()?;
body.push_str( body.push_str(
&serde_json::to_string(&json!({"delete":{"_index":physical_table,"_id":change.external_id}})) &serde_json::to_string(&json!({
.map_err(|error| RuntimeError::json("encode provider delete", error))?, "index": {
); "_index": physical_table,
body.push('\n'); "_id": change.external_id
} 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 immutable provider upsert", error))?,
.map_err(|error| RuntimeError::json("encode provider upsert", error))?, );
); body.push('\n');
body.push('\n'); body.push_str(
let mut payload = super::super::provider_payload(payload); &serde_json::to_string(payload)
if self.provider == "manticoresearch" { .map_err(|error| RuntimeError::json("encode immutable provider document", error))?,
prepare_manticore_payload(&mut payload, &token_ids)?; );
} body.push('\n');
body.push_str( }
&serde_json::to_string(&payload).map_err(|error| RuntimeError::json("encode provider document", error))?, if body.is_empty() {
); return Ok(());
body.push('\n');
}
} }
let path = if self.provider == "elasticsearch" {
"_bulk?refresh=wait_for"
} else {
"_bulk"
};
let response = self let response = self
.request(reqwest::Method::POST, path) .request(reqwest::Method::POST, "_bulk?refresh=wait_for")
.header("content-type", "application/x-ndjson") .header("content-type", "application/x-ndjson")
.body(body) .body(body)
.send() .send()
.await .await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?; .map_err(|_| RuntimeError::SearchProviderUnavailable)?;
if !response.status().is_success() { if !response.status().is_success() {
return Err(RuntimeError::SearchProviderUnavailable); return Err(provider_write_error(response.status().as_u16()));
} }
let value: Value = response let value: Value = response
.json() .json()
.await .await
.map_err(|error| RuntimeError::invalid_state(format!("invalid provider bulk response: {error}")))?; .map_err(|error| RuntimeError::invalid_state(format!("invalid provider bulk response: {error}")))?;
if value.get("errors").and_then(Value::as_bool) == Some(true) { validate_bulk_response(&value, changes.len())
return Err(RuntimeError::invalid_state("provider_apply_failed")); }
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 { fn request(&self, method: reqwest::Method, path: &str) -> reqwest::RequestBuilder {
@@ -257,48 +262,61 @@ impl RemoteProvider {
} }
request request
} }
}
async fn resolve_manticore_tokens( fn validate_delete_response(value: &Value) -> RuntimeResult<()> {
&self, if value.get("timed_out").and_then(Value::as_bool) == Some(true)
tokens: impl IntoIterator<Item = String>, || value
) -> RuntimeResult<std::collections::HashMap<String, i64>> { .get("failures")
let tokens = tokens.into_iter().collect::<std::collections::BTreeSet<_>>(); .and_then(Value::as_array)
if tokens.is_empty() { .is_some_and(|failures| !failures.is_empty())
return Ok(Default::default()); {
} return Err(RuntimeError::invalid_state("provider_apply_failed"));
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())
} }
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>> { async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>> {
@@ -322,13 +340,7 @@ async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>
Ok(bytes) Ok(bytes)
} }
fn normalize( fn normalize(value: Value) -> RuntimeResult<Value> {
value: Value,
manticore: bool,
_offset: u64,
_size: u64,
requested_fields: &[String],
) -> RuntimeResult<Value> {
let hits = value let hits = value
.pointer("/hits/hits") .pointer("/hits/hits")
.and_then(Value::as_array) .and_then(Value::as_array)
@@ -341,11 +353,7 @@ fn normalize(
let nodes = hits let nodes = hits
.iter() .iter()
.map(|hit| { .map(|hit| {
let fields = if manticore { let fields = hit.get("fields").cloned().unwrap_or_else(|| json!({}));
manticore_fields(hit.get("_source"), requested_fields)
} else {
hit.get("fields").cloned().unwrap_or_else(|| json!({}))
};
json!({ json!({
"id":hit.get("_id").and_then(Value::as_str).unwrap_or_default(), "id":hit.get("_id").and_then(Value::as_str).unwrap_or_default(),
"score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(), "score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(),
@@ -355,23 +363,16 @@ fn normalize(
}) })
}) })
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let next_cursor = if manticore { let next_cursor = hits
(!hits.is_empty()) .last()
.then(|| serde_json::to_string(&json!({"offset":_offset + _size}))) .and_then(|hit| hit.get("sort"))
.transpose() .map(serde_json::to_string)
.map_err(|error| RuntimeError::json("encode provider cursor", error))? .transpose()
} else { .map_err(|error| RuntimeError::json("encode provider cursor", error))?;
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})) 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 let buckets = value
.pointer("/aggregations/result/buckets") .pointer("/aggregations/result/buckets")
.and_then(Value::as_array) .and_then(Value::as_array)
@@ -386,12 +387,14 @@ fn normalize_aggregate(value: Value) -> RuntimeResult<Value> {
Ok(json!({ Ok(json!({
"key":bucket.get("key").cloned().unwrap_or(Value::Null), "key":bucket.get("key").cloned().unwrap_or(Value::Null),
"count":bucket.get("doc_count").cloned().unwrap_or(json!(0)), "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)), "hits":hits.iter().map(normalize_hit).collect::<Vec<_>>()
"nodes":hits.iter().map(normalize_hit).collect::<Vec<_>>()}
})) }))
}) })
.collect::<RuntimeResult<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 { fn normalize_hit(hit: &Value) -> Value {
@@ -403,3 +406,76 @@ fn normalize_hit(hit: &Value) -> Value {
"_source":hit.get("_source").cloned().unwrap_or_else(||json!({})), "_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");
}
}
@@ -1,10 +1,105 @@
use serde::Deserialize;
use serde_json::{Value, json}; use serde_json::{Value, json};
use super::types::{AggregateRequest, SearchOptions, SearchQuery, SearchRequest, SearchTable}; use super::{
use crate::runtime::{ AggregateOptions, AuthorizedSearchScope, DocReadScope, RuntimeAggregateRequest, RuntimeSearchQuery,
RuntimeError, RuntimeResult, RuntimeSearchRequest, SearchOptions, SearchTable,
backend_runtime::permission::{AuthorizedSearchScope, DocReadScope},
}; };
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> { pub(super) fn compile(request: &SearchRequest, scope: &AuthorizedSearchScope) -> RuntimeResult<Value> {
let query = compile_query(request.table, &request.query)?; 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)) .map(|field| validate_field(request.table, field).map(str::to_string))
.collect::<RuntimeResult<Vec<_>>>()?; .collect::<RuntimeResult<Vec<_>>>()?;
let mut dsl = json!({ let mut dsl = json!({
"_source":["workspace_id","doc_id"], "_source":["workspace_id","doc_id","source_version","permission_version"],
"fields":fields, "fields":fields,
"query":{"bool":{"must":must}}, "query":{"bool":{"must":must}},
"sort": stable_sort(request.table), "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 hit_dsl = compile(&search, scope)?;
let field = validate_field(request.table, &request.field)?; let field = validate_field(request.table, &request.field)?;
let limit = request.options.pagination.limit.unwrap_or(10); let limit = request.options.pagination.limit.unwrap_or(10);
if limit > 10_000 { let skip = request.options.pagination.skip.unwrap_or(0);
return Err(RuntimeError::invalid_input("aggregate limit exceeds 10000")); 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!({ Ok(json!({
"query":hit_dsl["query"], "query":hit_dsl["query"],
"from":request.options.pagination.skip.unwrap_or(0), "from":skip,
"size":0, "size":0,
"aggs":{"result":{"terms":{"field":field,"size":limit},"aggs":{"result":{"top_hits":{ "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!({})) "sort":hit_dsl["sort"],"highlight":hit_dsl.get("highlight").cloned().unwrap_or_else(||json!({}))
}}}}} }}}}}
})) }))
@@ -234,7 +333,6 @@ mod tests {
fn scope() -> AuthorizedSearchScope { fn scope() -> AuthorizedSearchScope {
AuthorizedSearchScope { AuthorizedSearchScope {
workspace_id: "workspace".to_string(), workspace_id: "workspace".to_string(),
permission_revision: 1,
docs: DocReadScope::All, docs: DocReadScope::All,
} }
} }
@@ -325,6 +423,10 @@ mod tests {
dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["highlight"]["fields"]["content"], dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["highlight"]["fields"]["content"],
json!({"pre_tags":["<b>"],"post_tags":["</b>"]}) 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; let mut invalid = aggregate;
invalid.field = "aclReadTokens".to_string(); invalid.field = "aclReadTokens".to_string();
@@ -332,5 +434,43 @@ mod tests {
invalid.field = "docId".to_string(); invalid.field = "docId".to_string();
invalid.options.pagination.limit = Some(10_001); invalid.options.pagination.limit = Some(10_001);
assert!(compile_aggregate(&invalid, &scope()).is_err()); 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());
} }
} }
@@ -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))
}
}
@@ -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());
}
}
File diff suppressed because it is too large. Load diff
@@ -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};
@@ -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(())
}
@@ -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)
}
@@ -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>,
}
@@ -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");
}
}
@@ -1,7 +1,5 @@
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::runtime::{RuntimeError, RuntimeResult};
#[napi_derive::napi(string_enum = "snake_case")] #[napi_derive::napi(string_enum = "snake_case")]
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")] #[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)] #[napi_derive::napi(object)]
#[derive(Clone, Debug, Default, Deserialize)] #[derive(Clone, Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")] #[serde(rename_all = "camelCase")]
@@ -89,23 +73,6 @@ pub struct AggregateOptions {
pub pagination: SearchPagination, 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)] #[napi_derive::napi(object)]
pub struct RuntimeSearchQuery { pub struct RuntimeSearchQuery {
pub query_type: String, pub query_type: String,
@@ -133,109 +100,3 @@ pub struct RuntimeAggregateRequest {
pub field: String, pub field: String,
pub options: AggregateOptions, 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());
}
}
@@ -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"]));
}
}
@@ -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,
));
}
}
@@ -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();
}
}
@@ -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,
};
@@ -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();
}
}
}
@@ -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();
}
}
@@ -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
@@ -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)
}
@@ -920,6 +920,16 @@ mod tests {
let enabled_without_provider: SearchRuntimeConfig = enabled_without_provider.indexer.unwrap().into(); let enabled_without_provider: SearchRuntimeConfig = enabled_without_provider.indexer.unwrap().into();
assert!(enabled_without_provider.enabled); assert!(enabled_without_provider.enabled);
assert_eq!(enabled_without_provider.provider, "embedded"); 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] #[test]
+20 -5
View File
@@ -21,15 +21,24 @@ pub(crate) enum RuntimeError {
#[error("search permission state unavailable")] #[error("search permission state unavailable")]
SearchPermissionUnavailable, 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")] #[error("search provider unavailable")]
SearchProviderUnavailable, SearchProviderUnavailable,
#[error("search query is not supported by the active provider")] #[error("search query is not supported by the active provider")]
SearchUnsupportedQuery, SearchUnsupportedQuery,
#[error("search stream replay gap")]
SearchReplayGap,
#[error("{context}: {source}")] #[error("{context}: {source}")]
Database { Database {
context: String, context: String,
@@ -111,9 +120,11 @@ impl RuntimeError {
} }
Self::SearchWorkspaceDenied Self::SearchWorkspaceDenied
| Self::SearchPermissionUnavailable | Self::SearchPermissionUnavailable
| Self::SearchIndexNotReady
| Self::SearchPermissionSyncing
| Self::SearchProviderUnavailable | Self::SearchProviderUnavailable
| Self::SearchUnsupportedQuery | Self::SearchUnsupportedQuery => false,
| Self::SearchReplayGap => false, Self::SearchIndexFailed(_) | Self::SearchSourceInvalid(_) => false,
_ => false, _ => false,
} }
} }
@@ -127,6 +138,10 @@ impl RuntimeError {
} if source.code().as_deref() == Some("40001") } 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 { pub(crate) fn to_napi_error(error: RuntimeError) -> Error {
+394 -19
View File
@@ -5,8 +5,7 @@ use super::{RuntimeError, RuntimeResult, types::EmbeddingHealth};
pub(crate) const RUNTIME_MIGRATIONS: &str = include_str!("sql/runtime_migrations.sql"); pub(crate) const RUNTIME_MIGRATIONS: &str = include_str!("sql/runtime_migrations.sql");
const EMBEDDING_MIGRATION: &str = include_str!("sql/embedding.sql"); const EMBEDDING_MIGRATION: &str = include_str!("sql/embedding.sql");
const SEARCH_MIGRATION: &str = include_str!("sql/search.sql"); const SEARCH_PROJECTION_MIGRATION: &str = include_str!("sql/search_projection.sql");
const SEARCH_ACL_TOKEN_MIGRATION: &str = include_str!("sql/search_acl_tokens.sql");
const EMBEDDING_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0046; const EMBEDDING_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0046;
const SEARCH_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0053; const SEARCH_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0053;
#[cfg(test)] #[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<()> { pub(crate) async fn migrate_search_tables(pool: &PgPool) -> RuntimeResult<()> {
migrate_component( migrate_component(
pool, pool,
"search", "search_projection",
SEARCH_ADVISORY_LOCK, SEARCH_ADVISORY_LOCK,
&[(1, &[SEARCH_MIGRATION]), (2, &[SEARCH_ACL_TOKEN_MIGRATION])], &[(1, &[SEARCH_PROJECTION_MIGRATION])],
) )
.await?; .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(()) Ok(())
} }
@@ -118,6 +113,10 @@ async fn migrate_component(
.begin() .begin()
.await .await
.map_err(|error| RuntimeError::database("Native migration transaction failed", error))?; .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)") sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(advisory_lock) .bind(advisory_lock)
.execute(&mut *transaction) .execute(&mut *transaction)
@@ -269,11 +268,31 @@ mod tests {
} }
#[test] #[test]
fn search_schema_uses_transactional_stream_heads() { fn search_schema_uses_terminal_control_plane() {
assert!(SEARCH_MIGRATION.contains("CREATE TABLE search_runtime_streams")); assert!(SEARCH_PROJECTION_MIGRATION.contains("DROP SCHEMA IF EXISTS search_projection CASCADE"));
assert!(SEARCH_MIGRATION.contains("PRIMARY KEY (table_key, stream_sequence)")); assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SCHEMA search_projection"));
assert!(!SEARCH_MIGRATION.contains("BIGSERIAL")); assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.generations"));
assert!(!SEARCH_MIGRATION.contains("CREATE SEQUENCE")); 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] #[tokio::test]
@@ -285,11 +304,367 @@ mod tests {
let (first, second) = tokio::join!(migrate_search_tables(&pool), migrate_search_tables(&pool)); let (first, second) = tokio::join!(migrate_search_tables(&pool), migrate_search_tables(&pool));
first.unwrap(); first.unwrap();
second.unwrap(); second.unwrap();
let versions: Vec<i32> = let versions: Vec<i32> = sqlx::query_scalar(
sqlx::query_scalar("SELECT version FROM native_schema_migrations WHERE component='search' ORDER BY version") "SELECT version FROM native_schema_migrations WHERE component='search_projection' ORDER BY version",
.fetch_all(&pool) )
.await .fetch_all(&pool)
.unwrap(); .await
assert_eq!(versions, vec![1, 2]); .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();
} }
} }
@@ -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();
@@ -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()
);
@@ -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();
@@ -6,9 +6,9 @@ use sqlx::{FromRow, PgPool, Postgres, Row, Transaction};
use uuid::Uuid; use uuid::Uuid;
use super::{ use super::{
CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupAckResult, RuntimeDocumentCleanupEffect, CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult,
RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult, RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult, StorageRuntime, load_workspace_live_doc_ids,
StorageRuntime, load_workspace_live_doc_ids, merge_current_doc, napi_error, merge_current_doc, napi_error,
}; };
#[derive(FromRow)] #[derive(FromRow)]
@@ -416,7 +416,6 @@ async fn delete_doc_rows(tx: &mut Transaction<'_, Postgres>, candidate: &Candida
"cleanupVersion": cleanup_version, "cleanupVersion": cleanup_version,
"commentAttachmentKeys": attachment_keys, "commentAttachmentKeys": attachment_keys,
"commentObjectsDone": false, "commentObjectsDone": false,
"searchDone": false,
})) }))
.execute(&mut **tx) .execute(&mut **tx)
.await .await
@@ -592,51 +591,43 @@ async fn execute_one(
Ok(Some((candidate, deleted_rows))) Ok(Some((candidate, deleted_rows)))
} }
fn payload_effect(effect: PendingEffect) -> RuntimeResult<RuntimeDocumentCleanupEffect> { fn payload_effect(effect: &PendingEffect) -> RuntimeResult<RuntimeDocumentCleanupEffect> {
let cleanup_version = effect let cleanup_version = effect
.cleanup_payload .cleanup_payload
.get("cleanupVersion") .get("cleanupVersion")
.and_then(Value::as_str) .and_then(Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("Document cleanup effect payload has no cleanupVersion"))?; .ok_or_else(|| RuntimeError::invalid_state("Document cleanup effect payload has no cleanupVersion"))?;
Ok(RuntimeDocumentCleanupEffect { Ok(RuntimeDocumentCleanupEffect {
workspace_id: effect.workspace_id, workspace_id: effect.workspace_id.clone(),
doc_id: effect.doc_id, doc_id: effect.doc_id.clone(),
cleanup_version: cleanup_version.to_string(), cleanup_version: cleanup_version.to_string(),
comment_objects_done: effect comment_objects_done: effect
.cleanup_payload .cleanup_payload
.get("commentObjectsDone") .get("commentObjectsDone")
.and_then(Value::as_bool) .and_then(Value::as_bool)
.unwrap_or(false), .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>, tx: &mut Transaction<'_, Postgres>,
workspace_id: &str, workspace_id: &str,
doc_id: &str, doc_id: &str,
cleanup_version: &str, cleanup_version: &str,
path: &str,
) -> RuntimeResult<bool> { ) -> RuntimeResult<bool> {
let completed = sqlx::query_scalar::<_, bool>( let completed = sqlx::query_scalar::<_, bool>(
r#" r#"
UPDATE document_cleanup_candidates 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 error = NULL, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending' WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending'
AND cleanup_payload->>'cleanupVersion' = $3 AND cleanup_payload->>'cleanupVersion' = $3
RETURNING COALESCE((cleanup_payload->>'commentObjectsDone')::boolean, false) RETURNING COALESCE((cleanup_payload->>'commentObjectsDone')::boolean, false)
AND COALESCE((cleanup_payload->>'searchDone')::boolean, false)
"#, "#,
) )
.bind(workspace_id) .bind(workspace_id)
.bind(doc_id) .bind(doc_id)
.bind(cleanup_version) .bind(cleanup_version)
.bind(path)
.fetch_optional(&mut **tx) .fetch_optional(&mut **tx)
.await .await
.map_err(|err| RuntimeError::database("Document cleanup effect completion failed", err))?; .map_err(|err| RuntimeError::database("Document cleanup effect completion failed", err))?;
@@ -658,7 +649,7 @@ async fn complete_effect(
Ok(completed) 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 let cleanup_version = effect
.cleanup_payload .cleanup_payload
.get("cleanupVersion") .get("cleanupVersion")
@@ -695,18 +686,11 @@ async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffec
.begin() .begin()
.await .await
.map_err(|err| RuntimeError::database("Document cleanup object effect transaction failed", err))?; .map_err(|err| RuntimeError::database("Document cleanup object effect transaction failed", err))?;
complete_effect( let completed = complete_comment_objects(&mut tx, &effect.workspace_id, &effect.doc_id, cleanup_version).await?;
&mut tx,
&effect.workspace_id,
&effect.doc_id,
cleanup_version,
"commentObjectsDone",
)
.await?;
tx.commit() tx.commit()
.await .await
.map_err(|err| RuntimeError::database("Document cleanup object effect commit failed", err))?; .map_err(|err| RuntimeError::database("Document cleanup object effect commit failed", err))?;
Ok(()) Ok(completed)
} }
async fn load_pending_effects( async fn load_pending_effects(
@@ -789,75 +773,33 @@ impl StorageRuntime {
} }
let effects = load_pending_effects(&pool, workspace_id.as_deref(), limit).await?; let effects = load_pending_effects(&pool, workspace_id.as_deref(), limit).await?;
for effect in effects { for effect in effects {
if let Err(err) = process_comment_objects(self, &effect).await { let mut result_effect = payload_effect(&effect)?;
result.failed += 1; match process_comment_objects(self, &effect).await {
sqlx::query( Ok(comment_objects_done) => {
r#" result_effect.comment_objects_done = comment_objects_done;
UPDATE document_cleanup_candidates result.effects.push(result_effect);
SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP }
WHERE workspace_id = $1 AND doc_id = $2 Err(err) => {
"#, result.failed += 1;
) result.effects.push(result_effect);
.bind(&effect.workspace_id) sqlx::query(
.bind(&effect.doc_id) r#"
.bind(err.to_string()) UPDATE document_cleanup_candidates
.execute(&pool) SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP
.await WHERE workspace_id = $1 AND doc_id = $2
.map_err(|db_err| RuntimeError::database("Document cleanup effect failure write failed", db_err))?; "#,
)
.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) 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)] #[cfg(test)]
@@ -870,7 +812,11 @@ mod tests {
use tokio::sync::Mutex; use tokio::sync::Mutex;
use super::*; 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)>> { async fn runtime_from_database_url() -> AnyResult<Option<(StorageRuntime, PgPool)>> {
let Ok(database_url) = std::env::var("DATABASE_URL") else { let Ok(database_url) = std::env::var("DATABASE_URL") else {
@@ -1233,10 +1179,17 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn document_cleanup_execute_postgres_semantics() -> AnyResult<()> { 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 { let Some((runtime, pool)) = runtime_from_database_url().await? else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set"); eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return Ok(()); 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 suffix = Uuid::new_v4().to_string();
let (user_id, workspace_id) = insert_user_workspace(&pool, &suffix).await?; let (user_id, workspace_id) = insert_user_workspace(&pool, &suffix).await?;
let object_root = tempfile::tempdir()?; let object_root = tempfile::tempdir()?;
@@ -1462,6 +1415,14 @@ mod tests {
.bind(doc_id) .bind(doc_id)
.execute(&pool) .execute(&pool)
.await?; .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 let executed = runtime
.execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10) .execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10)
.await .await
@@ -1470,7 +1431,16 @@ mod tests {
assert_eq!(executed.failed, 0); assert_eq!(executed.failed, 0);
assert_eq!(executed.effects.len(), 1); assert_eq!(executed.effects.len(), 1);
assert!(executed.effects[0].comment_objects_done); 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!( assert!(
runtime runtime
.head_object("blob".to_string(), attachment_object_key) .head_object("blob".to_string(), attachment_object_key)
@@ -1518,29 +1488,6 @@ mod tests {
assert_eq!(count, 0, "{table}.{column} should be nulled"); 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>( let candidate_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", "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_eq!(failed_object_delete.failed, 1);
assert!(!failed_object_delete.effects[0].comment_objects_done); assert!(!failed_object_delete.effects[0].comment_objects_done);
let retained = sqlx::query( let retained = sqlx::query(
"SELECT status, attempt_count, error, cleanup_payload->>'cleanupVersion' AS cleanup_version FROM \ "SELECT status, attempt_count, error FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
) )
.bind(&workspace_id) .bind(&workspace_id)
.bind(retry_doc_id) .bind(retry_doc_id)
@@ -1598,18 +1544,6 @@ mod tests {
assert_eq!(retained.get::<String, _>("status"), "effects_pending"); assert_eq!(retained.get::<String, _>("status"), "effects_pending");
assert_eq!(retained.get::<i32, _>("attempt_count"), 1); assert_eq!(retained.get::<i32, _>("attempt_count"), 1);
assert!(retained.get::<Option<String>, _>("error").is_some()); 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( sqlx::query(
"UPDATE document_cleanup_candidates SET cleanup_payload = jsonb_set(cleanup_payload, '{commentAttachmentKeys}', \ "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()))?; .map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(retried.executed, 0); assert_eq!(retried.executed, 0);
assert_eq!(retried.failed, 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>( let retry_candidate_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", "SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
) )
@@ -1635,6 +1570,10 @@ mod tests {
.await?; .await?;
assert_eq!(retry_candidate_count, 0); 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?; cleanup_workspace_fixture(&pool, &user_id, &workspace_id).await?;
Ok(()) Ok(())
} }
@@ -32,10 +32,10 @@ pub(super) use super::{
napi_error, to_napi_error, napi_error, to_napi_error,
types::{ types::{
RuntimeBlobCleanupExecuteResult, RuntimeBlobCleanupPlanResult, RuntimeBlobCleanupResult, RuntimeBlobCompleteResult, RuntimeBlobCleanupExecuteResult, RuntimeBlobCleanupPlanResult, RuntimeBlobCleanupResult, RuntimeBlobCompleteResult,
RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupAckResult, RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupEffect,
RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeMultipartUploadInit,
RuntimeMultipartUploadInit, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry, RuntimeObjectMetadata,
RuntimeObjectMetadata, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
}, },
}; };
@@ -476,7 +476,6 @@ pub struct RuntimeDocumentCleanupEffect {
pub doc_id: String, pub doc_id: String,
pub cleanup_version: String, pub cleanup_version: String,
pub comment_objects_done: bool, pub comment_objects_done: bool,
pub search_done: bool,
} }
#[napi_derive::napi(object)] #[napi_derive::napi(object)]
@@ -491,11 +490,6 @@ pub struct RuntimeDocumentCleanupExecuteResult {
pub effects: Vec<RuntimeDocumentCleanupEffect>, pub effects: Vec<RuntimeDocumentCleanupEffect>,
} }
#[napi_derive::napi(object)]
pub struct RuntimeDocumentCleanupAckResult {
pub completed: bool,
}
#[napi_derive::napi(object)] #[napi_derive::napi(object)]
pub struct RuntimeBlobCleanupPlanResult { pub struct RuntimeBlobCleanupPlanResult {
pub run_id: Option<String>, pub run_id: Option<String>,
@@ -527,34 +521,6 @@ pub struct RuntimeDocCompactionResult {
pub history_created: bool, 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)] #[napi_derive::napi(object)]
pub struct RuntimeEmbeddingProgress { pub struct RuntimeEmbeddingProgress {
pub total: i64, pub total: i64,
+375 -74
View File
@@ -3,12 +3,13 @@ mod query;
mod result; mod result;
mod schema; mod schema;
use std::sync::Arc; use std::{collections::HashMap, sync::Arc};
use memory_indexer::{MemoryIndex, Mutation, TermsAggregation}; use memory_indexer::{MemoryIndex, Mutation, TermsAggregation};
use napi::{Status, bindgen_prelude::Buffer}; use napi::Status;
use serde_json::Value as JsonValue; use serde_json::{Value as JsonValue, json};
use tokio::sync::RwLock; use tokio::sync::RwLock;
use uuid::Uuid;
use self::{ use self::{
document::compile_document, document::compile_document,
@@ -49,12 +50,12 @@ impl TableIndex {
} }
} }
struct IndexManager { struct GenerationIndex {
doc: TableIndex, doc: TableIndex,
block: TableIndex, block: TableIndex,
} }
impl IndexManager { impl GenerationIndex {
fn new() -> Self { fn new() -> Self {
Self { Self {
doc: TableIndex::new(TableSchema::doc()), doc: TableIndex::new(TableSchema::doc()),
@@ -71,9 +72,32 @@ impl IndexManager {
} }
} }
pub(crate) struct EmbeddedIndexCheckpoint { struct IndexManager {
pub sequence: i64, generations: RwLock<HashMap<Uuid, Arc<GenerationIndex>>>,
pub data: Buffer, }
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 { pub(crate) struct EmbeddedSearchIndex {
@@ -87,26 +111,52 @@ impl EmbeddedSearchIndex {
} }
} }
pub(crate) async fn restore(&self, table: String, checkpoint: Buffer) -> napi::Result<()> { pub(crate) async fn prepare_generation(&self, generation_id: Uuid) {
let table = self.manager.table(&table)?; self.manager.generation(generation_id, true).await;
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 reset(&self, table: String) -> napi::Result<()> { pub(crate) async fn has_generation(&self, generation_id: Uuid) -> bool {
let table = self.manager.table(&table)?; self.manager.generation(generation_id, false).await.is_some()
*table.index.write().await = MemoryIndex::new(table.schema.schema.clone());
Ok(())
} }
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<()> { 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 documents: Vec<JsonValue> = serde_json::from_str(&documents_json)?;
let generation_id = generation_id.to_string();
let documents = documents let documents = documents
.into_iter() .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<_>>>()?; .collect::<Result<Vec<_>>>()?;
table table
.index .index
@@ -117,14 +167,148 @@ impl EmbeddedSearchIndex {
Ok(()) Ok(())
} }
pub(crate) async fn delete(&self, table: String, id: String) -> napi::Result<()> { pub(crate) async fn gc_document_history_for_generation(
self.manager.table(&table)?.index.write().await.delete(&id); &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(()) 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> { pub(crate) async fn search(&self, table: String, dsl_json: String) -> napi::Result<String> {
let table = self.manager.table(&table)?; self.search_for_generation(Uuid::nil(), table, dsl_json).await
let dsl: JsonValue = serde_json::from_str(&dsl_json)?; }
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( let query = compile_query(
&table.schema, &table.schema,
dsl dsl
@@ -144,9 +328,27 @@ impl EmbeddedSearchIndex {
))?) ))?)
} }
#[cfg(test)]
pub(crate) async fn aggregate(&self, table: String, dsl_json: String) -> napi::Result<String> { pub(crate) async fn aggregate(&self, table: String, dsl_json: String) -> napi::Result<String> {
let table = self.manager.table(&table)?; self.aggregate_for_generation(Uuid::nil(), table, dsl_json).await
let dsl: JsonValue = serde_json::from_str(&dsl_json)?; }
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( let query = compile_query(
&table.schema, &table.schema,
dsl dsl
@@ -187,38 +389,23 @@ impl EmbeddedSearchIndex {
&highlight_tags(dsl.pointer("/aggs/result/aggs/result/top_hits").unwrap_or(&dsl)), &highlight_tags(dsl.pointer("/aggs/result/aggs/result/top_hits").unwrap_or(&dsl)),
))?) ))?)
} }
}
pub(crate) async fn checkpoint(&self, table: String) -> napi::Result<EmbeddedIndexCheckpoint> { fn ensure_projection_fields(dsl: &mut JsonValue) -> Result<()> {
let checkpoint = self let fields = dsl
.manager .as_object_mut()
.table(&table)? .ok_or_else(|| IndexError::InvalidInput("search DSL must be an object".into()))?
.index .entry("fields")
.read() .or_insert_with(|| json!([]));
.await let fields = fields
.checkpoint() .as_array_mut()
.map_err(IndexError::from)?; .ok_or_else(|| IndexError::InvalidInput("search fields must be an array".into()))?;
Ok(EmbeddedIndexCheckpoint { for field in ["doc_id", "source_version", "permission_version"] {
sequence: checkpoint.sequence as i64, if !fields.iter().any(|value| value.as_str() == Some(field)) {
data: checkpoint.bytes.into(), fields.push(json!(field));
}) }
}
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(())
} }
Ok(())
} }
fn highlight_tags(dsl: &JsonValue) -> HighlightTags { fn highlight_tags(dsl: &JsonValue) -> HighlightTags {
@@ -248,13 +435,17 @@ impl Default for EmbeddedSearchIndex {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use serde_json::{Value, json}; use serde_json::{Value, json};
use uuid::Uuid;
use super::EmbeddedSearchIndex; use super::EmbeddedSearchIndex;
fn doc(workspace: &str, id: &str, title: &str, updated_at: i64) -> Value { fn doc(workspace: &str, id: &str, title: &str, updated_at: i64) -> Value {
json!({ json!({
"generation_id": Uuid::nil().to_string(),
"workspace_id": workspace, "workspace_id": workspace,
"doc_id": id, "doc_id": id,
"source_version": 1,
"permission_version": 1,
"title": title, "title": title,
"summary": title, "summary": title,
"created_by_user_id": "user", "created_by_user_id": "user",
@@ -277,7 +468,7 @@ mod tests {
} }
#[tokio::test] #[tokio::test]
async fn exact_search_cursor_and_checkpoint_roundtrip() { async fn exact_search_cursor_and_aggregate() {
let index = EmbeddedSearchIndex::new(); let index = EmbeddedSearchIndex::new();
index index
.write( .write(
@@ -296,7 +487,7 @@ mod tests {
let first: Value = let first: Value =
serde_json::from_str(&index.search("doc".into(), search(query.clone(), None)).await.unwrap()).unwrap(); serde_json::from_str(&index.search("doc".into(), search(query.clone(), None)).await.unwrap()).unwrap();
assert_eq!(first["total"], 2); 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"])); assert_eq!(first["nodes"][0]["fields"]["doc_id"], json!(["two"]));
let cursor = first["nextCursor"].as_str().unwrap(); let cursor = first["nextCursor"].as_str().unwrap();
@@ -307,17 +498,7 @@ mod tests {
.unwrap(), .unwrap(),
) )
.unwrap(); .unwrap();
assert_eq!(second["nodes"][0]["id"], "workspace-1/one"); assert_eq!(second["nodes"][0]["id"], format!("{}/workspace-1/one/1/1", Uuid::nil()));
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);
let aggregate: Value = serde_json::from_str( let aggregate: Value = serde_json::from_str(
&index &index
@@ -343,12 +524,16 @@ mod tests {
} }
#[tokio::test] #[tokio::test]
async fn write_is_atomic_and_corrupt_checkpoint_is_rejected() { async fn write_is_atomic() {
let index = EmbeddedSearchIndex::new(); let index = EmbeddedSearchIndex::new();
index index
.write( .write(
"doc".into(), "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 .await
.unwrap(); .unwrap();
@@ -357,7 +542,6 @@ mod tests {
serde_json::from_str(&index.search("doc".into(), search(all.clone(), None)).await.unwrap()).unwrap(); serde_json::from_str(&index.search("doc".into(), search(all.clone(), None)).await.unwrap()).unwrap();
assert_eq!(result["total"], 1); assert_eq!(result["total"], 1);
index.reset("doc".into()).await.unwrap();
let documents = json!([ let documents = json!([
doc("workspace", "valid", "hello", 1), doc("workspace", "valid", "hello", 1),
{ "workspace_id": "workspace", "doc_id": "invalid", "unknown": true } { "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()); 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(); let result: Value = serde_json::from_str(&index.search("doc".into(), search(all, None)).await.unwrap()).unwrap();
assert_eq!(result["total"], 0); assert_eq!(result["total"], 1);
assert!(index.restore("doc".into(), vec![1, 2, 3].into()).await.is_err()); }
#[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"]));
} }
} }
@@ -16,6 +16,9 @@ impl TableSchema {
let mut fields = HashMap::new(); let mut fields = HashMap::new();
keyword(&mut builder, &mut fields, "workspace_id", true, false); keyword(&mut builder, &mut fields, "workspace_id", true, false);
keyword(&mut builder, &mut fields, "workspace_token", 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_id", true, true);
keyword(&mut builder, &mut fields, "doc_token", true, false); keyword(&mut builder, &mut fields, "doc_token", true, false);
fields.insert( fields.insert(
@@ -29,10 +32,19 @@ impl TableSchema {
keyword(&mut builder, &mut fields, "acl_read_tokens", true, false); keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable"); boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_member_default_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, "created_at", true, true);
integer(&mut builder, &mut fields, "updated_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 { pub fn block() -> Self {
@@ -58,6 +70,9 @@ impl TableSchema {
} }
keyword(&mut builder, &mut fields, "doc_id", true, true); keyword(&mut builder, &mut fields, "doc_id", true, true);
keyword(&mut builder, &mut fields, "workspace_token", 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_token", 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_id", true, true);
keyword(&mut builder, &mut fields, "block_token", true, false); 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); keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable"); boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_member_default_readable"); boolean(&mut builder, &mut fields, "acl_member_default_readable");
integer(&mut builder, &mut fields, "acl_revision", true, false); Self::finish(
Self::finish(builder, fields, &["workspace_id", "doc_id", "block_id"]) builder,
fields,
&[
"generation_id",
"workspace_id",
"doc_id",
"source_version",
"permission_version",
"block_id",
],
)
} }
fn finish( fn finish(
@@ -95,10 +120,14 @@ impl TableSchema {
.id_fields .id_fields
.iter() .iter()
.map(|field| { .map(|field| {
document let value = document
.get(*field) .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<_>>>() .collect::<Result<Vec<_>>>()
.map(|parts| parts.join("/")) .map(|parts| parts.join("/"))
-2
View File
@@ -11,5 +11,3 @@
# MAILER_USER="noreply@toeverything.info" # MAILER_USER="noreply@toeverything.info"
# MAILER_PASSWORD="affine" # MAILER_PASSWORD="affine"
# MAILER_SECURE=false # MAILER_SECURE=false
# AFFINE_INDEXER_ENABLED=true
@@ -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
```
@@ -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
@@ -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.
@@ -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
@@ -10,7 +10,7 @@ import {
type Config, type Config,
type EventBus, type EventBus,
type JobQueue, type JobQueue,
SearchProviderNotFound, SearchProviderUnavailable,
} from '../../base'; } from '../../base';
import { ServerFeature, type ServerService } from '../../core'; import { ServerFeature, type ServerService } from '../../core';
import type { DocReader } from '../../core/doc'; import type { DocReader } from '../../core/doc';
@@ -450,7 +450,7 @@ test('document tools enforce the user-selected hard scope', async t => {
readableAc, readableAc,
{ {
searchDocsByKeyword: async () => { searchDocsByKeyword: async () => {
throw new SearchProviderNotFound(); throw new SearchProviderUnavailable();
}, },
} as unknown as IndexerService, } as unknown as IndexerService,
vectorSearch, vectorSearch,
@@ -1,7 +1,6 @@
import { getCurrentUserQuery } from '@affine/graphql'; import { getCurrentUserQuery } from '@affine/graphql';
import { JobExecutor } from '../../../base/job/queue/executor'; import { JobExecutor } from '../../../base/job/queue/executor';
import { JobHandlerScanner } from '../../../base/job/queue/scanner';
import { DatabaseDocReader, DocReader } from '../../../core/doc'; import { DatabaseDocReader, DocReader } from '../../../core/doc';
import { RealtimeGateway } from '../../../core/realtime/gateway'; import { RealtimeGateway } from '../../../core/realtime/gateway';
import { createApp } from '../create-app'; import { createApp } from '../create-app';
@@ -43,7 +42,6 @@ e2e('should init worker service', async t => {
await withFlavor('worker', async app => { await withFlavor('worker', async app => {
const res = await app.GET('/info').expect(200); const res = await app.GET('/info').expect(200);
t.is(res.body.flavor, 'worker'); t.is(res.body.flavor, 'worker');
t.truthy(app.get(JobHandlerScanner).getHandler('indexer.indexDoc'));
t.throws(() => app.get(RealtimeGateway)); t.throws(() => app.get(RealtimeGateway));
await t.throwsAsync(app.gql({ query: getCurrentUserQuery })); await t.throwsAsync(app.gql({ query: getCurrentUserQuery }));
@@ -55,7 +53,6 @@ e2e('should init allinone service with worker handlers', async t => {
await withFlavor('allinone', async app => { await withFlavor('allinone', async app => {
const res = await app.GET('/info').expect(200); const res = await app.GET('/info').expect(200);
t.is(res.body.flavor, 'allinone'); t.is(res.body.flavor, 'allinone');
t.truthy(app.get(JobHandlerScanner).getHandler('indexer.indexDoc'));
}); });
}); });
@@ -1,16 +1,27 @@
import { serverConfigQuery, ServerFeature } from '@affine/graphql'; 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'; import { app, e2e } from '../test';
e2e('should indexer feature enabled by default', async t => { e2e(
const { serverConfig } = await app.gql({ query: serverConfigQuery }); '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( t.is(
serverConfig.features.includes(ServerFeature.Indexer), serverConfig.features.includes(ServerFeature.Indexer),
true, enabled,
JSON.stringify(serverConfig, null, 2) JSON.stringify(serverConfig, null, 2)
); );
}); }
);
e2e('should comment feature enabled by default', async t => { e2e('should comment feature enabled by default', async t => {
const { serverConfig } = await app.gql({ query: serverConfigQuery }); const { serverConfig } = await app.gql({ query: serverConfigQuery });
@@ -362,7 +362,9 @@ export async function createApp(
await app.init(); await app.init();
await app.get(BackendRuntimeProvider, { strict: false }).runMigrations(); await app.get(BackendRuntimeProvider, { strict: false }).runMigrations();
await app.get(StorageRuntimeProvider, { 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) { } catch (error) {
await app.close(); await app.close();
throw error; throw error;
@@ -4,14 +4,20 @@ import {
SearchTable, SearchTable,
} from '@affine/graphql'; } from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { createDocWithMarkdown } from '../../../native'; import { createDocWithMarkdown } from '../../../native';
import { IndexerService } from '../../../plugins/indexer/service';
import { Mockers } from '../../mocks'; import { Mockers } from '../../mocks';
import { app, e2e } from '../test'; 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 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 [ for (const [docId, markdown] of [
['doc-0', 'hello world\n\nhello again'], ['doc-0', 'hello world\n\nhello again'],
['doc-1', 'hello world'], ['doc-1', 'hello world'],
@@ -23,7 +29,7 @@ e2e('should aggregate by docId', async t => {
user: owner, user: owner,
blob: createDocWithMarkdown(docId, markdown, docId), blob: createDocWithMarkdown(docId, markdown, docId),
}); });
await app.get(IndexerService).indexDoc(workspace.id, docId); await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
} }
const result = await app.gql({ const result = await app.gql({
@@ -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);
}
);
@@ -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);
}
);
@@ -1,15 +1,19 @@
import { indexerSearchDocsQuery } from '@affine/graphql'; import { indexerSearchDocsQuery } from '@affine/graphql';
import { ConfigFactory } from '../../../base'; import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { createDocWithMarkdown } from '../../../native'; import { createDocWithMarkdown } from '../../../native';
import { SearchProviderType } from '../../../plugins/indexer/config';
import { IndexerService } from '../../../plugins/indexer/service';
import { Mockers } from '../../mocks'; import { Mockers } from '../../mocks';
import { app, e2e } from '../test'; 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 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']) { for (const docId of ['doc-0', 'doc-1', 'doc-2']) {
await app.create(Mockers.DocMeta, { workspaceId: workspace.id, docId }); await app.create(Mockers.DocMeta, { workspaceId: workspace.id, docId });
await app.create(Mockers.DocSnapshot, { await app.create(Mockers.DocSnapshot, {
@@ -18,33 +22,26 @@ e2e('should search docs by keyword', async t => {
user: owner, user: owner,
blob: createDocWithMarkdown(docId, `${docId} hello`, docId), blob: createDocWithMarkdown(docId, `${docId} hello`, docId),
}); });
await app.get(IndexerService).indexDoc(workspace.id, docId); await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
} }
const search = app.gql({ const search = app.gql({
query: indexerSearchDocsQuery, query: indexerSearchDocsQuery,
variables: { id: workspace.id, input: { keyword: 'hello', limit: 2 } }, 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; const result = await search;
t.is(result.workspace.searchDocs.length, 2); t.is(result.workspace.searchDocs.length, 2);
t.true(result.workspace.searchDocs.every(doc => doc.highlight.length > 0)); t.true(result.workspace.searchDocs.every(doc => doc.highlight.length > 0));
}); });
e2e( indexerE2e(
'should search docs by keyword failed when workspace is no permission', 'should search docs by keyword failed when workspace is no permission',
async t => { async t => {
const owner = await app.signup(); 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 app.signup();
await t.throwsAsync( await t.throwsAsync(
app.gql({ app.gql({
@@ -4,12 +4,15 @@ import {
SearchTable, SearchTable,
} from '@affine/graphql'; } from '@affine/graphql';
import { Config } from '../../../base';
import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { DocRole } from '../../../models'; import { DocRole } from '../../../models';
import { createDocWithMarkdown } from '../../../native'; import { createDocWithMarkdown } from '../../../native';
import { IndexerService } from '../../../plugins/indexer/service';
import { Mockers } from '../../mocks'; import { Mockers } from '../../mocks';
import { app, e2e } from '../test'; import { app, e2e } from '../test';
const indexerE2e = app.get(Config).indexer.enabled ? e2e : e2e.skip;
async function indexDoc( async function indexDoc(
workspaceId: string, workspaceId: string,
user: { id: string }, user: { id: string },
@@ -24,12 +27,15 @@ async function indexDoc(
user, user,
blob: createDocWithMarkdown(docId, markdown, docId), 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 owner = await app.signup();
const workspace = await app.create(Mockers.Workspace, { owner }); const workspace = await app.create(Mockers.Workspace, {
owner,
snapshot: true,
});
await indexDoc( await indexDoc(
workspace.id, workspace.id,
owner, owner,
@@ -114,100 +120,116 @@ e2e('should search with query', async t => {
); );
}); });
e2e('should filter no read permission docs on team workspace', async t => { indexerE2e(
const owner = await app.signup(); 'should filter no read permission docs on team workspace',
const workspace = await app.create(Mockers.Workspace, { owner }); async t => {
await app.create(Mockers.TeamWorkspace, { id: workspace.id }); const owner = await app.signup();
await indexDoc( const workspace = await app.create(Mockers.Workspace, {
workspace.id, owner,
owner, snapshot: true,
'private-doc', });
'team secret searchable', await app.create(Mockers.TeamWorkspace, { id: workspace.id });
DocRole.None await indexDoc(
); workspace.id,
owner,
'private-doc',
'team secret searchable',
DocRole.None
);
const member = await app.signup(); const member = await app.signup();
await app.create(Mockers.WorkspaceUser, { await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id, workspaceId: workspace.id,
userId: member.id, userId: member.id,
}); });
await app.get(IndexerService).reconcileWorkspace(workspace.id); await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const denied = await app.gql({ const denied = await app.gql({
query: indexerSearchQuery, query: indexerSearchQuery,
variables: { variables: {
id: workspace.id, id: workspace.id,
input: { input: {
table: SearchTable.block, table: SearchTable.block,
query: { query: {
type: SearchQueryType.match, type: SearchQueryType.match,
field: 'content', field: 'content',
match: 'secret', 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, { await app.create(Mockers.DocUser, {
workspaceId: workspace.id, workspaceId: workspace.id,
docId: 'private-doc', docId: 'private-doc',
userId: member.id, userId: member.id,
type: DocRole.Reader, type: DocRole.Reader,
}); });
await app.get(IndexerService).reconcileWorkspace(workspace.id); await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const allowed = await app.gql({ const allowed = await app.gql({
query: indexerSearchQuery, query: indexerSearchQuery,
variables: { variables: {
id: workspace.id, id: workspace.id,
input: { input: {
table: SearchTable.block, table: SearchTable.block,
query: { query: {
type: SearchQueryType.match, type: SearchQueryType.match,
field: 'content', field: 'content',
match: 'secret', 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.models.docUser.delete(workspace.id, 'private-doc', member.id);
await app.get(IndexerService).reconcileWorkspace(workspace.id); await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
const revoked = await app.gql({ const revoked = await app.gql({
query: indexerSearchQuery, query: indexerSearchQuery,
variables: { variables: {
id: workspace.id, id: workspace.id,
input: { input: {
table: SearchTable.block, table: SearchTable.block,
query: { query: {
type: SearchQueryType.match, type: SearchQueryType.match,
field: 'content', field: 'content',
match: 'secret', 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 => { indexerE2e(
const owner = await app.signup(); 'should return empty results when search not match any docs',
const workspace = await app.create(Mockers.Workspace, { owner }); async t => {
await app.get(IndexerService).reconcileWorkspace(workspace.id); const owner = await app.signup();
const result = await app.gql({ const workspace = await app.create(Mockers.Workspace, {
query: indexerSearchQuery, owner,
variables: { snapshot: true,
id: workspace.id, });
input: { await app.get(BackendRuntimeProvider).reconcileSearchProjection(1000);
table: SearchTable.doc, const result = await app.gql({
query: { type: SearchQueryType.match, field: 'title', match: 'absent' }, query: indexerSearchQuery,
options: { fields: ['docId'], pagination: { limit: 10 } }, 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.is(result.workspace.search.pagination.count, 0); t.deepEqual(result.workspace.search.nodes, []);
t.deepEqual(result.workspace.search.nodes, []); }
}); );
@@ -6,14 +6,18 @@ export const TEST_LOG_LEVEL: LogLevel =
(process.env.TEST_LOG_LEVEL as LogLevel) ?? 'fatal'; (process.env.TEST_LOG_LEVEL as LogLevel) ?? 'fatal';
async function flushDB(client: PrismaClient) { async function flushDB(client: PrismaClient) {
const result: { tablename: string }[] = const result: { schemaname: string; tablename: string }[] =
await client.$queryRaw`SELECT tablename await client.$queryRaw`SELECT schemaname, tablename
FROM pg_catalog.pg_tables FROM pg_catalog.pg_tables
WHERE schemaname != 'pg_catalog' WHERE schemaname != 'pg_catalog'
AND schemaname != 'information_schema'`; AND schemaname != 'information_schema'`;
const query = `TRUNCATE TABLE ${result const query = `TRUNCATE TABLE ${result
.map(({ tablename }) => tablename) .filter(({ tablename }) => !tablename.includes('migrations'))
.filter(name => !name.includes('migrations')) .map(({ schemaname, tablename }) =>
[schemaname, tablename]
.map(identifier => `"${identifier.replaceAll('"', '""')}"`)
.join('.')
)
.join(', ')}`; .join(', ')}`;
for (let attempt = 0; attempt < 3; attempt++) { for (let attempt = 0; attempt < 3; attempt++) {
+2 -3
View File
@@ -61,7 +61,7 @@ import { CalendarModule } from './plugins/calendar';
import { CaptchaModule } from './plugins/captcha'; import { CaptchaModule } from './plugins/captcha';
import { CopilotModule } from './plugins/copilot'; import { CopilotModule } from './plugins/copilot';
import { GCloudModule } from './plugins/gcloud'; import { GCloudModule } from './plugins/gcloud';
import { IndexerModule, IndexerWorkerModule } from './plugins/indexer'; import { IndexerModule } from './plugins/indexer';
import { LicenseModule } from './plugins/license'; import { LicenseModule } from './plugins/license';
import { OAuthModule } from './plugins/oauth'; import { OAuthModule } from './plugins/oauth';
import { PaymentModule } from './plugins/payment'; import { PaymentModule } from './plugins/payment';
@@ -169,9 +169,8 @@ export function buildAppModule(env: Env) {
.use(...FunctionalityModules) .use(...FunctionalityModules)
.useIf(() => !workerOnly, RealtimeGatewayModule) .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.isApi || env.isFrontend, IndexerModule)
.useIf(() => env.isWorker, IndexerWorkerModule)
// the worker owns doc consumers and schedulers // the worker owns doc consumers and schedulers
.useIf(() => env.isWorker, DocJobsModule) .useIf(() => env.isWorker, DocJobsModule)
@@ -14,6 +14,7 @@ export type UserFriendlyErrorBaseType =
| 'no_permission' | 'no_permission'
| 'quota_exceeded' | 'quota_exceeded'
| 'authentication_required' | 'authentication_required'
| 'service_unavailable'
| 'internal_server_error'; | 'internal_server_error';
type ErrorArgType = 'string' | 'number' | 'boolean'; type ErrorArgType = 'string' | 'number' | 'boolean';
@@ -36,6 +37,7 @@ const BaseTypeToHttpStatusMap: Record<UserFriendlyErrorBaseType, HttpStatus> = {
no_permission: HttpStatus.FORBIDDEN, no_permission: HttpStatus.FORBIDDEN,
quota_exceeded: HttpStatus.PAYMENT_REQUIRED, quota_exceeded: HttpStatus.PAYMENT_REQUIRED,
authentication_required: HttpStatus.UNAUTHORIZED, authentication_required: HttpStatus.UNAUTHORIZED,
service_unavailable: HttpStatus.SERVICE_UNAVAILABLE,
internal_server_error: HttpStatus.INTERNAL_SERVER_ERROR, internal_server_error: HttpStatus.INTERNAL_SERVER_ERROR,
}; };
@@ -969,6 +971,25 @@ export const USER_FRIENDLY_ERRORS = {
}, },
// indexer 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: { search_provider_not_found: {
type: 'resource_not_found', type: 'resource_not_found',
message: 'Search provider not found.', message: 'Search provider not found.',
@@ -1123,6 +1123,38 @@ export class InvalidAppConfigInput extends UserFriendlyError {
super('invalid_input', 'invalid_app_config_input', message, args); 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 { export class SearchProviderNotFound extends UserFriendlyError {
constructor(message?: string) { constructor(message?: string) {
@@ -1320,6 +1352,10 @@ export enum ErrorNames {
MENTION_USER_ONESELF_DENIED, MENTION_USER_ONESELF_DENIED,
INVALID_APP_CONFIG, INVALID_APP_CONFIG,
INVALID_APP_CONFIG_INPUT, INVALID_APP_CONFIG_INPUT,
SEARCH_INDEX_NOT_READY,
SEARCH_PERMISSION_SYNCING,
SEARCH_PROVIDER_UNAVAILABLE,
SEARCH_INDEX_FAILED,
SEARCH_PROVIDER_NOT_FOUND, SEARCH_PROVIDER_NOT_FOUND,
INVALID_SEARCH_PROVIDER_REQUEST, INVALID_SEARCH_PROVIDER_REQUEST,
INVALID_INDEXER_INPUT, INVALID_INDEXER_INPUT,
@@ -1335,5 +1371,5 @@ registerEnumType(ErrorNames, {
export const ErrorDataUnionType = createUnionType({ export const ErrorDataUnionType = createUnionType({
name: 'ErrorDataUnion', name: 'ErrorDataUnion',
types: () => 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,
}); });
@@ -71,14 +71,6 @@ defineModuleConfig('job', {
schema, schema,
}, },
'queues.indexer': {
desc: 'The config for indexer job queue',
default: {
concurrency: 1,
},
schema,
},
'queues.notification': { 'queues.notification': {
desc: 'The config for notification job queue', desc: 'The config for notification job queue',
default: { default: {
@@ -27,7 +27,6 @@ export enum Queue {
NOTIFICATION = 'notification', NOTIFICATION = 'notification',
DOC = 'doc', DOC = 'doc',
COPILOT = 'copilot', COPILOT = 'copilot',
INDEXER = 'indexer',
CALENDAR = 'calendar', CALENDAR = 'calendar',
BACKENDRUNTIME = 'backendRuntime', BACKENDRUNTIME = 'backendRuntime',
INVITE_ABUSE = 'inviteAbuse', INVITE_ABUSE = 'inviteAbuse',
@@ -1,11 +1,7 @@
import { ServerRole } from '../../../env'; import { ServerRole } from '../../../env';
import { Queue, QUEUES } from './def'; import { Queue, QUEUES } from './def';
export const WORKER_QUEUES = [ export const WORKER_QUEUES = [Queue.DOC, Queue.BACKENDRUNTIME] as const;
Queue.DOC,
Queue.INDEXER,
Queue.BACKENDRUNTIME,
] as const;
export function queuesForRole(role: ServerRole | undefined): Queue[] { export function queuesForRole(role: ServerRole | undefined): Queue[] {
switch (role) { switch (role) {
@@ -63,6 +63,7 @@ export type KnownMetricScopes =
| 'storage' | 'storage'
| 'process' | 'process'
| 'permission' | 'permission'
| 'search'
| 'workspace'; | 'workspace';
const metricCreators: MetricCreators = { const metricCreators: MetricCreators = {
@@ -5,6 +5,7 @@ import {
BackendRuntimeEmbeddingProducer, BackendRuntimeEmbeddingProducer,
BackendRuntimeEmbeddingService, BackendRuntimeEmbeddingService,
BackendRuntimeHousekeepingJob, BackendRuntimeHousekeepingJob,
BackendRuntimeSearchJob,
} from './job'; } from './job';
import { import {
BACKEND_RUNTIME_CONFIG_PATHS, BACKEND_RUNTIME_CONFIG_PATHS,
@@ -33,7 +34,11 @@ export class BackendRuntimeProducerModule {}
@Module({ @Module({
imports: [BackendRuntimeModule], imports: [BackendRuntimeModule],
providers: [BackendRuntimeEmbeddingJob, BackendRuntimeHousekeepingJob], providers: [
BackendRuntimeEmbeddingJob,
BackendRuntimeHousekeepingJob,
BackendRuntimeSearchJob,
],
}) })
export class BackendRuntimeWorkerModule {} export class BackendRuntimeWorkerModule {}
@@ -42,6 +47,7 @@ export {
BackendRuntimeEmbeddingProducer, BackendRuntimeEmbeddingProducer,
BackendRuntimeEmbeddingService, BackendRuntimeEmbeddingService,
BackendRuntimeHousekeepingJob, BackendRuntimeHousekeepingJob,
BackendRuntimeSearchJob,
} from './job'; } from './job';
export { export {
BACKEND_RUNTIME_CONFIG_PATHS, BACKEND_RUNTIME_CONFIG_PATHS,
@@ -2,11 +2,13 @@ import { Injectable, Logger } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule'; import { Cron, CronExpression } from '@nestjs/schedule';
import { import {
ConfigFactory,
CopilotSelectedSourcesFailed, CopilotSelectedSourcesFailed,
CopilotSelectedSourcesLimitExceeded, CopilotSelectedSourcesLimitExceeded,
CopilotSelectedSourcesProcessing, CopilotSelectedSourcesProcessing,
CopilotSelectedSourcesUnavailable, CopilotSelectedSourcesUnavailable,
JobQueue, JobQueue,
metrics,
OnEvent, OnEvent,
OnJob, OnJob,
} from '../../base'; } from '../../base';
@@ -30,6 +32,9 @@ declare global {
'backendRuntime.reconcileDocumentEmbeddings': { 'backendRuntime.reconcileDocumentEmbeddings': {
workspaceId: string; workspaceId: string;
}; };
'backendRuntime.reconcileSearchProjection': {
limit?: number;
};
} }
} }
@@ -257,3 +262,98 @@ export class BackendRuntimeHousekeepingJob {
return total; 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;
}
}
}
@@ -385,27 +385,9 @@ export class BackendRuntimeProvider
); );
} }
async indexSearchDocument(workspaceId: string, docId: string) { async reconcileSearchProjection(limit = 100) {
await this.measured('indexSearchDocument', runtime => return await this.measured('reconcileSearchProjection', runtime =>
runtime.indexSearchDocument(workspaceId, docId) runtime.reconcileSearchProjection(limit)
);
}
async deleteSearchDocument(workspaceId: string, docId: string) {
await this.measured('deleteSearchDocument', runtime =>
runtime.deleteSearchDocument(workspaceId, docId)
);
}
async reconcileSearchWorkspace(workspaceId: string) {
await this.measured('reconcileSearchWorkspace', runtime =>
runtime.reconcileSearchWorkspace(workspaceId)
);
}
async deleteSearchWorkspace(workspaceId: string) {
await this.measured('deleteSearchWorkspace', runtime =>
runtime.deleteSearchWorkspace(workspaceId)
); );
} }
@@ -7,6 +7,7 @@ import { createModule } from '../../../__tests__/create-module';
import { Mockers } from '../../../__tests__/mocks'; import { Mockers } from '../../../__tests__/mocks';
import { InvalidAppConfigInput } from '../../../base'; import { InvalidAppConfigInput } from '../../../base';
import { Models } from '../../../models'; import { Models } from '../../../models';
import { SearchProviderType } from '../../../plugins/indexer/config';
import { ServerService } from '../service'; import { ServerService } from '../service';
const module = await createModule({ const module = await createModule({
@@ -40,6 +41,19 @@ test('should update config', async t => {
t.is(service.getConfig().server.externalUrl, newValue); t.is(service.getConfig().server.externalUrl, newValue);
}); });
test('should enable the selected indexer provider', async t => {
await service.updateConfig(user.id, [
{
module: 'indexer',
key: 'provider.type',
value: SearchProviderType.Embedded,
},
]);
t.true(service.getConfig().indexer.enabled);
t.is(service.getConfig().indexer.provider.type, SearchProviderType.Embedded);
});
test('should validate config before update', async t => { test('should validate config before update', async t => {
await t.throwsAsync( await t.throwsAsync(
service.updateConfig(user.id, [ service.updateConfig(user.id, [
@@ -76,15 +76,7 @@ export class ServerService implements OnApplicationBootstrap {
const providerType = updates.find( const providerType = updates.find(
update => update.module === 'indexer' && update.key === 'provider.type' update => update.module === 'indexer' && update.key === 'provider.type'
); );
if (providerType?.value === 'embedded') { if (providerType) {
updates = updates.filter(update => update !== providerType);
updates = [
...updates.filter(
update => !(update.module === 'indexer' && update.key === 'enabled')
),
{ module: 'indexer', key: 'enabled', value: false },
];
} else if (providerType) {
updates = [ updates = [
...updates.filter( ...updates.filter(
update => !(update.module === 'indexer' && update.key === 'enabled') update => !(update.module === 'indexer' && update.key === 'enabled')
@@ -1,6 +1,13 @@
import test from 'ava'; import test from 'ava';
import { InternalServerError } from '../../../base'; import {
InternalServerError,
SearchIndexFailed,
SearchIndexNotReady,
SearchPermissionSyncing,
SearchProviderUnavailable,
SpaceAccessDenied,
} from '../../../base';
import { DocRole } from '../../../models'; import { DocRole } from '../../../models';
import { docLegacyBoundary } from '../context'; import { docLegacyBoundary } from '../context';
import { PermissionContextLoader } from '../context-loader'; import { PermissionContextLoader } from '../context-loader';
@@ -23,9 +30,11 @@ function createLoader() {
docPolicies: 0, docPolicies: 0,
docGrants: 0, docGrants: 0,
}; };
const queries: string[] = [];
const db = { const db = {
$queryRaw: async (strings: TemplateStringsArray) => { $queryRaw: async (strings: TemplateStringsArray) => {
const sql = strings.join(''); const sql = strings.join('');
queries.push(sql);
if (sql.includes('FROM workspace_members')) { if (sql.includes('FROM workspace_members')) {
calls.members += 1; calls.members += 1;
return [{ role: 'owner', state: 'active' }]; return [{ role: 'owner', state: 'active' }];
@@ -82,6 +91,7 @@ function createLoader() {
}; };
return { return {
calls, calls,
queries,
loader: new PermissionContextLoader(db as never, createCls() as never), loader: new PermissionContextLoader(db as never, createCls() as never),
}; };
} }
@@ -196,7 +206,7 @@ test('PermissionService supports anonymous preview without doc read', async t =>
}); });
test('PermissionContextLoader reads only terminal permission tables', async t => { test('PermissionContextLoader reads only terminal permission tables', async t => {
const { loader } = createLoader(); const { loader, queries } = createLoader();
const input = await loader.load({ const input = await loader.load({
userId: 'u1', userId: 'u1',
workspaceId: 'w1', workspaceId: 'w1',
@@ -212,6 +222,9 @@ test('PermissionContextLoader reads only terminal permission tables', async t =>
t.is(input.docs?.[0]?.explicitUserRole, 'manager'); t.is(input.docs?.[0]?.explicitUserRole, 'manager');
t.is(input.docs?.[0]?.memberDefaultRole, 'manager'); t.is(input.docs?.[0]?.memberDefaultRole, 'manager');
t.is(input.docs?.[1]?.publicRole, 'external'); t.is(input.docs?.[1]?.publicRole, 'external');
t.false(
queries.some(query => /workspace_permission_|search_runtime_/i.test(query))
);
}); });
test('PermissionContextLoader treats missing quota state as unknown and stale', async t => { test('PermissionContextLoader treats missing quota state as unknown and stale', async t => {
@@ -261,3 +274,11 @@ test('PermissionService maps native validation errors to internal errors', t =>
t.true(error instanceof InternalServerError); t.true(error instanceof InternalServerError);
}); });
test('search and permission boundaries expose stable HTTP contracts', t => {
t.is(new SearchIndexNotReady({ spaceId: 'w1' }).status, 503);
t.is(new SearchPermissionSyncing().status, 503);
t.is(new SearchProviderUnavailable().status, 503);
t.is(new SearchIndexFailed({ diagnosticId: 'w1' }).status, 503);
t.is(new SpaceAccessDenied({ spaceId: 'w1' }).status, 403);
});
@@ -287,17 +287,6 @@ export class StorageRuntimeProvider
); );
} }
async ackDocumentCleanupEffect(
workspaceId: string,
docId: string,
cleanupVersion: string,
effect: 'search' | 'copilot'
) {
return await this.measured('ackDocumentCleanupEffect', rt =>
rt.ackDocumentCleanupEffect(workspaceId, docId, cleanupVersion, effect)
);
}
async planUnreferencedWorkspaceBlobs( async planUnreferencedWorkspaceBlobs(
workspaceId: string, workspaceId: string,
gracePeriodDays: number, gracePeriodDays: number,
@@ -13,7 +13,6 @@ interface Context {
planUnreferencedWorkspaceBlobs: Sinon.SinonStub; planUnreferencedWorkspaceBlobs: Sinon.SinonStub;
executeBlobCleanupCandidates: Sinon.SinonStub; executeBlobCleanupCandidates: Sinon.SinonStub;
executeDocumentCleanupCandidates: Sinon.SinonStub; executeDocumentCleanupCandidates: Sinon.SinonStub;
ackDocumentCleanupEffect: Sinon.SinonStub;
}; };
event: { event: {
emitAsync: Sinon.SinonStub; emitAsync: Sinon.SinonStub;
@@ -51,7 +50,6 @@ test.beforeEach(t => {
planUnreferencedWorkspaceBlobs: Sinon.stub(), planUnreferencedWorkspaceBlobs: Sinon.stub(),
executeBlobCleanupCandidates: Sinon.stub(), executeBlobCleanupCandidates: Sinon.stub(),
executeDocumentCleanupCandidates: Sinon.stub(), executeDocumentCleanupCandidates: Sinon.stub(),
ackDocumentCleanupEffect: Sinon.stub(),
}; };
t.context.event = { t.context.event = {
emitAsync: Sinon.stub().resolves(undefined), emitAsync: Sinon.stub().resolves(undefined),
@@ -316,7 +314,7 @@ test('document projection worker drains metadata incrementally after a document
); );
}); });
test('document cleanup dispatches stable search effects', async t => { test('document cleanup emits blob updates after comment object cleanup', async t => {
t.context.runtime.executeDocumentCleanupCandidates.resolves({ t.context.runtime.executeDocumentCleanupCandidates.resolves({
scannedCandidates: 1, scannedCandidates: 1,
serializationRetries: 0, serializationRetries: 0,
@@ -331,22 +329,13 @@ test('document cleanup dispatches stable search effects', async t => {
docId: 'doc-1', docId: 'doc-1',
cleanupVersion: 'version-1', cleanupVersion: 'version-1',
commentObjectsDone: true, commentObjectsDone: true,
searchDone: false,
}, },
], ],
}); });
await t.context.job.executeDocumentCleanupCandidates({}); await t.context.job.executeDocumentCleanupCandidates({});
t.true( t.false(t.context.queue.add.called);
t.context.queue.add.calledWith(
'indexer.reconcileDocumentCleanup',
Sinon.match({ docId: 'doc-1' }),
{
jobId: 'document-cleanup:search:workspace-1:doc-1:version-1',
}
)
);
t.true( t.true(
t.context.event.emitAsync.calledWith('workspace.blobs.updated', { t.context.event.emitAsync.calledWith('workspace.blobs.updated', {
workspaceId: 'workspace-1', workspaceId: 'workspace-1',
@@ -354,22 +343,6 @@ test('document cleanup dispatches stable search effects', async t => {
); );
}); });
test('document cleanup effect ack delegates to storage runtime', async t => {
await t.context.job.ackDocumentCleanupEffect({
workspaceId: 'workspace-1',
docId: 'doc-1',
cleanupVersion: 'version-1',
effect: 'search',
});
t.deepEqual(t.context.runtime.ackDocumentCleanupEffect.firstCall.args, [
'workspace-1',
'doc-1',
'version-1',
'search',
]);
});
test('blob cleanup execution sweep drains marked runs and continues by page', async t => { test('blob cleanup execution sweep drains marked runs and continues by page', async t => {
t.context.db.$queryRaw t.context.db.$queryRaw
.onFirstCall() .onFirstCall()
@@ -33,12 +33,6 @@ declare global {
gracePeriodDays?: number; gracePeriodDays?: number;
limit?: number; limit?: number;
}; };
'backendRuntime.ackDocumentCleanupEffect': {
workspaceId: string;
docId: string;
cleanupVersion: string;
effect: 'search' | 'copilot';
};
'backendRuntime.planUnreferencedWorkspaceBlobs': { 'backendRuntime.planUnreferencedWorkspaceBlobs': {
workspaceId: string; workspaceId: string;
gracePeriodDays?: number; gracePeriodDays?: number;
@@ -321,11 +315,6 @@ export class StorageBlobJob {
.counter('document_cleanup_execute_failure_total') .counter('document_cleanup_execute_failure_total')
.add(result.failed); .add(result.failed);
for (const effect of result.effects) { for (const effect of result.effects) {
if (!effect.searchDone) {
await this.queue.add('indexer.reconcileDocumentCleanup', effect, {
jobId: `document-cleanup:search:${effect.workspaceId}:${effect.docId}:${effect.cleanupVersion}`,
});
}
if (effect.commentObjectsDone) { if (effect.commentObjectsDone) {
await this.event.emitAsync('workspace.blobs.updated', { await this.event.emitAsync('workspace.blobs.updated', {
workspaceId: effect.workspaceId, workspaceId: effect.workspaceId,
@@ -336,21 +325,6 @@ export class StorageBlobJob {
return result; return result;
} }
@OnJob('backendRuntime.ackDocumentCleanupEffect')
async ackDocumentCleanupEffect({
workspaceId,
docId,
cleanupVersion,
effect,
}: Jobs['backendRuntime.ackDocumentCleanupEffect']) {
await this.rt.ackDocumentCleanupEffect(
workspaceId,
docId,
cleanupVersion,
effect
);
}
@OnJob('backendRuntime.planUnreferencedWorkspaceBlobs') @OnJob('backendRuntime.planUnreferencedWorkspaceBlobs')
async planUnreferencedWorkspaceBlobs({ async planUnreferencedWorkspaceBlobs({
workspaceId, workspaceId,
-2
View File
@@ -71,7 +71,6 @@ import serverNativeModule, {
type RuntimeVerificationTokenRecord, type RuntimeVerificationTokenRecord,
type RuntimeWorkspaceArtifact, type RuntimeWorkspaceArtifact,
type RuntimeWorkspaceInviteLinkRecord, type RuntimeWorkspaceInviteLinkRecord,
type RuntimeWorkspaceStatsDailyRecalibrationResult,
type SafeFetchRequest, type SafeFetchRequest,
type SafeFetchResponse, type SafeFetchResponse,
type StorageProviderCapabilities, type StorageProviderCapabilities,
@@ -173,7 +172,6 @@ export type {
RuntimeVerificationTokenRecord, RuntimeVerificationTokenRecord,
RuntimeWorkspaceArtifact, RuntimeWorkspaceArtifact,
RuntimeWorkspaceInviteLinkRecord, RuntimeWorkspaceInviteLinkRecord,
RuntimeWorkspaceStatsDailyRecalibrationResult,
SafeFetchRequest, SafeFetchRequest,
SafeFetchResponse, SafeFetchResponse,
StorageProviderCapabilities, StorageProviderCapabilities,
@@ -1,6 +1,6 @@
import { Inject, Injectable } from '@nestjs/common'; import { Inject, Injectable } from '@nestjs/common';
import { SearchProviderNotFound } from '../../../base'; import { SearchProviderUnavailable } from '../../../base';
import { PermissionAccess } from '../../../core/permission'; import { PermissionAccess } from '../../../core/permission';
import type { DocVisibility } from '../../../core/utils/blocksuite'; import type { DocVisibility } from '../../../core/utils/blocksuite';
import { type DocChunkSimilarity, Models } from '../../../models'; import { type DocChunkSimilarity, Models } from '../../../models';
@@ -101,7 +101,7 @@ export class DocumentRetrievalService {
docIds, docIds,
}) })
.catch(error => { .catch(error => {
if (error instanceof SearchProviderNotFound) return null; if (error instanceof SearchProviderUnavailable) return null;
throw error; throw error;
}), }),
this.context.canEmbedding this.context.canEmbedding
@@ -1,77 +0,0 @@
import test from 'ava';
import { createModule } from '../../../__tests__/create-module';
import { JobQueue } from '../../../base';
import { ConfigModule } from '../../../base/config';
import { IndexerEvent } from '../event';
import { IndexerModule } from '../index';
import { IndexerScheduler } from '../scheduler';
const module = await createModule({
imports: [
IndexerModule,
ConfigModule.override({
indexer: {
enabled: true,
},
}),
],
});
const indexerEvent = module.get(IndexerEvent);
const indexerScheduler = new IndexerScheduler(module.get(JobQueue));
test.after.always(async () => {
await module.close();
});
test('should index workspace when root snapshot is updated', async t => {
// @ts-expect-error ignore missing fields
await indexerEvent.indexWorkspace({
workspaceId: 'test-workspace',
docId: 'test-workspace',
});
const { payload } = await module.queue.waitFor('indexer.indexWorkspace');
t.is(payload.workspaceId, 'test-workspace');
});
test('should not index workspace when non-root snapshot is updated', async t => {
const count = module.queue.count('indexer.indexWorkspace');
// @ts-expect-error ignore missing fields
await indexerEvent.indexWorkspace({
workspaceId: 'test-workspace',
docId: 'child-doc',
});
t.is(module.queue.count('indexer.indexWorkspace'), count);
});
test('should reindex documents after document access changes', async t => {
await indexerEvent.reindexDocOnGrantChange({
workspaceId: 'test-workspace',
docId: 'test-doc',
});
const { payload } = await module.queue.waitFor('indexer.indexDoc');
t.deepEqual(payload, {
workspaceId: 'test-workspace',
docId: 'test-doc',
});
});
test('should delete workspace', async t => {
// @ts-expect-error ignore missing fields
await indexerEvent.deleteUserWorkspaces({
ownedWorkspaces: ['test-workspace'],
});
const { payload } = await module.queue.waitFor('indexer.deleteWorkspace');
t.is(payload.workspaceId, 'test-workspace');
});
test('should schedule auto index workspaces', async t => {
await indexerScheduler.autoIndexWorkspaces();
const { payload } = await module.queue.waitFor('indexer.autoIndexWorkspaces');
t.is(payload.lastIndexedWorkspaceSid, undefined);
});
@@ -1,215 +0,0 @@
import { randomUUID } from 'node:crypto';
import test from 'ava';
import Sinon from 'sinon';
import { createModule } from '../../../__tests__/create-module';
import { Mockers } from '../../../__tests__/mocks';
import { JOB_SIGNAL } from '../../../base';
import { ConfigModule } from '../../../base/config';
import { ServerConfigModule } from '../../../core/config';
import { DocReader } from '../../../core/doc';
import { Models } from '../../../models';
import { addDocToRootDoc } from '../../../native';
import { IndexerModule, IndexerService, IndexerWorkerModule } from '../index';
import { IndexerJob } from '../job';
const module = await createModule({
imports: [
IndexerModule,
IndexerWorkerModule,
ServerConfigModule,
ConfigModule.override({
indexer: {
enabled: true,
},
}),
],
providers: [IndexerService],
});
const indexerService = module.get(IndexerService);
const indexerJob = module.get(IndexerJob);
const models = module.get(Models);
const docReader = module.get(DocReader);
const user = await module.create(Mockers.User);
const workspace = await module.create(Mockers.Workspace, {
snapshot: true,
owner: user,
});
test.after.always(async () => {
await module.close();
});
test.afterEach.always(() => {
Sinon.restore();
});
test('should handle indexer.indexDoc job', async t => {
const spy = Sinon.spy(indexerService, 'indexDoc');
await indexerJob.indexDoc({
workspaceId: workspace.id,
docId: randomUUID(),
});
t.is(spy.callCount, 1);
});
test('should handle indexer.deleteDoc job', async t => {
const spy = Sinon.spy(indexerService, 'deleteDoc');
await indexerJob.deleteDoc({
workspaceId: workspace.id,
docId: randomUUID(),
});
t.is(spy.callCount, 1);
});
test('should handle indexer.indexWorkspace job', async t => {
const spy = Sinon.stub(indexerService, 'reconcileWorkspace').resolves();
await indexerJob.indexWorkspace({
workspaceId: workspace.id,
});
t.true(spy.calledOnceWith(workspace.id));
// workspace should be indexed
const ws = await models.workspace.get(workspace.id);
t.is(ws!.indexed, true);
});
test('document cleanup reconcile deletes missing search state before ack', async t => {
const deleteSpy = Sinon.spy(indexerService, 'deleteDoc');
const indexSpy = Sinon.spy(indexerService, 'indexDoc');
const cleanupWorkspace = await module.create(Mockers.Workspace, {
owner: user,
});
await module.create(Mockers.DocSnapshot, {
workspaceId: cleanupWorkspace.id,
docId: cleanupWorkspace.id,
user,
blob: addDocToRootDoc(Buffer.from([0, 0]), 'live-doc', 'Live'),
});
await indexerJob.reconcileDocumentCleanup({
workspaceId: cleanupWorkspace.id,
docId: 'missing-doc',
cleanupVersion: 'version-1',
});
t.true(deleteSpy.calledOnceWith(cleanupWorkspace.id, 'missing-doc'));
t.false(indexSpy.called);
const { payload } = await module.queue.waitFor(
'backendRuntime.ackDocumentCleanupEffect'
);
t.deepEqual(payload, {
workspaceId: cleanupWorkspace.id,
docId: 'missing-doc',
cleanupVersion: 'version-1',
effect: 'search',
});
});
test('document cleanup reconcile reindexes restored doc before ack', async t => {
const deleteSpy = Sinon.spy(indexerService, 'deleteDoc');
const indexSpy = Sinon.spy(indexerService, 'indexDoc');
const cleanupWorkspace = await module.create(Mockers.Workspace, {
owner: user,
});
await module.create(Mockers.DocSnapshot, {
workspaceId: cleanupWorkspace.id,
docId: cleanupWorkspace.id,
user,
blob: addDocToRootDoc(Buffer.from([0, 0]), 'restored-doc', 'Restored'),
});
await module.create(Mockers.DocSnapshot, {
workspaceId: cleanupWorkspace.id,
docId: 'restored-doc',
user,
});
const getDocSpy = Sinon.spy(docReader, 'getDoc');
await indexerJob.reconcileDocumentCleanup({
workspaceId: cleanupWorkspace.id,
docId: 'restored-doc',
cleanupVersion: 'version-2',
});
t.true(indexSpy.calledOnceWith(cleanupWorkspace.id, 'restored-doc'));
t.false(deleteSpy.called);
t.true(getDocSpy.calledWith(cleanupWorkspace.id, cleanupWorkspace.id));
t.true(getDocSpy.calledWith(cleanupWorkspace.id, 'restored-doc'));
const { payload } = await module.queue.waitFor(
'backendRuntime.ackDocumentCleanupEffect'
);
t.deepEqual(payload, {
workspaceId: cleanupWorkspace.id,
docId: 'restored-doc',
cleanupVersion: 'version-2',
effect: 'search',
});
});
test('should handle indexer.deleteWorkspace job', async t => {
const spy = Sinon.spy(indexerService, 'deleteWorkspace');
await indexerJob.deleteWorkspace({
workspaceId: workspace.id,
});
t.is(spy.callCount, 1);
});
test('should handle indexer.autoIndexWorkspaces job', async t => {
const workspace = await module.create(Mockers.Workspace, {
snapshot: true,
});
const result = await indexerJob.autoIndexWorkspaces({
lastIndexedWorkspaceSid: workspace.sid - 1,
});
t.is(result, JOB_SIGNAL.Repeat);
const { payload } = await module.queue.waitFor('indexer.indexWorkspace');
t.is(payload.workspaceId, workspace.id);
// no new auto index job
const count = module.queue.count('indexer.autoIndexWorkspaces');
await indexerJob.autoIndexWorkspaces({
lastIndexedWorkspaceSid: workspace.sid,
});
t.is(module.queue.count('indexer.autoIndexWorkspaces'), count);
});
test('should not index workspace if it is not updated in 180 days', async t => {
const workspace = await module.create(Mockers.Workspace);
await module.create(Mockers.DocSnapshot, {
user,
workspaceId: workspace.id,
docId: workspace.id,
updatedAt: new Date(Date.now() - 180 * 24 * 60 * 60 * 1000 - 1),
});
const count = module.queue.count('indexer.indexWorkspace');
await indexerJob.autoIndexWorkspaces({
lastIndexedWorkspaceSid: workspace.sid - 1,
});
t.is(module.queue.count('indexer.indexWorkspace'), count);
});
test('should not index workspace if snapshot not exists', async t => {
// not create snapshot
const workspace = await module.create(Mockers.Workspace);
const count = module.queue.count('indexer.indexWorkspace');
await indexerJob.autoIndexWorkspaces({
lastIndexedWorkspaceSid: workspace.sid - 1,
});
t.is(module.queue.count('indexer.indexWorkspace'), count);
});
@@ -4,13 +4,18 @@ import Sinon from 'sinon';
import { import {
InternalServerError, InternalServerError,
InvalidIndexerInput, InvalidIndexerInput,
SearchProviderNotFound, SearchIndexFailed,
SearchIndexNotReady,
SearchPermissionSyncing,
SearchProviderUnavailable,
SpaceAccessDenied, SpaceAccessDenied,
WorkspacePermissionNotFound,
} from '../../../base'; } from '../../../base';
import { ConfigFactory } from '../../../base/config';
import { BackendRuntimeProvider } from '../../../core/backend-runtime'; import { BackendRuntimeProvider } from '../../../core/backend-runtime';
import { ServerService } from '../../../core/config'; import { BackendRuntimeSearchJob } from '../../../core/backend-runtime/job';
import { ServerFeature, ServerService } from '../../../core/config';
import { Models } from '../../../models'; import { Models } from '../../../models';
import { IndexerResolver } from '../resolver';
import { IndexerService } from '../service'; import { IndexerService } from '../service';
import { SearchQueryType, SearchTable } from '../types'; import { SearchQueryType, SearchTable } from '../types';
@@ -18,13 +23,47 @@ test.afterEach.always(() => {
Sinon.restore(); Sinon.restore();
}); });
function enabledServer() {
return {
getConfig: Sinon.stub().returns({ indexer: { enabled: true } }),
enableFeature: Sinon.stub(),
disableFeature: Sinon.stub(),
};
}
test('reflects native search readiness in the Node feature flag', async t => { test('reflects native search readiness in the Node feature flag', async t => {
const runtime = { const runtime = {
searchStatus: Sinon.stub(), searchStatus: Sinon.stub(),
searchAuthorized: Sinon.stub().resolves({
ok: true,
value: { total: 0, nodes: [] },
}),
}; };
runtime.searchStatus.onFirstCall().resolves({ ready: true }); runtime.searchStatus.onFirstCall().resolves({ ready: true });
runtime.searchStatus.onSecondCall().resolves({ ready: false }); runtime.searchStatus.onSecondCall().resolves({ ready: false });
runtime.searchStatus.onThirdCall().resolves({ ready: true });
const server = enabledServer();
const service = new IndexerService(
runtime as unknown as BackendRuntimeProvider,
{} as Models,
server as unknown as ServerService
);
await service.onApplicationBootstrap();
await service.onConfigChanged({ updates: { indexer: {} } } as never);
await service.search('actor', 'workspace', {} as never);
t.is(server.enableFeature.callCount, 2);
t.true(server.disableFeature.calledOnce);
t.is(runtime.searchStatus.callCount, 3);
});
test('does not query native search when the indexer is disabled', async t => {
const runtime = {
searchStatus: Sinon.stub(),
};
const server = { const server = {
getConfig: Sinon.stub().returns({ indexer: { enabled: false } }),
enableFeature: Sinon.stub(), enableFeature: Sinon.stub(),
disableFeature: Sinon.stub(), disableFeature: Sinon.stub(),
}; };
@@ -35,21 +74,50 @@ test('reflects native search readiness in the Node feature flag', async t => {
); );
await service.onApplicationBootstrap(); await service.onApplicationBootstrap();
await service.onConfigChanged({ updates: { indexer: {} } } as never);
t.true(server.enableFeature.calledOnce); t.false(runtime.searchStatus.called);
t.true(server.disableFeature.calledOnce); t.true(server.disableFeature.calledOnceWith(ServerFeature.Indexer));
t.is(runtime.searchStatus.callCount, 2); });
test('does not schedule or run native search reconciliation when disabled', async t => {
const runtime = {
reconcileSearchProjection: Sinon.stub(),
searchStatus: Sinon.stub(),
};
const queue = { add: Sinon.stub() };
const config = {
config: { indexer: { enabled: false } },
} as unknown as ConfigFactory;
const job = new BackendRuntimeSearchJob(
runtime as unknown as BackendRuntimeProvider,
queue as never,
config
);
await job.scheduleReconciliation();
t.is(queue.add.callCount, 0);
t.is(await job.reconcileProjection({ limit: 100 }), 0);
t.false(runtime.reconcileSearchProjection.called);
t.false(runtime.searchStatus.called);
config.config.indexer.enabled = true;
await job.scheduleReconciliation();
t.deepEqual(queue.add.firstCall.args[2], {
jobId: 'backend-runtime-search-reconciliation',
removeOnFail: true,
});
}); });
test('maps native search results and typed errors at the Node boundary', async t => { test('maps native search results and typed errors at the Node boundary', async t => {
const runtime = { const runtime = {
searchStatus: Sinon.stub().resolves({ ready: true }),
searchAuthorized: Sinon.stub(), searchAuthorized: Sinon.stub(),
aggregateAuthorized: Sinon.stub(),
}; };
const service = new IndexerService( const service = new IndexerService(
runtime as unknown as BackendRuntimeProvider, runtime as unknown as BackendRuntimeProvider,
{} as Models, {} as Models,
{} as ServerService enabledServer() as unknown as ServerService
); );
const input = { const input = {
table: SearchTable.block, table: SearchTable.block,
@@ -59,7 +127,7 @@ test('maps native search results and typed errors at the Node boundary', async t
runtime.searchAuthorized.resolves({ runtime.searchAuthorized.resolves({
ok: true, ok: true,
value: { value: {
total: 1, total: 99,
nodes: [ nodes: [
{ {
id: 'node', id: 'node',
@@ -85,12 +153,41 @@ test('maps native search results and typed errors at the Node boundary', async t
markdownPreview: ['<b>hello</b>'], markdownPreview: ['<b>hello</b>'],
}); });
const resolver = new IndexerResolver(service, {
user: Sinon.stub().returns({
workspace: Sinon.stub().returns({ assert: Sinon.stub().resolves() }),
}),
} as never);
const searchResult = await resolver.search(
{ id: 'actor' } as never,
{ id: 'workspace' } as never,
input
);
t.is(searchResult.pagination.count, 1);
runtime.aggregateAuthorized.resolves({
ok: true,
value: {
total: 99,
hasMore: true,
buckets: [{ key: 'doc', count: 1, hits: { nodes: [] } }],
},
});
const aggregateResult = await resolver.aggregate(
{ id: 'actor' } as never,
{ id: 'workspace' } as never,
{} as never
);
t.is(aggregateResult.pagination.count, 1);
for (const [errorCode, expected] of [ for (const [errorCode, expected] of [
['workspace_denied', SpaceAccessDenied], ['workspace_denied', SpaceAccessDenied],
['invalid_request', InvalidIndexerInput], ['invalid_request', InvalidIndexerInput],
['unsupported_query', InvalidIndexerInput], ['unsupported_query', InvalidIndexerInput],
['provider_unavailable', SearchProviderNotFound], ['provider_unavailable', SearchProviderUnavailable],
['permission_unavailable', WorkspacePermissionNotFound], ['index_not_ready', SearchIndexNotReady],
['permission_syncing', SearchPermissionSyncing],
['index_failed', SearchIndexFailed],
['unexpected', InternalServerError], ['unexpected', InternalServerError],
] as const) { ] as const) {
runtime.searchAuthorized.resolves({ ok: false, errorCode }); runtime.searchAuthorized.resolves({ ok: false, errorCode });
@@ -102,7 +199,29 @@ test('maps native search results and typed errors at the Node boundary', async t
}); });
test('searchDocs keeps filtering and enrichment in Node', async t => { test('searchDocs keeps filtering and enrichment in Node', async t => {
const blockNode = {
id: 'block',
score: 1,
fields: {
workspace_id: ['workspace'],
doc_id: ['doc'],
block_id: ['block'],
unit_id: ['unit'],
projection_version: [1],
source_hash: ['hash'],
visibility: ['visible'],
source_block_id: ['source-block'],
flavour: ['affine:paragraph'],
content: ['body'],
created_at: [2_000],
updated_at: [3_000],
created_by_user_id: ['creator'],
updated_by_user_id: ['updater'],
},
highlights: { content: ['<b>body</b>'] },
};
const runtime = { const runtime = {
searchStatus: Sinon.stub().resolves({ ready: true }),
aggregateAuthorized: Sinon.stub().resolves({ aggregateAuthorized: Sinon.stub().resolves({
ok: true, ok: true,
value: { value: {
@@ -113,34 +232,19 @@ test('searchDocs keeps filtering and enrichment in Node', async t => {
key: 'doc', key: 'doc',
count: 1, count: 1,
hits: { hits: {
nodes: [ nodes: [blockNode],
{
id: 'block',
score: 1,
fields: {
workspace_id: ['workspace'],
doc_id: ['doc'],
block_id: ['block'],
unit_id: ['unit'],
projection_version: [1],
source_hash: ['hash'],
visibility: ['visible'],
source_block_id: ['source-block'],
flavour: ['affine:paragraph'],
content: ['body'],
created_at: [2_000],
updated_at: [3_000],
created_by_user_id: ['creator'],
updated_by_user_id: ['updater'],
},
highlights: { content: ['<b>body</b>'] },
},
],
}, },
}, },
], ],
}, },
}), }),
searchAuthorized: Sinon.stub().resolves({
ok: true,
value: {
total: 2,
nodes: [blockNode, { ...blockNode, id: 'duplicate-block' }],
},
}),
}; };
const creator = { id: 'creator', name: 'Creator' }; const creator = { id: 'creator', name: 'Creator' };
const updater = { id: 'updater', name: 'Updater' }; const updater = { id: 'updater', name: 'Updater' };
@@ -162,9 +266,16 @@ test('searchDocs keeps filtering and enrichment in Node', async t => {
const service = new IndexerService( const service = new IndexerService(
runtime as unknown as BackendRuntimeProvider, runtime as unknown as BackendRuntimeProvider,
models as unknown as Models, models as unknown as Models,
{} as ServerService enabledServer() as unknown as ServerService
); );
for (const limit of [0, -1]) {
const error = await t.throwsAsync(
service.searchDocsByKeyword('actor', 'workspace', 'body', { limit })
);
t.true(error instanceof InvalidIndexerInput);
}
t.false(runtime.aggregateAuthorized.called);
t.deepEqual( t.deepEqual(
await service.searchDocsByKeyword('actor', 'workspace', 'body', { await service.searchDocsByKeyword('actor', 'workspace', 'body', {
docIds: [], docIds: [],
@@ -194,4 +305,22 @@ test('searchDocs keeps filtering and enrichment in Node', async t => {
} }
) )
); );
runtime.aggregateAuthorized.resolves({
ok: false,
errorCode: 'unsupported_query',
});
const basicDocs = await service.searchDocsByKeyword(
'actor',
'workspace',
'body',
{ limit: 5, docIds: ['doc'] }
);
t.is(basicDocs.length, 1);
t.is(basicDocs[0].docId, 'doc');
const basicRequest = runtime.searchAuthorized.firstCall.args[2];
t.true(basicRequest.options.fields.includes('docId'));
t.is(basicRequest.options.pagination?.limit, 20);
t.true(JSON.stringify(basicRequest.query).includes('doc'));
t.true(JSON.stringify(basicRequest.query).includes('workspace'));
}); });
@@ -4,8 +4,8 @@ import { defineModuleConfig } from '../../base';
export enum SearchProviderType { export enum SearchProviderType {
Embedded = 'embedded', Embedded = 'embedded',
Manticoresearch = 'manticoresearch',
Elasticsearch = 'elasticsearch', Elasticsearch = 'elasticsearch',
ManticoreSearch = 'manticoresearch',
} }
const SearchProviderTypeSchema = z.nativeEnum(SearchProviderType); const SearchProviderTypeSchema = z.nativeEnum(SearchProviderType);
@@ -21,9 +21,6 @@ declare global {
username: string; username: string;
password: string; password: string;
}; };
autoIndex: {
batchSize: number;
};
}; };
} }
} }
@@ -32,18 +29,15 @@ defineModuleConfig('indexer', {
enabled: { enabled: {
desc: 'Enable indexer plugin', desc: 'Enable indexer plugin',
default: false, default: false,
env: ['AFFINE_INDEXER_ENABLED', 'boolean'],
}, },
'provider.type': { 'provider.type': {
desc: 'Indexer search provider. Self-hosted uses the embedded provider by default; remote providers require an endpoint.', desc: 'Indexer search provider. Self-hosted uses the embedded provider by default; remote providers require an endpoint.',
default: SearchProviderType.Embedded, default: SearchProviderType.Embedded,
shape: SearchProviderTypeSchema, shape: SearchProviderTypeSchema,
env: ['AFFINE_INDEXER_SEARCH_PROVIDER', 'string'],
}, },
'provider.endpoint': { 'provider.endpoint': {
desc: 'Remote indexer endpoint. Not used by the embedded provider.', desc: 'Remote indexer endpoint. Not used by the embedded provider.',
default: '', default: '',
env: ['AFFINE_INDEXER_SEARCH_ENDPOINT', 'string'],
validate: val => { validate: val => {
// allow to be nullable and empty string // allow to be nullable and empty string
if (!val) { if (!val) {
@@ -54,25 +48,17 @@ defineModuleConfig('indexer', {
}, },
}, },
'provider.apiKey': { 'provider.apiKey': {
desc: 'Indexer search service api key. Optional for elasticsearch', desc: 'Indexer search service api key. Optional for remote providers',
link: 'https://www.elastic.co/guide/server/current/api-key.html', link: 'https://www.elastic.co/guide/server/current/api-key.html',
default: '', default: '',
env: ['AFFINE_INDEXER_SEARCH_API_KEY', 'string'],
}, },
'provider.username': { 'provider.username': {
desc: 'Indexer search service auth username, if not set, basic auth will be disabled. Optional for elasticsearch', desc: 'Indexer search service auth username, if not set, basic auth will be disabled. Optional for remote providers',
link: 'https://www.elastic.co/guide/en/elasticsearch/reference/current/http-clients.html', link: 'https://www.elastic.co/guide/en/elasticsearch/reference/current/http-clients.html',
default: '', default: '',
env: ['AFFINE_INDEXER_SEARCH_USERNAME', 'string'],
}, },
'provider.password': { 'provider.password': {
desc: 'Indexer search service auth password, if not set, basic auth will be disabled. Optional for elasticsearch', desc: 'Indexer search service auth password, if not set, basic auth will be disabled. Optional for remote providers',
default: '', default: '',
env: ['AFFINE_INDEXER_SEARCH_PASSWORD', 'string'],
},
'autoIndex.batchSize': {
desc: 'Number of workspaces automatically indexed per batch',
default: 10,
shape: z.number().int().positive().max(1000),
}, },
}); });
@@ -1,84 +0,0 @@
import { Injectable } from '@nestjs/common';
import { JobQueue, OnEvent } from '../../base';
@Injectable()
export class IndexerEvent {
constructor(private readonly queue: JobQueue) {}
@OnEvent('doc.grants.changed')
async reindexDocOnGrantChange({
workspaceId,
docId,
}: Events['doc.grants.changed']) {
await this.indexDoc({ workspaceId, docId });
}
@OnEvent('doc.owner.changed')
async reindexDocOnOwnerChange({
workspaceId,
docId,
}: Events['doc.owner.changed']) {
await this.indexDoc({ workspaceId, docId });
}
@OnEvent('doc.default_role.changed')
async reindexDocOnDefaultRoleChange({
workspaceId,
docId,
}: Events['doc.default_role.changed']) {
await this.indexDoc({ workspaceId, docId });
}
@OnEvent('doc.public_state.changed')
async reindexDocOnPublicStateChange({
workspaceId,
docId,
}: Events['doc.public_state.changed']) {
await this.indexDoc({ workspaceId, docId });
}
@OnEvent('doc.updated')
async indexDoc({ workspaceId, docId }: Events['doc.updated']) {
await this.queue.add(
'indexer.indexDoc',
{
workspaceId,
docId,
},
{
jobId: `indexDoc/${workspaceId}/${docId}`,
priority: 100,
}
);
}
@OnEvent('doc.snapshot.updated')
async indexWorkspace({ workspaceId, docId }: Events['doc.snapshot.updated']) {
if (workspaceId !== docId) {
return;
}
await this.queue.add(
'indexer.indexWorkspace',
{ workspaceId },
{ jobId: `indexWorkspace/${workspaceId}`, priority: 100 }
);
}
@OnEvent('user.deleted')
async deleteUserWorkspaces(payload: Events['user.deleted']) {
for (const workspace of payload.ownedWorkspaces) {
await this.queue.add(
'indexer.deleteWorkspace',
{
workspaceId: workspace,
},
{
jobId: `deleteWorkspace/${workspace}`,
priority: 0,
}
);
}
}
}
@@ -5,10 +5,7 @@ import { Module } from '@nestjs/common';
import { ServerConfigModule } from '../../core/config'; import { ServerConfigModule } from '../../core/config';
import { DocStorageModule } from '../../core/doc'; import { DocStorageModule } from '../../core/doc';
import { PermissionModule } from '../../core/permission'; import { PermissionModule } from '../../core/permission';
import { IndexerEvent } from './event';
import { IndexerJob } from './job';
import { IndexerResolver } from './resolver'; import { IndexerResolver } from './resolver';
import { IndexerScheduler } from './scheduler';
import { IndexerService } from './service'; import { IndexerService } from './service';
const INDEXER_SHARED_IMPORTS = [ const INDEXER_SHARED_IMPORTS = [
@@ -24,21 +21,9 @@ const INDEXER_SHARED_IMPORTS = [
}) })
export class IndexerServiceModule {} export class IndexerServiceModule {}
@Module({
imports: [IndexerServiceModule],
providers: [IndexerEvent],
})
export class IndexerProducerModule {}
@Module({ @Module({
imports: [IndexerServiceModule, DocStorageModule, PermissionModule], imports: [IndexerServiceModule, DocStorageModule, PermissionModule],
providers: [IndexerJob, IndexerScheduler], providers: [IndexerResolver],
})
export class IndexerWorkerModule {}
@Module({
imports: [IndexerServiceModule, DocStorageModule, PermissionModule],
providers: [IndexerResolver, IndexerEvent],
exports: [IndexerServiceModule], exports: [IndexerServiceModule],
}) })
export class IndexerModule {} export class IndexerModule {}
@@ -1,192 +0,0 @@
import { Injectable, Logger } from '@nestjs/common';
import { Config, JOB_SIGNAL, JobQueue, OnJob } from '../../base';
import { DocReader } from '../../core/doc';
import { readAllDocIdsFromWorkspaceSnapshot } from '../../core/utils/blocksuite';
import { Models } from '../../models';
import { IndexerService } from './service';
declare global {
interface Jobs {
'indexer.indexDoc': {
workspaceId: string;
docId: string;
};
'indexer.deleteDoc': {
workspaceId: string;
docId: string;
};
'indexer.indexWorkspace': {
workspaceId: string;
};
'indexer.deleteWorkspace': {
workspaceId: string;
};
'indexer.autoIndexWorkspaces': {
lastIndexedWorkspaceSid?: number;
};
'indexer.reconcileDocumentCleanup': {
workspaceId: string;
docId: string;
cleanupVersion: string;
};
}
}
@Injectable()
export class IndexerJob {
private readonly logger = new Logger(IndexerJob.name);
constructor(
private readonly models: Models,
private readonly service: IndexerService,
private readonly queue: JobQueue,
private readonly config: Config,
private readonly doc: DocReader
) {}
@OnJob('indexer.indexDoc')
async indexDoc({ workspaceId, docId }: Jobs['indexer.indexDoc']) {
// delete the 'indexer.deleteDoc' job from the queue
await this.queue.remove(
`deleteDoc/${workspaceId}/${docId}`,
'indexer.deleteDoc'
);
await this.service.indexDoc(workspaceId, docId);
await this.enqueueBlobRefProjection(workspaceId, docId);
}
@OnJob('indexer.deleteDoc')
async deleteDoc({ workspaceId, docId }: Jobs['indexer.deleteDoc']) {
// delete the 'indexer.updateDoc' job from the queue
await this.queue.remove(
`indexDoc/${workspaceId}/${docId}`,
'indexer.indexDoc'
);
await this.service.deleteDoc(workspaceId, docId);
}
@OnJob('indexer.reconcileDocumentCleanup')
async reconcileDocumentCleanup({
workspaceId,
docId,
cleanupVersion,
}: Jobs['indexer.reconcileDocumentCleanup']) {
const root = await this.doc.getDoc(workspaceId, workspaceId);
if (!root) {
throw new Error(`workspace root ${workspaceId} not found`);
}
const live = readAllDocIdsFromWorkspaceSnapshot(root.bin, true).includes(
docId
);
if (live) {
if (!(await this.doc.getDoc(workspaceId, docId))) {
throw new Error(`restored document ${workspaceId}/${docId} not found`);
}
await this.service.indexDoc(workspaceId, docId);
await this.enqueueBlobRefProjection(workspaceId, docId);
} else {
await this.service.deleteDoc(workspaceId, docId);
}
await this.queue.add('backendRuntime.ackDocumentCleanupEffect', {
workspaceId,
docId,
cleanupVersion,
effect: 'search',
});
}
@OnJob('indexer.indexWorkspace')
async indexWorkspace({ workspaceId }: Jobs['indexer.indexWorkspace']) {
await this.queue.remove(workspaceId, 'indexer.deleteWorkspace');
const workspace = await this.models.workspace.get(workspaceId);
if (!workspace) {
this.logger.warn(`workspace ${workspaceId} not found`);
return;
}
await this.service.reconcileWorkspace(workspaceId);
if (!workspace.indexed) {
await this.models.workspace.update(workspaceId, {
indexed: true,
});
}
this.logger.log(`reconciled workspace ${workspaceId}`);
}
@OnJob('indexer.deleteWorkspace')
async deleteWorkspace({ workspaceId }: Jobs['indexer.deleteWorkspace']) {
await this.queue.remove(
`indexWorkspace/${workspaceId}`,
'indexer.indexWorkspace'
);
await this.service.deleteWorkspace(workspaceId);
}
@OnJob('indexer.autoIndexWorkspaces')
async autoIndexWorkspaces(payload: Jobs['indexer.autoIndexWorkspaces']) {
const startSid = payload.lastIndexedWorkspaceSid ?? 0;
const workspaces = await this.models.workspace.list(
{ sid: { gt: startSid } },
{ id: true, indexed: true, sid: true },
this.config.indexer.autoIndex.batchSize
);
if (workspaces.length === 0) {
// Keep the current sid value when repeating
return JOB_SIGNAL.Repeat;
}
let addedCount = 0;
for (const workspace of workspaces) {
const snapshotMeta = await this.models.doc.getSnapshot(
workspace.id,
workspace.id,
{
select: {
updatedAt: true,
},
}
);
// ignore 180 days not updated workspaces
if (
!snapshotMeta?.updatedAt ||
Date.now() - snapshotMeta.updatedAt.getTime() >
180 * 24 * 60 * 60 * 1000
) {
continue;
}
await this.queue.add(
'indexer.indexWorkspace',
{ workspaceId: workspace.id },
{ jobId: `indexWorkspace/${workspace.id}` }
);
addedCount++;
}
const nextSid = workspaces[workspaces.length - 1].sid;
this.logger.log(
`Auto added ${addedCount} workspaces to queue, lastIndexedWorkspaceSid: ${startSid} -> ${nextSid}`
);
// update the lastIndexedWorkspaceSid in the payload and repeat the job after 30 seconds
payload.lastIndexedWorkspaceSid = nextSid;
return JOB_SIGNAL.Repeat;
}
private async enqueueBlobRefProjection(workspaceId: string, docId: string) {
const snapshot = await this.models.doc.getSnapshot(workspaceId, docId, {
select: { updatedAt: true },
});
if (!snapshot) {
return;
}
const sourceRevision = snapshot.updatedAt.getTime();
await this.queue.add(
'backendRuntime.projectWorkspaceDocBlobRefs',
{ workspaceId, docId, sourceRevision },
{
jobId: `doc:blob-ref-projection:${workspaceId}:${docId}:${sourceRevision}`,
priority: 100,
}
);
}
}
@@ -35,8 +35,8 @@ export class IndexerResolver {
return { return {
nodes: result.nodes, nodes: result.nodes,
pagination: { pagination: {
count: result.total, count: result.nodes.length,
hasMore: result.nodes.length > 0, hasMore: Boolean(result.nextCursor),
nextCursor: result.nextCursor, nextCursor: result.nextCursor,
}, },
}; };
@@ -56,7 +56,7 @@ export class IndexerResolver {
return { return {
buckets: result.buckets, buckets: result.buckets,
pagination: { pagination: {
count: result.total, count: result.buckets.length,
hasMore: result.hasMore, hasMore: result.hasMore,
}, },
}; };
@@ -3,6 +3,8 @@ import { camelCase, mapKeys } from 'lodash-es';
import { import {
AggregateInput, AggregateInput,
SearchDoc, SearchDoc,
SearchInput,
SearchQuery,
SearchQueryOccur, SearchQueryOccur,
SearchQueryType, SearchQueryType,
SearchTable, SearchTable,
@@ -147,6 +149,71 @@ export function buildSearchDocsInput(
}; };
} }
export function buildBasicSearchDocsInput(
workspaceId: string,
keyword: string,
options?: { limit?: number; docIds?: string[] }
): SearchInput {
const aggregate = buildSearchDocsInput(workspaceId, keyword, options);
const limit = options?.limit ?? 20;
const queries: SearchQuery[] = [
{
type: SearchQueryType.match as const,
field: 'workspaceId',
match: workspaceId,
},
{
type: SearchQueryType.match as const,
field: 'content',
match: keyword,
},
];
if (options?.docIds) {
queries.push({
type: SearchQueryType.boolean,
occur: SearchQueryOccur.should,
queries: options.docIds.map(docId => ({
type: SearchQueryType.match as const,
field: 'docId',
match: docId,
})),
});
}
return {
table: aggregate.table,
query: {
type: SearchQueryType.boolean,
occur: SearchQueryOccur.must,
queries,
},
options: {
fields: ['docId', ...aggregate.options.hits.fields],
highlights: aggregate.options.hits.highlights,
pagination: { limit: Math.min(Math.max(limit, 1) * 4, 10_000) },
},
};
}
export function collectBasicSearchDocs(
result: { nodes: SearchNodeWithMeta[] },
workspaceId: string,
limit: number
) {
const seen = new Set<string>();
const buckets: AggregateResult['buckets'] = [];
for (const node of result.nodes) {
const docId = node._source.docId;
if (seen.has(docId)) continue;
seen.add(docId);
buckets.push({ key: docId, count: 1, hits: { nodes: [node] } });
if (buckets.length === limit) break;
}
return collectSearchDocs(
{ total: buckets.length, hasMore: false, buckets },
workspaceId
);
}
export function collectSearchDocs( export function collectSearchDocs(
result: AggregateResult, result: AggregateResult,
workspaceId: string workspaceId: string
@@ -1,22 +0,0 @@
import { Injectable } from '@nestjs/common';
import { Cron, CronExpression } from '@nestjs/schedule';
import { JobQueue } from '../../base';
@Injectable()
export class IndexerScheduler {
constructor(private readonly queue: JobQueue) {}
@Cron(CronExpression.EVERY_30_SECONDS)
async autoIndexWorkspaces() {
await this.queue.add(
'indexer.autoIndexWorkspaces',
{},
{
// make sure only one job is running at a time
delay: 30 * 1000,
jobId: 'autoIndexWorkspaces',
}
);
}
}
@@ -4,16 +4,20 @@ import {
InternalServerError, InternalServerError,
InvalidIndexerInput, InvalidIndexerInput,
OnEvent, OnEvent,
SearchProviderNotFound, SearchIndexFailed,
SearchIndexNotReady,
SearchPermissionSyncing,
SearchProviderUnavailable,
SpaceAccessDenied, SpaceAccessDenied,
WorkspacePermissionNotFound,
} from '../../base'; } from '../../base';
import { BackendRuntimeProvider } from '../../core/backend-runtime'; import { BackendRuntimeProvider } from '../../core/backend-runtime';
import { ServerFeature, ServerService } from '../../core/config'; import { ServerFeature, ServerService } from '../../core/config';
import { Models } from '../../models'; import { Models } from '../../models';
import { import {
type AggregateResult, type AggregateResult,
buildBasicSearchDocsInput,
buildSearchDocsInput, buildSearchDocsInput,
collectBasicSearchDocs,
collectSearchDocs, collectSearchDocs,
formatSearchNodes, formatSearchNodes,
type SearchNode, type SearchNode,
@@ -50,12 +54,17 @@ export class IndexerService implements OnApplicationBootstrap {
} }
private async syncFeature() { private async syncFeature() {
if (!this.server.getConfig().indexer.enabled) {
this.server.disableFeature(ServerFeature.Indexer);
return;
}
const status = (await this.runtime.searchStatus()) as { ready: boolean }; const status = (await this.runtime.searchStatus()) as { ready: boolean };
if (status.ready) this.server.enableFeature(ServerFeature.Indexer); if (status.ready) this.server.enableFeature(ServerFeature.Indexer);
else this.server.disableFeature(ServerFeature.Indexer); else this.server.disableFeature(ServerFeature.Indexer);
} }
async search(actorUserId: string, workspaceId: string, input: SearchInput) { async search(actorUserId: string, workspaceId: string, input: SearchInput) {
await this.syncFeature();
const result = this.unwrap<SearchResult>( const result = this.unwrap<SearchResult>(
await this.runtime.searchAuthorized(actorUserId, workspaceId, input), await this.runtime.searchAuthorized(actorUserId, workspaceId, input),
workspaceId workspaceId
@@ -68,10 +77,15 @@ export class IndexerService implements OnApplicationBootstrap {
workspaceId: string, workspaceId: string,
input: AggregateInput input: AggregateInput
) { ) {
await this.syncFeature();
const result = this.unwrap<AggregateResult>( const result = this.unwrap<AggregateResult>(
await this.runtime.aggregateAuthorized(actorUserId, workspaceId, input), await this.runtime.aggregateAuthorized(actorUserId, workspaceId, input),
workspaceId workspaceId
); );
return this.formatAggregate(result);
}
private formatAggregate(result: AggregateResult) {
return { return {
...result, ...result,
buckets: result.buckets.map(bucket => ({ buckets: result.buckets.map(bucket => ({
@@ -84,38 +98,42 @@ export class IndexerService implements OnApplicationBootstrap {
}; };
} }
async indexDoc(workspaceId: string, docId: string) {
await this.runtime.indexSearchDocument(workspaceId, docId);
}
async deleteDoc(workspaceId: string, docId: string) {
await this.runtime.deleteSearchDocument(workspaceId, docId);
}
async reconcileWorkspace(workspaceId: string) {
await this.runtime.reconcileSearchWorkspace(workspaceId);
}
async deleteWorkspace(workspaceId: string) {
await this.runtime.deleteSearchWorkspace(workspaceId);
}
async searchDocsByKeyword( async searchDocsByKeyword(
actorUserId: string, actorUserId: string,
workspaceId: string, workspaceId: string,
keyword: string, keyword: string,
options?: { limit?: number; docIds?: string[] } options?: { limit?: number; docIds?: string[] }
): Promise<SearchDoc[]> { ): Promise<SearchDoc[]> {
await this.syncFeature();
if (options?.limit !== undefined && options.limit <= 0) {
throw new InvalidIndexerInput({
reason: 'searchDocs limit must be positive',
});
}
if (options?.docIds?.length === 0) return []; if (options?.docIds?.length === 0) return [];
const result = await this.aggregate( const aggregateOutput = await this.runtime.aggregateAuthorized(
actorUserId, actorUserId,
workspaceId, workspaceId,
buildSearchDocsInput(workspaceId, keyword, options) buildSearchDocsInput(workspaceId, keyword, options)
); );
const { docs, missingTitles, userIds } = collectSearchDocs( const collected =
result, !aggregateOutput.ok && aggregateOutput.errorCode === 'unsupported_query'
workspaceId ? collectBasicSearchDocs(
); await this.search(
actorUserId,
workspaceId,
buildBasicSearchDocsInput(workspaceId, keyword, options)
),
workspaceId,
options?.limit ?? 20
)
: collectSearchDocs(
this.formatAggregate(
this.unwrap<AggregateResult>(aggregateOutput, workspaceId)
),
workspaceId
);
const { docs, missingTitles, userIds } = collected;
if (missingTitles.length > 0) { if (missingTitles.length > 0) {
const metas = await this.models.doc.findMetas(missingTitles, { const metas = await this.models.doc.findMetas(missingTitles, {
select: { title: true }, select: { title: true },
@@ -146,9 +164,15 @@ export class IndexerService implements OnApplicationBootstrap {
case 'unsupported_query': case 'unsupported_query':
throw new InvalidIndexerInput({ reason: output.errorCode }); throw new InvalidIndexerInput({ reason: output.errorCode });
case 'provider_unavailable': case 'provider_unavailable':
throw new SearchProviderNotFound(); throw new SearchProviderUnavailable();
case 'permission_unavailable': case 'index_not_ready':
throw new WorkspacePermissionNotFound({ spaceId: workspaceId }); throw new SearchIndexNotReady({ spaceId: workspaceId });
case 'permission_syncing':
throw new SearchPermissionSyncing();
case 'index_failed':
throw new SearchIndexFailed({
diagnosticId: 'search_workspace_reconcile_failed',
});
default: default:
throw new InternalServerError(); throw new InternalServerError();
} }
@@ -296,13 +296,23 @@ export class SearchNodeObjectType {
@ObjectType() @ObjectType()
export class SearchResultPagination { export class SearchResultPagination {
@Field(() => Int) @Field(() => Int, {
description:
'Number of results returned in this response, not a global total',
})
count!: number; count!: number;
@Field(() => Boolean) @Field(() => Boolean, {
description:
'Whether the provider has more candidates; remaining visible results are not guaranteed',
})
hasMore!: boolean; hasMore!: boolean;
@Field(() => String, { nullable: true }) @Field(() => String, {
nullable: true,
description:
'Opaque provider candidate cursor; it does not guarantee complete visible-result pagination',
})
nextCursor?: string; nextCursor?: string;
} }
@@ -326,7 +336,9 @@ export class AggregateBucketObjectType {
@Field(() => String) @Field(() => String)
key!: string; key!: string;
@Field(() => Int) @Field(() => Int, {
description: 'Number of returned sample hits in this bucket',
})
count!: number; count!: number;
@Field(() => AggregateBucketHitsObjectType, { @Field(() => AggregateBucketHitsObjectType, {
+23 -1
View File
@@ -203,6 +203,7 @@ type AggregateBucketHitsObjectType {
} }
type AggregateBucketObjectType { type AggregateBucketObjectType {
"""Number of returned sample hits in this bucket"""
count: Int! count: Int!
"""The hits object""" """The hits object"""
@@ -925,7 +926,7 @@ type EditorType {
name: String! name: String!
} }
union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CalendarProviderRequestErrorDataType | CopilotDocNotFoundDataType | CopilotFailedToAddWorkspaceArtifactDataType | CopilotFailedToGenerateEmbeddingDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderNotSupportedDataType | CopilotProviderSideErrorDataType | DocActionDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | DocUpdateBlockedDataType | ExpectToGrantDocUserRolesDataType | ExpectToRevokeDocUserRolesDataType | ExpectToUpdateDocUserRoleDataType | GraphqlBadRequestDataType | HttpRequestErrorDataType | ImageFormatNotSupportedDataType | InvalidAppConfigDataType | InvalidAppConfigInputDataType | InvalidEmailDataType | InvalidHistoryTimestampDataType | InvalidIndexerInputDataType | InvalidLicenseToActivateDataType | InvalidLicenseUpdateParamsDataType | InvalidOauthCallbackCodeDataType | InvalidOauthResponseDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | InvalidSearchProviderRequestDataType | MemberNotFoundInSpaceDataType | MentionUserDocAccessDeniedDataType | MissingOauthQueryParameterDataType | NoCopilotProviderAvailableDataType | NoMoreSeatDataType | NotInSpaceDataType | QueryTooLongDataType | ResponseTooLargeErrorDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SpaceAccessDeniedDataType | SpaceNotFoundDataType | SpaceOwnerNotFoundDataType | SpaceShouldHaveOnlyOneOwnerDataType | SsrfBlockedErrorDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | UnsupportedClientVersionDataType | UnsupportedServerVersionDataType | UnsupportedSubscriptionPlanDataType | ValidationErrorDataType | VersionRejectedDataType | WorkspacePermissionNotFoundDataType | WrongSignInCredentialsDataType union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CalendarProviderRequestErrorDataType | CopilotDocNotFoundDataType | CopilotFailedToAddWorkspaceArtifactDataType | CopilotFailedToGenerateEmbeddingDataType | CopilotMessageNotFoundDataType | CopilotPromptNotFoundDataType | CopilotProviderNotSupportedDataType | CopilotProviderSideErrorDataType | DocActionDeniedDataType | DocHistoryNotFoundDataType | DocNotFoundDataType | DocUpdateBlockedDataType | ExpectToGrantDocUserRolesDataType | ExpectToRevokeDocUserRolesDataType | ExpectToUpdateDocUserRoleDataType | GraphqlBadRequestDataType | HttpRequestErrorDataType | ImageFormatNotSupportedDataType | InvalidAppConfigDataType | InvalidAppConfigInputDataType | InvalidEmailDataType | InvalidHistoryTimestampDataType | InvalidIndexerInputDataType | InvalidLicenseToActivateDataType | InvalidLicenseUpdateParamsDataType | InvalidOauthCallbackCodeDataType | InvalidOauthResponseDataType | InvalidPasswordLengthDataType | InvalidRuntimeConfigTypeDataType | InvalidSearchProviderRequestDataType | MemberNotFoundInSpaceDataType | MentionUserDocAccessDeniedDataType | MissingOauthQueryParameterDataType | NoCopilotProviderAvailableDataType | NoMoreSeatDataType | NotInSpaceDataType | QueryTooLongDataType | ResponseTooLargeErrorDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SearchIndexFailedDataType | SearchIndexNotReadyDataType | SpaceAccessDeniedDataType | SpaceNotFoundDataType | SpaceOwnerNotFoundDataType | SpaceShouldHaveOnlyOneOwnerDataType | SsrfBlockedErrorDataType | SubscriptionAlreadyExistsDataType | SubscriptionNotExistsDataType | SubscriptionPlanNotFoundDataType | UnknownOauthProviderDataType | UnsupportedClientVersionDataType | UnsupportedServerVersionDataType | UnsupportedSubscriptionPlanDataType | ValidationErrorDataType | VersionRejectedDataType | WorkspacePermissionNotFoundDataType | WrongSignInCredentialsDataType
enum ErrorNames { enum ErrorNames {
ACCESS_DENIED ACCESS_DENIED
@@ -1050,7 +1051,11 @@ enum ErrorNames {
RUNTIME_CONFIG_NOT_FOUND RUNTIME_CONFIG_NOT_FOUND
SAME_EMAIL_PROVIDED SAME_EMAIL_PROVIDED
SAME_SUBSCRIPTION_RECURRING SAME_SUBSCRIPTION_RECURRING
SEARCH_INDEX_FAILED
SEARCH_INDEX_NOT_READY
SEARCH_PERMISSION_SYNCING
SEARCH_PROVIDER_NOT_FOUND SEARCH_PROVIDER_NOT_FOUND
SEARCH_PROVIDER_UNAVAILABLE
SIGN_UP_FORBIDDEN SIGN_UP_FORBIDDEN
SPACE_ACCESS_DENIED SPACE_ACCESS_DENIED
SPACE_NOT_FOUND SPACE_NOT_FOUND
@@ -2141,6 +2146,14 @@ input SearchHighlight {
field: String! field: String!
} }
type SearchIndexFailedDataType {
diagnosticId: String!
}
type SearchIndexNotReadyDataType {
spaceId: String!
}
input SearchInput { input SearchInput {
options: SearchOptions! options: SearchOptions!
query: SearchQuery! query: SearchQuery!
@@ -2199,8 +2212,17 @@ type SearchResultObjectType {
} }
type SearchResultPagination { type SearchResultPagination {
"""Number of results returned in this response, not a global total"""
count: Int! count: Int!
"""
Whether the provider has more candidates; remaining visible results are not guaranteed
"""
hasMore: Boolean! hasMore: Boolean!
"""
Opaque provider candidate cursor; it does not guarantee complete visible-result pagination
"""
nextCursor: String nextCursor: String
} }
+20
View File
@@ -259,6 +259,7 @@ export interface AggregateBucketHitsObjectType {
export interface AggregateBucketObjectType { export interface AggregateBucketObjectType {
__typename?: 'AggregateBucketObjectType'; __typename?: 'AggregateBucketObjectType';
/** Number of returned sample hits in this bucket */
count: Scalars['Int']['output']; count: Scalars['Int']['output'];
/** The hits object */ /** The hits object */
hits: AggregateBucketHitsObjectType; hits: AggregateBucketHitsObjectType;
@@ -1108,6 +1109,8 @@ export type ErrorDataUnion =
| ResponseTooLargeErrorDataType | ResponseTooLargeErrorDataType
| RuntimeConfigNotFoundDataType | RuntimeConfigNotFoundDataType
| SameSubscriptionRecurringDataType | SameSubscriptionRecurringDataType
| SearchIndexFailedDataType
| SearchIndexNotReadyDataType
| SpaceAccessDeniedDataType | SpaceAccessDeniedDataType
| SpaceNotFoundDataType | SpaceNotFoundDataType
| SpaceOwnerNotFoundDataType | SpaceOwnerNotFoundDataType
@@ -1248,7 +1251,11 @@ export enum ErrorNames {
RUNTIME_CONFIG_NOT_FOUND = 'RUNTIME_CONFIG_NOT_FOUND', RUNTIME_CONFIG_NOT_FOUND = 'RUNTIME_CONFIG_NOT_FOUND',
SAME_EMAIL_PROVIDED = 'SAME_EMAIL_PROVIDED', SAME_EMAIL_PROVIDED = 'SAME_EMAIL_PROVIDED',
SAME_SUBSCRIPTION_RECURRING = 'SAME_SUBSCRIPTION_RECURRING', SAME_SUBSCRIPTION_RECURRING = 'SAME_SUBSCRIPTION_RECURRING',
SEARCH_INDEX_FAILED = 'SEARCH_INDEX_FAILED',
SEARCH_INDEX_NOT_READY = 'SEARCH_INDEX_NOT_READY',
SEARCH_PERMISSION_SYNCING = 'SEARCH_PERMISSION_SYNCING',
SEARCH_PROVIDER_NOT_FOUND = 'SEARCH_PROVIDER_NOT_FOUND', SEARCH_PROVIDER_NOT_FOUND = 'SEARCH_PROVIDER_NOT_FOUND',
SEARCH_PROVIDER_UNAVAILABLE = 'SEARCH_PROVIDER_UNAVAILABLE',
SIGN_UP_FORBIDDEN = 'SIGN_UP_FORBIDDEN', SIGN_UP_FORBIDDEN = 'SIGN_UP_FORBIDDEN',
SPACE_ACCESS_DENIED = 'SPACE_ACCESS_DENIED', SPACE_ACCESS_DENIED = 'SPACE_ACCESS_DENIED',
SPACE_NOT_FOUND = 'SPACE_NOT_FOUND', SPACE_NOT_FOUND = 'SPACE_NOT_FOUND',
@@ -2829,6 +2836,16 @@ export interface SearchHighlight {
field: Scalars['String']['input']; field: Scalars['String']['input'];
} }
export interface SearchIndexFailedDataType {
__typename?: 'SearchIndexFailedDataType';
diagnosticId: Scalars['String']['output'];
}
export interface SearchIndexNotReadyDataType {
__typename?: 'SearchIndexNotReadyDataType';
spaceId: Scalars['String']['output'];
}
export interface SearchInput { export interface SearchInput {
options: SearchOptions; options: SearchOptions;
query: SearchQuery; query: SearchQuery;
@@ -2889,8 +2906,11 @@ export interface SearchResultObjectType {
export interface SearchResultPagination { export interface SearchResultPagination {
__typename?: 'SearchResultPagination'; __typename?: 'SearchResultPagination';
/** Number of results returned in this response, not a global total */
count: Scalars['Int']['output']; count: Scalars['Int']['output'];
/** Whether the provider has more candidates; remaining visible results are not guaranteed */
hasMore: Scalars['Boolean']['output']; hasMore: Scalars['Boolean']['output'];
/** Opaque provider candidate cursor; it does not guarantee complete visible-result pagination */
nextCursor: Maybe<Scalars['String']['output']>; nextCursor: Maybe<Scalars['String']['output']>;
} }
+8 -22
View File
@@ -35,10 +35,6 @@
"type": "Object", "type": "Object",
"desc": "The config for doc job queue" "desc": "The config for doc job queue"
}, },
"queues.indexer": {
"type": "Object",
"desc": "The config for indexer job queue"
},
"queues.notification": { "queues.notification": {
"type": "Object", "type": "Object",
"desc": "The config for notification job queue" "desc": "The config for notification job queue"
@@ -398,39 +394,29 @@
"indexer": { "indexer": {
"enabled": { "enabled": {
"type": "Boolean", "type": "Boolean",
"desc": "Enable indexer plugin", "desc": "Enable indexer plugin"
"env": "AFFINE_INDEXER_ENABLED"
}, },
"provider.type": { "provider.type": {
"type": "String", "type": "String",
"desc": "Indexer search provider. Self-hosted uses the embedded provider by default; remote providers require an endpoint.", "desc": "Indexer search provider. Self-hosted uses the embedded provider by default; remote providers require an endpoint."
"env": "AFFINE_INDEXER_SEARCH_PROVIDER"
}, },
"provider.endpoint": { "provider.endpoint": {
"type": "String", "type": "String",
"desc": "Remote indexer endpoint. Not used by the embedded provider.", "desc": "Remote indexer endpoint. Not used by the embedded provider."
"env": "AFFINE_INDEXER_SEARCH_ENDPOINT"
}, },
"provider.apiKey": { "provider.apiKey": {
"type": "String", "type": "String",
"desc": "Indexer search service api key. Optional for elasticsearch", "desc": "Indexer search service api key. Optional for remote providers",
"link": "https://www.elastic.co/guide/server/current/api-key.html", "link": "https://www.elastic.co/guide/server/current/api-key.html"
"env": "AFFINE_INDEXER_SEARCH_API_KEY"
}, },
"provider.username": { "provider.username": {
"type": "String", "type": "String",
"desc": "Indexer search service auth username, if not set, basic auth will be disabled. Optional for elasticsearch", "desc": "Indexer search service auth username, if not set, basic auth will be disabled. Optional for remote providers",
"link": "https://www.elastic.co/guide/en/elasticsearch/reference/current/http-clients.html", "link": "https://www.elastic.co/guide/en/elasticsearch/reference/current/http-clients.html"
"env": "AFFINE_INDEXER_SEARCH_USERNAME"
}, },
"provider.password": { "provider.password": {
"type": "String", "type": "String",
"desc": "Indexer search service auth password, if not set, basic auth will be disabled. Optional for elasticsearch", "desc": "Indexer search service auth password, if not set, basic auth will be disabled. Optional for remote providers"
"env": "AFFINE_INDEXER_SEARCH_PASSWORD"
},
"autoIndex.batchSize": {
"type": "Number",
"desc": "Number of workspaces automatically indexed per batch"
} }
}, },
"oauth": { "oauth": {
@@ -170,14 +170,13 @@ export const KNOWN_CONFIG_GROUPS = [
{ {
key: 'provider.type', key: 'provider.type',
type: 'Enum', type: 'Enum',
options: ['embedded', 'manticoresearch', 'elasticsearch'], options: ['embedded', 'elasticsearch', 'manticoresearch'],
desc: 'Search provider. Embedded keeps external credentials for later reuse.', desc: 'Search provider. Embedded and Elasticsearch provide full search semantics; Manticore Search provides basic search semantics.',
}, },
'provider.endpoint', 'provider.endpoint',
'provider.apiKey', 'provider.apiKey',
'provider.username', 'provider.username',
'provider.password', 'provider.password',
'autoIndex.batchSize',
], ],
} as ConfigGroup<'indexer'>, } as ConfigGroup<'indexer'>,
]; ];
Loaded 100 of 107 files, more files were not shown because too many files have changed in this diff. Show more