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
2026-08-23 17:54:49 +08:00
committed by GitHub
parent a8eef53966
commit b530198a3b
107 changed files with 8612 additions and 5306 deletions
+1 -38
View File
@@ -56,10 +56,6 @@ export declare class BackendRuntime {
getWorkspaceInviteLinkById(inviteId: string): Promise<RuntimeWorkspaceInviteLinkRecord | null>
revokeWorkspaceInviteLink(workspaceId: string): Promise<boolean>
cleanupExpiredRuntimeStates(limit: number): Promise<number>
refreshWorkspaceAdminStatsDirty(batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRefreshResult>
recalibrateWorkspaceAdminStats(lastSid: number, batchLimit: number, owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsRecalibrationResult>
writeWorkspaceAdminStatsDailySnapshot(owner: string, leaseTtlMs: number): Promise<RuntimeWorkspaceStatsSnapshotResult>
recalibrateWorkspaceAdminStatsDaily(batchLimit: number, owner: string, leaseTtlMs: number, lockRetryTimes: number, lockRetryDelayMs: number): Promise<RuntimeWorkspaceStatsDailyRecalibrationResult>
constructor(privateKey?: string | undefined | null, configPaths?: Array<string> | undefined | null)
start(): Promise<void>
stop(): Promise<void>
@@ -68,10 +64,7 @@ export declare class BackendRuntime {
runMigrations(): Promise<void>
searchAuthorized(actorUserId: string, workspaceId: string, request: RuntimeSearchRequest): Promise<SearchOperationOutput>
aggregateAuthorized(actorUserId: string, workspaceId: string, request: RuntimeAggregateRequest): Promise<SearchOperationOutput>
indexSearchDocument(workspaceId: string, docId: string): Promise<void>
deleteSearchDocument(workspaceId: string, docId: string): Promise<void>
reconcileSearchWorkspace(workspaceId: string): Promise<void>
deleteSearchWorkspace(workspaceId: string): Promise<void>
reconcileSearchProjection(limit?: number | undefined | null): Promise<number>
filterReadableDocs(actorUserId: string, workspaceId: string, docIds: Array<string>): Promise<Array<string>>
searchStatus(): Promise<any>
embeddingHealth(): Promise<EmbeddingHealth>
@@ -113,7 +106,6 @@ export declare class StorageRuntime {
rebuildWorkspaceDocBlobRefs(workspaceId: string, limit: number): Promise<RuntimeDocBlobRefsResult>
reconcileWorkspaceDocuments(workspaceId: string): Promise<RuntimeDocumentCleanupReconcileResult>
executeDocumentCleanupCandidates(workspaceId: string | undefined | null, gracePeriodDays: number, limit: number): Promise<RuntimeDocumentCleanupExecuteResult>
ackDocumentCleanupEffect(workspaceId: string, docId: string, cleanupVersion: string, effect: string): Promise<RuntimeDocumentCleanupAckResult>
constructor()
start(): Promise<void>
configure(configJson: string): void
@@ -1261,16 +1253,11 @@ export interface RuntimeDocHistoryInput {
historyMaxAgeMs: number
}
export interface RuntimeDocumentCleanupAckResult {
completed: boolean
}
export interface RuntimeDocumentCleanupEffect {
workspaceId: string
docId: string
cleanupVersion: string
commentObjectsDone: boolean
searchDone: boolean
}
export interface RuntimeDocumentCleanupExecuteResult {
@@ -1551,30 +1538,6 @@ export interface RuntimeWorkspaceInviteQuotaUsage {
targetDomains: Array<RuntimeQuotaTargetDomainInput>
}
export interface RuntimeWorkspaceStatsDailyRecalibrationResult {
processed: number
lastSid: number
snapshotted: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsRecalibrationResult {
processed: number
lastSid: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsRefreshResult {
processed: number
backlog: number
skipped: boolean
}
export interface RuntimeWorkspaceStatsSnapshotResult {
snapshotted: number
skipped: boolean
}
export declare function safeFetch(request: SafeFetchRequest): Promise<SafeFetchResponse>
export type SafeFetchMethod = 'get'|
@@ -2,6 +2,3 @@ pub(super) const MAGIC_LINK_OTP_PURPOSE: &str = "magic_link_otp";
pub(super) const MAX_MAGIC_LINK_OTP_ATTEMPTS: i32 = 10;
pub(super) const WORKSPACE_INVITE_LINK_ID_PURPOSE: &str = "workspace_invite_link:id";
pub(super) const WORKSPACE_INVITE_LINK_WORKSPACE_PURPOSE: &str = "workspace_invite_link:workspace";
pub(super) const WORKSPACE_STATS_LEASE_KEY: &str = "workspace:admin-stats:refresh";
pub(super) const WORKSPACE_STATS_LOCK_NAMESPACE: i64 = 97_301;
pub(super) const WORKSPACE_STATS_REFRESH_LOCK_KEY: i64 = 1;
@@ -16,7 +16,6 @@ mod scope_compiler;
mod search;
#[cfg(test)]
mod tests;
mod workspace_stats;
use std::{
sync::{Arc, RwLock},
time::Duration,
@@ -26,6 +25,8 @@ use byok::LocalLeasePayload;
use copilot::{backend_provider, executable_protocol};
use embedding::register_artifact_source;
use napi::{Result, bindgen_prelude::Buffer};
#[cfg(test)]
pub(crate) use search::SEARCH_TEST_LOCK;
use search::SearchRuntime;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
@@ -66,10 +67,12 @@ fn search_operation_output(result: RuntimeResult<serde_json::Value>) -> SearchOp
error_code: Some(
match error {
RuntimeError::SearchWorkspaceDenied => "workspace_denied",
RuntimeError::SearchPermissionUnavailable => "permission_unavailable",
RuntimeError::SearchPermissionUnavailable => "permission_syncing",
RuntimeError::SearchIndexNotReady => "index_not_ready",
RuntimeError::SearchPermissionSyncing => "permission_syncing",
RuntimeError::SearchIndexFailed(_) => "index_failed",
RuntimeError::SearchProviderUnavailable => "provider_unavailable",
RuntimeError::SearchUnsupportedQuery => "unsupported_query",
RuntimeError::SearchReplayGap => "provider_unavailable",
RuntimeError::InvalidInput(_) | RuntimeError::Json { .. } => "invalid_request",
_ => "internal",
}
@@ -358,46 +361,14 @@ impl BackendRuntime {
}
#[napi]
pub async fn index_search_document(&self, workspace_id: String, doc_id: String) -> Result<()> {
let search = self.search_runtime().await?;
let result = if self.role.owns_background() {
search.index_document(&workspace_id, &doc_id).await
} else {
search.project_document_only(&workspace_id, &doc_id).await
};
result.map_err(to_napi_error)
}
#[napi]
pub async fn delete_search_document(&self, workspace_id: String, doc_id: String) -> Result<()> {
let search = self.search_runtime().await?;
let result = if self.role.owns_background() {
search.delete_document(&workspace_id, &doc_id).await
} else {
search.delete_document_only(&workspace_id, &doc_id).await
};
result.map_err(to_napi_error)
}
#[napi]
pub async fn reconcile_search_workspace(&self, workspace_id: String) -> Result<()> {
pub async fn reconcile_search_projection(&self, limit: Option<i32>) -> Result<i32> {
self.require_background()?;
self
.search_runtime()
.await?
.reconcile_workspace(permission::SystemSearchCapability::ReconcileIndex, &workspace_id)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn delete_search_workspace(&self, workspace_id: String) -> Result<()> {
self.require_background()?;
self
.search_runtime()
.await?
.delete_workspace(&workspace_id)
.reconcile_pending(limit.unwrap_or(100))
.await
.map(|count| count as i32)
.map_err(to_napi_error)
}
@@ -29,7 +29,7 @@ impl PermissionAuthorizer {
) -> RuntimeResult<AuthorizedSearchScope> {
match actor {
SearchActor::User { user_id } => {
let snapshot = self.store.search_snapshot(workspace_id, user_id).await?;
let snapshot = self.store.permission_snapshot(workspace_id, user_id, &[]).await?;
let owner_or_admin = matches!(snapshot.evaluation.workspace.role.as_deref(), Some("owner" | "admin"))
&& snapshot.evaluation.workspace.member_state.as_deref() == Some("active");
let decision = evaluate_permission(snapshot.evaluation)
@@ -54,17 +54,12 @@ impl PermissionAuthorizer {
};
Ok(AuthorizedSearchScope {
workspace_id: workspace_id.to_string(),
permission_revision: snapshot.revision,
docs,
})
}
}
}
pub(in crate::runtime::backend_runtime) async fn revision(&self, workspace_id: &str) -> RuntimeResult<i64> {
self.store.revision(workspace_id).await
}
pub(in crate::runtime::backend_runtime) async fn filter_readable_docs(
&self,
workspace_id: &str,
@@ -72,6 +67,9 @@ impl PermissionAuthorizer {
doc_ids: Vec<String>,
) -> RuntimeResult<BTreeSet<String>> {
let snapshot = self.store.permission_snapshot(workspace_id, user_id, &doc_ids).await?;
if snapshot.capability == DocAclCapability::Unknown {
return Err(RuntimeError::SearchPermissionUnavailable);
}
let output = evaluate_permission(snapshot.evaluation).map_err(|_| RuntimeError::SearchPermissionUnavailable)?;
Ok(
output
@@ -3,9 +3,7 @@ mod store;
mod types;
pub(super) use authorizer::PermissionAuthorizer;
#[cfg(test)]
pub(super) use types::AclPredicate;
pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor, SystemSearchCapability};
pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor};
#[cfg(test)]
mod tests;
@@ -20,10 +20,6 @@ impl PermissionStore {
Self { pool }
}
pub(super) async fn search_snapshot(&self, workspace_id: &str, user_id: &str) -> RuntimeResult<PermissionSnapshot> {
self.permission_snapshot(workspace_id, user_id, &[]).await
}
pub(super) async fn permission_snapshot(
&self,
workspace_id: &str,
@@ -40,12 +36,10 @@ impl PermissionStore {
.await
.map_err(|error| RuntimeError::database("configure permission snapshot", error))?;
let row = sqlx::query(
r#"SELECT revision.revision,
policy.visibility, coalesce(policy.sharing_enabled, true) AS sharing_enabled,
r#"SELECT policy.visibility, coalesce(policy.sharing_enabled, true) AS sharing_enabled,
coalesce(policy.member_default_doc_role, 'manager') AS member_default_doc_role,
member.role, member.state
FROM workspaces workspace
LEFT JOIN workspace_permission_revisions revision ON revision.workspace_id=workspace.id
LEFT JOIN workspace_access_policies policy ON policy.workspace_id=workspace.id
LEFT JOIN LATERAL (
SELECT role,state FROM workspace_members
@@ -60,10 +54,6 @@ impl PermissionStore {
.await
.map_err(|error| RuntimeError::database("load workspace permission facts", error))?
.ok_or_else(|| RuntimeError::invalid_input("workspace_not_found"))?;
let revision = row
.try_get::<Option<i64>, _>("revision")
.map_err(|error| RuntimeError::database("decode permission revision", error))?
.ok_or_else(|| RuntimeError::invalid_state("permission_state_unavailable"))?;
let sharing_enabled: bool = row
.try_get("sharing_enabled")
.map_err(|error| RuntimeError::database("decode workspace sharing", error))?;
@@ -134,7 +124,6 @@ impl PermissionStore {
member_state.as_deref() == Some("active") && matches!(role.as_deref(), Some("member" | "admin" | "owner"));
Ok(PermissionSnapshot {
revision,
capability,
evaluation: PermissionEvaluationInputV1 {
version: 1,
@@ -163,15 +152,6 @@ impl PermissionStore {
sharing_enabled,
})
}
pub(super) async fn revision(&self, workspace_id: &str) -> RuntimeResult<i64> {
sqlx::query_scalar("SELECT revision FROM workspace_permission_revisions WHERE workspace_id=$1")
.bind(workspace_id)
.fetch_optional(&self.pool)
.await
.map_err(|error| RuntimeError::database("read permission revision", error))?
.ok_or_else(|| RuntimeError::invalid_state("permission_state_unavailable"))
}
}
async fn load_doc_acl_capability(
@@ -7,7 +7,14 @@ static PERMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_
async fn setup() -> Option<(PgPool, String, String)> {
let database_url = std::env::var("DATABASE_URL").ok()?;
let pool = PgPool::connect(&database_url).await.unwrap();
crate::runtime::migrations::migrate_search_tables(&pool).await.unwrap();
let legacy_relations: Vec<Option<String>> = sqlx::query_scalar(
"SELECT to_regclass(name) FROM \
unnest(ARRAY['workspace_permission_revisions','workspace_permission_changes','search_runtime_generations']) name",
)
.fetch_all(&pool)
.await
.unwrap();
assert!(legacy_relations.into_iter().all(|relation| relation.is_none()));
let suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("search-permission-user-{suffix}");
let workspace_id = format!("search-permission-workspace-{suffix}");
@@ -67,7 +74,15 @@ async fn non_team_is_all_and_team_uses_projected_acl() {
};
assert_eq!(predicate.actor_user_id, user_id);
assert!(predicate.active_member);
assert!(team.permission_revision > free.permission_revision);
sqlx::query("UPDATE workspace_members SET role='admin' WHERE workspace_id=$1 AND user_id=$2")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
let admin = authorizer.authorize_search(&actor, &workspace_id).await.unwrap();
assert_eq!(admin.docs, DocReadScope::All);
}
#[tokio::test]
@@ -112,13 +127,20 @@ async fn inactive_member_is_denied_and_unknown_capability_fails_closed() {
}
#[tokio::test]
async fn fact_changes_advance_revision_and_write_ordered_change() {
async fn canonical_doc_acl_facts_are_evaluated_without_search_state() {
let _guard = PERMISSION_TEST_LOCK.lock().await;
let Some((pool, workspace_id, user_id)) = setup().await else {
return;
};
let authorizer = PermissionAuthorizer::new(pool.clone());
let before = authorizer.revision(&workspace_id).await.unwrap();
sqlx::query(
"INSERT INTO doc_access_policies(workspace_id,doc_id,visibility,member_default_role) \
VALUES($1,'doc','private','none'),($1,'hidden','private','none')",
)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO doc_grants(workspace_id,doc_id,principal_type,principal_id,role) VALUES($1,'doc','user',$2,'reader')",
)
@@ -127,39 +149,23 @@ async fn fact_changes_advance_revision_and_write_ordered_change() {
.execute(&pool)
.await
.unwrap();
let after = authorizer.revision(&workspace_id).await.unwrap();
assert_eq!(after, before + 1);
let change: (i64, Option<String>, String) = sqlx::query_as(
"SELECT revision,doc_id,scope FROM workspace_permission_changes WHERE workspace_id=$1 AND revision=$2",
)
.bind(&workspace_id)
.bind(after)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(change, (after, Some("doc".to_string()), "doc_grant".to_string()));
sqlx::query("UPDATE workspace_members SET updated_at=now() WHERE workspace_id=$1 AND user_id=$2")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
let readable = authorizer
.filter_readable_docs(&workspace_id, &user_id, vec!["doc".to_string(), "hidden".to_string()])
.await
.unwrap();
assert_eq!(authorizer.revision(&workspace_id).await.unwrap(), after);
assert_eq!(readable, ["doc".to_string()].into_iter().collect());
let moved_workspace_id = format!("{workspace_id}-moved");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&moved_workspace_id)
.execute(&pool)
.await
.unwrap();
assert_eq!(authorizer.revision(&moved_workspace_id).await.unwrap(), 0);
sqlx::query("UPDATE doc_grants SET workspace_id=$1 WHERE workspace_id=$2 AND doc_id='doc'")
.bind(&moved_workspace_id)
sqlx::query("UPDATE doc_access_policies SET member_default_role='reader' WHERE workspace_id=$1 AND doc_id='hidden'")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
assert_eq!(authorizer.revision(&workspace_id).await.unwrap(), after + 1);
assert_eq!(authorizer.revision(&moved_workspace_id).await.unwrap(), 1);
let readable = authorizer
.filter_readable_docs(&workspace_id, &user_id, vec!["doc".to_string(), "hidden".to_string()])
.await
.unwrap();
assert_eq!(
readable,
["doc".to_string(), "hidden".to_string()].into_iter().collect()
);
}
@@ -1,10 +1,5 @@
use crate::permission::PermissionEvaluationInputV1;
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum SystemSearchCapability {
ReconcileIndex,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum SearchActor {
User { user_id: String },
@@ -33,12 +28,10 @@ pub(in crate::runtime::backend_runtime) enum DocReadScope {
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) struct AuthorizedSearchScope {
pub(in crate::runtime::backend_runtime) workspace_id: String,
pub(in crate::runtime::backend_runtime) permission_revision: i64,
pub(in crate::runtime::backend_runtime) docs: DocReadScope,
}
pub(super) struct PermissionSnapshot {
pub(super) revision: i64,
pub(super) capability: DocAclCapability,
pub(super) evaluation: PermissionEvaluationInputV1,
pub(super) actor_user_id: String,
@@ -250,7 +250,15 @@ mod tests {
};
let _guard = crate::runtime::migrations::EMBEDDING_TEST_LOCK.lock().await;
let pool = PgPool::connect(&database_url).await.unwrap();
crate::runtime::migrations::migrate_search_tables(&pool).await.unwrap();
let legacy_relations: Vec<Option<String>> = sqlx::query_scalar(
"SELECT to_regclass(name) FROM \
unnest(ARRAY['workspace_permission_revisions','workspace_permission_changes','search_runtime_generations']) \
name",
)
.fetch_all(&pool)
.await
.unwrap();
assert!(legacy_relations.into_iter().all(|relation| relation.is_none()));
let suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("scope-user-{suffix}");
let collaborator_id = format!("scope-collaborator-{suffix}");
@@ -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 uuid::Uuid;
use super::{SCHEMA_FINGERPRINT, provider::RemoteProvider, types::SearchTable};
use super::{SCHEMA_FINGERPRINT, SearchProvider, SearchTable};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
#[derive(Clone)]
@@ -22,82 +22,120 @@ impl ActiveGeneration {
}
}
pub(super) async fn prepare(
pub(super) async fn ensure(
pool: &PgPool,
config: &SearchRuntimeConfig,
remote: Option<&RemoteProvider>,
remote: Option<&SearchProvider>,
rebuild_embedded: bool,
) -> RuntimeResult<ActiveGeneration> {
let fingerprint = config_fingerprint(config);
let expected_config_hash = config_hash(config);
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search generation", error))?;
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-runtime-generation', 0))")
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("lock search generation", error))?;
let existing = sqlx::query(
r#"SELECT generation_id,provider,manifest FROM search_runtime_generations
WHERE state IN ('active','pending') AND provider=$1 AND config_fingerprint=$2 AND schema_fingerprint=$3
ORDER BY (state='active') DESC LIMIT 1"#,
r#"SELECT id,provider,config_hash,schema_version,manifest,state
FROM search_projection.generations
WHERE state='building'
ORDER BY created_at DESC
LIMIT 1"#,
)
.bind(&config.provider)
.bind(&fingerprint)
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load search generation", error))?;
let generation = if let Some(row) = existing {
decode(row)?
} else {
let pending: bool =
sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM search_runtime_generations WHERE state='pending')")
.fetch_one(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("check pending search generation", error))?;
if pending {
return Err(RuntimeError::invalid_state("search_generation_change_in_progress"));
}
let generation_id = Uuid::new_v4();
let suffix = generation_id.simple().to_string();
let manifest = if config.provider == "embedded" {
json!({"doc":"doc","block":"block"})
let provider: String = row
.try_get("provider")
.map_err(|error| RuntimeError::database("decode search generation provider", error))?;
let existing_config_hash: Vec<u8> = row
.try_get("config_hash")
.map_err(|error| RuntimeError::database("decode search generation config hash", error))?;
let schema_version: i32 = row
.try_get("schema_version")
.map_err(|error| RuntimeError::database("decode search generation schema version", error))?;
if provider != config.provider
|| existing_config_hash != expected_config_hash
|| schema_version != SCHEMA_FINGERPRINT
{
let id: Uuid = row
.try_get("id")
.map_err(|error| RuntimeError::database("decode superseded search generation id", error))?;
sqlx::query(
"UPDATE search_projection.generations SET state='failed', last_error='search generation superseded by \
configuration' WHERE id=$1 AND state='building'",
)
.bind(id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("fail superseded search generation", error))?;
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
} else {
json!({
"doc":format!("affine_search_doc_{suffix}"),
"block":format!("affine_search_block_{suffix}"),
})
};
sqlx::query(
r#"INSERT INTO search_runtime_generations
(generation_id,provider,state,config_fingerprint,schema_fingerprint,manifest)
VALUES ($1,$2,'pending',$3,$4,$5)"#,
)
.bind(generation_id)
.bind(&config.provider)
.bind(&fingerprint)
.bind(SCHEMA_FINGERPRINT)
.bind(&manifest)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("create pending search generation", error))?;
for table in [SearchTable::Doc, SearchTable::Block] {
sqlx::query("INSERT INTO search_runtime_provider_cursors(generation_id,table_key) VALUES ($1,$2)")
.bind(generation_id)
.bind(table.as_str())
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("initialize search generation cursor", error))?;
let state: String = row
.try_get("state")
.map_err(|error| RuntimeError::database("decode search generation state", error))?;
if rebuild_embedded {
if state == "building" {
sqlx::query(
"UPDATE search_projection.generations SET state='failed', last_error='embedded generation lost on \
restart' WHERE id=$1 AND state='building'",
)
.bind(
row
.try_get::<Uuid, _>("id")
.map_err(|error| RuntimeError::database("decode search generation id", error))?,
)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("fail stale embedded generation", error))?;
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
} else {
decode(row)?
}
} else {
decode(row)?
}
}
ActiveGeneration {
id: generation_id,
manifest,
} else {
let active = sqlx::query(
r#"SELECT id,provider,config_hash,schema_version,manifest
FROM search_projection.generations
WHERE state='active'
ORDER BY activated_at DESC NULLS LAST LIMIT 1"#,
)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load active search generation", error))?;
if let Some(row) = active {
let provider: String = row
.try_get("provider")
.map_err(|error| RuntimeError::database("decode active search provider", error))?;
let existing_config_hash: Vec<u8> = row
.try_get("config_hash")
.map_err(|error| RuntimeError::database("decode active search config hash", error))?;
let schema_version: i32 = row
.try_get("schema_version")
.map_err(|error| RuntimeError::database("decode active search schema version", error))?;
if !rebuild_embedded
&& provider == config.provider
&& existing_config_hash == expected_config_hash
&& schema_version == SCHEMA_FINGERPRINT
{
decode(row)?
} else {
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
}
} else {
create_generation(&mut transaction, config, expected_config_hash.clone()).await?
}
};
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit pending search generation", error))?;
.map_err(|error| RuntimeError::database("commit search generation", error))?;
if let Some(remote) = remote {
for table in [SearchTable::Doc, SearchTable::Block] {
@@ -114,14 +152,14 @@ pub(super) async fn load_active(
pool: &PgPool,
config: &SearchRuntimeConfig,
) -> RuntimeResult<Option<ActiveGeneration>> {
let fingerprint = config_fingerprint(config);
let row = sqlx::query(
r#"SELECT generation_id,provider,manifest FROM search_runtime_generations
WHERE state='active' AND provider=$1 AND config_fingerprint=$2 AND schema_fingerprint=$3
r#"SELECT id,provider,config_hash,manifest
FROM search_projection.generations
WHERE state='active' AND provider=$1 AND config_hash=$2 AND schema_version=$3
ORDER BY activated_at DESC NULLS LAST LIMIT 1"#,
)
.bind(&config.provider)
.bind(&fingerprint)
.bind(config_hash(config))
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(pool)
.await
@@ -129,28 +167,122 @@ pub(super) async fn load_active(
row.map(decode).transpose()
}
pub(super) async fn fail(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
sqlx::query("UPDATE search_runtime_generations SET state='failed' WHERE generation_id=$1 AND state='pending'")
.bind(generation.id)
.execute(pool)
.await
.map_err(|error| RuntimeError::database("fail pending search generation", error))?;
async fn fail(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
sqlx::query(
"UPDATE search_projection.generations SET state='failed', last_error=$2 WHERE id=$1 AND state NOT IN \
('active','draining')",
)
.bind(generation.id)
.bind("search generation build failed")
.execute(pool)
.await
.map_err(|error| RuntimeError::database("fail search generation", error))?;
Ok(())
}
pub(super) async fn activate(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
pub(super) async fn activate(
pool: &PgPool,
generation: &ActiveGeneration,
config: &SearchRuntimeConfig,
) -> RuntimeResult<()> {
let expected_config_hash = config_hash(config);
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin search generation activation", error))?;
sqlx::query("UPDATE search_runtime_generations SET state='draining' WHERE state='active' AND generation_id<>$1")
.bind(generation.id)
sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended('search-projection-generation', 0))")
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("drain previous search generation", error))?;
.map_err(|error| RuntimeError::database("lock search generation activation", error))?;
let row = sqlx::query(
r#"SELECT generation.state,generation.provider,generation.config_hash,generation.schema_version,
generation.scan_high_water_sid IS NOT NULL
AND generation.scan_cursor_sid IS NOT NULL
AND generation.scan_cursor_sid >= generation.scan_high_water_sid AS scan_complete,
NOT EXISTS (
SELECT 1
FROM search_projection.workspace_states state
WHERE state.generation_id=generation.id
AND (
NOT state.covered
OR state.required_permission_version > state.applied_permission_version
OR state.pending_scope <> 'none'
OR state.last_error IS NOT NULL
)
) AS workspaces_covered,
NOT EXISTS (
SELECT 1
FROM workspaces workspace
WHERE workspace.sid <= COALESCE(generation.scan_high_water_sid,0)
AND NOT EXISTS (
SELECT 1
FROM search_projection.workspace_states state
WHERE state.generation_id=generation.id AND state.workspace_id=workspace.id
)
) AS workspaces_seeded,
NOT EXISTS (
SELECT 1 FROM search_projection.document_states state
WHERE state.generation_id=generation.id
AND (state.target_source_version <> state.published_source_version
OR state.target_source_exists <> state.published_source_exists
OR state.target_permission_version <> state.published_permission_version)
) AS publications_complete
FROM search_projection.generations generation
WHERE generation.id=$1
FOR UPDATE"#,
)
.bind(generation.id)
.fetch_optional(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load search generation activation state", error))?;
let Some(row) = row else {
return Err(RuntimeError::SearchIndexNotReady);
};
let state: String = row
.try_get("state")
.map_err(|error| RuntimeError::database("decode search generation activation state", error))?;
let provider: String = row
.try_get("provider")
.map_err(|error| RuntimeError::database("decode search generation activation provider", error))?;
let row_config_hash: Vec<u8> = row
.try_get("config_hash")
.map_err(|error| RuntimeError::database("decode search generation activation config hash", error))?;
let schema_version: i32 = row
.try_get("schema_version")
.map_err(|error| RuntimeError::database("decode search generation activation schema", error))?;
let scan_complete: bool = row
.try_get("scan_complete")
.map_err(|error| RuntimeError::database("decode search generation scan state", error))?;
let workspaces_covered: bool = row
.try_get("workspaces_covered")
.map_err(|error| RuntimeError::database("decode search generation workspace coverage", error))?;
let workspaces_seeded: bool = row
.try_get("workspaces_seeded")
.map_err(|error| RuntimeError::database("decode search generation workspace seed", error))?;
let publications_complete: bool = row
.try_get("publications_complete")
.map_err(|error| RuntimeError::database("decode search generation publication state", error))?;
if state != "building"
|| provider != config.provider
|| row_config_hash != expected_config_hash
|| schema_version != SCHEMA_FINGERPRINT
|| !scan_complete
|| !workspaces_covered
|| !workspaces_seeded
|| !publications_complete
{
return Err(RuntimeError::SearchIndexNotReady);
}
sqlx::query(
"UPDATE search_runtime_generations SET state='active', activated_at=coalesce(activated_at,now()) WHERE \
generation_id=$1 AND state IN ('pending','active')",
"UPDATE search_projection.generations SET state='draining', drained_at=now() WHERE state='active' AND id<>$1",
)
.bind(generation.id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("drain previous search generation", error))?;
sqlx::query(
"UPDATE search_projection.generations SET state='active', activated_at=coalesce(activated_at,now()) WHERE id=$1 \
AND state='building'",
)
.bind(generation.id)
.execute(&mut *transaction)
@@ -162,7 +294,7 @@ pub(super) async fn activate(pool: &PgPool, generation: &ActiveGeneration) -> Ru
.map_err(|error| RuntimeError::database("commit search generation activation", error))
}
fn config_fingerprint(config: &SearchRuntimeConfig) -> String {
pub(super) fn config_hash(config: &SearchRuntimeConfig) -> Vec<u8> {
let mut hash = Sha256::new();
for value in [
&config.provider,
@@ -174,16 +306,108 @@ fn config_fingerprint(config: &SearchRuntimeConfig) -> String {
hash.update(value.as_bytes());
hash.update([0]);
}
hash.finalize().iter().map(|byte| format!("{byte:02x}")).collect()
hash.finalize().to_vec()
}
fn decode(row: sqlx::postgres::PgRow) -> RuntimeResult<ActiveGeneration> {
Ok(ActiveGeneration {
id: row
.try_get("generation_id")
.try_get("id")
.map_err(|error| RuntimeError::database("decode search generation id", error))?,
manifest: row
.try_get("manifest")
.map_err(|error| RuntimeError::database("decode search generation manifest", error))?,
})
}
async fn create_generation(
transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
config: &SearchRuntimeConfig,
config_hash: Vec<u8>,
) -> RuntimeResult<ActiveGeneration> {
let generation_id = Uuid::new_v4();
let suffix = generation_id.simple().to_string();
let manifest = if config.provider == "embedded" {
json!({"doc":"doc","block":"block"})
} else {
json!({
"doc":format!("affine_search_doc_{suffix}"),
"block":format!("affine_search_block_{suffix}")
})
};
sqlx::query(
r#"INSERT INTO search_projection.generations
(id,provider,state,config_hash,schema_version,manifest)
VALUES ($1,$2,'building',$3,$4,$5)"#,
)
.bind(generation_id)
.bind(&config.provider)
.bind(&config_hash)
.bind(SCHEMA_FINGERPRINT)
.bind(&manifest)
.execute(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("create search generation", error))?;
Ok(ActiveGeneration {
id: generation_id,
manifest,
})
}
#[cfg(test)]
mod tests {
use sqlx::PgPool;
use uuid::Uuid;
use super::ensure;
use crate::runtime::{
SearchRuntimeConfig, backend_runtime::search::SEARCH_TEST_LOCK, migrations::migrate_search_tables,
};
#[tokio::test]
async fn configuration_change_supersedes_an_incomplete_candidate() {
let _guard = SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await
.unwrap();
let stale_id = Uuid::new_v4();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(stale_id)
.execute(&pool)
.await
.unwrap();
let generation = ensure(&pool, &SearchRuntimeConfig::default(), None, false)
.await
.unwrap();
assert_ne!(generation.id, stale_id);
let stale: (String, Option<String>) =
sqlx::query_as("SELECT state,last_error FROM search_projection.generations WHERE id=$1")
.bind(stale_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stale.0, "failed");
assert_eq!(
stale.1.as_deref(),
Some("search generation superseded by configuration")
);
sqlx::query("DELETE FROM search_projection.generations WHERE id IN ($1,$2)")
.bind(stale_id)
.bind(generation.id)
.execute(&pool)
.await
.unwrap();
}
}
@@ -1,19 +1,32 @@
mod checkpoint;
mod generation;
mod projection;
mod provider;
mod query;
mod query_runtime;
mod result_filter;
mod runtime;
mod store;
mod types;
mod worker;
use generation::{ActiveGeneration, activate, config_hash, ensure, load_active};
use projection::{ProjectionInput, project_document};
use provider::{SearchChange, SearchProvider, projection_external_id};
use query::{compile, compile_aggregate};
use result_filter::{candidates, retain_visible_nodes};
pub(super) use runtime::SearchRuntime;
use types::{AggregateOptions, RuntimeSearchQuery, SearchOptions, SearchTable};
pub(super) use types::{RuntimeAggregateRequest, RuntimeSearchRequest};
use worker::{reconcile_workspace, sweep_generation_orphans};
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 {
use sha2::{Digest, Sha256};
@@ -32,4 +45,19 @@ fn provider_payload(payload: &serde_json::Value) -> serde_json::Value {
}
#[cfg(test)]
mod tests;
mod tests {
use super::SCHEMA_FINGERPRINT;
#[test]
fn terminal_schema_uses_one_generation_version() {
assert_eq!(SCHEMA_FINGERPRINT, 1);
}
#[test]
fn external_document_ids_keep_document_and_block_tables_distinct() {
let cases = [("workspace/doc", false), ("workspace/doc/block", true)];
for (external_id, block) in cases {
assert_eq!(external_id.matches('/').count() >= 2, block);
}
}
}
@@ -1,32 +1,42 @@
use serde_json::{Value, json};
use sqlx::{PgPool, Row};
use super::store::ProjectionInput;
use crate::{
permission::doc_role_allows,
runtime::{RuntimeError, RuntimeResult, storage_runtime::load_current_doc},
};
#[derive(Clone, Debug, PartialEq)]
pub(super) struct ProjectionInput {
pub(super) workspace_id: String,
pub(super) doc_id: String,
pub(super) payload: Value,
}
pub(super) async fn project_document(
pool: &PgPool,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<Option<(ProjectionInput, Vec<ProjectionInput>)>> {
let Some(current) = load_current_doc(pool, workspace_id, doc_id).await? else {
let Some(current) = load_current_doc(pool, workspace_id, doc_id)
.await
.map_err(|error| match error {
RuntimeError::InvalidState(message) => RuntimeError::SearchSourceInvalid(message),
error => error,
})?
else {
return Ok(None);
};
let revision = current.updated_at.timestamp_millis();
let projection =
affine_doc_loader::project_document_search(current.blob, doc_id.to_string(), revision.to_string())
.map_err(|error| RuntimeError::invalid_state(format!("search document projection failed: {error}")))?;
let metadata = sqlx::query(
.map_err(|error| RuntimeError::SearchSourceInvalid(format!("document projection failed: {error}")))?;
let Some(metadata) = sqlx::query(
r#"SELECT snapshot.created_at,snapshot.updated_at,snapshot.created_by,snapshot.updated_by,
revision.revision AS acl_revision,
coalesce(doc_policy.visibility,'private') AS visibility,
doc_policy.public_role,
coalesce(doc_policy.member_default_role,workspace_policy.member_default_doc_role,'manager') AS member_default_role
FROM snapshots snapshot
LEFT JOIN workspace_permission_revisions revision ON revision.workspace_id=snapshot.workspace_id
LEFT JOIN workspace_access_policies workspace_policy ON workspace_policy.workspace_id=snapshot.workspace_id
LEFT JOIN doc_access_policies doc_policy
ON doc_policy.workspace_id=snapshot.workspace_id AND doc_policy.doc_id=snapshot.guid
@@ -37,11 +47,9 @@ pub(super) async fn project_document(
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load search projection metadata", error))?
.ok_or_else(|| RuntimeError::invalid_state("search snapshot metadata unavailable"))?;
let acl_revision = metadata
.try_get::<Option<i64>, _>("acl_revision")
.map_err(|error| RuntimeError::database("decode search ACL revision", error))?
.ok_or_else(|| RuntimeError::invalid_state("permission_state_unavailable"))?;
else {
return Ok(None);
};
let visibility: String = metadata
.try_get("visibility")
.map_err(|error| RuntimeError::database("decode search doc visibility", error))?;
@@ -51,14 +59,18 @@ pub(super) async fn project_document(
let member_default_role: String = metadata
.try_get("member_default_role")
.map_err(|error| RuntimeError::database("decode search member default role", error))?;
let acl_public_readable = visibility == "public"
&& public_role
.as_deref()
.is_some_and(|role| doc_role_allows(role, "Doc.Read").unwrap_or(false));
let acl_member_default_readable = doc_role_allows(&member_default_role, "Doc.Read")
.map_err(|_| RuntimeError::invalid_state("permission_state_unavailable"))?;
let acl_public_readable = if visibility == "public" {
match public_role.as_deref() {
Some(role) => readable_role(role, "public")?,
None => false,
}
} else {
false
};
let acl_member_default_readable = readable_role(&member_default_role, "member default")?;
let grants = sqlx::query(
"SELECT principal_id,role FROM doc_grants WHERE workspace_id=$1 AND doc_id=$2 AND principal_type='user'",
"SELECT principal_id,role FROM doc_grants WHERE workspace_id=$1 AND doc_id=$2 AND principal_type='user' ORDER BY \
principal_id,role",
)
.bind(workspace_id)
.bind(doc_id)
@@ -74,7 +86,7 @@ pub(super) async fn project_document(
let role: String = row
.try_get("role")
.map_err(|error| RuntimeError::database("decode grant role", error))?;
Ok(doc_role_allows(&role, "Doc.Read").unwrap_or(false).then_some(id))
Ok(readable_role(&role, "grant")?.then_some(id))
})
.collect::<RuntimeResult<Vec<_>>>()?
.into_iter()
@@ -96,7 +108,6 @@ pub(super) async fn project_document(
public_readable: acl_public_readable,
member_default_readable: acl_member_default_readable,
read_user_ids: acl_read_user_ids,
revision: acl_revision,
};
let document_payload = with_acl(
json!({
@@ -113,14 +124,7 @@ pub(super) async fn project_document(
}),
&acl,
);
let document = input(
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}"),
revision,
document_payload,
&acl,
);
let document = input(workspace_id, doc_id, document_payload);
let blocks = projection
.units
.into_iter()
@@ -142,24 +146,21 @@ pub(super) async fn project_document(
}),
&acl,
);
input(
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}/{block_id}"),
revision,
payload,
&acl,
)
input(workspace_id, doc_id, payload)
})
.collect();
Ok(Some((document, blocks)))
}
fn readable_role(role: &str, source: &str) -> RuntimeResult<bool> {
doc_role_allows(role, "Doc.Read")
.map_err(|_| RuntimeError::SearchSourceInvalid(format!("invalid {source} document role")))
}
struct AclFields {
public_readable: bool,
member_default_readable: bool,
read_user_ids: Vec<String>,
revision: i64,
}
fn with_acl(mut payload: Value, acl: &AclFields) -> Value {
@@ -181,27 +182,25 @@ fn with_acl(mut payload: Value, acl: &AclFields) -> Value {
tokens.push("public".to_string());
}
object.insert("acl_read_tokens".to_string(), json!(tokens));
object.insert("acl_revision".to_string(), json!(acl.revision));
payload
}
fn input(
workspace_id: &str,
doc_id: &str,
external_id: &str,
revision: i64,
payload: Value,
acl: &AclFields,
) -> ProjectionInput {
fn input(workspace_id: &str, doc_id: &str, payload: Value) -> ProjectionInput {
ProjectionInput {
external_id: external_id.to_string(),
workspace_id: workspace_id.to_string(),
doc_id: doc_id.to_string(),
revision,
payload,
acl_public_readable: acl.public_readable,
acl_member_default_readable: acl.member_default_readable,
acl_read_user_ids: acl.read_user_ids.clone(),
acl_revision: acl.revision,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn invalid_acl_role_is_a_permanent_search_source_error() {
let error = readable_role("future-role", "grant").unwrap_err();
assert!(matches!(error, RuntimeError::SearchSourceInvalid(_)));
assert!(error.is_permanent_search_source());
}
}
@@ -1,326 +1,690 @@
use serde_json::{Value, json};
use std::time::Duration;
use crate::runtime::{RuntimeError, RuntimeResult};
use reqwest::{Client, redirect::Policy};
use serde_json::{Map, Value, json};
use sha2::{Digest, Sha256};
pub(super) fn prepare_manticore_payload(
payload: &mut Value,
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<()> {
let object = payload.as_object_mut().expect("search payload is an object");
object.remove("acl_read_user_ids");
if let Some(Value::Array(tokens)) = object.remove("acl_read_tokens") {
object.insert(
"acl_read_token_ids".to_string(),
Value::Array(
tokens
.iter()
.filter_map(Value::as_str)
.map(|token| {
token_ids
.get(token)
.copied()
.map(Value::from)
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))
})
.collect::<RuntimeResult<Vec<_>>>()?,
),
);
}
if let Some(Value::Array(tokens)) = object.get("ref_doc_id").cloned() {
object.insert(
"ref_doc_token_ids".to_string(),
Value::Array(
tokens
.iter()
.filter_map(Value::as_str)
.map(|token| {
token_ids
.get(token)
.copied()
.map(Value::from)
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))
})
.collect::<RuntimeResult<Vec<_>>>()?,
),
);
}
for field in ["created_at", "updated_at"] {
if let Some(value) = object.get_mut(field)
&& let Some(milliseconds) = value.as_i64()
{
*value = json!(milliseconds / 1_000);
}
}
for (field, value) in object.iter_mut() {
if let Value::Array(values) = value {
*value = if matches!(field.as_str(), "acl_read_token_ids" | "ref_doc_token_ids") {
continue;
} else if field == "content" {
Value::String(values.iter().filter_map(Value::as_str).collect::<Vec<_>>().join(" "))
} else {
Value::String(
serde_json::to_string(values).map_err(|error| RuntimeError::json("encode manticore array", error))?,
)
};
} else if value.is_null() {
*value = Value::String(String::new());
}
}
Ok(())
use super::{SearchChange, SearchTable, provider_write_error, webpki_tls_config};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
pub(in crate::runtime::backend_runtime::search) struct ManticoreSearchProvider {
client: Client,
endpoint: String,
api_key: String,
username: String,
password: String,
}
pub(super) fn prepare_manticore_search(
dsl: &mut Value,
cursor: Option<Value>,
size: u64,
initial_offset: u64,
requested_fields: &[String],
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<u64> {
normalize_manticore_terms(dsl, token_ids)?;
let object = dsl.as_object_mut().expect("search DSL is an object");
let mut source = object
.get("_source")
.and_then(Value::as_array)
impl ManticoreSearchProvider {
pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
let endpoint = config.endpoint.trim_end_matches('/');
let url = url::Url::parse(endpoint).map_err(|_| RuntimeError::config("invalid search provider endpoint"))?;
if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return Err(RuntimeError::config("invalid search provider endpoint"));
}
let client = Client::builder()
.tls_backend_preconfigured(
webpki_tls_config()
.map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?,
)
.redirect(Policy::none())
.timeout(Duration::from_secs(30))
.build()
.map_err(|error| RuntimeError::invalid_state(format!("search HTTP client failed: {error}")))?;
Ok(Self {
client,
endpoint: endpoint.to_string(),
api_key: config.api_key.clone(),
username: config.username.clone(),
password: config.password.clone(),
})
}
pub(super) async fn search(&self, physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
let request = translate_search_request(physical_table, dsl)?;
let response = self
.request(reqwest::Method::POST, "search")
.json(&request)
.send()
.await;
let response = response.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(if status.as_u16() == 400 {
RuntimeError::SearchUnsupportedQuery
} else {
RuntimeError::SearchProviderUnavailable
});
}
let value: Value =
serde_json::from_slice(&body).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize(value, request.get("offset").and_then(Value::as_u64).unwrap_or_default())
}
pub(super) async fn aggregate(&self, _physical_table: &str, _dsl: Value) -> RuntimeResult<Value> {
Err(RuntimeError::SearchUnsupportedQuery)
}
pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
let response = self
.request(reqwest::Method::POST, "sql?mode=raw")
.header("content-type", "application/x-www-form-urlencoded")
.body(create_table_sql(physical_table, table))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid search provider schema response", error))?;
if value.get("error").is_some() {
return Err(RuntimeError::invalid_state("provider_schema_failed"));
}
Ok(())
}
pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
if changes.is_empty() {
return Ok(());
}
let mut body = String::new();
for change in changes {
let payload = change.upsert_payload()?;
let mut document = manticore_document(payload);
document.insert("external_id".to_string(), json!(change.external_id));
body.push_str(
&serde_json::to_string(&json!({
"replace": {
"table": physical_table,
"id": document_id(&change.external_id),
"doc": document,
}
}))
.map_err(|error| RuntimeError::json("encode manticore document", error))?,
);
body.push('\n');
}
if body.is_empty() {
return Ok(());
}
let response = self
.request(reqwest::Method::POST, "bulk")
.header("content-type", "application/x-ndjson")
.body(body)
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(provider_write_error(status.as_u16()));
}
let value: Value =
serde_json::from_slice(&body).map_err(|error| RuntimeError::json("invalid manticore bulk response", error))?;
validate_bulk_response(&value, changes.len())
}
pub(super) async fn gc_document_history(
&self,
physical_table: &str,
workspace_id: &str,
doc_id: &str,
source_version: i64,
permission_version: i64,
limit: usize,
) -> RuntimeResult<()> {
let query = json!({
"bool": {
"must": [
{"equals":{"workspace_id":workspace_id}},
{"equals":{"doc_id":doc_id}},
],
"must_not": [{"bool":{"must":[
{"equals":{"source_version":source_version}},
{"equals":{"permission_version":permission_version}}
]}}]
}
});
self
.delete_bounded(physical_table, query, limit.max(1))
.await
.map(|_| ())
}
pub(super) async fn gc_workspace(
&self,
physical_table: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> RuntimeResult<bool> {
let limit = limit.max(1);
let query = json!({"bool":{"must":[
{"equals":{"workspace_id":workspace_id}},
{"range":{"source_version":{"lte":source_version_high_water}}}
]}});
let deleted = self.delete_bounded(physical_table, query, limit).await?;
Ok(deleted == limit)
}
async fn delete_bounded(&self, physical_table: &str, query: Value, limit: usize) -> RuntimeResult<usize> {
let response = self
.request(reqwest::Method::POST, "search")
.json(&json!({
"table":physical_table,
"query":query,
"limit":limit,
"sort":[{"id":"asc"}]
}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid manticore GC search response", error))?;
let ids = value
.pointer("/hits/hits")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid manticore GC search response"))?
.iter()
.map(|hit| {
let id = hit
.get("_id")
.or_else(|| hit.get("id"))
.and_then(|id| id.as_u64().or_else(|| id.as_str().and_then(|id| id.parse().ok())))
.ok_or_else(|| RuntimeError::invalid_state("invalid manticore GC search response"))?;
Ok(json!(id))
})
.collect::<RuntimeResult<Vec<_>>>()?;
if ids.is_empty() {
return Ok(0);
}
let response = self
.request(reqwest::Method::POST, "delete")
.json(&json!({"table":physical_table,"id":ids}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value =
serde_json::from_slice(&body).map_err(|error| RuntimeError::json("invalid manticore delete response", error))?;
if value.get("error").is_some() {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(ids.len())
}
fn request(&self, method: reqwest::Method, path: &str) -> reqwest::RequestBuilder {
let mut request = self.client.request(method, format!("{}/{path}", self.endpoint));
if !self.api_key.is_empty() {
request = request.header("Authorization", format!("ApiKey {}", self.api_key));
} else if !self.username.is_empty() {
request = request.basic_auth(&self.username, Some(&self.password));
}
request
}
}
fn create_table_sql(physical_table: &str, table: SearchTable) -> String {
let text_field = table.text_field();
let mut columns = vec![format!("{text_field} text")];
for field in [
"external_id",
"workspace_id",
"workspace_token",
"generation_id",
"doc_id",
"doc_token",
"created_by_user_id",
"updated_by_user_id",
] {
columns.push(format!("{field} string"));
}
for field in ["source_version", "permission_version", "created_at", "updated_at"] {
columns.push(format!("{field} bigint"));
}
columns.push("acl_public_readable integer".to_string());
columns.push("acl_member_default_readable integer".to_string());
columns.push("acl_read_tokens text".to_string());
let extra_fields = if table == SearchTable::Doc {
vec!["summary string", "journal string"]
} else {
vec![
"block_id string",
"block_token string",
"unit_id string",
"projection_version integer",
"source_hash string",
"visibility string",
"element_id string",
"frame_id string",
"source_block_id string",
"flavour string",
"blob string",
"ref_doc_id string",
"ref string",
"parent_flavour string",
"parent_block_id string",
"additional string",
"markdown_preview string",
]
};
columns.extend(extra_fields.into_iter().map(str::to_string));
format!("CREATE TABLE IF NOT EXISTS {physical_table} ({})", columns.join(", "))
}
fn manticore_document(payload: &Value) -> Map<String, Value> {
payload
.as_object()
.into_iter()
.flatten()
.filter_map(Value::as_str)
.map(str::to_string)
.collect::<std::collections::BTreeSet<_>>();
source.extend(requested_fields.iter().cloned());
object.insert("_source".to_string(), json!(source));
object.remove("fields");
if let Some(highlight) = object.get_mut("highlight")
&& let Some(options) = highlight
.get("fields")
.filter(|(_, value)| !value.is_null())
.map(|(key, value)| {
let value = match value {
Value::Bool(value) => json!(i32::from(*value)),
Value::Array(values) => json!(
values
.iter()
.filter(|value| !value.is_null())
.map(|value| value.as_str().map(str::to_string).unwrap_or_else(|| value.to_string()))
.collect::<Vec<_>>()
.join(" ")
),
value => value.clone(),
};
(key.clone(), value)
})
.collect()
}
fn document_id(external_id: &str) -> u64 {
let digest = Sha256::digest(external_id.as_bytes());
let mut bytes = [0; 8];
bytes.copy_from_slice(&digest[..8]);
u64::from_be_bytes(bytes) | 1
}
fn validate_bulk_response(value: &Value, expected_changes: usize) -> RuntimeResult<()> {
let items = value
.get("items")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid manticore bulk response"))?;
let mut affected = 0_u64;
for item in items {
let result = item
.as_object()
.and_then(|item| item.values().next())
.and_then(Value::as_object)
.and_then(|fields| fields.values().next())
.cloned()
{
*highlight = options;
}
let offset = if let Some(Value::String(cursor)) = cursor {
let offset = serde_json::from_str::<Value>(&cursor)
.map_err(|error| RuntimeError::json("invalid search cursor", error))?
.get("offset")
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
let status = result
.get("status")
.and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
if offset.saturating_add(size) > 10_000 {
return Err(RuntimeError::invalid_input("search cursor exceeds 10000"));
.and_then(|status| u16::try_from(status).ok())
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
if result.get("error").is_some() || !(200..300).contains(&status) {
return Err(provider_write_error(status));
}
object.insert("from".to_string(), json!(offset));
offset
} else if cursor.is_some() {
return Err(RuntimeError::invalid_input("invalid search cursor"));
} else {
initial_offset
};
Ok(offset)
}
pub(super) fn manticore_fields(source: Option<&Value>, requested_fields: &[String]) -> Value {
let source = source.and_then(Value::as_object);
Value::Object(
requested_fields
.iter()
.filter_map(|field| {
let mut value = source?.get(field)?.clone();
if matches!(field.as_str(), "created_at" | "updated_at")
&& let Some(seconds) = value.as_i64()
{
value = json!(seconds * 1_000);
} else if let Some(encoded) = value.as_str()
&& encoded.starts_with('[')
&& let Ok(decoded) = serde_json::from_str(encoded)
{
value = decoded;
}
if !value.is_array() {
value = Value::Array(vec![value]);
}
Some((field.clone(), value))
})
.collect(),
)
}
fn normalize_manticore_terms(
value: &mut Value,
token_ids: &std::collections::HashMap<String, i64>,
) -> RuntimeResult<()> {
if let Some(term) = manticore_term(value, token_ids)? {
*value = term;
return Ok(());
let numeric = ["created", "updated", "deleted"]
.into_iter()
.filter_map(|field| result.get(field).and_then(Value::as_u64))
.sum::<u64>();
affected += if numeric > 0 {
numeric
} else if result.get("created").and_then(Value::as_bool).is_some()
|| result
.get("result")
.and_then(Value::as_str)
.is_some_and(|result| matches!(result, "created" | "updated"))
{
1
} else {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
};
}
match value {
Value::Object(object) => {
if let Some(Value::Object(boolean)) = object.get_mut("bool")
&& boolean.get("boost").and_then(Value::as_f64) == Some(1.0)
{
boolean.remove("boost");
}
if let Some(Value::Object(terms)) = object.get_mut("terms") {
terms.entry("order").or_insert_with(|| json!({"_count":"desc"}));
}
for child in object.values_mut() {
normalize_manticore_terms(child, token_ids)?;
}
}
Value::Array(array) => {
for child in array {
normalize_manticore_terms(child, token_ids)?;
}
}
_ => {}
if value.get("errors").and_then(Value::as_bool) != Some(false) || affected != expected_changes as u64 {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(())
}
fn manticore_term(value: &Value, token_ids: &std::collections::HashMap<String, i64>) -> RuntimeResult<Option<Value>> {
let Some(term) = value.get("term").and_then(Value::as_object) else {
return Ok(None);
fn translate_search_request(physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
let object = dsl
.as_object()
.ok_or_else(|| RuntimeError::invalid_input("invalid search request"))?;
let size = object.get("size").and_then(Value::as_u64).unwrap_or(10);
let offset = if let Some(cursor) = object.get("cursor") {
let cursor = cursor
.as_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
serde_json::from_str::<Value>(cursor)
.ok()
.and_then(|cursor| cursor.get("offset").and_then(Value::as_u64))
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?
} else {
object.get("from").and_then(Value::as_u64).unwrap_or_default()
};
if term.len() != 1 {
return Ok(None);
let query = object.get("query").cloned().unwrap_or_else(|| json!({"match_all":{}}));
let query = translate_query(query)?;
let mut request = json!({
"table":physical_table,
"query":query,
"limit":size,
"offset":offset,
});
if let Some(sort) = object.get("sort") {
request["sort"] = translate_sort(sort)?;
}
let Some((field, clause)) = term.iter().next() else {
return Ok(None);
};
let value = clause.get("value").unwrap_or(clause);
Ok(match value {
Value::String(value) => {
if matches!(field.as_str(), "acl_read_tokens" | "ref_doc_id") {
let token_id = token_ids
.get(value)
.copied()
.ok_or_else(|| RuntimeError::invalid_state("Manticore exact token mapping is incomplete"))?;
let field = if field == "acl_read_tokens" {
"acl_read_token_ids"
} else {
"ref_doc_token_ids"
let mut source = Vec::new();
if let Some(fields) = object.get("_source").and_then(Value::as_array) {
source.extend(fields.iter().filter_map(Value::as_str).map(str::to_string));
}
if let Some(fields) = object.get("fields").and_then(Value::as_array) {
source.extend(fields.iter().filter_map(Value::as_str).map(str::to_string));
}
source.push("external_id".to_string());
source.push("doc_id".to_string());
source.push("source_version".to_string());
source.push("permission_version".to_string());
source.sort();
source.dedup();
request["_source"] = json!(source);
if let Some(highlight) = object.get("highlight") {
request["highlight"] = translate_highlight(highlight)?;
}
Ok(request)
}
fn translate_highlight(highlight: &Value) -> RuntimeResult<Value> {
let fields = highlight
.get("fields")
.and_then(Value::as_object)
.ok_or(RuntimeError::SearchUnsupportedQuery)?;
let mut request = json!({"fields":fields.keys().collect::<Vec<_>>()});
let tags = fields.values().filter_map(Value::as_object).next();
if let Some(pre_tag) = tags
.and_then(|options| options.get("pre_tags"))
.and_then(Value::as_array)
.and_then(|tags| tags.first())
.and_then(Value::as_str)
{
request["pre_tags"] = json!(pre_tag);
}
if let Some(post_tag) = tags
.and_then(|options| options.get("post_tags"))
.and_then(Value::as_array)
.and_then(|tags| tags.first())
.and_then(Value::as_str)
{
request["post_tags"] = json!(post_tag);
}
Ok(request)
}
fn translate_sort(sort: &Value) -> RuntimeResult<Value> {
let values = sort
.as_array()
.ok_or_else(|| RuntimeError::invalid_input("invalid search sort"))?
.iter()
.filter_map(|value| match value {
Value::String(field) => match field.as_str() {
"_score" => None,
"_id" => Some(json!({"id":"asc"})),
field => Some(json!({field:"asc"})),
},
Value::Object(object) => {
let (field, direction) = object.iter().next()?;
let field = match field.as_str() {
"_id" => "id",
"_score" => return None,
field => field,
};
return Ok(Some(json!({"equals":{field:token_id}})));
}
let (field, value) = match field.as_str() {
"workspace_id" => ("workspace_token", super::super::exact_token(value)),
"doc_id" => ("doc_token", super::super::exact_token(value)),
"block_id" => ("block_token", super::super::exact_token(value)),
_ => (field.as_str(), value.clone()),
};
if let Some(boost) = clause.get("boost").and_then(Value::as_f64) {
Some(json!({"match":{field:{"query":value,"boost":boost}}}))
} else {
Some(json!({"equals":{field:value}}))
Some(json!({field:direction}))
}
_ => None,
})
.collect::<Vec<_>>();
Ok(json!(values))
}
fn translate_query(query: Value) -> RuntimeResult<Value> {
let Some(object) = query.as_object() else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
if object.contains_key("match_all") {
return Ok(json!({"match_all":{}}));
}
if let Some(match_query) = object.get("match") {
let Some((field, value)) = match_query.as_object().and_then(|object| object.iter().next()) else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
let value = value
.get("query")
.cloned()
.or_else(|| value.as_str().map(|value| json!(value)))
.ok_or(RuntimeError::SearchUnsupportedQuery)?;
return Ok(json!({"match":{field:value}}));
}
if let Some(term_query) = object.get("term") {
let Some((field, value)) = term_query.as_object().and_then(|object| object.iter().next()) else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
let value = value
.get("value")
.cloned()
.ok_or(RuntimeError::SearchUnsupportedQuery)?;
if field == "acl_read_tokens" {
let value = value.as_str().ok_or(RuntimeError::SearchUnsupportedQuery)?;
return Ok(json!({"match":{field:value}}));
}
Value::Bool(value) => Some(json!({"equals":{field:u8::from(*value)}})),
Value::Number(value) => Some(json!({"equals":{field:value}})),
_ => None,
return Ok(json!({"equals":{field:manticore_scalar(value)?}}));
}
if let Some(bool_query) = object.get("bool") {
let Some(bool_query) = bool_query.as_object() else {
return Err(RuntimeError::SearchUnsupportedQuery);
};
let mut translated = Map::new();
for occurrence in ["must", "should", "must_not"] {
let Some(clauses) = bool_query.get(occurrence) else {
continue;
};
let clauses = if let Some(array) = clauses.as_array() {
array
.iter()
.map(|clause| translate_query(clause.clone()))
.collect::<RuntimeResult<Vec<_>>>()?
} else {
vec![translate_query(clauses.clone())?]
};
translated.insert(occurrence.to_string(), json!(clauses));
}
return Ok(json!({"bool":translated}));
}
if object.contains_key("boost") {
return translate_query(object.get("boost").cloned().unwrap_or_default());
}
Err(RuntimeError::SearchUnsupportedQuery)
}
fn manticore_scalar(value: Value) -> RuntimeResult<Value> {
Ok(match value {
Value::Bool(value) => json!(i32::from(value)),
Value::String(value) => json!(value),
Value::Number(value) => Value::Number(value),
_ => return Err(RuntimeError::SearchUnsupportedQuery),
})
}
pub(super) fn manticore_exact_tokens(value: &Value) -> Vec<String> {
let mut tokens = Vec::new();
collect_exact_tokens(value, &mut tokens);
tokens
async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>> {
if response
.content_length()
.is_some_and(|length| length > MAX_RESPONSE_BYTES as u64)
{
return Err(RuntimeError::invalid_state("provider_response_too_large"));
}
let mut bytes = Vec::new();
while let Some(chunk) = response
.chunk()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?
{
if bytes.len() + chunk.len() > MAX_RESPONSE_BYTES {
return Err(RuntimeError::invalid_state("provider_response_too_large"));
}
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
fn collect_exact_tokens(value: &Value, tokens: &mut Vec<String>) {
match value {
Value::Object(object) => {
if let Some(token) = object
.get("term")
.and_then(|term| term.get("acl_read_tokens").or_else(|| term.get("ref_doc_id")))
.and_then(|clause| clause.get("value").unwrap_or(clause).as_str())
{
tokens.push(token.to_string());
}
for field in ["acl_read_tokens", "ref_doc_id"] {
if let Some(values) = object.get(field).and_then(Value::as_array) {
tokens.extend(values.iter().filter_map(Value::as_str).map(str::to_string));
}
}
for child in object.values() {
collect_exact_tokens(child, tokens);
}
}
Value::Array(values) => {
for child in values {
collect_exact_tokens(child, tokens);
}
}
_ => {}
}
fn normalize(value: Value, offset: u64) -> RuntimeResult<Value> {
let hits = value
.pointer("/hits/hits")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider response"))?;
let total = value
.pointer("/hits/total")
.and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_state("inexact provider total"))?;
let nodes = hits
.iter()
.map(|hit| {
let source = hit.get("_source").cloned().unwrap_or_else(|| json!({}));
json!({
"id":source.get("external_id").and_then(Value::as_str).unwrap_or_default(),
"score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(),
"fields":fields_from_source(&source),
"highlights":hit.get("highlight").cloned().unwrap_or_else(||json!({})),
"_source":source,
})
})
.collect::<Vec<_>>();
let next_cursor = (offset.saturating_add(nodes.len() as u64) < total)
.then(|| json!({"offset":offset.saturating_add(nodes.len() as u64)}).to_string());
Ok(json!({"total":total,"nodes":nodes,"nextCursor":next_cursor}))
}
fn fields_from_source(source: &Value) -> Value {
let Some(source) = source.as_object() else {
return json!({});
};
source
.iter()
.map(|(field, value)| (field.clone(), json!([value])))
.collect::<Map<_, _>>()
.into()
}
#[cfg(test)]
mod tests {
use super::{super::super::exact_token, *};
use serde_json::json;
use super::{
SearchTable, document_id, manticore_document, translate_query, translate_search_request, translate_sort,
validate_bulk_response,
};
use crate::runtime::RuntimeError;
#[test]
fn payload_and_fields_preserve_terminal_types() {
let mut payload = json!({
"content":["hello","world"],
"ref_doc_id":["doc-a","doc-b"],
"summary":null,
"created_at":2_000,
"updated_at":3_000,
"acl_read_tokens":["member"]
});
prepare_manticore_payload(
&mut payload,
&[
("member".to_string(), 7),
("doc-a".to_string(), 8),
("doc-b".to_string(), 9),
]
.into_iter()
.collect(),
)
.unwrap();
assert_eq!(payload["content"], "hello world");
assert_eq!(payload["ref_doc_id"], "[\"doc-a\",\"doc-b\"]");
assert_eq!(payload["summary"], "");
assert_eq!(payload["created_at"], 2);
assert_eq!(payload["acl_read_token_ids"], json!([7]));
assert_eq!(payload["ref_doc_token_ids"], json!([8, 9]));
let fields = manticore_fields(
Some(&payload),
&[
"ref_doc_id".to_string(),
"summary".to_string(),
"updated_at".to_string(),
],
fn translates_the_shared_basic_query_subset() {
assert_eq!(
translate_query(json!({
"bool":{"must":[
{"term":{"workspace_id":{"value":"workspace"}}},
{"match":{"content":{"query":"hello","boost":1.0}}}
]}
}))
.unwrap(),
json!({
"bool":{"must":[
{"equals":{"workspace_id":"workspace"}},
{"match":{"content":"hello"}}
]}
})
);
assert_eq!(
translate_search_request(
"blocks",
json!({
"query":{"match":{"content":{"query":"hello"}}},
"highlight":{"fields":{"content":{"pre_tags":["<b>"],"post_tags":["</b>"]}}}
})
)
.unwrap()["highlight"],
json!({"fields":["content"],"pre_tags":"<b>","post_tags":"</b>"})
);
assert_eq!(fields["ref_doc_id"], json!(["doc-a", "doc-b"]));
assert_eq!(fields["summary"], json!([""]));
assert_eq!(fields["updated_at"], json!([3_000]));
}
#[test]
fn nested_terms_use_exact_identity_and_acl_tokens() {
let mut dsl = json!({"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":"workspace","boost":2.0}}},
{"bool":{"must_not":[{"term":{"doc_id":{"value":"doc"}}}]}},
{"term":{"acl_read_tokens":{"value":"member"}}},
{"term":{"ref_doc_id":{"value":"ref-doc"}}}
],"boost":1.0}}});
normalize_manticore_terms(
&mut dsl,
&[("member".to_string(), 9), ("ref-doc".to_string(), 10)]
.into_iter()
.collect(),
)
.unwrap();
fn converts_projection_arrays_and_booleans_to_basic_attributes() {
assert_eq!(
dsl,
json!({"query":{"bool":{"must":[
{"match":{"workspace_token":{"query":exact_token("workspace"),"boost":2.0}}},
{"bool":{"must_not":[{"equals":{"doc_token":exact_token("doc")}}]}},
{"equals":{"acl_read_token_ids":9}},
{"equals":{"ref_doc_token_ids":10}}
]}}})
manticore_document(&json!({"acl_read_tokens":["user","member"],"acl_public_readable":false})),
json!({"acl_read_tokens":"user member","acl_public_readable":0})
.as_object()
.unwrap()
.clone()
);
assert!(
validate_bulk_response(
&json!({"errors":false,"items":[{"bulk":{"status":201,"created":2,"updated":0,"deleted":0}}]}),
2,
)
.is_ok()
);
assert!(
validate_bulk_response(
&json!({"errors":false,"items":[
{"replace":{"status":201,"created":true}},
{"replace":{"status":200,"created":false,"result":"updated"}}
]}),
2,
)
.is_ok()
);
assert!(
validate_bulk_response(
&json!({"errors":true,"items":[{"replace":{"status":400,"error":{"reason":"invalid"}}}]}),
1,
)
.is_err()
);
assert!(matches!(
validate_bulk_response(
&json!({"errors":true,"items":[{"replace":{"status":400,"error":{"reason":"invalid"}}}]}),
1
),
Err(RuntimeError::SearchSourceInvalid(_))
));
}
#[test]
fn uses_stable_nonzero_document_ids() {
assert_ne!(document_id("workspace/doc/block"), 0);
assert_eq!(document_id("workspace/doc/block"), document_id("workspace/doc/block"));
assert_ne!(
document_id("generation/workspace/doc/1/1/block"),
document_id("generation/workspace/doc/2/1/block")
);
}
#[test]
fn drops_score_sort_but_keeps_basic_attribute_order() {
assert_eq!(
translate_sort(&json!(["_score",{"updated_at":"desc"},{"_id":"asc"}])).unwrap(),
json!([{ "updated_at":"desc" }, { "id":"asc" }])
);
assert_eq!(SearchTable::Doc.text_field(), "title");
}
}
@@ -1,40 +1,188 @@
mod manticore;
mod remote;
pub(super) use remote::RemoteProvider;
use serde_json::{Value, json};
use super::types::SearchTable;
use self::{manticore::ManticoreSearchProvider, remote::RemoteProvider};
use super::{SearchTable, webpki_tls_config};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
pub(super) fn mapping(table: SearchTable, provider: &str) -> Value {
fn provider_write_error(status: u16) -> RuntimeError {
match status {
408 | 429 | 500..=599 => RuntimeError::SearchProviderUnavailable,
400..=499 => RuntimeError::SearchSourceInvalid("search provider rejected projection".to_string()),
_ => RuntimeError::invalid_state("provider_apply_failed"),
}
}
#[derive(Clone, Debug, PartialEq)]
pub(super) struct SearchChange {
pub(super) table: SearchTable,
pub(super) external_id: String,
pub(super) workspace_id: String,
pub(super) doc_id: String,
pub(super) source_version: i64,
pub(super) permission_version: i64,
pub(super) payload: Value,
}
impl SearchChange {
fn upsert_payload(&self) -> RuntimeResult<&Value> {
let payload = &self.payload;
if payload.get("workspace_id").and_then(Value::as_str) != Some(&self.workspace_id)
|| payload.get("doc_id").and_then(Value::as_str) != Some(&self.doc_id)
|| payload.get("source_version").and_then(Value::as_i64) != Some(self.source_version)
|| payload.get("permission_version").and_then(Value::as_i64) != Some(self.permission_version)
{
return Err(RuntimeError::invalid_input(
"provider projection tuple does not match payload",
));
}
let generation_id = payload
.get("generation_id")
.and_then(Value::as_str)
.ok_or_else(|| RuntimeError::invalid_input("provider projection generation is required"))?;
if projection_external_id(
self.table,
generation_id,
&self.workspace_id,
&self.doc_id,
payload.get("block_id").and_then(Value::as_str),
self.source_version,
self.permission_version,
)? != self.external_id
{
return Err(RuntimeError::invalid_input(
"provider external id does not match projection tuple",
));
}
Ok(payload)
}
}
pub(super) enum SearchProvider {
/// Elasticsearch-compatible provider with the RFC6 shared projection
/// contract.
Elasticsearch(RemoteProvider),
/// Manticore Search provides candidate retrieval. Canonical permission facts
/// filter every ACL-scoped result before it leaves the runtime.
ManticoreSearch(ManticoreSearchProvider),
}
impl SearchProvider {
pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
match config.provider.as_str() {
"elasticsearch" => RemoteProvider::new(config).map(Self::Elasticsearch),
"manticoresearch" => ManticoreSearchProvider::new(config).map(Self::ManticoreSearch),
_ => Err(crate::runtime::RuntimeError::config("unsupported search provider")),
}
}
pub(super) async fn search(&self, physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
match self {
Self::Elasticsearch(provider) => provider.search(physical_table, dsl).await,
Self::ManticoreSearch(provider) => provider.search(physical_table, dsl).await,
}
}
pub(super) async fn aggregate(&self, physical_table: &str, dsl: Value) -> RuntimeResult<Value> {
match self {
Self::Elasticsearch(provider) => provider.aggregate(physical_table, dsl).await,
Self::ManticoreSearch(provider) => provider.aggregate(physical_table, dsl).await,
}
}
pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
match self {
Self::Elasticsearch(provider) => provider.provision(physical_table, table).await,
Self::ManticoreSearch(provider) => provider.provision(physical_table, table).await,
}
}
pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
match self {
Self::Elasticsearch(provider) => provider.apply(physical_table, changes).await,
Self::ManticoreSearch(provider) => provider.apply(physical_table, changes).await,
}
}
pub(super) async fn gc_document_history(
&self,
physical_table: &str,
workspace_id: &str,
doc_id: &str,
source_version: i64,
permission_version: i64,
limit: usize,
) -> RuntimeResult<()> {
match self {
Self::Elasticsearch(provider) => {
provider
.gc_document_history(
physical_table,
workspace_id,
doc_id,
source_version,
permission_version,
limit,
)
.await
}
Self::ManticoreSearch(provider) => {
provider
.gc_document_history(
physical_table,
workspace_id,
doc_id,
source_version,
permission_version,
limit,
)
.await
}
}
}
pub(super) async fn gc_workspace(
&self,
physical_table: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> RuntimeResult<bool> {
match self {
Self::Elasticsearch(provider) => {
provider
.gc_workspace(physical_table, workspace_id, source_version_high_water, limit)
.await
}
Self::ManticoreSearch(provider) => {
provider
.gc_workspace(physical_table, workspace_id, source_version_high_water, limit)
.await
}
}
}
}
fn mapping(table: SearchTable) -> Value {
let text_field = table.text_field();
let mut properties = serde_json::Map::from_iter([
("workspace_id".into(), json!({"type":"keyword"})),
("workspace_token".into(), json!({"type":"keyword"})),
("generation_id".into(), json!({"type":"keyword"})),
("source_version".into(), json!({"type":"long"})),
("permission_version".into(), json!({"type":"long"})),
("doc_id".into(), json!({"type":"keyword"})),
("doc_token".into(), json!({"type":"keyword"})),
(text_field.into(), json!({"type":"text"})),
(
"created_at".into(),
json!({"type":if provider == "manticoresearch" { "long" } else { "date" }}),
),
(
"updated_at".into(),
json!({"type":if provider == "manticoresearch" { "long" } else { "date" }}),
),
("created_at".into(), json!({"type":"date"})),
("updated_at".into(), json!({"type":"date"})),
("created_by_user_id".into(), json!({"type":"keyword"})),
("updated_by_user_id".into(), json!({"type":"keyword"})),
("acl_public_readable".into(), json!({"type":"boolean"})),
("acl_member_default_readable".into(), json!({"type":"boolean"})),
(
"acl_read_tokens".into(),
if provider == "manticoresearch" {
json!({"type":"keyword","mva":true})
} else {
json!({"type":"keyword"})
},
),
("acl_revision".into(), json!({"type":"long"})),
("acl_read_tokens".into(), json!({"type":"keyword"})),
]);
if table == SearchTable::Block {
for field in [
@@ -65,61 +213,76 @@ pub(super) fn mapping(table: SearchTable, provider: &str) -> Value {
json!({"mappings":{"properties":properties}})
}
pub(super) fn manticore_schema(table: SearchTable, physical_table: &str) -> String {
let common = r#"
workspace_id string attribute indexed,
workspace_token string attribute indexed,
doc_id string attribute indexed,
doc_token string attribute indexed,"#;
let fields = match table {
SearchTable::Doc => format!(
r#"{common}
title text,
summary string stored,
journal string stored,
created_by_user_id string attribute indexed,
updated_by_user_id string attribute indexed,
created_at timestamp,
updated_at timestamp,
acl_public_readable bool,
acl_member_default_readable bool,
acl_read_token_ids multi64,
acl_revision bigint"#,
),
SearchTable::Block => format!(
r#"{common}
block_id string attribute indexed,
block_token string attribute indexed,
unit_id string attribute indexed,
projection_version bigint,
source_hash string attribute indexed,
visibility string attribute indexed,
element_id string attribute indexed,
frame_id string attribute indexed,
source_block_id string attribute indexed,
content text,
flavour string attribute indexed,
blob string attribute indexed,
ref_doc_id string attribute indexed,
ref_doc_token_ids multi64,
ref string stored,
parent_flavour string attribute indexed,
parent_block_id string attribute indexed,
additional string stored,
markdown_preview string stored,
created_by_user_id string attribute indexed,
updated_by_user_id string attribute indexed,
created_at timestamp,
updated_at timestamp,
acl_public_readable bool,
acl_member_default_readable bool,
acl_read_token_ids multi64,
acl_revision bigint"#,
),
};
format!(
"CREATE TABLE IF NOT EXISTS {physical_table} ({fields}) charset_table='non_cjk, chinese' ngram_len='1' \
ngram_chars='U+1100..U+11FF, U+3130..U+318F, U+A960..U+A97F, U+AC00..U+D7AF, U+D7B0..U+D7FF, U+3040..U+30FF, \
U+0E00..U+0E7F' index_field_lengths='1'"
)
pub(super) fn projection_external_id(
table: SearchTable,
generation_id: &str,
workspace_id: &str,
doc_id: &str,
block_id: Option<&str>,
source_version: i64,
permission_version: i64,
) -> RuntimeResult<String> {
let mut id = format!("{generation_id}/{workspace_id}/{doc_id}/{source_version}/{permission_version}");
if table == SearchTable::Block {
let block_id = block_id
.filter(|block_id| !block_id.is_empty())
.ok_or_else(|| RuntimeError::invalid_input("block projection id requires block_id"))?;
id.push('/');
id.push_str(block_id);
}
Ok(id)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{SearchChange, SearchTable, projection_external_id};
#[test]
fn projection_ids_are_stable_and_keep_block_identity_last() {
let doc = projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 7, 3).unwrap();
assert_eq!(doc, "generation/workspace/doc/7/3");
assert_eq!(
projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 7, 3).unwrap(),
doc
);
assert_eq!(
projection_external_id(
SearchTable::Block,
"generation",
"workspace",
"doc",
Some("block"),
7,
3,
)
.unwrap(),
"generation/workspace/doc/7/3/block"
);
assert_ne!(
projection_external_id(SearchTable::Doc, "generation", "workspace", "doc", None, 8, 3).unwrap(),
doc
);
}
#[test]
fn change_rejects_an_external_id_outside_its_tuple() {
let change = SearchChange {
table: SearchTable::Doc,
external_id: "generation/workspace/doc/6/3".into(),
workspace_id: "workspace".into(),
doc_id: "doc".into(),
source_version: 7,
permission_version: 3,
payload: json!({
"generation_id":"generation",
"workspace_id":"workspace",
"doc_id":"doc",
"source_version":7,
"permission_version":3
}),
};
assert!(change.upsert_payload().is_err());
}
}
@@ -2,12 +2,8 @@ use std::time::Duration;
use reqwest::{Client, redirect::Policy};
use serde_json::{Value, json};
use sqlx::PgPool;
use super::{
super::{store::SearchChange, types::SearchTable, webpki_tls_config},
manticore::{manticore_exact_tokens, manticore_fields, prepare_manticore_payload, prepare_manticore_search},
};
use super::{SearchChange, SearchTable, provider_write_error, webpki_tls_config};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
@@ -15,68 +11,42 @@ const MAX_RESPONSE_BYTES: usize = 50 * 1024 * 1024;
pub(in crate::runtime::backend_runtime::search) struct RemoteProvider {
client: Client,
endpoint: String,
provider: String,
api_key: String,
username: String,
password: String,
pool: PgPool,
}
impl RemoteProvider {
pub(in crate::runtime::backend_runtime::search) fn new(
config: &SearchRuntimeConfig,
pool: PgPool,
) -> RuntimeResult<Self> {
pub(super) fn new(config: &SearchRuntimeConfig) -> RuntimeResult<Self> {
let endpoint = config.endpoint.trim_end_matches('/');
let url = url::Url::parse(endpoint).map_err(|_| RuntimeError::config("invalid search provider endpoint"))?;
if !matches!(url.scheme(), "http" | "https") || url.host_str().is_none() {
return Err(RuntimeError::config("invalid search provider endpoint"));
}
let mut client = Client::builder()
let client = Client::builder()
.tls_backend_preconfigured(
webpki_tls_config()
.map_err(|error| RuntimeError::invalid_state(format!("search TLS config failed: {error}")))?,
)
.redirect(Policy::none())
.timeout(Duration::from_secs(30));
if config.provider == "manticoresearch" {
client = client.pool_max_idle_per_host(0);
}
let client = client
.build()
.map_err(|error| RuntimeError::invalid_state(format!("search HTTP client failed: {error}")))?;
Ok(Self {
client,
endpoint: endpoint.to_string(),
provider: config.provider.clone(),
api_key: config.api_key.clone(),
username: config.username.clone(),
password: config.password.clone(),
pool,
})
}
pub(in crate::runtime::backend_runtime::search) async fn search(
&self,
physical_table: &str,
mut dsl: Value,
) -> RuntimeResult<Value> {
let requested_fields = dsl
.get("fields")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_str)
.map(str::to_string)
.collect::<Vec<_>>();
pub(super) async fn search(&self, physical_table: &str, mut dsl: Value) -> RuntimeResult<Value> {
ensure_projection_source(&mut dsl);
dsl["track_total_hits"] = json!(true);
let size = dsl.get("size").and_then(Value::as_u64).unwrap_or(10);
let mut offset = dsl.get("from").and_then(Value::as_u64).unwrap_or(0);
let cursor = dsl.as_object_mut().and_then(|object| object.remove("cursor"));
if self.provider == "manticoresearch" {
let token_ids = self.resolve_manticore_tokens(manticore_exact_tokens(&dsl)).await?;
offset = prepare_manticore_search(&mut dsl, cursor, size, offset, &requested_fields, &token_ids)?;
} else if let Some(cursor) = cursor {
if let Some(cursor) = cursor {
let cursor = cursor
.as_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
@@ -107,22 +77,20 @@ impl RemoteProvider {
}
let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize(
value,
self.provider == "manticoresearch",
offset,
size,
&requested_fields,
)
normalize(value)
}
pub(in crate::runtime::backend_runtime::search) async fn aggregate(
&self,
physical_table: &str,
mut dsl: Value,
) -> RuntimeResult<Value> {
if self.provider == "manticoresearch" {
return Err(RuntimeError::SearchUnsupportedQuery);
pub(super) async fn aggregate(&self, physical_table: &str, mut dsl: Value) -> RuntimeResult<Value> {
if let Some(top_hits) = dsl.pointer_mut("/aggs/result/aggs/result/top_hits") {
ensure_projection_source(top_hits);
}
let limit = dsl
.pointer("/aggs/result/terms/size")
.and_then(Value::as_u64)
.unwrap_or(10);
let skip = dsl.get("from").and_then(Value::as_u64).unwrap_or_default();
if let Some(size) = dsl.pointer_mut("/aggs/result/terms/size") {
*size = json!(skip.saturating_add(limit).saturating_add(1));
}
dsl["track_total_hits"] = json!(true);
let response = self
@@ -134,32 +102,18 @@ impl RemoteProvider {
let status = response.status();
let bytes = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchUnsupportedQuery);
return Err(if status.as_u16() == 400 {
RuntimeError::SearchUnsupportedQuery
} else {
RuntimeError::SearchProviderUnavailable
});
}
let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize_aggregate(value)
normalize_aggregate(value, skip as usize, limit as usize)
}
pub(in crate::runtime::backend_runtime::search) async fn provision(
&self,
physical_table: &str,
table: SearchTable,
) -> RuntimeResult<()> {
if self.provider == "manticoresearch" {
let response = self
.request(reqwest::Method::POST, "cli")
.header("content-type", "text/plain")
.body(super::manticore_schema(table, physical_table))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
return response
.status()
.is_success()
.then_some(())
.ok_or(RuntimeError::SearchProviderUnavailable);
}
pub(super) async fn provision(&self, physical_table: &str, table: SearchTable) -> RuntimeResult<()> {
if self
.request(reqwest::Method::HEAD, physical_table)
.send()
@@ -170,7 +124,7 @@ impl RemoteProvider {
}
let response = self
.request(reqwest::Method::PUT, physical_table)
.json(&super::mapping(table, &self.provider))
.json(&super::mapping(table))
.send()
.await;
match response {
@@ -179,73 +133,124 @@ impl RemoteProvider {
}
}
pub(in crate::runtime::backend_runtime::search) async fn apply(
&self,
physical_table: &str,
changes: &[SearchChange],
) -> RuntimeResult<()> {
pub(super) async fn apply(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
if changes.is_empty() {
return Ok(());
}
let token_ids = if self.provider == "manticoresearch" {
self
.resolve_manticore_tokens(
changes
.iter()
.filter_map(|change| change.payload.as_ref())
.flat_map(manticore_exact_tokens),
)
.await?
} else {
Default::default()
};
self.apply_elasticsearch(physical_table, changes).await
}
async fn apply_elasticsearch(&self, physical_table: &str, changes: &[SearchChange]) -> RuntimeResult<()> {
let mut body = String::new();
for change in changes {
if change.operation == "delete" {
body.push_str(
&serde_json::to_string(&json!({"delete":{"_index":physical_table,"_id":change.external_id}}))
.map_err(|error| RuntimeError::json("encode provider delete", error))?,
);
body.push('\n');
} else if let Some(payload) = &change.payload {
body.push_str(
&serde_json::to_string(&json!({"index":{"_index":physical_table,"_id":change.external_id}}))
.map_err(|error| RuntimeError::json("encode provider upsert", error))?,
);
body.push('\n');
let mut payload = super::super::provider_payload(payload);
if self.provider == "manticoresearch" {
prepare_manticore_payload(&mut payload, &token_ids)?;
}
body.push_str(
&serde_json::to_string(&payload).map_err(|error| RuntimeError::json("encode provider document", error))?,
);
body.push('\n');
}
let payload = change.upsert_payload()?;
body.push_str(
&serde_json::to_string(&json!({
"index": {
"_index": physical_table,
"_id": change.external_id
}
}))
.map_err(|error| RuntimeError::json("encode immutable provider upsert", error))?,
);
body.push('\n');
body.push_str(
&serde_json::to_string(payload)
.map_err(|error| RuntimeError::json("encode immutable provider document", error))?,
);
body.push('\n');
}
if body.is_empty() {
return Ok(());
}
let path = if self.provider == "elasticsearch" {
"_bulk?refresh=wait_for"
} else {
"_bulk"
};
let response = self
.request(reqwest::Method::POST, path)
.request(reqwest::Method::POST, "_bulk?refresh=wait_for")
.header("content-type", "application/x-ndjson")
.body(body)
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
if !response.status().is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
return Err(provider_write_error(response.status().as_u16()));
}
let value: Value = response
.json()
.await
.map_err(|error| RuntimeError::invalid_state(format!("invalid provider bulk response: {error}")))?;
if value.get("errors").and_then(Value::as_bool) == Some(true) {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
validate_bulk_response(&value, changes.len())
}
pub(super) async fn gc_document_history(
&self,
physical_table: &str,
workspace_id: &str,
doc_id: &str,
source_version: i64,
permission_version: i64,
limit: usize,
) -> RuntimeResult<()> {
let query = json!({"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"term":{"doc_id":{"value":doc_id}}}
],"must_not":[{"bool":{"must":[
{"term":{"source_version":{"value":source_version}}},
{"term":{"permission_version":{"value":permission_version}}}
]}}]}});
let response = self
.request(
reqwest::Method::POST,
&format!("{physical_table}/_delete_by_query?conflicts=proceed&refresh=true"),
)
.json(&json!({"query":query,"max_docs":limit.max(1)}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
Ok(())
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid provider history GC response", error))?;
validate_delete_response(&value)
}
pub(super) async fn gc_workspace(
&self,
physical_table: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> RuntimeResult<bool> {
let limit = limit.max(1);
let response = self
.request(
reqwest::Method::POST,
&format!("{physical_table}/_delete_by_query?conflicts=proceed&refresh=true"),
)
.json(&json!({
"query":{"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"range":{"source_version":{"lte":source_version_high_water}}}
]}},
"max_docs":limit
}))
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let body = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
let value: Value = serde_json::from_slice(&body)
.map_err(|error| RuntimeError::json("invalid provider workspace GC response", error))?;
validate_delete_response(&value)?;
let deleted = value
.get("deleted")
.and_then(Value::as_u64)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider workspace GC response"))?;
Ok(deleted >= limit as u64)
}
fn request(&self, method: reqwest::Method, path: &str) -> reqwest::RequestBuilder {
@@ -257,48 +262,61 @@ impl RemoteProvider {
}
request
}
}
async fn resolve_manticore_tokens(
&self,
tokens: impl IntoIterator<Item = String>,
) -> RuntimeResult<std::collections::HashMap<String, i64>> {
let tokens = tokens.into_iter().collect::<std::collections::BTreeSet<_>>();
if tokens.is_empty() {
return Ok(Default::default());
}
let tokens = tokens.into_iter().collect::<Vec<_>>();
let mut transaction = self
.pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin Manticore exact token resolution", error))?;
sqlx::query(
r#"INSERT INTO search_runtime_acl_tokens(token)
SELECT candidate.token FROM unnest($1::text[]) candidate(token)
LEFT JOIN search_runtime_acl_tokens existing USING(token)
WHERE existing.token IS NULL ON CONFLICT DO NOTHING"#,
)
.bind(&tokens)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("allocate Manticore exact token IDs", error))?;
let rows: Vec<(String, i64)> =
sqlx::query_as("SELECT token,token_id FROM search_runtime_acl_tokens WHERE token=ANY($1)")
.bind(&tokens)
.fetch_all(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load Manticore exact token IDs", error))?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit Manticore exact token resolution", error))?;
if rows.len() != tokens.len() {
return Err(RuntimeError::invalid_state(
"Manticore exact token mapping is incomplete",
));
}
Ok(rows.into_iter().collect())
fn validate_delete_response(value: &Value) -> RuntimeResult<()> {
if value.get("timed_out").and_then(Value::as_bool) == Some(true)
|| value
.get("failures")
.and_then(Value::as_array)
.is_some_and(|failures| !failures.is_empty())
{
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(())
}
fn validate_bulk_response(value: &Value, expected_items: usize) -> RuntimeResult<()> {
let items = value
.get("items")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider bulk response"))?;
if items.len() != expected_items {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
for item in items {
let result = item
.as_object()
.and_then(|item| item.values().next())
.and_then(Value::as_object)
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
let status = result
.get("status")
.and_then(Value::as_u64)
.and_then(|status| u16::try_from(status).ok())
.ok_or_else(|| RuntimeError::invalid_state("provider_apply_failed"))?;
if result.get("error").is_some() || !(200..300).contains(&status) {
return Err(provider_write_error(status));
}
}
if value.get("errors").and_then(Value::as_bool) != Some(false) {
return Err(RuntimeError::invalid_state("provider_apply_failed"));
}
Ok(())
}
fn ensure_projection_source(dsl: &mut Value) {
let mut source = dsl
.get("_source")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
for field in ["doc_id", "source_version", "permission_version"] {
if !source.iter().any(|value| value.as_str() == Some(field)) {
source.push(json!(field));
}
}
dsl["_source"] = json!(source);
}
async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>> {
@@ -322,13 +340,7 @@ async fn read_response(mut response: reqwest::Response) -> RuntimeResult<Vec<u8>
Ok(bytes)
}
fn normalize(
value: Value,
manticore: bool,
_offset: u64,
_size: u64,
requested_fields: &[String],
) -> RuntimeResult<Value> {
fn normalize(value: Value) -> RuntimeResult<Value> {
let hits = value
.pointer("/hits/hits")
.and_then(Value::as_array)
@@ -341,11 +353,7 @@ fn normalize(
let nodes = hits
.iter()
.map(|hit| {
let fields = if manticore {
manticore_fields(hit.get("_source"), requested_fields)
} else {
hit.get("fields").cloned().unwrap_or_else(|| json!({}))
};
let fields = hit.get("fields").cloned().unwrap_or_else(|| json!({}));
json!({
"id":hit.get("_id").and_then(Value::as_str).unwrap_or_default(),
"score":hit.get("_score").and_then(Value::as_f64).unwrap_or_default(),
@@ -355,23 +363,16 @@ fn normalize(
})
})
.collect::<Vec<_>>();
let next_cursor = if manticore {
(!hits.is_empty())
.then(|| serde_json::to_string(&json!({"offset":_offset + _size})))
.transpose()
.map_err(|error| RuntimeError::json("encode provider cursor", error))?
} else {
hits
.last()
.and_then(|hit| hit.get("sort"))
.map(serde_json::to_string)
.transpose()
.map_err(|error| RuntimeError::json("encode provider cursor", error))?
};
let next_cursor = hits
.last()
.and_then(|hit| hit.get("sort"))
.map(serde_json::to_string)
.transpose()
.map_err(|error| RuntimeError::json("encode provider cursor", error))?;
Ok(json!({"total":total,"nodes":nodes,"nextCursor":next_cursor}))
}
fn normalize_aggregate(value: Value) -> RuntimeResult<Value> {
fn normalize_aggregate(value: Value, skip: usize, limit: usize) -> RuntimeResult<Value> {
let buckets = value
.pointer("/aggregations/result/buckets")
.and_then(Value::as_array)
@@ -386,12 +387,14 @@ fn normalize_aggregate(value: Value) -> RuntimeResult<Value> {
Ok(json!({
"key":bucket.get("key").cloned().unwrap_or(Value::Null),
"count":bucket.get("doc_count").cloned().unwrap_or(json!(0)),
"hits":{"total":bucket.pointer("/result/hits/total/value").or_else(||bucket.pointer("/result/hits/total")).cloned().unwrap_or(json!(0)),
"nodes":hits.iter().map(normalize_hit).collect::<Vec<_>>()}
"hits":hits.iter().map(normalize_hit).collect::<Vec<_>>()
}))
})
.collect::<RuntimeResult<Vec<_>>>()?;
Ok(json!({"total":nodes.len(),"hasMore":false,"buckets":nodes}))
let total = nodes.len();
let has_more = total > skip.saturating_add(limit);
let nodes = nodes.into_iter().skip(skip).take(limit).collect::<Vec<_>>();
Ok(json!({"total":total,"hasMore":has_more,"buckets":nodes}))
}
fn normalize_hit(hit: &Value) -> Value {
@@ -403,3 +406,76 @@ fn normalize_hit(hit: &Value) -> Value {
"_source":hit.get("_source").cloned().unwrap_or_else(||json!({})),
})
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{normalize_aggregate, validate_bulk_response, validate_delete_response};
use crate::runtime::RuntimeError;
#[test]
fn delete_response_rejects_partial_provider_failures() {
assert!(validate_delete_response(&json!({"failures":[]})).is_ok());
assert!(validate_delete_response(&json!({"timed_out":true})).is_err());
assert!(validate_delete_response(&json!({"failures":[{"reason":"stale"}]})).is_err());
}
#[test]
fn bulk_response_requires_every_item_to_succeed() {
for (response, expected_items, valid) in [
(
json!({"errors":false,"items":[{"index":{"status":201}},{"index":{"status":200}}]}),
2,
true,
),
(json!({"errors":false,"items":[{"index":{"status":404}}]}), 1, false),
(
json!({"errors":true,"items":[{"index":{"status":201}},{"index":{"status":400,"error":{}}}]}),
2,
false,
),
(json!({"errors":false}), 1, false),
] {
assert_eq!(validate_bulk_response(&response, expected_items).is_ok(), valid);
}
assert!(matches!(
validate_bulk_response(&json!({"errors":true,"items":[{"index":{"status":400,"error":{}}}]}), 1),
Err(RuntimeError::SearchSourceInvalid(_))
));
assert!(matches!(
validate_bulk_response(&json!({"errors":true,"items":[{"index":{"status":429,"error":{}}}]}), 1),
Err(RuntimeError::SearchProviderUnavailable)
));
}
#[test]
fn aggregate_normalization_reports_extra_buckets() {
let value = json!({
"aggregations":{"result":{"buckets":[
{"key":"one","doc_count":2,"result":{"hits":{"hits":[]}}},
{"key":"two","doc_count":1,"result":{"hits":{"hits":[]}}}
]}}
});
let result = normalize_aggregate(value, 0, 1).unwrap();
assert_eq!(result["total"], 2);
assert_eq!(result["hasMore"], true);
assert_eq!(result["buckets"].as_array().unwrap().len(), 1);
assert!(result["buckets"][0]["hits"].is_array());
}
#[test]
fn aggregate_normalization_applies_bucket_skip_before_limit() {
let value = json!({
"aggregations":{"result":{"buckets":[
{"key":"one","doc_count":2,"result":{"hits":{"hits":[]}}},
{"key":"two","doc_count":1,"result":{"hits":{"hits":[]}}},
{"key":"three","doc_count":1,"result":{"hits":{"hits":[]}}}
]}}
});
let result = normalize_aggregate(value, 1, 1).unwrap();
assert_eq!(result["total"], 3);
assert_eq!(result["hasMore"], true);
assert_eq!(result["buckets"][0]["key"], "two");
}
}
@@ -1,10 +1,105 @@
use serde::Deserialize;
use serde_json::{Value, json};
use super::types::{AggregateRequest, SearchOptions, SearchQuery, SearchRequest, SearchTable};
use crate::runtime::{
RuntimeError, RuntimeResult,
backend_runtime::permission::{AuthorizedSearchScope, DocReadScope},
use super::{
AggregateOptions, AuthorizedSearchScope, DocReadScope, RuntimeAggregateRequest, RuntimeSearchQuery,
RuntimeSearchRequest, SearchOptions, SearchTable,
};
use crate::runtime::{RuntimeError, RuntimeResult};
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct AggregateRequest {
pub(super) table: SearchTable,
query: SearchQuery,
field: String,
options: AggregateOptions,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchRequest {
pub(super) table: SearchTable,
query: SearchQuery,
options: SearchOptions,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchQuery {
#[serde(rename = "type")]
query_type: String,
field: Option<String>,
#[serde(rename = "match")]
match_value: Option<String>,
query: Option<Box<SearchQuery>>,
queries: Option<Vec<SearchQuery>>,
occur: Option<String>,
boost: Option<f64>,
}
impl RuntimeSearchRequest {
pub(super) fn into_search_request(self) -> RuntimeResult<SearchRequest> {
let mut decoded_nodes = 0;
Ok(SearchRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
options: self.options,
})
}
}
impl RuntimeAggregateRequest {
pub(super) fn into_aggregate_request(self) -> RuntimeResult<AggregateRequest> {
let mut decoded_nodes = 0;
Ok(AggregateRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
field: self.field,
options: self.options,
})
}
}
const MAX_QUERY_GRAPH_NODES: usize = 100;
const MAX_QUERY_DEPTH: usize = 100;
const MAX_DECODED_QUERY_NODES: usize = 1_000;
fn decode_query(
nodes: &[RuntimeSearchQuery],
index: u32,
depth: usize,
decoded_nodes: &mut usize,
) -> RuntimeResult<SearchQuery> {
if nodes.len() > MAX_QUERY_GRAPH_NODES || depth > MAX_QUERY_DEPTH || *decoded_nodes >= MAX_DECODED_QUERY_NODES {
return Err(RuntimeError::invalid_input("search query is too complex"));
}
*decoded_nodes += 1;
let node = nodes
.get(index as usize)
.ok_or_else(|| RuntimeError::invalid_input("invalid search query node"))?;
Ok(SearchQuery {
query_type: node.query_type.clone(),
field: node.field.clone(),
match_value: node.match_value.clone(),
query: node
.query
.map(|index| decode_query(nodes, index, depth + 1, decoded_nodes).map(Box::new))
.transpose()?,
queries: node
.queries
.as_ref()
.map(|indices| {
indices
.iter()
.map(|index| decode_query(nodes, *index, depth + 1, decoded_nodes))
.collect()
})
.transpose()?,
occur: node.occur.clone(),
boost: node.boost,
})
}
pub(super) fn compile(request: &SearchRequest, scope: &AuthorizedSearchScope) -> RuntimeResult<Value> {
let query = compile_query(request.table, &request.query)?;
@@ -26,7 +121,7 @@ pub(super) fn compile(request: &SearchRequest, scope: &AuthorizedSearchScope) ->
.map(|field| validate_field(request.table, field).map(str::to_string))
.collect::<RuntimeResult<Vec<_>>>()?;
let mut dsl = json!({
"_source":["workspace_id","doc_id"],
"_source":["workspace_id","doc_id","source_version","permission_version"],
"fields":fields,
"query":{"bool":{"must":must}},
"sort": stable_sort(request.table),
@@ -73,15 +168,19 @@ pub(super) fn compile_aggregate(request: &AggregateRequest, scope: &AuthorizedSe
let hit_dsl = compile(&search, scope)?;
let field = validate_field(request.table, &request.field)?;
let limit = request.options.pagination.limit.unwrap_or(10);
if limit > 10_000 {
return Err(RuntimeError::invalid_input("aggregate limit exceeds 10000"));
let skip = request.options.pagination.skip.unwrap_or(0);
if skip.saturating_add(limit) > 10_000 {
return Err(RuntimeError::invalid_input("aggregate pagination exceeds 10000"));
}
if request.options.pagination.cursor.is_some() {
return Err(RuntimeError::invalid_input("aggregate cursor is unsupported"));
}
Ok(json!({
"query":hit_dsl["query"],
"from":request.options.pagination.skip.unwrap_or(0),
"from":skip,
"size":0,
"aggs":{"result":{"terms":{"field":field,"size":limit},"aggs":{"result":{"top_hits":{
"size":hit_dsl["size"],"_source":hit_dsl["_source"],"fields":hit_dsl["fields"],
"size":hit_dsl["size"],"from":hit_dsl.get("from").cloned().unwrap_or_else(||json!(0)),"_source":hit_dsl["_source"],"fields":hit_dsl["fields"],
"sort":hit_dsl["sort"],"highlight":hit_dsl.get("highlight").cloned().unwrap_or_else(||json!({}))
}}}}}
}))
@@ -234,7 +333,6 @@ mod tests {
fn scope() -> AuthorizedSearchScope {
AuthorizedSearchScope {
workspace_id: "workspace".to_string(),
permission_revision: 1,
docs: DocReadScope::All,
}
}
@@ -325,6 +423,10 @@ mod tests {
dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["highlight"]["fields"]["content"],
json!({"pre_tags":["<b>"],"post_tags":["</b>"]})
);
let mut hit_skip = aggregate.clone();
hit_skip.options.hits.pagination.skip = Some(1);
let hit_skip_dsl = compile_aggregate(&hit_skip, &scope()).unwrap();
assert_eq!(hit_skip_dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["from"], 1);
let mut invalid = aggregate;
invalid.field = "aclReadTokens".to_string();
@@ -332,5 +434,43 @@ mod tests {
invalid.field = "docId".to_string();
invalid.options.pagination.limit = Some(10_001);
assert!(compile_aggregate(&invalid, &scope()).is_err());
invalid.options.pagination.limit = Some(10_000);
invalid.options.pagination.skip = Some(1);
assert!(compile_aggregate(&invalid, &scope()).is_err());
invalid.options.pagination.limit = Some(10);
invalid.options.pagination.skip = None;
invalid.options.pagination.cursor = Some("candidate".to_string());
assert!(compile_aggregate(&invalid, &scope()).is_err());
}
fn runtime_node(query_type: &str) -> RuntimeSearchQuery {
RuntimeSearchQuery {
query_type: query_type.to_string(),
field: None,
match_value: None,
query: None,
queries: None,
occur: None,
boost: None,
}
}
#[test]
fn rejects_invalid_or_overly_complex_query_graphs() {
assert!(decode_query(&[runtime_node("all")], 1, 0, &mut 0).is_err());
let mut oversized = (0..101).map(|_| runtime_node("all")).collect::<Vec<_>>();
oversized[0].query = Some(1);
assert!(decode_query(&oversized, 0, 0, &mut 0).is_err());
let mut recursive = vec![runtime_node("boost")];
recursive[0].query = Some(0);
assert!(decode_query(&recursive, 0, 0, &mut 0).is_err());
let mut shared_child = (0..100).map(|_| runtime_node("boolean")).collect::<Vec<_>>();
for (index, node) in shared_child.iter_mut().enumerate().take(99) {
node.queries = Some(vec![(index + 1) as u32, (index + 1) as u32]);
}
assert!(decode_query(&shared_child, 0, 0, &mut 0).is_err());
}
}
@@ -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 crate::runtime::{RuntimeError, RuntimeResult};
#[napi_derive::napi(string_enum = "snake_case")]
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
@@ -26,20 +24,6 @@ impl SearchTable {
}
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchQuery {
#[serde(rename = "type")]
pub(super) query_type: String,
pub(super) field: Option<String>,
#[serde(rename = "match")]
pub(super) match_value: Option<String>,
pub(super) query: Option<Box<SearchQuery>>,
pub(super) queries: Option<Vec<SearchQuery>>,
pub(super) occur: Option<String>,
pub(super) boost: Option<f64>,
}
#[napi_derive::napi(object)]
#[derive(Clone, Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -89,23 +73,6 @@ pub struct AggregateOptions {
pub pagination: SearchPagination,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct AggregateRequest {
pub(super) table: SearchTable,
pub(super) query: SearchQuery,
pub(super) field: String,
pub(super) options: AggregateOptions,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct SearchRequest {
pub(super) table: SearchTable,
pub(super) query: SearchQuery,
pub(super) options: SearchOptions,
}
#[napi_derive::napi(object)]
pub struct RuntimeSearchQuery {
pub query_type: String,
@@ -133,109 +100,3 @@ pub struct RuntimeAggregateRequest {
pub field: String,
pub options: AggregateOptions,
}
#[cfg(test)]
impl SearchRequest {
pub(super) fn parse(value: serde_json::Value) -> RuntimeResult<Self> {
serde_json::from_value(value).map_err(|error| RuntimeError::json("invalid search request", error))
}
}
impl RuntimeSearchRequest {
pub(super) fn into_search_request(self) -> RuntimeResult<SearchRequest> {
let mut decoded_nodes = 0;
Ok(SearchRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
options: self.options,
})
}
}
impl RuntimeAggregateRequest {
pub(super) fn into_aggregate_request(self) -> RuntimeResult<AggregateRequest> {
let mut decoded_nodes = 0;
Ok(AggregateRequest {
table: self.table,
query: decode_query(&self.queries, self.root_query, 0, &mut decoded_nodes)?,
field: self.field,
options: self.options,
})
}
}
const MAX_QUERY_GRAPH_NODES: usize = 100;
const MAX_QUERY_DEPTH: usize = 100;
const MAX_DECODED_QUERY_NODES: usize = 1_000;
fn decode_query(
nodes: &[RuntimeSearchQuery],
index: u32,
depth: usize,
decoded_nodes: &mut usize,
) -> RuntimeResult<SearchQuery> {
if nodes.len() > MAX_QUERY_GRAPH_NODES || depth > MAX_QUERY_DEPTH || *decoded_nodes >= MAX_DECODED_QUERY_NODES {
return Err(RuntimeError::invalid_input("search query is too complex"));
}
*decoded_nodes += 1;
let node = nodes
.get(index as usize)
.ok_or_else(|| RuntimeError::invalid_input("invalid search query node"))?;
Ok(SearchQuery {
query_type: node.query_type.clone(),
field: node.field.clone(),
match_value: node.match_value.clone(),
query: node
.query
.map(|index| decode_query(nodes, index, depth + 1, decoded_nodes).map(Box::new))
.transpose()?,
queries: node
.queries
.as_ref()
.map(|indices| {
indices
.iter()
.map(|index| decode_query(nodes, *index, depth + 1, decoded_nodes))
.collect()
})
.transpose()?,
occur: node.occur.clone(),
boost: node.boost,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn node(query_type: &str) -> RuntimeSearchQuery {
RuntimeSearchQuery {
query_type: query_type.to_string(),
field: None,
match_value: None,
query: None,
queries: None,
occur: None,
boost: None,
}
}
#[test]
fn rejects_invalid_or_overly_complex_query_graphs() {
assert!(decode_query(&[node("all")], 1, 0, &mut 0).is_err());
let mut oversized = (0..101).map(|_| node("all")).collect::<Vec<_>>();
oversized[0].query = Some(1);
assert!(decode_query(&oversized, 0, 0, &mut 0).is_err());
let mut recursive = vec![node("boost")];
recursive[0].query = Some(0);
assert!(decode_query(&recursive, 0, 0, &mut 0).is_err());
let mut shared_child = (0..100).map(|_| node("boolean")).collect::<Vec<_>>();
for (index, node) in shared_child.iter_mut().enumerate().take(99) {
node.queries = Some(vec![(index + 1) as u32, (index + 1) as u32]);
}
assert!(decode_query(&shared_child, 0, 0, &mut 0).is_err());
}
}
@@ -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();
assert!(enabled_without_provider.enabled);
assert_eq!(enabled_without_provider.provider, "embedded");
let manticore = app_config_from_flat_overrides([
("indexer.enabled", serde_json::json!(true)),
("indexer.provider.type", serde_json::json!("manticoresearch")),
("indexer.provider.endpoint", serde_json::json!("http://localhost:9308")),
])
.unwrap();
let manticore: SearchRuntimeConfig = manticore.indexer.unwrap().into();
assert!(manticore.enabled);
assert_eq!(manticore.provider, "manticoresearch");
}
#[test]
+20 -5
View File
@@ -21,15 +21,24 @@ pub(crate) enum RuntimeError {
#[error("search permission state unavailable")]
SearchPermissionUnavailable,
#[error("search index is not ready")]
SearchIndexNotReady,
#[error("search permission projection is syncing")]
SearchPermissionSyncing,
#[error("search index failed: {0}")]
SearchIndexFailed(String),
#[error("search source is invalid: {0}")]
SearchSourceInvalid(String),
#[error("search provider unavailable")]
SearchProviderUnavailable,
#[error("search query is not supported by the active provider")]
SearchUnsupportedQuery,
#[error("search stream replay gap")]
SearchReplayGap,
#[error("{context}: {source}")]
Database {
context: String,
@@ -111,9 +120,11 @@ impl RuntimeError {
}
Self::SearchWorkspaceDenied
| Self::SearchPermissionUnavailable
| Self::SearchIndexNotReady
| Self::SearchPermissionSyncing
| Self::SearchProviderUnavailable
| Self::SearchUnsupportedQuery
| Self::SearchReplayGap => false,
| Self::SearchUnsupportedQuery => false,
Self::SearchIndexFailed(_) | Self::SearchSourceInvalid(_) => false,
_ => false,
}
}
@@ -127,6 +138,10 @@ impl RuntimeError {
} if source.code().as_deref() == Some("40001")
)
}
pub(crate) fn is_permanent_search_source(&self) -> bool {
matches!(self, Self::SearchSourceInvalid(_))
}
}
pub(crate) fn to_napi_error(error: RuntimeError) -> Error {
+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");
const EMBEDDING_MIGRATION: &str = include_str!("sql/embedding.sql");
const SEARCH_MIGRATION: &str = include_str!("sql/search.sql");
const SEARCH_ACL_TOKEN_MIGRATION: &str = include_str!("sql/search_acl_tokens.sql");
const SEARCH_PROJECTION_MIGRATION: &str = include_str!("sql/search_projection.sql");
const EMBEDDING_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0046;
const SEARCH_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0053;
#[cfg(test)]
@@ -70,15 +69,11 @@ pub(crate) async fn embedding_schema_health(pool: &PgPool) -> RuntimeResult<Embe
pub(crate) async fn migrate_search_tables(pool: &PgPool) -> RuntimeResult<()> {
migrate_component(
pool,
"search",
"search_projection",
SEARCH_ADVISORY_LOCK,
&[(1, &[SEARCH_MIGRATION]), (2, &[SEARCH_ACL_TOKEN_MIGRATION])],
&[(1, &[SEARCH_PROJECTION_MIGRATION])],
)
.await?;
sqlx::query("INSERT INTO search_runtime_streams(table_key) VALUES ('doc'), ('block') ON CONFLICT DO NOTHING")
.execute(pool)
.await
.map_err(|error| RuntimeError::database("Repair search runtime streams", error))?;
Ok(())
}
@@ -118,6 +113,10 @@ async fn migrate_component(
.begin()
.await
.map_err(|error| RuntimeError::database("Native migration transaction failed", error))?;
transaction
.execute("SET LOCAL lock_timeout = '5s'")
.await
.map_err(|error| RuntimeError::database("Native migration lock timeout failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(advisory_lock)
.execute(&mut *transaction)
@@ -269,11 +268,31 @@ mod tests {
}
#[test]
fn search_schema_uses_transactional_stream_heads() {
assert!(SEARCH_MIGRATION.contains("CREATE TABLE search_runtime_streams"));
assert!(SEARCH_MIGRATION.contains("PRIMARY KEY (table_key, stream_sequence)"));
assert!(!SEARCH_MIGRATION.contains("BIGSERIAL"));
assert!(!SEARCH_MIGRATION.contains("CREATE SEQUENCE"));
fn search_schema_uses_terminal_control_plane() {
assert!(SEARCH_PROJECTION_MIGRATION.contains("DROP SCHEMA IF EXISTS search_projection CASCADE"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SCHEMA search_projection"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.generations"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.workspace_states"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_projection.document_states"));
assert!(
SEARCH_PROJECTION_MIGRATION.contains("CREATE SEQUENCE IF NOT EXISTS search_projection.source_mutation_version")
);
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SEQUENCE IF NOT EXISTS search_projection.permission_version"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE SEQUENCE IF NOT EXISTS search_projection.claim_fence"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TRIGGER search_projection_snapshot_capture"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("CREATE TRIGGER search_projection_membership_capture"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("ON CONFLICT (generation_id, workspace_id, doc_id) DO UPDATE"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("WHERE state IN ('building', 'active')"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("target_source_version <> published_source_version"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("gc_table TEXT NOT NULL DEFAULT 'doc'"));
assert!(SEARCH_PROJECTION_MIGRATION.contains("gc_cursor TEXT"));
assert!(
SEARCH_PROJECTION_MIGRATION
.contains("NEW.state = 'active' AND (TG_OP = 'INSERT' OR OLD.state IS DISTINCT FROM NEW.state)")
);
assert!(!SEARCH_PROJECTION_MIGRATION.contains("clock_timestamp()"));
assert!(!SEARCH_PROJECTION_MIGRATION.contains("CREATE TABLE search_runtime_projections"));
assert!(!SEARCH_PROJECTION_MIGRATION.contains("payload JSONB NOT NULL"));
}
#[tokio::test]
@@ -285,11 +304,367 @@ mod tests {
let (first, second) = tokio::join!(migrate_search_tables(&pool), migrate_search_tables(&pool));
first.unwrap();
second.unwrap();
let versions: Vec<i32> =
sqlx::query_scalar("SELECT version FROM native_schema_migrations WHERE component='search' ORDER BY version")
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(versions, vec![1, 2]);
let versions: Vec<i32> = sqlx::query_scalar(
"SELECT version FROM native_schema_migrations WHERE component='search_projection' ORDER BY version",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(versions, vec![1]);
}
#[tokio::test]
async fn permission_triggers_route_and_coalesce_document_and_workspace_scopes() {
let _guard = crate::runtime::backend_runtime::SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await
.unwrap();
let suffix = uuid::Uuid::new_v4().simple().to_string();
let generation_id = uuid::Uuid::new_v4();
let workspace_id = format!("trigger-workspace-{suffix}");
let doc_id = format!("trigger-doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO search_projection.workspace_states(generation_id,workspace_id,covered,pending_scope) \
VALUES($1,$2,true,'none')",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('00','hex'),now())")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO doc_access_policies(workspace_id,doc_id,visibility,member_default_role) \
VALUES($1,$2,'private','none')",
)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("UPDATE doc_access_policies SET member_default_role='reader' WHERE workspace_id=$1 AND doc_id=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let document_scope: (String, i64) = sqlx::query_as(
r#"SELECT state.pending_scope,count(task.doc_id)
FROM search_projection.workspace_states state
LEFT JOIN search_projection.document_states task
ON task.generation_id=state.generation_id AND task.workspace_id=state.workspace_id
WHERE state.generation_id=$1 AND state.workspace_id=$2
GROUP BY state.pending_scope"#,
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(document_scope, ("none".to_string(), 1));
let moved_workspace_id = format!("trigger-moved-workspace-{suffix}");
let moved_doc_id = format!("trigger-moved-doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&moved_workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('00','hex'),now())")
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.execute(&pool)
.await
.unwrap();
let old_version: i64 = sqlx::query_scalar(
"SELECT target_permission_version FROM search_projection.document_states WHERE generation_id=$1 AND \
workspace_id=$2 AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
let moved_version: i64 = sqlx::query_scalar(
"SELECT target_permission_version FROM search_projection.document_states WHERE generation_id=$1 AND \
workspace_id=$2 AND doc_id=$3",
)
.bind(generation_id)
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.fetch_one(&pool)
.await
.unwrap();
sqlx::query("UPDATE doc_access_policies SET workspace_id=$3,doc_id=$4 WHERE workspace_id=$1 AND doc_id=$2")
.bind(&workspace_id)
.bind(&doc_id)
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.execute(&pool)
.await
.unwrap();
let versions: (i64, i64) = sqlx::query_as(
r#"SELECT
(SELECT target_permission_version FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3),
(SELECT target_permission_version FROM search_projection.document_states
WHERE generation_id=$1 AND workspace_id=$4 AND doc_id=$5)"#,
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.bind(&moved_workspace_id)
.bind(&moved_doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(versions.0 > old_version);
assert!(versions.1 > moved_version);
sqlx::query("INSERT INTO workspace_access_policies(workspace_id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let pending_scope: String = sqlx::query_scalar(
"SELECT pending_scope FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2",
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(pending_scope, "workspace");
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&moved_workspace_id)
.execute(&pool)
.await
.unwrap();
}
#[tokio::test]
async fn source_mutations_use_unique_versions_and_rollback_only_leaves_a_sequence_gap() {
let _guard = crate::runtime::backend_runtime::SEARCH_TEST_LOCK.lock().await;
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
migrate_search_tables(&pool).await.unwrap();
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await
.unwrap();
let suffix = uuid::Uuid::new_v4().simple().to_string();
let generation_id = uuid::Uuid::new_v4();
let workspace_id = format!("source-version-workspace-{suffix}");
let doc_id = format!("source-version-doc-{suffix}");
sqlx::query("INSERT INTO workspaces(id) VALUES($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO search_projection.workspace_states(generation_id,workspace_id,pending_scope) VALUES($1,$2,'none')",
)
.bind(generation_id)
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('00','hex'),'2026-01-01 UTC')",
)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let inserted: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
sqlx::query("UPDATE snapshots SET blob=decode('01','hex') WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let updated: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(updated > inserted);
sqlx::query("DELETE FROM snapshots WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let (deleted, target_exists, published_exists): (i64, bool, bool) = sqlx::query_as(
"SELECT target_source_version,target_source_exists,published_source_exists FROM \
search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(deleted > updated);
assert!(!target_exists);
assert!(!published_exists);
sqlx::query(
"INSERT INTO snapshots(workspace_id,guid,blob,updated_at) VALUES($1,$2,decode('02','hex'),'2026-01-01 UTC')",
)
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let recreated: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(recreated > deleted);
let mut transaction = pool.begin().await.unwrap();
sqlx::query("UPDATE snapshots SET blob=decode('03','hex') WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&mut *transaction)
.await
.unwrap();
let rolled_back: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&mut *transaction)
.await
.unwrap();
transaction.rollback().await.unwrap();
let after_rollback: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(after_rollback, recreated);
assert!(rolled_back > recreated);
sqlx::query("UPDATE snapshots SET blob=decode('04','hex') WHERE workspace_id=$1 AND guid=$2")
.bind(&workspace_id)
.bind(&doc_id)
.execute(&pool)
.await
.unwrap();
let after_gap: i64 = sqlx::query_scalar(
"SELECT target_source_version FROM search_projection.document_states WHERE generation_id=$1 AND workspace_id=$2 \
AND doc_id=$3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_one(&pool)
.await
.unwrap();
assert!(after_gap > rolled_back);
sqlx::query("DELETE FROM workspaces WHERE id=$1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let (remaining_states, delete_gc_scheduled): (i64, bool) = sqlx::query_as(
r#"SELECT count(*),bool_and(progress->>'kind'='deleted')
FROM search_projection.workspace_states WHERE generation_id=$1 AND workspace_id=$2"#,
)
.bind(generation_id)
.bind(&workspace_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(remaining_states, 1);
assert!(delete_gc_scheduled);
sqlx::query("DELETE FROM search_projection.generations WHERE id=$1")
.bind(generation_id)
.execute(&pool)
.await
.unwrap();
}
}
@@ -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 super::{
CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupAckResult, RuntimeDocumentCleanupEffect,
RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult,
StorageRuntime, load_workspace_live_doc_ids, merge_current_doc, napi_error,
CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult,
RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult, StorageRuntime, load_workspace_live_doc_ids,
merge_current_doc, napi_error,
};
#[derive(FromRow)]
@@ -416,7 +416,6 @@ async fn delete_doc_rows(tx: &mut Transaction<'_, Postgres>, candidate: &Candida
"cleanupVersion": cleanup_version,
"commentAttachmentKeys": attachment_keys,
"commentObjectsDone": false,
"searchDone": false,
}))
.execute(&mut **tx)
.await
@@ -592,51 +591,43 @@ async fn execute_one(
Ok(Some((candidate, deleted_rows)))
}
fn payload_effect(effect: PendingEffect) -> RuntimeResult<RuntimeDocumentCleanupEffect> {
fn payload_effect(effect: &PendingEffect) -> RuntimeResult<RuntimeDocumentCleanupEffect> {
let cleanup_version = effect
.cleanup_payload
.get("cleanupVersion")
.and_then(Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("Document cleanup effect payload has no cleanupVersion"))?;
Ok(RuntimeDocumentCleanupEffect {
workspace_id: effect.workspace_id,
doc_id: effect.doc_id,
workspace_id: effect.workspace_id.clone(),
doc_id: effect.doc_id.clone(),
cleanup_version: cleanup_version.to_string(),
comment_objects_done: effect
.cleanup_payload
.get("commentObjectsDone")
.and_then(Value::as_bool)
.unwrap_or(false),
search_done: effect
.cleanup_payload
.get("searchDone")
.and_then(Value::as_bool)
.unwrap_or(false),
})
}
async fn complete_effect(
async fn complete_comment_objects(
tx: &mut Transaction<'_, Postgres>,
workspace_id: &str,
doc_id: &str,
cleanup_version: &str,
path: &str,
) -> RuntimeResult<bool> {
let completed = sqlx::query_scalar::<_, bool>(
r#"
UPDATE document_cleanup_candidates
SET cleanup_payload = jsonb_set(cleanup_payload, ARRAY[$4], 'true'),
SET cleanup_payload = jsonb_set(cleanup_payload, '{commentObjectsDone}', 'true'),
error = NULL, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending'
AND cleanup_payload->>'cleanupVersion' = $3
RETURNING COALESCE((cleanup_payload->>'commentObjectsDone')::boolean, false)
AND COALESCE((cleanup_payload->>'searchDone')::boolean, false)
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(cleanup_version)
.bind(path)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup effect completion failed", err))?;
@@ -658,7 +649,7 @@ async fn complete_effect(
Ok(completed)
}
async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffect) -> RuntimeResult<()> {
async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffect) -> RuntimeResult<bool> {
let cleanup_version = effect
.cleanup_payload
.get("cleanupVersion")
@@ -695,18 +686,11 @@ async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffec
.begin()
.await
.map_err(|err| RuntimeError::database("Document cleanup object effect transaction failed", err))?;
complete_effect(
&mut tx,
&effect.workspace_id,
&effect.doc_id,
cleanup_version,
"commentObjectsDone",
)
.await?;
let completed = complete_comment_objects(&mut tx, &effect.workspace_id, &effect.doc_id, cleanup_version).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("Document cleanup object effect commit failed", err))?;
Ok(())
Ok(completed)
}
async fn load_pending_effects(
@@ -789,75 +773,33 @@ impl StorageRuntime {
}
let effects = load_pending_effects(&pool, workspace_id.as_deref(), limit).await?;
for effect in effects {
if let Err(err) = process_comment_objects(self, &effect).await {
result.failed += 1;
sqlx::query(
r#"
UPDATE document_cleanup_candidates
SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2
"#,
)
.bind(&effect.workspace_id)
.bind(&effect.doc_id)
.bind(err.to_string())
.execute(&pool)
.await
.map_err(|db_err| RuntimeError::database("Document cleanup effect failure write failed", db_err))?;
let mut result_effect = payload_effect(&effect)?;
match process_comment_objects(self, &effect).await {
Ok(comment_objects_done) => {
result_effect.comment_objects_done = comment_objects_done;
result.effects.push(result_effect);
}
Err(err) => {
result.failed += 1;
result.effects.push(result_effect);
sqlx::query(
r#"
UPDATE document_cleanup_candidates
SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP
WHERE workspace_id = $1 AND doc_id = $2
"#,
)
.bind(&effect.workspace_id)
.bind(&effect.doc_id)
.bind(err.to_string())
.execute(&pool)
.await
.map_err(|db_err| RuntimeError::database("Document cleanup effect failure write failed", db_err))?;
}
}
}
result.effects = load_pending_effects(&pool, workspace_id.as_deref(), limit)
.await?
.into_iter()
.map(payload_effect)
.collect::<RuntimeResult<_>>()?;
Ok(result)
}
#[napi]
pub async fn ack_document_cleanup_effect(
&self,
workspace_id: String,
doc_id: String,
cleanup_version: String,
effect: String,
) -> napi::Result<RuntimeDocumentCleanupAckResult> {
let path = match effect.as_str() {
"search" => "searchDone",
_ => return Err(napi_error("document cleanup effect must be search")),
};
let pool = self.pool().await?;
let mut tx = pool
.begin()
.await
.map_err(|err| RuntimeError::database("Document cleanup ack transaction failed", err))?;
let current_version = sqlx::query_scalar::<_, String>(
r#"
SELECT cleanup_payload->>'cleanupVersion'
FROM document_cleanup_candidates
WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending'
"#,
)
.bind(&workspace_id)
.bind(&doc_id)
.fetch_optional(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup ack candidate load failed", err))?;
let Some(current_version) = current_version else {
tx.rollback()
.await
.map_err(|err| RuntimeError::database("Document cleanup duplicate ack rollback failed", err))?;
return Ok(RuntimeDocumentCleanupAckResult { completed: true });
};
if current_version != cleanup_version {
return Err(napi_error("document cleanup effect candidate version mismatch"));
}
let completed = complete_effect(&mut tx, &workspace_id, &doc_id, &cleanup_version, path).await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("Document cleanup ack commit failed", err))?;
Ok(RuntimeDocumentCleanupAckResult { completed })
}
}
#[cfg(test)]
@@ -870,7 +812,11 @@ mod tests {
use tokio::sync::Mutex;
use super::*;
use crate::runtime::{migrations::migrate_runtime_tables, storage_runtime::StorageRuntimeConfig};
use crate::runtime::{
backend_runtime::SEARCH_TEST_LOCK,
migrations::{migrate_runtime_tables, migrate_search_tables},
storage_runtime::StorageRuntimeConfig,
};
async fn runtime_from_database_url() -> AnyResult<Option<(StorageRuntime, PgPool)>> {
let Ok(database_url) = std::env::var("DATABASE_URL") else {
@@ -1233,10 +1179,17 @@ mod tests {
#[tokio::test]
async fn document_cleanup_execute_postgres_semantics() -> AnyResult<()> {
let _search_guard = SEARCH_TEST_LOCK.lock().await;
let Some((runtime, pool)) = runtime_from_database_url().await? else {
eprintln!("skipping postgres integration test: DATABASE_URL is not set");
return Ok(());
};
migrate_search_tables(&pool)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
sqlx::query("DELETE FROM search_projection.generations WHERE state='building'")
.execute(&pool)
.await?;
let suffix = Uuid::new_v4().to_string();
let (user_id, workspace_id) = insert_user_workspace(&pool, &suffix).await?;
let object_root = tempfile::tempdir()?;
@@ -1462,6 +1415,14 @@ mod tests {
.bind(doc_id)
.execute(&pool)
.await?;
let generation_id = Uuid::new_v4();
sqlx::query(
r#"INSERT INTO search_projection.generations(id,provider,state,config_hash,schema_version)
VALUES($1,'embedded','building',decode(repeat('00',32),'hex'),1)"#,
)
.bind(generation_id)
.execute(&pool)
.await?;
let executed = runtime
.execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10)
.await
@@ -1470,7 +1431,16 @@ mod tests {
assert_eq!(executed.failed, 0);
assert_eq!(executed.effects.len(), 1);
assert!(executed.effects[0].comment_objects_done);
assert!(!executed.effects[0].search_done);
let search_delete_tasks = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM search_projection.document_states WHERE generation_id = $1 AND workspace_id = $2 AND \
doc_id = $3",
)
.bind(generation_id)
.bind(&workspace_id)
.bind(doc_id)
.fetch_one(&pool)
.await?;
assert_eq!(search_delete_tasks, 1);
assert!(
runtime
.head_object("blob".to_string(), attachment_object_key)
@@ -1518,29 +1488,6 @@ mod tests {
assert_eq!(count, 0, "{table}.{column} should be nulled");
}
let effect = &executed.effects[0];
let search = runtime
.ack_document_cleanup_effect(
workspace_id.clone(),
doc_id.to_string(),
effect.cleanup_version.clone(),
"search".to_string(),
)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert!(search.completed);
// the copilot effect was removed; acknowledging it must be rejected
assert!(
runtime
.ack_document_cleanup_effect(
workspace_id.clone(),
doc_id.to_string(),
effect.cleanup_version.clone(),
"copilot".to_string(),
)
.await
.is_err()
);
let candidate_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
)
@@ -1588,8 +1535,7 @@ mod tests {
assert_eq!(failed_object_delete.failed, 1);
assert!(!failed_object_delete.effects[0].comment_objects_done);
let retained = sqlx::query(
"SELECT status, attempt_count, error, cleanup_payload->>'cleanupVersion' AS cleanup_version FROM \
document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
"SELECT status, attempt_count, error FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
)
.bind(&workspace_id)
.bind(retry_doc_id)
@@ -1598,18 +1544,6 @@ mod tests {
assert_eq!(retained.get::<String, _>("status"), "effects_pending");
assert_eq!(retained.get::<i32, _>("attempt_count"), 1);
assert!(retained.get::<Option<String>, _>("error").is_some());
let retry_cleanup_version = retained.get::<String, _>("cleanup_version");
let ack = runtime
.ack_document_cleanup_effect(
workspace_id.clone(),
retry_doc_id.to_string(),
retry_cleanup_version.clone(),
"search".to_string(),
)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert!(!ack.completed);
sqlx::query(
"UPDATE document_cleanup_candidates SET cleanup_payload = jsonb_set(cleanup_payload, '{commentAttachmentKeys}', \
@@ -1625,7 +1559,8 @@ mod tests {
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(retried.executed, 0);
assert_eq!(retried.failed, 0);
assert!(retried.effects.is_empty());
assert_eq!(retried.effects.len(), 1);
assert!(retried.effects[0].comment_objects_done);
let retry_candidate_count = sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2",
)
@@ -1635,6 +1570,10 @@ mod tests {
.await?;
assert_eq!(retry_candidate_count, 0);
sqlx::query("DELETE FROM search_projection.generations WHERE id = $1")
.bind(generation_id)
.execute(&pool)
.await?;
cleanup_workspace_fixture(&pool, &user_id, &workspace_id).await?;
Ok(())
}
@@ -32,10 +32,10 @@ pub(super) use super::{
napi_error, to_napi_error,
types::{
RuntimeBlobCleanupExecuteResult, RuntimeBlobCleanupPlanResult, RuntimeBlobCleanupResult, RuntimeBlobCompleteResult,
RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupAckResult,
RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult,
RuntimeMultipartUploadInit, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry,
RuntimeObjectMetadata, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupEffect,
RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeMultipartUploadInit,
RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry, RuntimeObjectMetadata,
RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest,
},
};
@@ -476,7 +476,6 @@ pub struct RuntimeDocumentCleanupEffect {
pub doc_id: String,
pub cleanup_version: String,
pub comment_objects_done: bool,
pub search_done: bool,
}
#[napi_derive::napi(object)]
@@ -491,11 +490,6 @@ pub struct RuntimeDocumentCleanupExecuteResult {
pub effects: Vec<RuntimeDocumentCleanupEffect>,
}
#[napi_derive::napi(object)]
pub struct RuntimeDocumentCleanupAckResult {
pub completed: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeBlobCleanupPlanResult {
pub run_id: Option<String>,
@@ -527,34 +521,6 @@ pub struct RuntimeDocCompactionResult {
pub history_created: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsRefreshResult {
pub processed: i64,
pub backlog: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsRecalibrationResult {
pub processed: i64,
pub last_sid: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsSnapshotResult {
pub snapshotted: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeWorkspaceStatsDailyRecalibrationResult {
pub processed: i64,
pub last_sid: i64,
pub snapshotted: i64,
pub skipped: bool,
}
#[napi_derive::napi(object)]
pub struct RuntimeEmbeddingProgress {
pub total: i64,
+375 -74
View File
@@ -3,12 +3,13 @@ mod query;
mod result;
mod schema;
use std::sync::Arc;
use std::{collections::HashMap, sync::Arc};
use memory_indexer::{MemoryIndex, Mutation, TermsAggregation};
use napi::{Status, bindgen_prelude::Buffer};
use serde_json::Value as JsonValue;
use napi::Status;
use serde_json::{Value as JsonValue, json};
use tokio::sync::RwLock;
use uuid::Uuid;
use self::{
document::compile_document,
@@ -49,12 +50,12 @@ impl TableIndex {
}
}
struct IndexManager {
struct GenerationIndex {
doc: TableIndex,
block: TableIndex,
}
impl IndexManager {
impl GenerationIndex {
fn new() -> Self {
Self {
doc: TableIndex::new(TableSchema::doc()),
@@ -71,9 +72,32 @@ impl IndexManager {
}
}
pub(crate) struct EmbeddedIndexCheckpoint {
pub sequence: i64,
pub data: Buffer,
struct IndexManager {
generations: RwLock<HashMap<Uuid, Arc<GenerationIndex>>>,
}
impl IndexManager {
fn new() -> Self {
Self {
generations: RwLock::new(HashMap::new()),
}
}
async fn generation(&self, generation_id: Uuid, create: bool) -> Option<Arc<GenerationIndex>> {
if let Some(generation) = self.generations.read().await.get(&generation_id).cloned() {
return Some(generation);
}
if !create {
return None;
}
let mut generations = self.generations.write().await;
Some(
generations
.entry(generation_id)
.or_insert_with(|| Arc::new(GenerationIndex::new()))
.clone(),
)
}
}
pub(crate) struct EmbeddedSearchIndex {
@@ -87,26 +111,52 @@ impl EmbeddedSearchIndex {
}
}
pub(crate) async fn restore(&self, table: String, checkpoint: Buffer) -> napi::Result<()> {
let table = self.manager.table(&table)?;
let index =
MemoryIndex::from_checkpoint(table.schema.schema.clone(), checkpoint.as_ref()).map_err(IndexError::from)?;
*table.index.write().await = index;
Ok(())
pub(crate) async fn prepare_generation(&self, generation_id: Uuid) {
self.manager.generation(generation_id, true).await;
}
pub(crate) async fn reset(&self, table: String) -> napi::Result<()> {
let table = self.manager.table(&table)?;
*table.index.write().await = MemoryIndex::new(table.schema.schema.clone());
Ok(())
pub(crate) async fn has_generation(&self, generation_id: Uuid) -> bool {
self.manager.generation(generation_id, false).await.is_some()
}
pub(crate) async fn retain_generation(&self, generation_id: Uuid) {
self
.manager
.generations
.write()
.await
.retain(|id, _| *id == generation_id);
}
#[cfg(test)]
pub(crate) async fn write(&self, table: String, documents_json: String) -> napi::Result<()> {
let table = self.manager.table(&table)?;
self.write_for_generation(Uuid::nil(), table, documents_json).await
}
pub(crate) async fn write_for_generation(
&self,
generation_id: Uuid,
table: String,
documents_json: String,
) -> napi::Result<()> {
let generation = self
.manager
.generation(generation_id, true)
.await
.expect("created embedded search generation");
let table = generation.table(&table)?;
let documents: Vec<JsonValue> = serde_json::from_str(&documents_json)?;
let generation_id = generation_id.to_string();
let documents = documents
.into_iter()
.map(|document| compile_document(&table.schema, document))
.map(|document| {
if document.get("generation_id").and_then(JsonValue::as_str) != Some(generation_id.as_str()) {
return Err(IndexError::InvalidInput(
"embedded projection generation does not match target index".into(),
));
}
compile_document(&table.schema, document)
})
.collect::<Result<Vec<_>>>()?;
table
.index
@@ -117,14 +167,148 @@ impl EmbeddedSearchIndex {
Ok(())
}
pub(crate) async fn delete(&self, table: String, id: String) -> napi::Result<()> {
self.manager.table(&table)?.index.write().await.delete(&id);
pub(crate) async fn gc_document_history_for_generation(
&self,
generation_id: Uuid,
table_name: &str,
workspace_id: &str,
doc_id: &str,
published_tuple: (i64, i64),
limit: usize,
) -> napi::Result<()> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(table_name)?;
let source_version_field = table.schema.field("source_version")?;
let permission_version_field = table.schema.field("permission_version")?;
let query = compile_query(
&table.schema,
&json!({
"bool":{"must":[
{"term":{"workspace_id":{"value":workspace_id}}},
{"term":{"doc_id":{"value":doc_id}}}
]}
}),
)?;
let mut index = table.index.write().await;
let matches = index
.search(
&query,
memory_indexer::SearchOptions {
limit: limit.max(1),
offset: 0,
after: None,
sort: vec![memory_indexer::Sort::DocumentId],
stored_fields: vec![source_version_field, permission_version_field],
highlight_fields: Vec::new(),
},
)
.map_err(IndexError::from)?;
for hit in matches.hits {
let mut hit_source_version = None;
let mut hit_permission_version = None;
for (field, values) in hit.fields {
let value = values.first().and_then(|value| match value {
memory_indexer::Value::I64(value) => Some(*value),
_ => None,
});
if field == source_version_field {
hit_source_version = value;
} else if field == permission_version_field {
hit_permission_version = value;
}
}
if (hit_source_version, hit_permission_version) != (Some(published_tuple.0), Some(published_tuple.1)) {
index.delete(&hit.id);
}
}
Ok(())
}
pub(crate) async fn gc_workspace_for_generation(
&self,
generation_id: Uuid,
table_name: &str,
workspace_id: &str,
source_version_high_water: i64,
limit: usize,
) -> napi::Result<bool> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(table_name)?;
let source_version_field = table.schema.field("source_version")?;
let query = compile_query(&table.schema, &json!({"term":{"workspace_id":{"value":workspace_id}}}))?;
let mut index = table.index.write().await;
let matches = index
.search(
&query,
memory_indexer::SearchOptions {
limit: limit.max(1),
offset: 0,
after: None,
sort: vec![
memory_indexer::Sort::Field {
field: source_version_field,
order: memory_indexer::SortOrder::Asc,
},
memory_indexer::Sort::DocumentId,
],
stored_fields: vec![source_version_field],
highlight_fields: Vec::new(),
},
)
.map_err(IndexError::from)?;
let may_have_more = matches.hits.len() == limit.max(1)
&& matches.hits.iter().any(|hit| {
hit.fields.iter().any(|(field, values)| {
*field == source_version_field
&& values.first().is_none_or(
|value| !matches!(value, memory_indexer::Value::I64(version) if *version > source_version_high_water),
)
})
});
for hit in matches.hits {
let source_version = hit
.fields
.iter()
.find(|(field, _)| *field == source_version_field)
.and_then(|(_, values)| values.first())
.and_then(|value| match value {
memory_indexer::Value::I64(value) => Some(*value),
_ => None,
});
if source_version.is_none_or(|source_version| source_version <= source_version_high_water) {
index.delete(&hit.id);
}
}
Ok(may_have_more)
}
#[cfg(test)]
pub(crate) async fn search(&self, table: String, dsl_json: String) -> napi::Result<String> {
let table = self.manager.table(&table)?;
let dsl: JsonValue = serde_json::from_str(&dsl_json)?;
self.search_for_generation(Uuid::nil(), table, dsl_json).await
}
pub(crate) async fn search_for_generation(
&self,
generation_id: Uuid,
table: String,
dsl_json: String,
) -> napi::Result<String> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(&table)?;
let mut dsl: JsonValue = serde_json::from_str(&dsl_json)?;
ensure_projection_fields(&mut dsl)?;
let query = compile_query(
&table.schema,
dsl
@@ -144,9 +328,27 @@ impl EmbeddedSearchIndex {
))?)
}
#[cfg(test)]
pub(crate) async fn aggregate(&self, table: String, dsl_json: String) -> napi::Result<String> {
let table = self.manager.table(&table)?;
let dsl: JsonValue = serde_json::from_str(&dsl_json)?;
self.aggregate_for_generation(Uuid::nil(), table, dsl_json).await
}
pub(crate) async fn aggregate_for_generation(
&self,
generation_id: Uuid,
table: String,
dsl_json: String,
) -> napi::Result<String> {
let generation = self
.manager
.generation(generation_id, false)
.await
.ok_or_else(|| napi::Error::from_reason("embedded search generation is not initialized"))?;
let table = generation.table(&table)?;
let mut dsl: JsonValue = serde_json::from_str(&dsl_json)?;
if let Some(top_hits) = dsl.pointer_mut("/aggs/result/aggs/result/top_hits") {
ensure_projection_fields(top_hits)?;
}
let query = compile_query(
&table.schema,
dsl
@@ -187,38 +389,23 @@ impl EmbeddedSearchIndex {
&highlight_tags(dsl.pointer("/aggs/result/aggs/result/top_hits").unwrap_or(&dsl)),
))?)
}
}
pub(crate) async fn checkpoint(&self, table: String) -> napi::Result<EmbeddedIndexCheckpoint> {
let checkpoint = self
.manager
.table(&table)?
.index
.read()
.await
.checkpoint()
.map_err(IndexError::from)?;
Ok(EmbeddedIndexCheckpoint {
sequence: checkpoint.sequence as i64,
data: checkpoint.bytes.into(),
})
}
pub(crate) async fn optimize(&self, table: String) -> napi::Result<()> {
self.manager.table(&table)?.index.write().await.optimize();
Ok(())
}
pub(crate) async fn mark_checkpoint_persisted(&self, table: String, sequence: i64) -> napi::Result<()> {
self
.manager
.table(&table)?
.index
.write()
.await
.mark_checkpoint_persisted(sequence as u64)
.map_err(IndexError::from)?;
Ok(())
fn ensure_projection_fields(dsl: &mut JsonValue) -> Result<()> {
let fields = dsl
.as_object_mut()
.ok_or_else(|| IndexError::InvalidInput("search DSL must be an object".into()))?
.entry("fields")
.or_insert_with(|| json!([]));
let fields = fields
.as_array_mut()
.ok_or_else(|| IndexError::InvalidInput("search fields must be an array".into()))?;
for field in ["doc_id", "source_version", "permission_version"] {
if !fields.iter().any(|value| value.as_str() == Some(field)) {
fields.push(json!(field));
}
}
Ok(())
}
fn highlight_tags(dsl: &JsonValue) -> HighlightTags {
@@ -248,13 +435,17 @@ impl Default for EmbeddedSearchIndex {
#[cfg(test)]
mod tests {
use serde_json::{Value, json};
use uuid::Uuid;
use super::EmbeddedSearchIndex;
fn doc(workspace: &str, id: &str, title: &str, updated_at: i64) -> Value {
json!({
"generation_id": Uuid::nil().to_string(),
"workspace_id": workspace,
"doc_id": id,
"source_version": 1,
"permission_version": 1,
"title": title,
"summary": title,
"created_by_user_id": "user",
@@ -277,7 +468,7 @@ mod tests {
}
#[tokio::test]
async fn exact_search_cursor_and_checkpoint_roundtrip() {
async fn exact_search_cursor_and_aggregate() {
let index = EmbeddedSearchIndex::new();
index
.write(
@@ -296,7 +487,7 @@ mod tests {
let first: Value =
serde_json::from_str(&index.search("doc".into(), search(query.clone(), None)).await.unwrap()).unwrap();
assert_eq!(first["total"], 2);
assert_eq!(first["nodes"][0]["id"], "workspace-1/two");
assert_eq!(first["nodes"][0]["id"], format!("{}/workspace-1/two/1/1", Uuid::nil()));
assert_eq!(first["nodes"][0]["fields"]["doc_id"], json!(["two"]));
let cursor = first["nextCursor"].as_str().unwrap();
@@ -307,17 +498,7 @@ mod tests {
.unwrap(),
)
.unwrap();
assert_eq!(second["nodes"][0]["id"], "workspace-1/one");
let checkpoint = index.checkpoint("doc".into()).await.unwrap();
index.reset("doc".into()).await.unwrap();
let empty: Value =
serde_json::from_str(&index.search("doc".into(), search(query.clone(), None)).await.unwrap()).unwrap();
assert_eq!(empty["total"], 0);
index.restore("doc".into(), checkpoint.data).await.unwrap();
let restored: Value =
serde_json::from_str(&index.search("doc".into(), search(query, None)).await.unwrap()).unwrap();
assert_eq!(restored["total"], 2);
assert_eq!(second["nodes"][0]["id"], format!("{}/workspace-1/one/1/1", Uuid::nil()));
let aggregate: Value = serde_json::from_str(
&index
@@ -343,12 +524,16 @@ mod tests {
}
#[tokio::test]
async fn write_is_atomic_and_corrupt_checkpoint_is_rejected() {
async fn write_is_atomic() {
let index = EmbeddedSearchIndex::new();
index
.write(
"doc".into(),
json!([{ "workspace_id": "workspace", "doc_id": "null-values", "summary": [null] }]).to_string(),
json!([{
"generation_id":Uuid::nil().to_string(),"workspace_id":"workspace","doc_id":"null-values",
"source_version":1,"permission_version":1,"summary":[null]
}])
.to_string(),
)
.await
.unwrap();
@@ -357,7 +542,6 @@ mod tests {
serde_json::from_str(&index.search("doc".into(), search(all.clone(), None)).await.unwrap()).unwrap();
assert_eq!(result["total"], 1);
index.reset("doc".into()).await.unwrap();
let documents = json!([
doc("workspace", "valid", "hello", 1),
{ "workspace_id": "workspace", "doc_id": "invalid", "unknown": true }
@@ -365,7 +549,124 @@ mod tests {
assert!(index.write("doc".into(), documents.to_string()).await.is_err());
let result: Value = serde_json::from_str(&index.search("doc".into(), search(all, None)).await.unwrap()).unwrap();
assert_eq!(result["total"], 0);
assert!(index.restore("doc".into(), vec![1, 2, 3].into()).await.is_err());
assert_eq!(result["total"], 1);
}
#[tokio::test]
async fn immutable_tuples_do_not_overwrite_each_other() {
let index = EmbeddedSearchIndex::new();
let mut old = doc("workspace", "doc", "old", 1);
old["source_version"] = json!(1);
let mut current = doc("workspace", "doc", "current", 1);
current["source_version"] = json!(2);
index
.write("doc".into(), json!([old, current]).to_string())
.await
.unwrap();
let result: Value = serde_json::from_str(
&index
.search(
"doc".into(),
json!({"query":{"term":{"doc_id":{"value":"doc"}}},"fields":["title","source_version"],"size":10})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 2);
assert_ne!(result["nodes"][0]["id"], result["nodes"][1]["id"]);
let retired = Uuid::new_v4();
index.prepare_generation(retired).await;
index.retain_generation(Uuid::nil()).await;
assert!(index.has_generation(Uuid::nil()).await);
assert!(!index.has_generation(retired).await);
}
#[tokio::test]
async fn history_gc_keeps_the_published_tuple() {
let index = EmbeddedSearchIndex::new();
let mut stale = doc("workspace", "doc", "old", 1);
stale["source_version"] = json!(1);
let mut current = doc("workspace", "doc", "current", 1);
current["source_version"] = json!(2);
index
.write("doc".into(), json!([stale, current]).to_string())
.await
.unwrap();
index
.gc_document_history_for_generation(Uuid::nil(), "doc", "workspace", "doc", (2, 1), 10)
.await
.unwrap();
let result: Value = serde_json::from_str(
&index
.search(
"doc".into(),
json!({"query":{"term":{"doc_id":{"value":"doc"}}},"fields":["title"],"size":10}).to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(result["nodes"][0]["fields"]["title"], json!(["current"]));
index
.write(
"doc".into(),
json!([
doc("deleted-workspace", "one", "one", 1),
doc("deleted-workspace", "two", "two", 1)
])
.to_string(),
)
.await
.unwrap();
assert!(
index
.gc_workspace_for_generation(Uuid::nil(), "doc", "deleted-workspace", i64::MAX, 1)
.await
.unwrap()
);
assert!(
index
.gc_workspace_for_generation(Uuid::nil(), "doc", "deleted-workspace", i64::MAX, 1)
.await
.unwrap()
);
assert!(
!index
.gc_workspace_for_generation(Uuid::nil(), "doc", "deleted-workspace", i64::MAX, 1)
.await
.unwrap()
);
let old = doc("recreated-workspace", "old", "old", 1);
let mut recreated = doc("recreated-workspace", "new", "new", 1);
recreated["source_version"] = json!(3);
index
.write("doc".into(), json!([old, recreated]).to_string())
.await
.unwrap();
assert!(
!index
.gc_workspace_for_generation(Uuid::nil(), "doc", "recreated-workspace", 2, 10)
.await
.unwrap()
);
let result: Value = serde_json::from_str(
&index
.search(
"doc".into(),
json!({"query":{"term":{"workspace_id":{"value":"recreated-workspace"}}},"fields":["doc_id"],"size":10})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(result["total"], 1);
assert_eq!(result["nodes"][0]["fields"]["doc_id"], json!(["new"]));
}
}
@@ -16,6 +16,9 @@ impl TableSchema {
let mut fields = HashMap::new();
keyword(&mut builder, &mut fields, "workspace_id", true, false);
keyword(&mut builder, &mut fields, "workspace_token", true, false);
keyword(&mut builder, &mut fields, "generation_id", true, false);
integer(&mut builder, &mut fields, "source_version", true, true);
integer(&mut builder, &mut fields, "permission_version", true, false);
keyword(&mut builder, &mut fields, "doc_id", true, true);
keyword(&mut builder, &mut fields, "doc_token", true, false);
fields.insert(
@@ -29,10 +32,19 @@ impl TableSchema {
keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_member_default_readable");
integer(&mut builder, &mut fields, "acl_revision", true, false);
integer(&mut builder, &mut fields, "created_at", true, true);
integer(&mut builder, &mut fields, "updated_at", true, true);
Self::finish(builder, fields, &["workspace_id", "doc_id"])
Self::finish(
builder,
fields,
&[
"generation_id",
"workspace_id",
"doc_id",
"source_version",
"permission_version",
],
)
}
pub fn block() -> Self {
@@ -58,6 +70,9 @@ impl TableSchema {
}
keyword(&mut builder, &mut fields, "doc_id", true, true);
keyword(&mut builder, &mut fields, "workspace_token", true, false);
keyword(&mut builder, &mut fields, "generation_id", true, false);
integer(&mut builder, &mut fields, "source_version", true, true);
integer(&mut builder, &mut fields, "permission_version", true, false);
keyword(&mut builder, &mut fields, "doc_token", true, false);
keyword(&mut builder, &mut fields, "block_id", true, true);
keyword(&mut builder, &mut fields, "block_token", true, false);
@@ -74,8 +89,18 @@ impl TableSchema {
keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_member_default_readable");
integer(&mut builder, &mut fields, "acl_revision", true, false);
Self::finish(builder, fields, &["workspace_id", "doc_id", "block_id"])
Self::finish(
builder,
fields,
&[
"generation_id",
"workspace_id",
"doc_id",
"source_version",
"permission_version",
"block_id",
],
)
}
fn finish(
@@ -95,10 +120,14 @@ impl TableSchema {
.id_fields
.iter()
.map(|field| {
document
let value = document
.get(*field)
.and_then(serde_json::Value::as_str)
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} is required")))
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} is required")))?;
value
.as_str()
.map(str::to_string)
.or_else(|| value.as_i64().map(|value| value.to_string()))
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} has invalid identity type")))
})
.collect::<Result<Vec<_>>>()
.map(|parts| parts.join("/"))