Jayden/AFFiNE-Mirror
0
0
Fork 0
mirror of https://github.com/toeverything/AFFiNE.git synced 2026-08-31 21:59:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(server): indexer & worker & sync perf (#15504)

Browse Source
This commit is contained in:
DarkSky
2026-08-21 08:11:37 +08:00
committed by GitHub
parent 8c9aad9a9b
commit c57004ea2c
259 changed files with 16530 additions and 17793 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/backend/native/src/entitlement.rs
+16
View File
@@ -28,6 +28,22 @@ const ONE_GB: i64 = 1024 * ONE_MB;
const ONE_DAY_SECONDS: i64 = 24 * 60 * 60;
const MAX_SEAT_QUANTITY: i32 = 100_000;
pub(crate) fn entitlement_priority(status: &str, plan: &str) -> i32 {
let status = match status {
"active" => 200,
"grace" => 100,
_ => 0,
};
let plan = match plan {
"team" | "selfhost_team" => 40,
"lifetime_pro" => 30,
"pro" => 20,
"ai" => 10,
_ => 0,
};
status + plan
}
#[napi(object)]
pub struct ResolveEntitlementInput {
pub deployment_type: String,
packages/backend/native/src/lib.rs
+1
View File
@@ -14,6 +14,7 @@ pub mod llm;
pub mod permission;
pub mod runtime;
pub mod safe_fetch;
pub(crate) mod search_index;
pub mod tiktoken;
mod userdata_acl;
mod utils;
packages/backend/native/src/permission/candidates.rs
+1 -1
View File
@@ -23,7 +23,7 @@ pub(super) fn parse_workspace_role(role: &str) -> anyhow::Result<WorkspaceRole>
}
}
fn parse_doc_role(role: &str) -> anyhow::Result<DocRole> {
pub(super) fn parse_doc_role(role: &str) -> anyhow::Result<DocRole> {
match role {
"none" => Ok(DocRole::None),
"external" => Ok(DocRole::External),
packages/backend/native/src/permission/mod.rs
+5
View File
@@ -10,6 +10,11 @@ use napi_derive::napi;
use serde_json::Value;
pub use types::*;
pub(crate) fn doc_role_allows(role: &str, action: &str) -> anyhow::Result<bool> {
let role = candidates::parse_doc_role(role)?;
Ok(actions::doc_actions_for_role(role).contains(action))
}
#[napi]
pub fn evaluate_permission_v1(input: Value) -> Result<Value> {
let input = serde_json::from_value::<PermissionEvaluationInputV1>(input)
packages/backend/native/src/runtime/backend_runtime/byok/profile.rs
+2 -116
View File
@@ -47,10 +47,7 @@ pub(in super::super) async fn list(pool: &PgPool, workspace_id: &str) -> Runtime
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("list BYOK profiles failed", error))?;
// Rows written by the previous release while it shares the database carry
// only the database-default definition and fail to parse; skip them until
// that release is retired.
Ok(rows.into_iter().filter_map(|row| profile_output(row).ok()).collect())
rows.into_iter().map(profile_output).collect()
}
pub(in super::super) async fn create(
@@ -103,21 +100,7 @@ pub(in super::super) async fn create(
definition, sort_order, enabled, created_by, updated_by, created_at, updated_at
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $10, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
ON CONFLICT (workspace_id, provider, name) DO UPDATE
SET id = EXCLUDED.id,
description = EXCLUDED.description,
encrypted_api_key = EXCLUDED.encrypted_api_key,
definition = EXCLUDED.definition,
sort_order = EXCLUDED.sort_order,
enabled = EXCLUDED.enabled,
revision = 1,
credential_generation = 1,
validation = NULL,
created_by = EXCLUDED.created_by,
updated_by = EXCLUDED.updated_by,
created_at = EXCLUDED.created_at,
updated_at = EXCLUDED.updated_at
WHERE ai_workspace_byok_configs.definition = '{}'::jsonb
ON CONFLICT (workspace_id, provider, name) DO NOTHING
RETURNING id, workspace_id, provider, name, description, encrypted_api_key,
definition, sort_order, enabled, revision, credential_generation, validation
"#,
@@ -591,100 +574,3 @@ pub(super) fn require_text(value: &str, field: &'static str) -> RuntimeResult<()
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::{ByokPolicy, PgPool, Uuid, create, list};
use crate::{
llm::{
ByokCapabilityInput, ByokEndpointInput, ByokModelDeclarationInput, ByokProfileDefinitionInput,
CreateByokProfileInput, Deployment,
},
runtime::config::CopilotByokRuntimeConfig,
};
#[tokio::test]
async fn list_skips_rows_with_unparseable_legacy_definition() {
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
let workspace_id = format!("byok-legacy-{}", Uuid::new_v4());
sqlx::query("INSERT INTO workspaces (id) VALUES ($1)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
// a row as written by the previous release while it shares the database:
// definition is left at the database default and cannot be parsed
for (id, name, definition) in [
(Uuid::new_v4().to_string(), "legacy", "{}"),
(
Uuid::new_v4().to_string(),
"valid",
r#"{"endpoint":{"kind":"provider_default"},"models":[]}"#,
),
] {
sqlx::query(
"INSERT INTO ai_workspace_byok_configs (id, workspace_id, provider, name, encrypted_api_key, definition, \
created_at, updated_at) VALUES ($1, $2, 'openai', $3, 'x', $4::jsonb, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)",
)
.bind(&id)
.bind(&workspace_id)
.bind(name)
.bind(definition)
.execute(&pool)
.await
.unwrap();
}
let profiles = list(&pool, &workspace_id).await.unwrap();
assert_eq!(profiles.len(), 1);
assert_eq!(profiles[0].name, "valid");
let input = || CreateByokProfileInput {
workspace_id: workspace_id.clone(),
provider: "openai".to_string(),
name: "legacy".to_string(),
description: None,
credential: "replacement-key".to_string(),
definition: ByokProfileDefinitionInput {
endpoint: ByokEndpointInput {
kind: "provider_default".to_string(),
url: None,
dialect: None,
},
models: vec![ByokModelDeclarationInput {
model_id: "gpt-4o-mini".to_string(),
enabled: true,
capabilities: vec![ByokCapabilityInput {
input: vec!["text".to_string()],
output: vec!["text".to_string()],
features: vec![],
attachment_kinds: vec![],
attachment_sources: vec![],
}],
}],
},
enabled: true,
actor_user_id: "user-1".to_string(),
};
let policy = ByokPolicy::from(Deployment::Cloud, &CopilotByokRuntimeConfig::default());
create(&pool, &[7; 32], &policy, input()).await.unwrap();
let profiles = list(&pool, &workspace_id).await.unwrap();
assert_eq!(profiles.len(), 2);
assert!(profiles.iter().any(|profile| profile.name == "legacy"));
assert!(create(&pool, &[7; 32], &policy, input()).await.is_err());
sqlx::query("DELETE FROM ai_workspace_byok_configs WHERE workspace_id = $1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("DELETE FROM workspaces WHERE id = $1")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
}
}
packages/backend/native/src/runtime/backend_runtime/copilot/context.rs
+7 -11
View File
@@ -80,16 +80,11 @@ async fn load_server_profiles(
.map_err(|error| RuntimeError::database("load authorized BYOK profiles failed", error))?;
rows
.into_iter()
.filter_map(|row| {
// Rows written by the previous release while it shares the database
// carry only the database-default definition; skip them until that
// release is retired instead of failing the whole profile load.
let definition = match serde_json::from_value::<ByokProfileDefinition>(row.definition) {
Ok(definition) => definition,
Err(_) => return None,
};
.map(|row| {
let definition = serde_json::from_value::<ByokProfileDefinition>(row.definition)
.map_err(|error| RuntimeError::json("invalid stored BYOK definition", error))?;
if !policy.allows(&row.provider, &definition.endpoint) {
return None;
return Ok(None);
}
let aad = server_aad(
&row.workspace_id,
@@ -97,7 +92,7 @@ async fn load_server_profiles(
&row.provider,
definition.endpoint_identity(),
);
Some(Ok(authorized_byok_profile(
Ok(Some(authorized_byok_profile(
row.id,
ProfileSource::Server,
row.provider,
@@ -110,7 +105,8 @@ async fn load_server_profiles(
},
)))
})
.collect()
.collect::<RuntimeResult<Vec<_>>>()
.map(|profiles| profiles.into_iter().flatten().collect())
}
async fn load_local_profiles(
packages/backend/native/src/runtime/backend_runtime/embedding/index.rs
+2 -2
View File
@@ -160,7 +160,7 @@ mod tests {
.unwrap();
assert_eq!(first.active_index_id, repeated.active_index_id);
assert_eq!(first.index_epoch, repeated.index_epoch);
let failed_probe = super::super::store::claim_index_probe(&pool, "probe-a")
let failed_probe = super::super::store::claim_index_probe_for_workspace(&pool, "probe-a", &workspace_id)
.await
.unwrap()
.unwrap();
@@ -178,7 +178,7 @@ mod tests {
.execute(&pool)
.await
.unwrap();
let recovered_probe = super::super::store::claim_index_probe(&pool, "probe-b")
let recovered_probe = super::super::store::claim_index_probe_for_workspace(&pool, "probe-b", &workspace_id)
.await
.unwrap()
.unwrap();
packages/backend/native/src/runtime/backend_runtime/embedding/mod.rs
+30 -20
View File
@@ -12,7 +12,7 @@ use std::{
};
use sqlx::PgPool;
use tokio::sync::{Mutex, Notify};
use tokio::sync::Notify;
pub(super) use types::EmbeddingTarget;
use types::*;
@@ -36,10 +36,31 @@ pub(super) struct EmbeddingService {
object_storage: RwLock<Arc<ObjectStorageService>>,
provider: BackgroundEmbeddingProvider,
wake: Notify,
worker: Mutex<Option<worker::WorkerHandle>>,
candidate_cancellations: StdMutex<HashMap<String, Option<tokio::sync::watch::Sender<bool>>>>,
}
pub(super) struct EmbeddingWorker {
handle: Option<worker::WorkerHandle>,
}
impl EmbeddingWorker {
pub(super) fn start(service: Arc<EmbeddingService>) -> Self {
Self {
handle: Some(worker::start(service)),
}
}
pub(super) async fn stop(mut self) {
if let Some(handle) = self.handle.take() {
handle.stop().await;
}
}
pub(super) fn is_running(&self) -> bool {
self.handle.is_some()
}
}
impl EmbeddingService {
pub(super) fn new(
pool: PgPool,
@@ -51,28 +72,10 @@ impl EmbeddingService {
object_storage: RwLock::new(object_storage),
provider,
wake: Notify::new(),
worker: Mutex::new(None),
candidate_cancellations: StdMutex::new(HashMap::new()),
})
}
pub(super) async fn start(self: &Arc<Self>) {
let mut worker = self.worker.lock().await;
if worker.is_none() {
*worker = Some(worker::start(Arc::clone(self)));
}
}
pub(super) async fn stop(&self) {
if let Some(worker) = self.worker.lock().await.take() {
worker.stop().await;
}
}
pub(super) async fn is_running(&self) -> bool {
self.worker.lock().await.is_some()
}
fn wake(&self) {
self.wake.notify_one();
}
@@ -237,6 +240,13 @@ pub(in crate::runtime::backend_runtime) async fn register_artifact_source(
pool: &PgPool,
artifact: &crate::runtime::types::RuntimeWorkspaceArtifact,
) -> RuntimeResult<()> {
let schema_ready: bool = sqlx::query_scalar("SELECT to_regclass('embedding_sources') IS NOT NULL")
.fetch_one(pool)
.await
.map_err(|error| RuntimeError::database("Embedding source schema health check failed", error))?;
if !schema_ready {
return Ok(());
}
uuid::Uuid::parse_str(&artifact.id).map_err(|_| RuntimeError::invalid_input("artifact_id_invalid"))?;
source::register_artifact(pool, artifact).await
}
packages/backend/native/src/runtime/backend_runtime/embedding/store.rs
+28
View File
@@ -106,6 +106,34 @@ pub(super) async fn claim_index_probe(pool: &PgPool, owner: &str) -> RuntimeResu
.map_err(|error| RuntimeError::database("claim embedding index probe failed", error))
}
#[cfg(test)]
pub(super) async fn claim_index_probe_for_workspace(
pool: &PgPool,
owner: &str,
workspace_id: &str,
) -> RuntimeResult<Option<IndexProbeClaim>> {
sqlx::query_as(
r#"WITH candidate AS(
SELECT index_fact.id FROM embedding_indexes index_fact
JOIN embedding_workspace_states state ON state.active_index_id=index_fact.id
WHERE state.workspace_id=$2 AND state.runtime_state='active' AND(
index_fact.health_status='pending'
OR index_fact.health_status='retry_wait' AND index_fact.next_probe_at<=clock_timestamp()
OR index_fact.probe_lease_until<=clock_timestamp())
ORDER BY index_fact.next_probe_at NULLS FIRST,index_fact.updated_at
FOR UPDATE OF index_fact SKIP LOCKED LIMIT 1
) UPDATE embedding_indexes index_fact SET probe_lease_owner=$1,
probe_lease_until=clock_timestamp()+interval '2 minutes',updated_at=now()
FROM candidate WHERE index_fact.id=candidate.id
RETURNING index_fact.id,index_fact.workspace_id,index_fact.fingerprint,index_fact.probe_lease_owner"#,
)
.bind(owner)
.bind(workspace_id)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("claim embedding index probe for workspace failed", error))
}
pub(super) async fn complete_index_probe(pool: &PgPool, claim: &IndexProbeClaim) -> RuntimeResult<()> {
sqlx::query(
"UPDATE embedding_indexes SET \
packages/backend/native/src/runtime/backend_runtime/mod.rs
+249 -25
View File
@@ -8,9 +8,12 @@ mod doc_storage;
mod embedding;
mod gate;
mod housekeeping;
mod permission;
mod role;
mod rolling_quota;
mod runtime_state;
mod scope_compiler;
mod search;
#[cfg(test)]
mod tests;
mod workspace_stats;
@@ -23,16 +26,21 @@ use byok::LocalLeasePayload;
use copilot::{backend_provider, executable_protocol};
use embedding::register_artifact_source;
use napi::{Result, bindgen_prelude::Buffer};
use search::SearchRuntime;
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row, postgres::PgPoolOptions};
use tokio::sync::Mutex;
use self::types::{BackendRuntimeHealth, EmbeddingHealth};
use self::{
role::ServerRole,
search::{RuntimeAggregateRequest, RuntimeSearchRequest},
types::{BackendRuntimeHealth, EmbeddingHealth, SearchOperationOutput},
};
use super::object_storage::ObjectStorageService;
pub(crate) use super::types;
pub(super) use super::{
BackendRuntimeConfig, ConfigSource, InviteQuotaConfig, RuntimeError, RuntimeResult,
migrations::{migrate_embedding_tables, migrate_runtime_tables},
migrations::{embedding_schema_health, migrate_all_tables},
napi_error, to_napi_error,
};
use crate::llm::{
@@ -45,15 +53,45 @@ pub(super) fn token_hash(token: &str) -> String {
hex::encode(Sha256::digest(token.as_bytes()))
}
fn search_operation_output(result: RuntimeResult<serde_json::Value>) -> SearchOperationOutput {
match result {
Ok(value) => SearchOperationOutput {
ok: true,
value: Some(value),
error_code: None,
},
Err(error) => SearchOperationOutput {
ok: false,
value: None,
error_code: Some(
match error {
RuntimeError::SearchWorkspaceDenied => "workspace_denied",
RuntimeError::SearchPermissionUnavailable => "permission_unavailable",
RuntimeError::SearchProviderUnavailable => "provider_unavailable",
RuntimeError::SearchUnsupportedQuery => "unsupported_query",
RuntimeError::SearchReplayGap => "provider_unavailable",
RuntimeError::InvalidInput(_) | RuntimeError::Json { .. } => "invalid_request",
_ => "internal",
}
.to_string(),
),
},
}
}
#[napi_derive::napi]
pub struct BackendRuntime {
config_source: ConfigSource,
role: ServerRole,
script_mode: bool,
config: Arc<RwLock<Arc<BackendRuntimeConfig>>>,
config_reload: Mutex<()>,
pool: Mutex<Option<PgPool>>,
embedding_health: RwLock<EmbeddingHealth>,
object_storage: RwLock<Arc<ObjectStorageService>>,
embedding: Mutex<Option<Arc<embedding::EmbeddingService>>>,
embedding_worker: Mutex<Option<embedding::EmbeddingWorker>>,
search: Mutex<Option<Arc<SearchRuntime>>>,
managed_token_providers: Arc<copilot::ManagedTokenProviderCache>,
}
@@ -62,16 +100,21 @@ impl BackendRuntime {
#[napi(constructor)]
pub fn new(private_key: Option<String>, config_paths: Option<Vec<String>>) -> Result<Self> {
let config_source = ConfigSource::new(config_paths);
let (role, script_mode) = ServerRole::from_environment().map_err(napi_error)?;
let config = BackendRuntimeConfig::from_config_source(private_key, &config_source).map_err(to_napi_error)?;
let object_storage = ObjectStorageService::from_config_source(&config_source).map_err(to_napi_error)?;
Ok(Self {
config_source,
role,
script_mode,
config: Arc::new(RwLock::new(Arc::new(config))),
config_reload: Mutex::new(()),
pool: Mutex::new(None),
embedding_health: RwLock::new(EmbeddingHealth::disabled("runtime_not_started", None)),
object_storage: RwLock::new(Arc::new(object_storage)),
embedding: Mutex::new(None),
embedding_worker: Mutex::new(None),
search: Mutex::new(None),
managed_token_providers: Arc::new(Default::default()),
})
}
@@ -109,7 +152,30 @@ impl BackendRuntime {
.write()
.map_err(|_| RuntimeError::invalid_state("object storage service lock poisoned"))? = Arc::new(object_storage);
let mut embedding_health = migrate_embedding_tables(&pool).await;
let embedding_health = if self.script_mode {
EmbeddingHealth::disabled("script_runtime", None)
} else {
let config = self.config()?;
if config.search.enabled {
if config.search.provider == "embedded" && !self.role.allows_embedded_search() {
return Err(RuntimeError::config(format!(
"embedded search is only available for the allinone role (current role: {})",
self.role.as_str()
)));
}
let search = Arc::new(SearchRuntime::new(pool.clone(), config.search.clone())?);
if self.role.owns_background() {
search.initialize().await?;
}
*self.search.lock().await = Some(search);
} else {
*self.search.lock().await = None;
}
embedding_schema_health(&pool).await?
};
if self.script_mode {
*self.search.lock().await = None;
}
if embedding_health.enabled {
let provider = copilot::BackgroundEmbeddingProvider::new(
pool.clone(),
@@ -117,18 +183,30 @@ impl BackendRuntime {
Arc::clone(&self.managed_token_providers),
);
let embedding = embedding::EmbeddingService::new(pool.clone(), self.object_storage()?, provider);
if std::env::var("NODE_ENV").as_deref() != Ok("test")
|| std::env::var("AFFINE_EMBEDDING_WORKER").as_deref() == Ok("1")
if self.role.owns_background()
&& (std::env::var("NODE_ENV").as_deref() != Ok("test")
|| std::env::var("AFFINE_EMBEDDING_WORKER").as_deref() == Ok("1"))
{
embedding.start().await;
*self.embedding_worker.lock().await = Some(embedding::EmbeddingWorker::start(Arc::clone(&embedding)));
}
embedding_health.worker_running = embedding.is_running().await;
let mut embedding_health = embedding_health;
embedding_health.worker_running = self
.embedding_worker
.lock()
.await
.as_ref()
.is_some_and(embedding::EmbeddingWorker::is_running);
*self.embedding.lock().await = Some(embedding);
*self
.embedding_health
.write()
.map_err(|_| RuntimeError::invalid_state("embedding health lock poisoned"))? = embedding_health;
} else {
*self
.embedding_health
.write()
.map_err(|_| RuntimeError::invalid_state("embedding health lock poisoned"))? = embedding_health;
}
*self
.embedding_health
.write()
.map_err(|_| RuntimeError::invalid_state("embedding health lock poisoned"))? = embedding_health;
*guard = Some(pool);
Ok(())
@@ -136,9 +214,11 @@ impl BackendRuntime {
#[napi]
pub async fn stop(&self) -> Result<()> {
if let Some(embedding) = self.embedding.lock().await.take() {
embedding.stop().await;
self.search.lock().await.take();
if let Some(worker) = self.embedding_worker.lock().await.take() {
worker.stop().await;
}
self.embedding.lock().await.take();
let pool = self.pool.lock().await.take();
if let Some(pool) = pool {
pool.close().await;
@@ -168,6 +248,26 @@ impl BackendRuntime {
.await
.map_err(to_napi_error)?;
self.update_config(config).map_err(to_napi_error)?;
if !self.script_mode {
let config = self.config().map_err(to_napi_error)?;
if config.search.enabled {
if config.search.provider == "embedded" && !self.role.allows_embedded_search() {
return Err(napi_error(format!(
"embedded search is only available for the allinone role (current role: {})",
self.role.as_str()
)));
}
let search = Arc::new(SearchRuntime::new(pool.clone(), config.search.clone()).map_err(to_napi_error)?);
if self.role.owns_background() {
search.initialize().await.map_err(to_napi_error)?;
}
*self.search.lock().await = Some(search);
} else {
*self.search.lock().await = None;
}
} else {
*self.search.lock().await = None;
}
let object_storage = Arc::new(object_storage);
*self
.object_storage
@@ -176,17 +276,19 @@ impl BackendRuntime {
if let Some(embedding) = self.embedding.lock().await.as_ref() {
embedding.reload_object_storage(object_storage).map_err(to_napi_error)?;
}
let workspace_ids = sqlx::query_scalar::<_, String>("SELECT id FROM workspaces")
.fetch_all(&pool)
.await
.map_err(|error| {
to_napi_error(RuntimeError::database(
"load workspaces for embedding reconciliation failed",
error,
))
})?;
for workspace_id in workspace_ids {
self.reconcile_embedding_workspace(&workspace_id).await?;
if self.role.owns_background() {
let workspace_ids = sqlx::query_scalar::<_, String>("SELECT id FROM workspaces")
.fetch_all(&pool)
.await
.map_err(|error| {
to_napi_error(RuntimeError::database(
"load workspaces for embedding reconciliation failed",
error,
))
})?;
for workspace_id in workspace_ids {
self.reconcile_embedding_workspace(&workspace_id).await?;
}
}
Ok(())
}
@@ -217,7 +319,106 @@ impl BackendRuntime {
#[napi]
pub async fn run_migrations(&self) -> Result<()> {
let pool = self.pool().await?;
migrate_runtime_tables(&pool).await.map_err(to_napi_error)
let embedding_health = migrate_all_tables(&pool).await.map_err(to_napi_error)?;
*self
.embedding_health
.write()
.map_err(|_| napi_error("embedding health lock poisoned"))? = embedding_health;
Ok(())
}
#[napi]
pub async fn search_authorized(
&self,
actor_user_id: String,
workspace_id: String,
request: RuntimeSearchRequest,
) -> Result<SearchOperationOutput> {
let result = self
.search_runtime()
.await?
.search_authorized(&actor_user_id, &workspace_id, request)
.await;
Ok(search_operation_output(result))
}
#[napi]
pub async fn aggregate_authorized(
&self,
actor_user_id: String,
workspace_id: String,
request: RuntimeAggregateRequest,
) -> Result<SearchOperationOutput> {
let result = self
.search_runtime()
.await?
.aggregate_authorized(&actor_user_id, &workspace_id, request)
.await;
Ok(search_operation_output(result))
}
#[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<()> {
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)
.await
.map_err(to_napi_error)
}
#[napi]
pub async fn filter_readable_docs(
&self,
actor_user_id: String,
workspace_id: String,
doc_ids: Vec<String>,
) -> Result<Vec<String>> {
let authorizer = permission::PermissionAuthorizer::new(self.pool().await?);
authorizer
.filter_readable_docs(&workspace_id, &actor_user_id, doc_ids)
.await
.map(|ids| ids.into_iter().collect())
.map_err(to_napi_error)
}
#[napi]
pub async fn search_status(&self) -> Result<serde_json::Value> {
self.search_runtime().await?.status().await.map_err(to_napi_error)
}
#[napi]
@@ -370,6 +571,7 @@ impl BackendRuntime {
#[napi]
pub async fn reconcile_embedding_workspaces(&self) -> Result<i64> {
self.require_background()?;
let workspace_ids = sqlx::query_scalar::<_, String>("SELECT id FROM workspaces")
.fetch_all(&self.pool().await?)
.await
@@ -405,6 +607,7 @@ impl BackendRuntime {
#[napi]
pub async fn cleanup_unreferenced_artifacts(&self, limit: i64) -> Result<i64> {
self.require_background()?;
if limit <= 0 {
return Err(napi_error("artifact cleanup limit must be positive"));
}
@@ -572,6 +775,27 @@ impl BackendRuntime {
.ok_or_else(|| RuntimeError::invalid_state("BackendRuntime must be started before using postgres operations"))
}
fn require_background(&self) -> Result<()> {
if self.role.owns_background() {
Ok(())
} else {
Err(napi_error(format!(
"backend runtime role {} does not own background work",
self.role.as_str()
)))
}
}
async fn search_runtime(&self) -> Result<Arc<SearchRuntime>> {
self
.search
.lock()
.await
.as_ref()
.cloned()
.ok_or_else(|| napi_error("search_provider_not_ready"))
}
async fn reconcile_embedding_workspace(&self, workspace_id: &str) -> Result<()> {
let enabled = sqlx::query_scalar::<_, bool>("SELECT enable_doc_embedding FROM workspaces WHERE id=$1")
.bind(workspace_id)
packages/backend/native/src/runtime/backend_runtime/permission/authorizer.rs
+90
View File
@@ -0,0 +1,90 @@
use std::collections::BTreeSet;
use sqlx::PgPool;
use super::{
store::PermissionStore,
types::{AclPredicate, AuthorizedSearchScope, DocAclCapability, DocReadScope, SearchActor},
};
use crate::{
permission::evaluate_permission,
runtime::{RuntimeError, RuntimeResult},
};
pub(in crate::runtime::backend_runtime) struct PermissionAuthorizer {
store: PermissionStore,
}
impl PermissionAuthorizer {
pub(in crate::runtime::backend_runtime) fn new(pool: PgPool) -> Self {
Self {
store: PermissionStore::new(pool),
}
}
pub(in crate::runtime::backend_runtime) async fn authorize_search(
&self,
actor: &SearchActor,
workspace_id: &str,
) -> RuntimeResult<AuthorizedSearchScope> {
match actor {
SearchActor::User { user_id } => {
let snapshot = self.store.search_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)
.map_err(|_| RuntimeError::SearchPermissionUnavailable)?
.workspace
.decisions
.into_iter()
.find(|decision| decision.action == "Workspace.Read")
.ok_or(RuntimeError::SearchPermissionUnavailable)?;
if !decision.allowed {
return Err(RuntimeError::SearchWorkspaceDenied);
}
let docs = match snapshot.capability {
DocAclCapability::Disabled => DocReadScope::All,
DocAclCapability::Unknown => return Err(RuntimeError::SearchPermissionUnavailable),
DocAclCapability::Enabled if owner_or_admin => DocReadScope::All,
DocAclCapability::Enabled => DocReadScope::ProjectedAcl(AclPredicate {
actor_user_id: snapshot.actor_user_id,
active_member: snapshot.active_member,
sharing_enabled: snapshot.sharing_enabled,
}),
};
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,
user_id: &str,
doc_ids: Vec<String>,
) -> RuntimeResult<BTreeSet<String>> {
let snapshot = self.store.permission_snapshot(workspace_id, user_id, &doc_ids).await?;
let output = evaluate_permission(snapshot.evaluation).map_err(|_| RuntimeError::SearchPermissionUnavailable)?;
Ok(
output
.docs
.into_iter()
.filter(|doc| {
doc
.decisions
.iter()
.any(|decision| decision.action == "Doc.Read" && decision.allowed)
})
.map(|doc| doc.doc_id)
.collect(),
)
}
}
packages/backend/native/src/runtime/backend_runtime/permission/mod.rs
+11
View File
@@ -0,0 +1,11 @@
mod authorizer;
mod store;
mod types;
pub(super) use authorizer::PermissionAuthorizer;
#[cfg(test)]
pub(super) use types::AclPredicate;
pub(super) use types::{AuthorizedSearchScope, DocReadScope, SearchActor, SystemSearchCapability};
#[cfg(test)]
mod tests;
packages/backend/native/src/runtime/backend_runtime/permission/store.rs
+225
View File
@@ -0,0 +1,225 @@
use chrono::{DateTime, Utc};
use sqlx::{PgPool, Row};
use super::types::{DocAclCapability, PermissionSnapshot};
use crate::{
entitlement::entitlement_priority,
permission::{
PermissionDocInputV1, PermissionEvaluationInputV1, PermissionRuntimeInputV1, PermissionSubjectInputV1,
PermissionWorkspaceInputV1,
},
runtime::{RuntimeError, RuntimeResult},
};
pub(super) struct PermissionStore {
pool: PgPool,
}
impl PermissionStore {
pub(super) fn new(pool: PgPool) -> Self {
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,
user_id: &str,
doc_ids: &[String],
) -> RuntimeResult<PermissionSnapshot> {
let mut transaction = self
.pool
.begin()
.await
.map_err(|error| RuntimeError::database("begin permission snapshot", error))?;
sqlx::query("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ, READ ONLY")
.execute(&mut *transaction)
.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,
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
WHERE workspace_id=workspace.id AND user_id=$2
ORDER BY (state='active') DESC, updated_at DESC LIMIT 1
) member ON true
WHERE workspace.id=$1"#,
)
.bind(workspace_id)
.bind(user_id)
.fetch_optional(&mut *transaction)
.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))?;
let role: Option<String> = row
.try_get("role")
.map_err(|error| RuntimeError::database("decode workspace member role", error))?;
let member_state: Option<String> = row
.try_get("state")
.map_err(|error| RuntimeError::database("decode workspace member state", error))?;
let visibility: Option<String> = row
.try_get("visibility")
.map_err(|error| RuntimeError::database("decode workspace visibility", error))?;
let member_default_doc_role: String = row
.try_get("member_default_doc_role")
.map_err(|error| RuntimeError::database("decode member default doc role", error))?;
let capability = load_doc_acl_capability(&mut transaction, workspace_id).await?;
let docs = if doc_ids.is_empty() {
Vec::new()
} else {
sqlx::query(
r#"SELECT candidate.doc_id, policy.visibility, policy.public_role,
coalesce(policy.member_default_role, $3) AS member_default_role,
grant_fact.role AS explicit_user_role
FROM unnest($4::text[]) candidate(doc_id)
LEFT JOIN doc_access_policies policy
ON policy.workspace_id=$1 AND policy.doc_id=candidate.doc_id
LEFT JOIN doc_grants grant_fact
ON grant_fact.workspace_id=$1 AND grant_fact.doc_id=candidate.doc_id
AND grant_fact.principal_type='user' AND grant_fact.principal_id=$2"#,
)
.bind(workspace_id)
.bind(user_id)
.bind(&member_default_doc_role)
.bind(doc_ids)
.fetch_all(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("load document permission facts", error))?
.into_iter()
.map(|row| {
Ok(PermissionDocInputV1 {
doc_id: row
.try_get("doc_id")
.map_err(|error| RuntimeError::database("decode permission doc id", error))?,
actions: vec!["Doc.Read".to_string()],
explicit_user_role: row
.try_get("explicit_user_role")
.map_err(|error| RuntimeError::database("decode explicit doc role", error))?,
member_default_role: row
.try_get("member_default_role")
.map_err(|error| RuntimeError::database("decode member default role", error))?,
public_role: row
.try_get("public_role")
.map_err(|error| RuntimeError::database("decode public doc role", error))?,
visibility: row
.try_get("visibility")
.map_err(|error| RuntimeError::database("decode doc visibility", error))?,
sharing_enabled: Some(sharing_enabled),
..Default::default()
})
})
.collect::<RuntimeResult<Vec<_>>>()?
};
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit permission snapshot", error))?;
let active_member =
member_state.as_deref() == Some("active") && matches!(role.as_deref(), Some("member" | "admin" | "owner"));
Ok(PermissionSnapshot {
revision,
capability,
evaluation: PermissionEvaluationInputV1 {
version: 1,
legacy_compat_mode: false,
subject: PermissionSubjectInputV1 {
user_id: Some(user_id.to_string()),
..Default::default()
},
runtime: PermissionRuntimeInputV1 {
known: true,
sharing_enabled: Some(sharing_enabled),
..Default::default()
},
workspace: PermissionWorkspaceInputV1 {
role,
member_state,
public: visibility.as_deref() == Some("public"),
sharing_enabled: Some(sharing_enabled),
..Default::default()
},
workspace_actions: vec!["Workspace.Read".to_string()],
docs,
},
actor_user_id: user_id.to_string(),
active_member,
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(
transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
workspace_id: &str,
) -> RuntimeResult<DocAclCapability> {
let rows = sqlx::query(
r#"SELECT plan,status,expires_at,grace_until,validated_at,source,signed_payload
FROM entitlements
WHERE target_type='workspace' AND target_id=$1
AND ((status='active' AND (expires_at IS NULL OR expires_at>now()))
OR (status='grace' AND grace_until>now()))"#,
)
.bind(workspace_id)
.fetch_all(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("load workspace entitlement facts", error))?;
let mut best: Option<(i32, String)> = None;
for row in rows {
let plan: String = row
.try_get("plan")
.map_err(|error| RuntimeError::database("decode entitlement plan", error))?;
let status: String = row
.try_get("status")
.map_err(|error| RuntimeError::database("decode entitlement status", error))?;
let source: String = row
.try_get("source")
.map_err(|error| RuntimeError::database("decode entitlement source", error))?;
let validated_at: Option<DateTime<Utc>> = row
.try_get("validated_at")
.map_err(|error| RuntimeError::database("decode entitlement validation", error))?;
let signed_payload: Option<Vec<u8>> = row
.try_get("signed_payload")
.map_err(|error| RuntimeError::database("decode entitlement payload", error))?;
if source == "selfhost_license" && (validated_at.is_none() || signed_payload.is_none()) {
continue;
}
let priority = entitlement_priority(&status, &plan);
if best.as_ref().is_none_or(|(current, _)| priority > *current) {
best = Some((priority, plan));
}
}
match best.map(|(_, plan)| plan) {
Some(plan) if matches!(plan.as_str(), "team" | "selfhost_team") => Ok(DocAclCapability::Enabled),
Some(plan) if matches!(plan.as_str(), "free" | "pro" | "lifetime_pro" | "ai" | "selfhost_free") => {
Ok(DocAclCapability::Disabled)
}
Some(_) => Ok(DocAclCapability::Unknown),
None => Ok(DocAclCapability::Disabled),
}
}
packages/backend/native/src/runtime/backend_runtime/permission/tests.rs
+165
View File
@@ -0,0 +1,165 @@
use sqlx::PgPool;
use super::{DocReadScope, PermissionAuthorizer, SearchActor};
static PERMISSION_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
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 suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("search-permission-user-{suffix}");
let workspace_id = format!("search-permission-workspace-{suffix}");
sqlx::query(
"INSERT INTO users(id,name,email,registered,email_verified,disabled) VALUES($1,'Search Permission \
User',$2,true,now(),false)",
)
.bind(&user_id)
.bind(format!("search-permission-{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)")
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
sqlx::query("INSERT INTO workspace_members(workspace_id,user_id,role,state) VALUES($1,$2,'member','active')")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
Some((pool, workspace_id, user_id))
}
#[tokio::test]
async fn non_team_is_all_and_team_uses_projected_acl() {
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 actor = SearchActor::User {
user_id: user_id.clone(),
};
let free = authorizer.authorize_search(&actor, &workspace_id).await.unwrap();
assert_eq!(free.docs, DocReadScope::All);
sqlx::query(
"INSERT INTO entitlements(id,target_type,target_id,source,plan,status,validated_at) \
VALUES($1,'workspace',$2,'admin_grant','team','active',now())",
)
.bind(format!("search-permission-entitlement-{workspace_id}"))
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let team = authorizer.authorize_search(&actor, &workspace_id).await.unwrap();
let DocReadScope::ProjectedAcl(predicate) = team.docs else {
panic!("team member must use projected ACL");
};
assert_eq!(predicate.actor_user_id, user_id);
assert!(predicate.active_member);
assert!(team.permission_revision > free.permission_revision);
}
#[tokio::test]
async fn inactive_member_is_denied_and_unknown_capability_fails_closed() {
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 actor = SearchActor::User {
user_id: user_id.clone(),
};
sqlx::query("UPDATE workspace_members SET state='suspended' WHERE workspace_id=$1 AND user_id=$2")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
let error = authorizer.authorize_search(&actor, &workspace_id).await.unwrap_err();
assert!(matches!(error, crate::runtime::RuntimeError::SearchWorkspaceDenied));
sqlx::query("UPDATE workspace_members SET state='active' WHERE workspace_id=$1 AND user_id=$2")
.bind(&workspace_id)
.bind(&user_id)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO entitlements(id,target_type,target_id,source,plan,status,validated_at) \
VALUES($1,'workspace',$2,'admin_grant','future_plan','active',now())",
)
.bind(format!("search-permission-entitlement-{workspace_id}"))
.bind(&workspace_id)
.execute(&pool)
.await
.unwrap();
let error = authorizer.authorize_search(&actor, &workspace_id).await.unwrap_err();
assert!(matches!(
error,
crate::runtime::RuntimeError::SearchPermissionUnavailable
));
}
#[tokio::test]
async fn fact_changes_advance_revision_and_write_ordered_change() {
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_grants(workspace_id,doc_id,principal_type,principal_id,role) VALUES($1,'doc','user',$2,'reader')",
)
.bind(&workspace_id)
.bind(&user_id)
.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)
.await
.unwrap();
assert_eq!(authorizer.revision(&workspace_id).await.unwrap(), after);
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)
.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);
}
packages/backend/native/src/runtime/backend_runtime/permission/types.rs
+47
View File
@@ -0,0 +1,47 @@
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 },
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum DocAclCapability {
Enabled,
Disabled,
Unknown,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) struct AclPredicate {
pub(in crate::runtime::backend_runtime) actor_user_id: String,
pub(in crate::runtime::backend_runtime) active_member: bool,
pub(in crate::runtime::backend_runtime) sharing_enabled: bool,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(in crate::runtime::backend_runtime) enum DocReadScope {
All,
ProjectedAcl(AclPredicate),
}
#[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,
pub(super) active_member: bool,
pub(super) sharing_enabled: bool,
}
packages/backend/native/src/runtime/backend_runtime/role.rs
+98
View File
@@ -0,0 +1,98 @@
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum ServerRole {
Frontend,
Api,
Worker,
AllInOne,
}
impl ServerRole {
pub(super) fn from_environment() -> Result<(Self, bool), String> {
let script_mode = matches!(std::env::var("SERVER_FLAVOR").as_deref(), Ok("script"));
if let Ok(value) = std::env::var("AFFINE_SERVER_ROLE") {
return Self::parse(&value).map(|role| (role, script_mode));
}
match std::env::var("SERVER_FLAVOR") {
Err(std::env::VarError::NotPresent) => Ok((Self::AllInOne, false)),
Ok(value) => Self::from_flavor(&value),
Err(std::env::VarError::NotUnicode(_)) => Err("backend runtime role source is not valid unicode".to_string()),
}
}
fn from_flavor(value: &str) -> Result<(Self, bool), String> {
match value {
"allinone" => Ok((Self::AllInOne, false)),
"front" => Ok((Self::Frontend, false)),
"graphql" => Ok((Self::Api, false)),
"worker" => Ok((Self::Worker, false)),
"sync" | "renderer" => Ok((Self::Frontend, false)),
// The CLI uses the BackendRuntime only for database/object-storage work.
// It is not one of the four server roles and must not initialize search.
"script" => Ok((Self::Frontend, true)),
value => Err(format!("unsupported backend runtime role source value: {value}")),
}
}
fn parse(value: &str) -> Result<Self, String> {
match value {
"frontend" => Ok(Self::Frontend),
"api" => Ok(Self::Api),
"worker" => Ok(Self::Worker),
"allinone" => Ok(Self::AllInOne),
_ => Err(format!(
"unsupported backend runtime role: {value}; expected frontend, api, worker, or allinone"
)),
}
}
pub(super) fn owns_background(self) -> bool {
matches!(self, Self::Worker | Self::AllInOne)
}
pub(super) fn allows_embedded_search(self) -> bool {
matches!(self, Self::AllInOne)
}
pub(super) fn as_str(self) -> &'static str {
match self {
Self::Frontend => "frontend",
Self::Api => "api",
Self::Worker => "worker",
Self::AllInOne => "allinone",
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn role_parser_is_closed() {
assert_eq!(ServerRole::parse("frontend"), Ok(ServerRole::Frontend));
assert_eq!(ServerRole::parse("api"), Ok(ServerRole::Api));
assert_eq!(ServerRole::parse("worker"), Ok(ServerRole::Worker));
assert_eq!(ServerRole::parse("allinone"), Ok(ServerRole::AllInOne));
assert!(ServerRole::parse("graphql").is_err());
assert!(ServerRole::from_flavor("doc").is_err());
assert_eq!(ServerRole::from_flavor("script"), Ok((ServerRole::Frontend, true)));
assert_eq!(ServerRole::from_flavor("worker"), Ok((ServerRole::Worker, false)));
}
#[test]
fn only_all_in_one_allows_embedded_search() {
assert!(!ServerRole::Frontend.allows_embedded_search());
assert!(!ServerRole::Api.allows_embedded_search());
assert!(!ServerRole::Worker.allows_embedded_search());
assert!(ServerRole::AllInOne.allows_embedded_search());
}
#[test]
fn only_worker_compositions_own_background() {
assert!(!ServerRole::Frontend.owns_background());
assert!(!ServerRole::Api.owns_background());
assert!(ServerRole::Worker.owns_background());
assert!(ServerRole::AllInOne.owns_background());
}
}
packages/backend/native/src/runtime/backend_runtime/scope_compiler.rs
+10 -44
View File
@@ -6,18 +6,22 @@ use affine_doc_loader::{
use chrono::Utc;
use sqlx::{PgPool, Row};
use super::{RuntimeError, RuntimeResult, types};
use super::{RuntimeError, RuntimeResult, permission::PermissionAuthorizer, types};
use crate::{runtime::storage_runtime::load_current_doc, userdata_acl};
const REQUIRED_DOCUMENT_LIMIT: usize = 64;
pub(super) struct ScopeCompiler {
pool: PgPool,
authorizer: PermissionAuthorizer,
}
impl ScopeCompiler {
pub(super) fn new(pool: PgPool) -> Self {
Self { pool }
Self {
authorizer: PermissionAuthorizer::new(pool.clone()),
pool,
}
}
pub(super) async fn compile(
@@ -63,10 +67,11 @@ impl ScopeCompiler {
.await?;
let readable = self
.readable_doc_ids(
.authorizer
.filter_readable_docs(
&input.workspace_id,
&input.user_id,
facts.documents.iter().map(|doc| doc.id.as_str()),
facts.documents.iter().map(|doc| doc.id.clone()).collect(),
)
.await?;
let mut required_docs = BTreeSet::new();
@@ -154,46 +159,6 @@ impl ScopeCompiler {
Ok(())
}
async fn readable_doc_ids<'a>(
&self,
workspace_id: &str,
user_id: &str,
doc_ids: impl Iterator<Item = &'a str>,
) -> RuntimeResult<BTreeSet<String>> {
let doc_ids = doc_ids.map(str::to_string).collect::<Vec<_>>();
let rows = sqlx::query(
r#"SELECT candidate.doc_id FROM unnest($3::text[]) candidate(doc_id)
WHERE EXISTS (
SELECT 1 FROM workspace_access_policies workspace_policy
LEFT JOIN doc_access_policies doc_policy
ON doc_policy.workspace_id=workspace_policy.workspace_id AND doc_policy.doc_id=candidate.doc_id
LEFT JOIN workspace_members member
ON member.workspace_id=workspace_policy.workspace_id AND member.user_id=$2 AND member.state='active'
LEFT JOIN doc_grants grant_fact
ON grant_fact.workspace_id=workspace_policy.workspace_id AND grant_fact.doc_id=candidate.doc_id
AND grant_fact.principal_type='user' AND grant_fact.principal_id=$2
WHERE workspace_policy.workspace_id=$1 AND (
member.id IS NOT NULL AND grant_fact.role=ANY(ARRAY['owner','manager','editor','commenter','reader']::text[])
OR member.id IS NULL AND workspace_policy.sharing_enabled
AND grant_fact.role=ANY(ARRAY['owner','manager','editor','commenter','reader']::text[])
OR member.role=ANY(ARRAY['owner','admin']::text[])
OR member.id IS NOT NULL AND grant_fact.principal_id IS NULL
AND coalesce(doc_policy.member_default_role,workspace_policy.member_default_doc_role)
=ANY(ARRAY['owner','manager','editor','commenter','reader']::text[])
OR workspace_policy.sharing_enabled AND doc_policy.visibility='public'
AND doc_policy.public_role=ANY(ARRAY['owner','manager','editor','commenter','reader','external']::text[])
)
)"#,
)
.bind(workspace_id)
.bind(user_id)
.bind(doc_ids)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("filter scope document permissions failed", error))?;
Ok(rows.into_iter().map(|row| row.get("doc_id")).collect())
}
async fn artifact_is_readable(&self, workspace_id: &str, user_id: &str, artifact_id: &str) -> RuntimeResult<bool> {
let id = artifact_id
.parse::<uuid::Uuid>()
@@ -285,6 +250,7 @@ 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 suffix = uuid::Uuid::new_v4().simple().to_string();
let user_id = format!("scope-user-{suffix}");
let collaborator_id = format!("scope-collaborator-{suffix}");
packages/backend/native/src/runtime/backend_runtime/search/checkpoint.rs
+193
View File
@@ -0,0 +1,193 @@
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()
}
packages/backend/native/src/runtime/backend_runtime/search/generation.rs
+189
View File
@@ -0,0 +1,189 @@
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use sqlx::{PgPool, Row};
use uuid::Uuid;
use super::{SCHEMA_FINGERPRINT, provider::RemoteProvider, types::SearchTable};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
#[derive(Clone)]
pub(super) struct ActiveGeneration {
pub(super) id: Uuid,
pub(super) manifest: Value,
}
impl ActiveGeneration {
pub(super) fn physical_table(&self, table: SearchTable) -> RuntimeResult<&str> {
self
.manifest
.get(table.as_str())
.and_then(Value::as_str)
.ok_or_else(|| RuntimeError::invalid_state("search generation manifest is incomplete"))
}
}
pub(super) async fn prepare(
pool: &PgPool,
config: &SearchRuntimeConfig,
remote: Option<&RemoteProvider>,
) -> RuntimeResult<ActiveGeneration> {
let fingerprint = config_fingerprint(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))")
.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"#,
)
.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"})
} 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))?;
}
ActiveGeneration {
id: generation_id,
manifest,
}
};
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit pending search generation", error))?;
if let Some(remote) = remote {
for table in [SearchTable::Doc, SearchTable::Block] {
if let Err(error) = remote.provision(generation.physical_table(table)?, table).await {
fail(pool, &generation).await?;
return Err(error);
}
}
}
Ok(generation)
}
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
ORDER BY activated_at DESC NULLS LAST LIMIT 1"#,
)
.bind(&config.provider)
.bind(&fingerprint)
.bind(SCHEMA_FINGERPRINT)
.fetch_optional(pool)
.await
.map_err(|error| RuntimeError::database("load active search generation", error))?;
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))?;
Ok(())
}
pub(super) async fn activate(pool: &PgPool, generation: &ActiveGeneration) -> RuntimeResult<()> {
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)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("drain previous search generation", error))?;
sqlx::query(
"UPDATE search_runtime_generations SET state='active', activated_at=coalesce(activated_at,now()) WHERE \
generation_id=$1 AND state IN ('pending','active')",
)
.bind(generation.id)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("activate search generation", error))?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("commit search generation activation", error))
}
fn config_fingerprint(config: &SearchRuntimeConfig) -> String {
let mut hash = Sha256::new();
for value in [
&config.provider,
&config.endpoint,
&config.api_key,
&config.username,
&config.password,
] {
hash.update(value.as_bytes());
hash.update([0]);
}
hash.finalize().iter().map(|byte| format!("{byte:02x}")).collect()
}
fn decode(row: sqlx::postgres::PgRow) -> RuntimeResult<ActiveGeneration> {
Ok(ActiveGeneration {
id: row
.try_get("generation_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))?,
})
}
packages/backend/native/src/runtime/backend_runtime/search/mod.rs
+33
View File
@@ -0,0 +1,33 @@
mod checkpoint;
mod generation;
mod projection;
mod provider;
mod query;
mod runtime;
mod store;
mod types;
mod worker;
pub(super) use runtime::SearchRuntime;
pub(super) use types::{RuntimeAggregateRequest, RuntimeSearchRequest};
const SCHEMA_FINGERPRINT: &str = "search-runtime-v5";
fn exact_token(value: &str) -> String {
use sha2::{Digest, Sha256};
Sha256::digest(value.as_bytes())
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
fn provider_payload(payload: &serde_json::Value) -> serde_json::Value {
let mut payload = payload.clone();
if let Some(object) = payload.as_object_mut() {
object.remove("acl_read_user_ids");
}
payload
}
#[cfg(test)]
mod tests;
packages/backend/native/src/runtime/backend_runtime/search/projection.rs
+207
View File
@@ -0,0 +1,207 @@
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},
};
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 {
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(
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
WHERE snapshot.workspace_id=$1 AND snapshot.guid=$2"#,
)
.bind(workspace_id)
.bind(doc_id)
.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"))?;
let visibility: String = metadata
.try_get("visibility")
.map_err(|error| RuntimeError::database("decode search doc visibility", error))?;
let public_role: Option<String> = metadata
.try_get("public_role")
.map_err(|error| RuntimeError::database("decode search public role", error))?;
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 grants = sqlx::query(
"SELECT principal_id,role FROM doc_grants WHERE workspace_id=$1 AND doc_id=$2 AND principal_type='user'",
)
.bind(workspace_id)
.bind(doc_id)
.fetch_all(pool)
.await
.map_err(|error| RuntimeError::database("load search doc grants", error))?;
let acl_read_user_ids = grants
.into_iter()
.map(|row| {
let id: String = row
.try_get("principal_id")
.map_err(|error| RuntimeError::database("decode grant principal", error))?;
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))
})
.collect::<RuntimeResult<Vec<_>>>()?
.into_iter()
.flatten()
.collect::<Vec<_>>();
let created_at: chrono::DateTime<chrono::Utc> = metadata
.try_get("created_at")
.map_err(|error| RuntimeError::database("decode search created time", error))?;
let updated_at: chrono::DateTime<chrono::Utc> = metadata
.try_get("updated_at")
.map_err(|error| RuntimeError::database("decode search updated time", error))?;
let created_by: Option<String> = metadata
.try_get("created_by")
.map_err(|error| RuntimeError::database("decode search creator", error))?;
let updated_by: Option<String> = metadata
.try_get("updated_by")
.map_err(|error| RuntimeError::database("decode search updater", error))?;
let acl = AclFields {
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!({
"workspace_id": workspace_id,
"workspace_token": super::exact_token(workspace_id),
"doc_id": doc_id,
"doc_token": super::exact_token(doc_id),
"title": projection.title,
"summary": projection.units.iter().map(|unit| unit.text.as_str()).collect::<Vec<_>>().join("\n").chars().take(1000).collect::<String>(),
"created_by_user_id": created_by.clone().unwrap_or_default(),
"updated_by_user_id": updated_by.clone().unwrap_or_default(),
"created_at": created_at.timestamp_millis(),
"updated_at": updated_at.timestamp_millis(),
}),
&acl,
);
let document = input(
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}"),
revision,
document_payload,
&acl,
);
let blocks = projection
.units
.into_iter()
.map(|unit| {
let block_id = unit.block_id.clone().unwrap_or_else(|| unit.unit_id.clone());
let payload = with_acl(
json!({
"workspace_id":workspace_id,"workspace_token":super::exact_token(workspace_id),
"doc_id":doc_id,"doc_token":super::exact_token(doc_id),
"block_id":block_id,"block_token":super::exact_token(&block_id),
"unit_id":unit.unit_id,"projection_version":projection.version,
"source_hash":projection.source_hash,"visibility":serde_json::to_value(unit.visibility).unwrap_or(Value::Null),
"element_id":unit.element_id,"frame_id":unit.frame_id,"source_block_id":unit.block_id,
"blob":unit.blob_id,"ref_doc_id":unit.ref_doc_ids,"ref":unit.refs,"content":unit.text,
"flavour":format!("affine:{}",unit.unit_type),"parent_flavour":unit.parent_flavour,
"parent_block_id":unit.parent_block_id,"additional":unit.additional,
"created_by_user_id":created_by.clone().unwrap_or_default(),"updated_by_user_id":updated_by.clone().unwrap_or_default(),
"created_at":created_at.timestamp_millis(),"updated_at":updated_at.timestamp_millis(),
}),
&acl,
);
input(
workspace_id,
doc_id,
&format!("{workspace_id}/{doc_id}/{block_id}"),
revision,
payload,
&acl,
)
})
.collect();
Ok(Some((document, blocks)))
}
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 {
let object = payload.as_object_mut().expect("search projection payload is an object");
object.insert("acl_public_readable".to_string(), json!(acl.public_readable));
object.insert(
"acl_member_default_readable".to_string(),
json!(acl.member_default_readable),
);
let mut tokens = acl
.read_user_ids
.iter()
.map(|user_id| super::exact_token(user_id))
.collect::<Vec<_>>();
if acl.member_default_readable {
tokens.push("member".to_string());
}
if acl.public_readable {
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 {
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,
}
}
packages/backend/native/src/runtime/backend_runtime/search/provider/manticore.rs
+326
View File
@@ -0,0 +1,326 @@
use serde_json::{Value, json};
use crate::runtime::{RuntimeError, RuntimeResult};
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(())
}
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)
.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")
.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")
.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"));
}
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(());
}
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)?;
}
}
_ => {}
}
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);
};
if term.len() != 1 {
return Ok(None);
}
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"
};
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}}))
}
}
Value::Bool(value) => Some(json!({"equals":{field:u8::from(*value)}})),
Value::Number(value) => Some(json!({"equals":{field:value}})),
_ => None,
})
}
pub(super) fn manticore_exact_tokens(value: &Value) -> Vec<String> {
let mut tokens = Vec::new();
collect_exact_tokens(value, &mut tokens);
tokens
}
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);
}
}
_ => {}
}
}
#[cfg(test)]
mod tests {
use super::{super::super::exact_token, *};
#[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(),
],
);
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();
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}}
]}}})
);
}
}
packages/backend/native/src/runtime/backend_runtime/search/provider/mod.rs
+125
View File
@@ -0,0 +1,125 @@
mod manticore;
mod remote;
pub(super) use remote::RemoteProvider;
use serde_json::{Value, json};
use super::types::SearchTable;
pub(super) fn mapping(table: SearchTable, provider: &str) -> 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"})),
("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_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"})),
]);
if table == SearchTable::Block {
for field in [
"block_id",
"block_token",
"unit_id",
"source_hash",
"visibility",
"element_id",
"frame_id",
"source_block_id",
"flavour",
"blob",
"ref_doc_id",
"parent_flavour",
"parent_block_id",
] {
properties.insert(field.into(), json!({"type":"keyword"}));
}
properties.insert("projection_version".into(), json!({"type":"integer"}));
for field in ["ref", "additional", "markdown_preview"] {
properties.insert(field.into(), json!({"type":"text","index":false}));
}
} else {
properties.insert("summary".into(), json!({"type":"text","index":false}));
properties.insert("journal".into(), json!({"type":"keyword"}));
}
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'"
)
}
packages/backend/native/src/runtime/backend_runtime/search/provider/remote.rs
+401
View File
@@ -0,0 +1,401 @@
use std::time::Duration;
use reqwest::{Client, redirect::Policy};
use serde_json::{Value, json};
use sqlx::PgPool;
use super::{
super::{store::SearchChange, types::SearchTable},
manticore::{manticore_exact_tokens, manticore_fields, prepare_manticore_payload, prepare_manticore_search},
};
use crate::runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig};
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> {
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()
.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<_>>();
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 {
let cursor = cursor
.as_str()
.ok_or_else(|| RuntimeError::invalid_input("invalid search cursor"))?;
dsl["search_after"] =
serde_json::from_str(cursor).map_err(|error| RuntimeError::json("invalid search cursor", error))?;
}
let mut request = self
.client
.post(format!("{}/{physical_table}/_search", self.endpoint))
.json(&dsl);
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));
}
let response = request
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let bytes = 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(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize(
value,
self.provider == "manticoresearch",
offset,
size,
&requested_fields,
)
}
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);
}
dsl["track_total_hits"] = json!(true);
let response = self
.request(reqwest::Method::POST, &format!("{physical_table}/_search"))
.json(&dsl)
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
let status = response.status();
let bytes = read_response(response).await?;
if !status.is_success() {
return Err(RuntimeError::SearchUnsupportedQuery);
}
let value: Value =
serde_json::from_slice(&bytes).map_err(|error| RuntimeError::json("invalid search provider response", error))?;
normalize_aggregate(value)
}
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);
}
if self
.request(reqwest::Method::HEAD, physical_table)
.send()
.await
.is_ok_and(|response| response.status().is_success())
{
return Ok(());
}
let response = self
.request(reqwest::Method::PUT, physical_table)
.json(&super::mapping(table, &self.provider))
.send()
.await;
match response {
Ok(response) if response.status().is_success() => Ok(()),
_ => Err(RuntimeError::SearchProviderUnavailable),
}
}
pub(in crate::runtime::backend_runtime::search) 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()
};
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 path = if self.provider == "elasticsearch" {
"_bulk?refresh=wait_for"
} else {
"_bulk"
};
let response = self
.request(reqwest::Method::POST, path)
.header("content-type", "application/x-ndjson")
.body(body)
.send()
.await
.map_err(|_| RuntimeError::SearchProviderUnavailable)?;
if !response.status().is_success() {
return Err(RuntimeError::SearchProviderUnavailable);
}
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"));
}
Ok(())
}
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
}
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())
}
}
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 normalize(
value: Value,
manticore: bool,
_offset: u64,
_size: u64,
requested_fields: &[String],
) -> 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/value")
.or_else(|| 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 fields = if manticore {
manticore_fields(hit.get("_source"), requested_fields)
} else {
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(),
"fields":fields,
"highlights":hit.get("highlight").cloned().unwrap_or_else(||json!({})),
"_source":hit.get("_source").cloned().unwrap_or_else(||json!({})),
})
})
.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))?
};
Ok(json!({"total":total,"nodes":nodes,"nextCursor":next_cursor}))
}
fn normalize_aggregate(value: Value) -> RuntimeResult<Value> {
let buckets = value
.pointer("/aggregations/result/buckets")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider aggregate response"))?;
let nodes = buckets
.iter()
.map(|bucket| {
let hits = bucket
.pointer("/result/hits/hits")
.and_then(Value::as_array)
.ok_or_else(|| RuntimeError::invalid_state("invalid provider aggregate hits"))?;
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<_>>()}
}))
})
.collect::<RuntimeResult<Vec<_>>>()?;
Ok(json!({"total":nodes.len(),"hasMore":false,"buckets":nodes}))
}
fn normalize_hit(hit: &Value) -> Value {
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(),
"fields":hit.get("fields").cloned().unwrap_or_else(||json!({})),
"highlights":hit.get("highlight").cloned().unwrap_or_else(||json!({})),
"_source":hit.get("_source").cloned().unwrap_or_else(||json!({})),
})
}
packages/backend/native/src/runtime/backend_runtime/search/query.rs
+336
View File
@@ -0,0 +1,336 @@
use serde_json::{Value, json};
use super::types::{AggregateRequest, SearchOptions, SearchQuery, SearchRequest, SearchTable};
use crate::runtime::{
RuntimeError, RuntimeResult,
backend_runtime::permission::{AuthorizedSearchScope, DocReadScope},
};
pub(super) fn compile(request: &SearchRequest, scope: &AuthorizedSearchScope) -> RuntimeResult<Value> {
let query = compile_query(request.table, &request.query)?;
let mut must = vec![json!({"term":{"workspace_id":{"value":scope.workspace_id}}}), query];
if let DocReadScope::ProjectedAcl(predicate) = &scope.docs {
let mut should = vec![json!({"term":{"acl_read_tokens":{"value":super::exact_token(&predicate.actor_user_id)}}})];
if predicate.active_member {
should.push(json!({"term":{"acl_read_tokens":{"value":"member"}}}));
}
if predicate.sharing_enabled {
should.push(json!({"term":{"acl_read_tokens":{"value":"public"}}}));
}
must.push(json!({"bool":{"should":should}}));
}
let fields = request
.options
.fields
.iter()
.map(|field| validate_field(request.table, field).map(str::to_string))
.collect::<RuntimeResult<Vec<_>>>()?;
let mut dsl = json!({
"_source":["workspace_id","doc_id"],
"fields":fields,
"query":{"bool":{"must":must}},
"sort": stable_sort(request.table),
});
let pagination = &request.options.pagination;
if pagination.limit.unwrap_or(10) > 10_000 {
return Err(RuntimeError::invalid_input("search limit exceeds 10000"));
}
dsl["size"] = json!(pagination.limit.unwrap_or(10));
if let Some(skip) = pagination.skip {
if skip.saturating_add(pagination.limit.unwrap_or(10)) > 10_000 {
return Err(RuntimeError::invalid_input("search offset exceeds 10000"));
}
dsl["from"] = json!(skip);
}
if let Some(cursor) = &pagination.cursor {
dsl["cursor"] = json!(cursor);
}
if !request.options.highlights.is_empty() {
let mut highlights = serde_json::Map::new();
for highlight in &request.options.highlights {
let field = validate_field(request.table, &highlight.field)?;
highlights.insert(
field.to_string(),
json!({"pre_tags":[highlight.before],"post_tags":[highlight.end]}),
);
}
dsl["highlight"] = json!({"fields":highlights});
}
Ok(dsl)
}
pub(super) fn compile_aggregate(request: &AggregateRequest, scope: &AuthorizedSearchScope) -> RuntimeResult<Value> {
let hits = SearchOptions {
fields: request.options.hits.fields.clone(),
highlights: request.options.hits.highlights.clone(),
pagination: request.options.hits.pagination.clone(),
};
let search = SearchRequest {
table: request.table,
query: request.query.clone(),
options: hits,
};
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"));
}
Ok(json!({
"query":hit_dsl["query"],
"from":request.options.pagination.skip.unwrap_or(0),
"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"],
"sort":hit_dsl["sort"],"highlight":hit_dsl.get("highlight").cloned().unwrap_or_else(||json!({}))
}}}}}
}))
}
fn compile_query(table: SearchTable, query: &SearchQuery) -> RuntimeResult<Value> {
let boost = query.boost.unwrap_or(1.0);
if !boost.is_finite() || boost <= 0.0 {
return Err(RuntimeError::invalid_input("invalid search boost"));
}
match query.query_type.as_str() {
"match" => {
let field = validate_field(
table,
query
.field
.as_deref()
.ok_or_else(|| RuntimeError::invalid_input("match field is required"))?,
)?;
let value = query
.match_value
.as_deref()
.filter(|value| !value.is_empty())
.ok_or_else(|| RuntimeError::invalid_input("match value is required"))?;
if field == table.text_field() {
Ok(json!({"match":{field:{"query":value,"boost":boost}}}))
} else {
Ok(json!({"term":{field:{"value":value,"boost":boost}}}))
}
}
"boolean" => {
let occur = query
.occur
.as_deref()
.filter(|occur| matches!(*occur, "must" | "should" | "must_not"))
.ok_or_else(|| RuntimeError::invalid_input("invalid boolean occurrence"))?;
let clauses = query
.queries
.as_deref()
.ok_or_else(|| RuntimeError::invalid_input("boolean queries are required"))?
.iter()
.map(|query| compile_query(table, query))
.collect::<RuntimeResult<Vec<_>>>()?;
Ok(json!({"bool":{occur:clauses,"boost":boost}}))
}
"exists" => {
let field = validate_field(
table,
query
.field
.as_deref()
.ok_or_else(|| RuntimeError::invalid_input("exists field is required"))?,
)?;
Ok(json!({"exists":{"field":field,"boost":boost}}))
}
"all" => Ok(json!({"match_all":{"boost":boost}})),
"boost" => {
let mut nested = query
.query
.as_deref()
.ok_or_else(|| RuntimeError::invalid_input("boost query is required"))?
.clone();
nested.boost = Some(boost);
compile_query(table, &nested)
}
_ => Err(RuntimeError::invalid_input("unsupported search query")),
}
}
fn stable_sort(table: SearchTable) -> Value {
match table {
SearchTable::Doc => json!(["_score", {"updated_at":"desc"}, "doc_id"]),
SearchTable::Block => json!(["_score", {"updated_at":"desc"}, "doc_id", "block_id"]),
}
}
fn validate_field(table: SearchTable, field: &str) -> RuntimeResult<&'static str> {
let normalized = match field {
"workspaceId" => "workspace_id",
"docId" => "doc_id",
"blockId" => "block_id",
"createdByUserId" => "created_by_user_id",
"updatedByUserId" => "updated_by_user_id",
"createdAt" => "created_at",
"updatedAt" => "updated_at",
"refDocId" => "ref_doc_id",
"parentFlavour" => "parent_flavour",
"parentBlockId" => "parent_block_id",
"unitId" => "unit_id",
"projectionVersion" => "projection_version",
"sourceHash" => "source_hash",
"elementId" => "element_id",
"frameId" => "frame_id",
"sourceBlockId" => "source_block_id",
"markdownPreview" => "markdown_preview",
value => value,
};
let allowed = match table {
SearchTable::Doc => [
"workspace_id",
"doc_id",
"title",
"summary",
"journal",
"created_by_user_id",
"updated_by_user_id",
"created_at",
"updated_at",
]
.as_slice(),
SearchTable::Block => [
"workspace_id",
"doc_id",
"block_id",
"unit_id",
"projection_version",
"source_hash",
"visibility",
"element_id",
"frame_id",
"source_block_id",
"content",
"flavour",
"blob",
"ref_doc_id",
"ref",
"parent_flavour",
"parent_block_id",
"additional",
"markdown_preview",
"created_by_user_id",
"updated_by_user_id",
"created_at",
"updated_at",
]
.as_slice(),
};
allowed
.iter()
.find(|candidate| **candidate == normalized)
.copied()
.ok_or_else(|| RuntimeError::invalid_input("unknown or internal search field"))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::backend_runtime::permission::{AuthorizedSearchScope, DocReadScope};
fn scope() -> AuthorizedSearchScope {
AuthorizedSearchScope {
workspace_id: "workspace".to_string(),
permission_revision: 1,
docs: DocReadScope::All,
}
}
fn request(query: Value) -> SearchRequest {
serde_json::from_value(json!({
"table": "block",
"query": query,
"options": {
"fields": ["docId", "createdAt"],
"highlights": [{"field": "content", "before": "<b>", "end": "</b>"}],
"pagination": {"limit": 20, "skip": 5}
}
}))
.unwrap()
}
#[test]
fn compiles_supported_query_variants_and_options() {
let cases = [
(json!({"type":"all"}), json!({"match_all":{"boost":1.0}})),
(
json!({"type":"exists","field":"refDocId"}),
json!({"exists":{"field":"ref_doc_id","boost":1.0}}),
),
(
json!({"type":"boost","boost":2.5,"query":{"type":"match","field":"content","match":"hello"}}),
json!({"match":{"content":{"query":"hello","boost":2.5}}}),
),
(
json!({"type":"boolean","occur":"must_not","queries":[{"type":"match","field":"docId","match":"doc"}]}),
json!({"bool":{"must_not":[{"term":{"doc_id":{"value":"doc","boost":1.0}}}],"boost":1.0}}),
),
];
for (query, expected) in cases {
let dsl = compile(&request(query), &scope()).unwrap();
assert_eq!(dsl["query"]["bool"]["must"][1], expected);
assert_eq!(dsl["fields"], json!(["doc_id", "created_at"]));
assert_eq!(dsl["from"], 5);
assert_eq!(dsl["size"], 20);
assert_eq!(
dsl["highlight"]["fields"]["content"],
json!({"pre_tags":["<b>"],"post_tags":["</b>"]})
);
}
}
#[test]
fn validates_query_fields_and_pagination_limits() {
for query in [
json!({"type":"match","field":"aclReadTokens","match":"member"}),
json!({"type":"exists","field":"unknown"}),
json!({"type":"boolean","occur":"invalid","queries":[]}),
json!({"type":"boost","boost":0,"query":{"type":"all"}}),
] {
assert!(compile(&request(query), &scope()).is_err());
}
let mut oversized = request(json!({"type":"all"}));
oversized.options.pagination.limit = Some(10_001);
assert!(compile(&oversized, &scope()).is_err());
oversized.options.pagination.limit = Some(10_000);
oversized.options.pagination.skip = Some(1);
assert!(compile(&oversized, &scope()).is_err());
}
#[test]
fn compiles_aggregate_contract_and_rejects_invalid_fields() {
let aggregate: AggregateRequest = serde_json::from_value(json!({
"table":"block",
"query":{"type":"match","field":"content","match":"hello"},
"field":"docId",
"options":{
"hits":{
"fields":["docId","content"],
"highlights":[{"field":"content","before":"<b>","end":"</b>"}],
"pagination":{"limit":2}
},
"pagination":{"limit":50,"skip":3}
}
}))
.unwrap();
let dsl = compile_aggregate(&aggregate, &scope()).unwrap();
assert_eq!(dsl["from"], 3);
assert_eq!(dsl["aggs"]["result"]["terms"], json!({"field":"doc_id","size":50}));
assert_eq!(dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["size"], 2);
assert_eq!(
dsl["aggs"]["result"]["aggs"]["result"]["top_hits"]["highlight"]["fields"]["content"],
json!({"pre_tags":["<b>"],"post_tags":["</b>"]})
);
let mut invalid = aggregate;
invalid.field = "aclReadTokens".to_string();
assert!(compile_aggregate(&invalid, &scope()).is_err());
invalid.field = "docId".to_string();
invalid.options.pagination.limit = Some(10_001);
assert!(compile_aggregate(&invalid, &scope()).is_err());
}
}
packages/backend/native/src/runtime/backend_runtime/search/runtime.rs
+470
View File
@@ -0,0 +1,470 @@
use sqlx::PgPool;
use tokio::sync::{Mutex, RwLock};
use super::{
super::permission::{PermissionAuthorizer, SearchActor, SystemSearchCapability},
generation::{self, ActiveGeneration},
projection::project_document,
provider::RemoteProvider,
query,
store::{SearchStore, SearchTable},
types::{RuntimeAggregateRequest, RuntimeSearchRequest},
};
use crate::{
runtime::{RuntimeError, RuntimeResult, SearchRuntimeConfig},
search_index::EmbeddedSearchIndex,
};
pub(in crate::runtime::backend_runtime) struct SearchRuntime {
pool: PgPool,
store: SearchStore,
authorizer: PermissionAuthorizer,
pub(super) embedded: EmbeddedSearchIndex,
remote: Option<RemoteProvider>,
config: SearchRuntimeConfig,
generation: RwLock<Option<ActiveGeneration>>,
embedded_cursors: RwLock<[i64; 2]>,
embedded_permission_cursors: RwLock<std::collections::HashMap<String, i64>>,
sync_lock: Mutex<()>,
}
impl SearchRuntime {
pub(in crate::runtime::backend_runtime) fn new(pool: PgPool, config: SearchRuntimeConfig) -> RuntimeResult<Self> {
if !matches!(
config.provider.as_str(),
"embedded" | "elasticsearch" | "manticoresearch"
) {
return Err(RuntimeError::config("unsupported search provider"));
}
let remote = (config.provider != "embedded")
.then(|| RemoteProvider::new(&config, pool.clone()))
.transpose()?;
Ok(Self {
store: SearchStore::new(pool.clone()),
authorizer: PermissionAuthorizer::new(pool.clone()),
embedded: EmbeddedSearchIndex::new(),
remote,
config,
generation: RwLock::new(None),
embedded_cursors: RwLock::new([0; 2]),
embedded_permission_cursors: RwLock::new(std::collections::HashMap::new()),
sync_lock: Mutex::new(()),
pool,
})
}
pub(in crate::runtime::backend_runtime) async fn initialize(&self) -> RuntimeResult<()> {
let stream_count =
sqlx::query_scalar::<_, i64>("SELECT count(*) FROM search_runtime_streams WHERE table_key IN ('doc', 'block')")
.fetch_one(&self.pool)
.await
.map_err(|error| RuntimeError::database("check search runtime streams", error))?;
if stream_count != SearchTable::ORDERED.len() as i64 {
return Ok(());
}
let active = generation::prepare(&self.pool, &self.config, self.remote.as_ref()).await?;
if let Err(error) = super::worker::rebuild(
&self.pool,
&self.store,
&self.embedded,
self.remote.as_ref(),
&active,
&self.embedded_cursors,
true,
)
.await
{
generation::fail(&self.pool, &active).await?;
return Err(error);
}
generation::activate(&self.pool, &active).await?;
*self.generation.write().await = Some(active);
self.refresh_all_permission_cursors().await?;
Ok(())
}
async fn project_document(&self, workspace_id: &str, doc_id: &str) -> RuntimeResult<()> {
match project_document(&self.pool, workspace_id, doc_id).await? {
Some((document, blocks)) => self.store.replace_document(document, blocks).await?,
None => {
let revision = chrono::Utc::now().timestamp_millis();
self.store.delete_document(workspace_id, doc_id, revision).await?;
}
}
Ok(())
}
pub(in crate::runtime::backend_runtime) async fn project_document_only(
&self,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<()> {
self.project_document(workspace_id, doc_id).await
}
pub(in crate::runtime::backend_runtime) async fn index_document(
&self,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<()> {
self.project_document(workspace_id, doc_id).await?;
self.sync().await?;
self.refresh_permission_cursor(workspace_id).await
}
pub(in crate::runtime::backend_runtime) async fn delete_document_only(
&self,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<()> {
self
.store
.delete_document(workspace_id, doc_id, chrono::Utc::now().timestamp_millis())
.await?;
Ok(())
}
pub(in crate::runtime::backend_runtime) async fn delete_document(
&self,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<()> {
self.delete_document_only(workspace_id, doc_id).await?;
self.sync().await?;
self.refresh_permission_cursor(workspace_id).await
}
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?;
self
.check_permission_revision(workspace_id, scope.permission_revision)
.await?;
let generation = self.active_generation().await?;
self.ensure_query_ready(&generation).await?;
let dsl = query::compile(&request, &scope)?;
let result = if let Some(remote) = &self.remote {
remote.search(generation.physical_table(request.table)?, dsl).await?
} else {
let result = self
.embedded
.search(
request.table.as_str().to_string(),
serde_json::to_string(&dsl).map_err(|error| RuntimeError::json("encode embedded search", error))?,
)
.await?;
serde_json::from_str(&result).map_err(|error| RuntimeError::json("decode embedded search", error))?
};
if self.authorizer.revision(workspace_id).await? == scope.permission_revision {
return Ok(result);
}
if attempt == 1 {
return Err(RuntimeError::SearchPermissionUnavailable);
}
}
unreachable!()
}
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?;
self
.check_permission_revision(workspace_id, scope.permission_revision)
.await?;
let generation = self.active_generation().await?;
self.ensure_query_ready(&generation).await?;
let dsl = query::compile_aggregate(&request, &scope)?;
let result = if let Some(remote) = &self.remote {
remote.aggregate(generation.physical_table(request.table)?, dsl).await?
} else {
let result = self
.embedded
.aggregate(
request.table.as_str().to_string(),
serde_json::to_string(&dsl).map_err(|error| RuntimeError::json("encode embedded aggregate", error))?,
)
.await?;
let mut value: serde_json::Value =
serde_json::from_str(&result).map_err(|error| RuntimeError::json("decode embedded aggregate", error))?;
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});
}
}
value
};
if self.authorizer.revision(workspace_id).await? == scope.permission_revision {
return Ok(result);
}
if attempt == 1 {
return Err(RuntimeError::SearchPermissionUnavailable);
}
}
unreachable!()
}
async fn active_generation(&self) -> RuntimeResult<ActiveGeneration> {
if let Some(active) = self.generation.read().await.clone() {
return Ok(active);
}
if let Some(active) = generation::load_active(&self.pool, &self.config).await? {
*self.generation.write().await = Some(active.clone());
return Ok(active);
}
Err(RuntimeError::invalid_state("search_runtime_not_ready"))
}
pub(super) async fn sync(&self) -> RuntimeResult<()> {
let _guard = self.sync_lock.lock().await;
let generation = self.active_generation().await?;
let result = super::worker::sync(
&self.pool,
&self.store,
&self.embedded,
self.remote.as_ref(),
&generation,
&self.embedded_cursors,
)
.await;
if matches!(result, Err(RuntimeError::SearchReplayGap)) && self.remote.is_none() {
return super::worker::rebuild(
&self.pool,
&self.store,
&self.embedded,
None,
&generation,
&self.embedded_cursors,
false,
)
.await;
}
result
}
async fn ensure_query_ready(&self, generation: &ActiveGeneration) -> RuntimeResult<()> {
if self.remote.is_none() {
let heads = sqlx::query_as::<_, (String, i64)>(
"SELECT table_key,head FROM search_runtime_streams WHERE table_key IN ('doc','block')",
)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load embedded search stream heads", error))?;
let cursors = *self.embedded_cursors.read().await;
for (table, head) in heads {
let cursor = cursors[if table == "doc" { 0 } else { 1 }];
if cursor != head {
return Err(RuntimeError::SearchProviderUnavailable);
}
}
return Ok(());
}
let rows = sqlx::query_as::<_, (String, i64, i64)>(
"SELECT streams.table_key,streams.head,cursors.source_cursor
FROM search_runtime_streams streams
JOIN search_runtime_provider_cursors cursors
ON cursors.table_key=streams.table_key AND cursors.generation_id=$1
WHERE streams.table_key IN ('doc','block')",
)
.bind(generation.id)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load remote search readiness", error))?;
if rows.len() != SearchTable::ORDERED.len() || rows.iter().any(|(_, head, cursor)| head != cursor) {
return Err(RuntimeError::SearchProviderUnavailable);
}
Ok(())
}
async fn check_permission_revision(&self, workspace_id: &str, revision: i64) -> RuntimeResult<()> {
let generation = self.active_generation().await?;
let applied = if self.remote.is_none() {
self.embedded_permission_cursors.read().await.get(workspace_id).copied()
} else {
sqlx::query_scalar(
"SELECT permission_revision FROM search_runtime_permission_cursors 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 permission cursor", error))?
};
let applied = applied.unwrap_or(-1);
if applied >= revision {
return Ok(());
}
Err(RuntimeError::SearchPermissionUnavailable)
}
async fn refresh_permission_cursor(&self, workspace_id: &str) -> RuntimeResult<()> {
let revision: i64 = sqlx::query_scalar(
"SELECT coalesce(max(revision),0)::bigint FROM workspace_permission_changes WHERE workspace_id=$1",
)
.bind(workspace_id)
.fetch_one(&self.pool)
.await
.map_err(|error| RuntimeError::database("load search permission revision", error))?;
let generation = self.active_generation().await?;
if self.remote.is_none() {
let mut cursors = self.embedded_permission_cursors.write().await;
let applied = cursors.entry(workspace_id.to_string()).or_insert(revision);
*applied = (*applied).max(revision);
} else {
sqlx::query(
r#"INSERT INTO search_runtime_permission_cursors(generation_id,workspace_id,permission_revision)
VALUES ($1,$2,$3) ON CONFLICT (generation_id,workspace_id) DO UPDATE SET
permission_revision=GREATEST(search_runtime_permission_cursors.permission_revision,EXCLUDED.permission_revision), updated_at=now()"#,
)
.bind(generation.id)
.bind(workspace_id)
.bind(revision)
.execute(&self.pool)
.await
.map_err(|error| RuntimeError::database("advance search permission cursor", error))?;
}
Ok(())
}
async fn refresh_all_permission_cursors(&self) -> RuntimeResult<()> {
let workspace_ids: Vec<String> = sqlx::query_scalar("SELECT id FROM workspaces")
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load search permission workspaces", error))?;
for workspace_id in workspace_ids {
self.refresh_permission_cursor(&workspace_id).await?;
}
Ok(())
}
pub(in crate::runtime::backend_runtime) async fn reconcile_workspace(
&self,
capability: SystemSearchCapability,
workspace_id: &str,
) -> RuntimeResult<()> {
match capability {
SystemSearchCapability::ReconcileIndex => {}
}
let doc_ids: Vec<String> = sqlx::query_scalar("SELECT page_id FROM workspace_pages WHERE workspace_id=$1")
.bind(workspace_id)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load search workspace documents", error))?;
let indexed_doc_ids: Vec<String> =
sqlx::query_scalar("SELECT DISTINCT doc_id FROM search_runtime_projections WHERE workspace_id=$1")
.bind(workspace_id)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load indexed workspace documents", error))?;
let live_doc_ids = doc_ids.iter().cloned().collect::<std::collections::BTreeSet<_>>();
for doc_id in doc_ids {
self.project_document(workspace_id, &doc_id).await?;
}
let deletion_revision = chrono::Utc::now().timestamp_millis();
for doc_id in indexed_doc_ids {
if !live_doc_ids.contains(&doc_id) {
self
.store
.delete_document(workspace_id, &doc_id, deletion_revision)
.await?;
}
}
self.sync().await?;
self.refresh_permission_cursor(workspace_id).await
}
pub(in crate::runtime::backend_runtime) async fn delete_workspace(&self, workspace_id: &str) -> RuntimeResult<()> {
let doc_ids: Vec<String> =
sqlx::query_scalar("SELECT DISTINCT doc_id FROM search_runtime_projections WHERE workspace_id=$1")
.bind(workspace_id)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load indexed workspace documents", error))?;
for doc_id in doc_ids {
self
.store
.delete_document(workspace_id, &doc_id, chrono::Utc::now().timestamp_millis())
.await?;
}
self.sync().await?;
self.refresh_permission_cursor(workspace_id).await
}
pub(in crate::runtime::backend_runtime) async fn status(&self) -> RuntimeResult<serde_json::Value> {
let generation = match self.active_generation().await {
Ok(generation) => generation,
Err(RuntimeError::InvalidState(message)) if message == "search_runtime_not_ready" => {
return Ok(serde_json::json!({
"ready": false,
"provider": self.config.provider,
"tables": [],
}));
}
Err(error) => return Err(error),
};
use sqlx::Row;
let heads = sqlx::query(
"SELECT table_key,head FROM search_runtime_streams WHERE table_key IN ('doc','block') ORDER BY table_key",
)
.fetch_all(&self.pool)
.await
.map_err(|error| RuntimeError::database("load search runtime heads", error))?;
let local_cursors = *self.embedded_cursors.read().await;
let mut tables = Vec::with_capacity(heads.len());
for row in heads {
let table_key: String = row.get("table_key");
let head: i64 = row.get("head");
let cursor = if self.remote.is_none() {
local_cursors[if table_key == "doc" { 0 } else { 1 }]
} else {
sqlx::query_scalar(
"SELECT source_cursor FROM search_runtime_provider_cursors WHERE generation_id=$1 AND table_key=$2",
)
.bind(generation.id)
.bind(&table_key)
.fetch_one(&self.pool)
.await
.map_err(|error| RuntimeError::database("load search provider cursor", error))?
};
tables.push(serde_json::json!({"table":table_key,"head":head,"cursor":cursor,"lag":head-cursor}));
}
Ok(serde_json::json!({
"ready":tables.len()==2 && tables.iter().all(|table|table["lag"]==0),
"generationId":generation.id.to_string(),
"provider":self.config.provider,
"tables":tables,
}))
}
}
packages/backend/native/src/runtime/backend_runtime/search/store/mod.rs
+6
View File
@@ -0,0 +1,6 @@
mod projection;
pub(super) mod stream;
mod types;
pub(super) use projection::SearchStore;
pub(super) use types::{ProjectionInput, SearchChange, SearchSnapshot, SearchTable};
packages/backend/native/src/runtime/backend_runtime/search/store/projection.rs
+450
View File
@@ -0,0 +1,450 @@
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(())
}
packages/backend/native/src/runtime/backend_runtime/search/store/stream.rs
+26
View File
@@ -0,0 +1,26 @@
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)
}
packages/backend/native/src/runtime/backend_runtime/search/store/types.rs
+54
View File
@@ -0,0 +1,54 @@
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>,
}
packages/backend/native/src/runtime/backend_runtime/search/tests.rs
+635
View File
@@ -0,0 +1,635 @@
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");
}
}
packages/backend/native/src/runtime/backend_runtime/search/types.rs
+241
View File
@@ -0,0 +1,241 @@
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")]
pub enum SearchTable {
Doc,
Block,
}
impl SearchTable {
pub(super) fn as_str(self) -> &'static str {
match self {
Self::Doc => "doc",
Self::Block => "block",
}
}
pub(super) fn text_field(self) -> &'static str {
match self {
Self::Doc => "title",
Self::Block => "content",
}
}
}
#[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")]
pub struct SearchPagination {
pub limit: Option<u32>,
pub skip: Option<u32>,
pub cursor: Option<String>,
}
#[napi_derive::napi(object)]
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchHighlight {
pub field: String,
pub before: String,
pub end: String,
}
#[napi_derive::napi(object)]
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchOptions {
pub fields: Vec<String>,
#[serde(default)]
pub highlights: Vec<SearchHighlight>,
#[serde(default)]
pub pagination: SearchPagination,
}
#[napi_derive::napi(object)]
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AggregateHitsOptions {
pub fields: Vec<String>,
#[serde(default)]
pub highlights: Vec<SearchHighlight>,
#[serde(default)]
pub pagination: SearchPagination,
}
#[napi_derive::napi(object)]
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AggregateOptions {
pub hits: AggregateHitsOptions,
#[serde(default)]
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,
pub field: Option<String>,
pub match_value: Option<String>,
pub query: Option<u32>,
pub queries: Option<Vec<u32>>,
pub occur: Option<String>,
pub boost: Option<f64>,
}
#[napi_derive::napi(object)]
pub struct RuntimeSearchRequest {
pub table: SearchTable,
pub queries: Vec<RuntimeSearchQuery>,
pub root_query: u32,
pub options: SearchOptions,
}
#[napi_derive::napi(object)]
pub struct RuntimeAggregateRequest {
pub table: SearchTable,
pub queries: Vec<RuntimeSearchQuery>,
pub root_query: u32,
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());
}
}
packages/backend/native/src/runtime/backend_runtime/search/worker.rs
+272
View File
@@ -0,0 +1,272 @@
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"]));
}
}
packages/backend/native/src/runtime/backend_runtime/tests.rs
+18
View File
@@ -24,6 +24,7 @@ fn migrations_include_runtime_tables_without_worker_heartbeats() {
assert!(RUNTIME_MIGRATIONS.contains("storage_reconciliation_checkpoints"));
assert!(RUNTIME_MIGRATIONS.contains("document_cleanup_candidates"));
assert!(RUNTIME_MIGRATIONS.contains("doc_blob_refs"));
assert!(RUNTIME_MIGRATIONS.contains("doc_blob_ref_projections"));
assert!(RUNTIME_MIGRATIONS.contains("blob_cleanup_candidates"));
assert!(!RUNTIME_MIGRATIONS.contains("runtime_worker_heartbeats"));
}
@@ -100,12 +101,15 @@ async fn runtime_from_database_url() -> AnyResult<Option<BackendRuntime>> {
Ok(Some(BackendRuntime {
config_source: Default::default(),
role: ServerRole::AllInOne,
script_mode: false,
config: Arc::new(RwLock::new(Arc::new(BackendRuntimeConfig {
database_url,
invite_quota: Default::default(),
private_key: Arc::new(zeroize::Zeroizing::new("test-private-key".to_string())),
deployment: crate::llm::Deployment::Cloud,
copilot: Default::default(),
search: Default::default(),
}))),
config_reload: Mutex::new(()),
pool: Mutex::new(Some(pool)),
@@ -114,6 +118,8 @@ async fn runtime_from_database_url() -> AnyResult<Option<BackendRuntime>> {
crate::runtime::object_storage::ObjectStorageService::from_config_files()?,
)),
embedding: Mutex::new(None),
embedding_worker: Mutex::new(None),
search: Mutex::new(None),
managed_token_providers: Arc::new(Default::default()),
}))
}
@@ -260,12 +266,16 @@ async fn runtime_gate_sql_semantics_are_atomic_and_ttl_bound() {
for _ in 0..16 {
let runtime = BackendRuntime {
config_source: Default::default(),
role: ServerRole::AllInOne,
script_mode: false,
config: Arc::new(RwLock::new(runtime.config().unwrap())),
config_reload: Mutex::new(()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
embedding_health: RwLock::new(super::EmbeddingHealth::disabled("test", None)),
object_storage: RwLock::new(runtime.object_storage().unwrap()),
embedding: Mutex::new(None),
embedding_worker: Mutex::new(None),
search: Mutex::new(None),
managed_token_providers: Arc::new(Default::default()),
};
tasks.push(tokio::spawn(async move {
@@ -597,12 +607,16 @@ async fn coordination_lease_sql_semantics_are_fenced_and_ttl_bound() {
for index in 0..16 {
let runtime = BackendRuntime {
config_source: Default::default(),
role: ServerRole::AllInOne,
script_mode: false,
config: Arc::new(RwLock::new(runtime.config().unwrap())),
config_reload: Mutex::new(()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
embedding_health: RwLock::new(super::EmbeddingHealth::disabled("test", None)),
object_storage: RwLock::new(runtime.object_storage().unwrap()),
embedding: Mutex::new(None),
embedding_worker: Mutex::new(None),
search: Mutex::new(None),
managed_token_providers: Arc::new(Default::default()),
};
tasks.push(tokio::spawn(async move {
@@ -807,12 +821,16 @@ async fn verification_token_sql_state_machine_handles_keep_verify_and_cleanup()
for _ in 0..16 {
let runtime = BackendRuntime {
config_source: Default::default(),
role: ServerRole::AllInOne,
script_mode: false,
config: Arc::new(RwLock::new(runtime.config().unwrap())),
config_reload: Mutex::new(()),
pool: Mutex::new(Some(runtime.pool().await.unwrap())),
embedding_health: RwLock::new(super::EmbeddingHealth::disabled("test", None)),
object_storage: RwLock::new(runtime.object_storage().unwrap()),
embedding: Mutex::new(None),
embedding_worker: Mutex::new(None),
search: Mutex::new(None),
managed_token_providers: Arc::new(Default::default()),
};
let token = concurrent_token.clone();
packages/backend/native/src/runtime/config.rs
+95
View File
@@ -20,6 +20,30 @@ pub(crate) struct BackendRuntimeConfig {
pub(crate) private_key: Arc<Zeroizing<String>>,
pub(crate) deployment: Deployment,
pub(crate) copilot: CopilotRuntimeConfig,
pub(crate) search: SearchRuntimeConfig,
}
#[derive(Clone, Debug)]
pub(crate) struct SearchRuntimeConfig {
pub(crate) enabled: bool,
pub(crate) provider: String,
pub(crate) endpoint: String,
pub(crate) api_key: String,
pub(crate) username: String,
pub(crate) password: String,
}
impl Default for SearchRuntimeConfig {
fn default() -> Self {
Self {
enabled: false,
provider: "embedded".to_string(),
endpoint: String::new(),
api_key: String::new(),
username: String::new(),
password: String::new(),
}
}
}
#[derive(Clone, Debug)]
@@ -348,6 +372,7 @@ impl BackendRuntimeConfig {
.map(TryInto::try_into)
.transpose()?
.unwrap_or_default(),
search: app_config.indexer.map(Into::into).unwrap_or_default(),
}
.validated()
}
@@ -385,6 +410,10 @@ impl BackendRuntimeConfig {
.map(TryInto::try_into)
.transpose()?
.unwrap_or_else(|| self.copilot.clone()),
search: app_config
.indexer
.map(Into::into)
.unwrap_or_else(|| self.search.clone()),
}
.validated()
}
@@ -453,6 +482,42 @@ struct AppConfigFile {
db: Option<DbConfigFile>,
crypto: Option<CryptoConfigFile>,
copilot: Option<CopilotRuntimeConfigFile>,
indexer: Option<SearchRuntimeConfigFile>,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase", default)]
struct SearchRuntimeConfigFile {
enabled: bool,
provider: SearchProviderConfigFile,
}
#[derive(Default, Deserialize)]
#[serde(rename_all = "camelCase", default)]
struct SearchProviderConfigFile {
#[serde(rename = "type")]
provider: String,
endpoint: String,
api_key: String,
username: String,
password: String,
}
impl From<SearchRuntimeConfigFile> for SearchRuntimeConfig {
fn from(value: SearchRuntimeConfigFile) -> Self {
Self {
enabled: value.enabled,
provider: if value.provider.provider.is_empty() {
"embedded".to_string()
} else {
value.provider.provider
},
endpoint: value.provider.endpoint,
api_key: value.provider.api_key,
username: value.provider.username,
password: value.provider.password,
}
}
}
#[derive(Default, Deserialize)]
@@ -815,6 +880,35 @@ mod tests {
assert!(copilot.byok.allow_custom_endpoint);
}
#[test]
fn search_config_keeps_disabled_state_separate_from_embedded_provider() {
let disabled = app_config_from_flat_overrides([
("indexer.enabled", serde_json::json!(false)),
("indexer.provider.type", serde_json::json!("embedded")),
])
.unwrap();
let disabled: SearchRuntimeConfig = disabled.indexer.unwrap().into();
assert!(!disabled.enabled);
assert_eq!(disabled.provider, "embedded");
let enabled = app_config_from_flat_overrides([
("indexer.enabled", serde_json::json!(true)),
("indexer.provider.type", serde_json::json!("elasticsearch")),
])
.unwrap();
let enabled: SearchRuntimeConfig = enabled.indexer.unwrap().into();
assert!(enabled.enabled);
assert_eq!(enabled.provider, "elasticsearch");
let enabled_without_provider = app_config_from_module_json(serde_json::json!({
"indexer": { "enabled": true }
}))
.unwrap();
let enabled_without_provider: SearchRuntimeConfig = enabled_without_provider.indexer.unwrap().into();
assert!(enabled_without_provider.enabled);
assert_eq!(enabled_without_provider.provider, "embedded");
}
#[test]
fn partial_database_config_preserves_file_config_siblings() {
let mut file_config = expand_module_config_paths(serde_json::json!({
@@ -874,6 +968,7 @@ mod tests {
private_key: Arc::new(Zeroizing::new("active-private-key".to_string())),
deployment: Deployment::Cloud,
copilot: CopilotRuntimeConfig::default(),
search: SearchRuntimeConfig::default(),
};
let empty = serde_json::Value::Object(Map::new());
packages/backend/native/src/runtime/error.rs
+20
View File
@@ -15,6 +15,21 @@ pub(crate) enum RuntimeError {
#[error("{0}")]
InvalidState(String),
#[error("workspace access denied")]
SearchWorkspaceDenied,
#[error("search permission state unavailable")]
SearchPermissionUnavailable,
#[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,
@@ -94,6 +109,11 @@ impl RuntimeError {
| Self::NapiBoundary(message) => {
message.contains("NoSuchKey") || message.contains("NotFound") || message.contains("not found")
}
Self::SearchWorkspaceDenied
| Self::SearchPermissionUnavailable
| Self::SearchProviderUnavailable
| Self::SearchUnsupportedQuery
| Self::SearchReplayGap => false,
_ => false,
}
}
packages/backend/native/src/runtime/migrations.rs
+138 -35
View File
@@ -5,7 +5,10 @@ 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 EMBEDDING_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0046;
const SEARCH_ADVISORY_LOCK: i64 = 0x4146_4649_4e45_0053;
#[cfg(test)]
pub(crate) static EMBEDDING_TEST_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
@@ -17,6 +20,14 @@ pub(crate) async fn migrate_runtime_tables(pool: &PgPool) -> RuntimeResult<()> {
Ok(())
}
pub(crate) async fn migrate_all_tables(pool: &PgPool) -> RuntimeResult<EmbeddingHealth> {
migrate_runtime_tables(pool).await?;
let embedding = migrate_embedding_tables_inner(pool).await?;
migrate_search_tables(pool).await?;
Ok(embedding)
}
#[cfg(test)]
pub(crate) async fn migrate_embedding_tables(pool: &PgPool) -> EmbeddingHealth {
match migrate_embedding_tables_inner(pool).await {
Ok(health) => health,
@@ -24,7 +35,7 @@ pub(crate) async fn migrate_embedding_tables(pool: &PgPool) -> EmbeddingHealth {
}
}
async fn migrate_embedding_tables_inner(pool: &PgPool) -> RuntimeResult<EmbeddingHealth> {
pub(crate) async fn embedding_schema_health(pool: &PgPool) -> RuntimeResult<EmbeddingHealth> {
let Some(version) = pgvector_version(pool).await? else {
return Ok(EmbeddingHealth::disabled("pgvector_unavailable", None));
};
@@ -32,33 +43,19 @@ async fn migrate_embedding_tables_inner(pool: &PgPool) -> RuntimeResult<Embeddin
return Ok(EmbeddingHealth::disabled("pgvector_version_unsupported", Some(version)));
}
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("Embedding migration transaction failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(EMBEDDING_ADVISORY_LOCK)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("Embedding migration lock failed", error))?;
transaction
.execute(
r#"CREATE TABLE IF NOT EXISTS native_schema_migrations (
component TEXT NOT NULL,
version INTEGER NOT NULL,
checksum TEXT NOT NULL,
applied_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (component, version)
)"#,
)
.await
.map_err(|error| RuntimeError::database("Embedding migration ledger failed", error))?;
apply_migration(&mut transaction, 1, &[EMBEDDING_MIGRATION]).await?;
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("Embedding migration commit failed", error))?;
let schema_ready: bool = sqlx::query_scalar(
"SELECT to_regclass('embedding_workspace_states') IS NOT NULL
AND to_regclass('embedding_indexes') IS NOT NULL
AND to_regclass('embedding_sources') IS NOT NULL
AND to_regclass('embedding_projections') IS NOT NULL
AND to_regclass('embedding_chunks') IS NOT NULL",
)
.fetch_one(pool)
.await
.map_err(|error| RuntimeError::database("Embedding schema health check failed", error))?;
if !schema_ready {
return Ok(EmbeddingHealth::disabled("schema_not_migrated", Some(version)));
}
Ok(EmbeddingHealth {
enabled: true,
@@ -70,23 +67,103 @@ async fn migrate_embedding_tables_inner(pool: &PgPool) -> RuntimeResult<Embeddin
})
}
pub(crate) async fn migrate_search_tables(pool: &PgPool) -> RuntimeResult<()> {
migrate_component(
pool,
"search",
SEARCH_ADVISORY_LOCK,
&[(1, &[SEARCH_MIGRATION]), (2, &[SEARCH_ACL_TOKEN_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(())
}
async fn migrate_embedding_tables_inner(pool: &PgPool) -> RuntimeResult<EmbeddingHealth> {
let Some(version) = pgvector_version(pool).await? else {
return Ok(EmbeddingHealth::disabled("pgvector_unavailable", None));
};
if !pgvector_at_least_0_8(&version) {
return Ok(EmbeddingHealth::disabled("pgvector_version_unsupported", Some(version)));
}
migrate_component(
pool,
"embedding",
EMBEDDING_ADVISORY_LOCK,
&[(1, &[EMBEDDING_MIGRATION])],
)
.await?;
Ok(EmbeddingHealth {
enabled: true,
state: "ready".to_string(),
reason: None,
pgvector_version: Some(version),
schema_version: Some(1),
worker_running: false,
})
}
async fn migrate_component(
pool: &PgPool,
component: &str,
advisory_lock: i64,
migrations: &[(i32, &[&str])],
) -> RuntimeResult<()> {
let mut transaction = pool
.begin()
.await
.map_err(|error| RuntimeError::database("Native migration transaction failed", error))?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(advisory_lock)
.execute(&mut *transaction)
.await
.map_err(|error| RuntimeError::database("Native migration lock failed", error))?;
transaction
.execute(
r#"CREATE TABLE IF NOT EXISTS native_schema_migrations (
component TEXT NOT NULL,
version INTEGER NOT NULL,
checksum TEXT NOT NULL,
applied_at TIMESTAMPTZ NOT NULL DEFAULT now(),
PRIMARY KEY (component, version)
)"#,
)
.await
.map_err(|error| RuntimeError::database("Native migration ledger failed", error))?;
for (version, statements) in migrations {
apply_migration(&mut transaction, component, *version, statements).await?;
}
transaction
.commit()
.await
.map_err(|error| RuntimeError::database("Native migration commit failed", error))
}
async fn apply_migration(
transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
component: &str,
version: i32,
statements: &[&str],
) -> RuntimeResult<()> {
let checksum = migration_checksum(statements);
let applied = sqlx::query("SELECT checksum FROM native_schema_migrations WHERE component='embedding' AND version=$1")
let applied = sqlx::query("SELECT checksum FROM native_schema_migrations WHERE component=$1 AND version=$2")
.bind(component)
.bind(version)
.fetch_optional(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("Embedding migration ledger read failed", error))?;
.map_err(|error| RuntimeError::database("Native migration ledger read failed", error))?;
if let Some(applied) = applied {
let stored: String = applied
.try_get("checksum")
.map_err(|error| RuntimeError::database("Embedding migration checksum decode failed", error))?;
.map_err(|error| RuntimeError::database("Native migration checksum decode failed", error))?;
if stored != checksum {
return Err(RuntimeError::invalid_state("Embedding migration checksum mismatch"));
return Err(RuntimeError::invalid_state("Native migration checksum mismatch"));
}
return Ok(());
}
@@ -94,14 +171,15 @@ async fn apply_migration(
transaction
.execute(*statement)
.await
.map_err(|error| RuntimeError::database("Embedding migration failed", error))?;
.map_err(|error| RuntimeError::database("Native component migration failed", error))?;
}
sqlx::query("INSERT INTO native_schema_migrations(component,version,checksum) VALUES('embedding',$1,$2)")
sqlx::query("INSERT INTO native_schema_migrations(component,version,checksum) VALUES($1,$2,$3)")
.bind(component)
.bind(version)
.bind(checksum)
.execute(&mut **transaction)
.await
.map_err(|error| RuntimeError::database("Embedding migration record failed", error))?;
.map_err(|error| RuntimeError::database("Native migration record failed", error))?;
Ok(())
}
@@ -189,4 +267,29 @@ mod tests {
.unwrap();
assert_eq!(dimensions, 1024);
}
#[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"));
}
#[tokio::test]
async fn concurrent_search_migration_is_idempotent() {
let Ok(database_url) = std::env::var("DATABASE_URL") else {
return;
};
let pool = PgPool::connect(&database_url).await.unwrap();
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]);
}
}
packages/backend/native/src/runtime/mod.rs
+1 -1
View File
@@ -10,7 +10,7 @@ pub(crate) mod types;
pub(crate) use config::{
BackendRuntimeConfig, ConfigSource, CopilotManagedProfileConfig, CopilotManagedProfileConfigFile,
CopilotRuntimeConfig, CopilotRuntimeConfigFile, InviteQuotaConfig,
CopilotRuntimeConfig, CopilotRuntimeConfigFile, InviteQuotaConfig, SearchRuntimeConfig,
};
use config::{SUPPORTED_BYOK_PROVIDERS, validate_copilot_config};
pub use config_descriptor::{AppConfigDescriptor, app_config_descriptors, validate_app_config_value};
packages/backend/native/src/runtime/sql/runtime_migrations.sql
+22
View File
@@ -113,6 +113,28 @@ CREATE INDEX IF NOT EXISTS doc_blob_refs_workspace_blob_idx
CREATE INDEX IF NOT EXISTS doc_blob_refs_workspace_status_idx
ON doc_blob_refs (workspace_id, status);
CREATE TABLE IF NOT EXISTS doc_blob_ref_projections (
workspace_id TEXT NOT NULL,
doc_id TEXT NOT NULL,
source_revision TIMESTAMPTZ(3),
parser_version INTEGER NOT NULL,
status TEXT NOT NULL CHECK (status IN ('pending', 'running', 'fresh', 'failed', 'missing')),
indexed_at TIMESTAMPTZ(3),
error_code TEXT,
error_summary TEXT,
attempt_count INTEGER NOT NULL DEFAULT 0 CHECK (attempt_count >= 0),
updated_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (workspace_id, doc_id),
CHECK (status <> 'fresh' OR (source_revision IS NOT NULL AND indexed_at IS NOT NULL)),
CHECK (error_summary IS NULL OR octet_length(error_summary) <= 512)
);
CREATE INDEX IF NOT EXISTS doc_blob_ref_projections_workspace_status_idx
ON doc_blob_ref_projections (workspace_id, status, updated_at DESC);
CREATE INDEX IF NOT EXISTS doc_blob_ref_projections_workspace_revision_idx
ON doc_blob_ref_projections (workspace_id, source_revision);
CREATE TABLE IF NOT EXISTS blob_cleanup_candidates (
workspace_id TEXT NOT NULL,
blob_key TEXT NOT NULL,
packages/backend/native/src/runtime/sql/search.sql
+252
View File
@@ -0,0 +1,252 @@
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();
packages/backend/native/src/runtime/sql/search_acl_tokens.sql
+5
View File
@@ -0,0 +1,5 @@
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()
);
packages/backend/native/src/runtime/storage_runtime/blob_cleanup.rs
+82 -16
View File
@@ -5,7 +5,7 @@ use sqlx::{FromRow, PgPool};
use super::{
RuntimeBlobCleanupExecuteResult, RuntimeBlobCleanupPlanResult, RuntimeError, RuntimeResult, StorageRuntime,
napi_error,
doc_blob_refs::PARSER_VERSION, load_workspace_canonical_doc_ids, napi_error,
};
#[derive(FromRow)]
@@ -83,6 +83,70 @@ async fn projection_is_stale(pool: &PgPool, workspace_id: &str) -> RuntimeResult
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Blob cleanup retention activity check failed", err))?;
if sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM updates WHERE workspace_id = $1)")
.bind(workspace_id)
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Blob cleanup pending update check failed", err))?
{
return Ok(true);
}
let mut current_doc_ids = match load_workspace_canonical_doc_ids(pool, workspace_id).await {
Ok(ids) => ids,
Err(_) => return Ok(true),
};
current_doc_ids.push(workspace_id.to_string());
current_doc_ids.extend(
sqlx::query_scalar::<_, String>(
"SELECT doc_id FROM document_cleanup_candidates WHERE workspace_id = $1 AND status IN ('marked', 'failed')",
)
.bind(workspace_id)
.fetch_all(pool)
.await
.map_err(|err| RuntimeError::database("Blob cleanup retained document load failed", err))?,
);
current_doc_ids.sort();
current_doc_ids.dedup();
let has_nonfresh_projection = sqlx::query_scalar::<_, bool>(
r#"
SELECT EXISTS(
SELECT 1
FROM unnest($2::text[]) AS ids(doc_id)
LEFT JOIN snapshots s
ON s.workspace_id = $1 AND s.guid = ids.doc_id
LEFT JOIN doc_blob_ref_projections p
ON p.workspace_id = $1 AND p.doc_id = ids.doc_id
WHERE s.guid IS NULL
OR p.doc_id IS NULL
OR p.status <> 'fresh'
OR p.parser_version <> $3
OR p.source_revision IS DISTINCT FROM s.updated_at
)
OR EXISTS(
SELECT 1 FROM doc_blob_ref_projections
WHERE workspace_id = $1 AND status <> 'fresh'
)
OR EXISTS(
SELECT 1
FROM doc_blob_refs r
LEFT JOIN doc_blob_ref_projections p
ON p.workspace_id = r.workspace_id AND p.doc_id = r.doc_id
WHERE r.workspace_id = $1
AND (
p.doc_id IS NULL
OR p.status <> 'fresh'
OR r.parser_version <> p.parser_version
OR r.snapshot_updated_at IS DISTINCT FROM p.source_revision
)
)
"#,
)
.bind(workspace_id)
.bind(&current_doc_ids)
.bind(PARSER_VERSION)
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Blob cleanup projection state check failed", err))?;
let has_stale_rows = sqlx::query_scalar::<_, bool>(
"SELECT EXISTS(SELECT 1 FROM doc_blob_refs WHERE workspace_id = $1 AND status <> 'fresh')",
)
@@ -90,7 +154,7 @@ async fn projection_is_stale(pool: &PgPool, workspace_id: &str) -> RuntimeResult
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Blob cleanup projection freshness check failed", err))?;
Ok(activity_after_checkpoint || has_stale_rows)
Ok(activity_after_checkpoint || has_nonfresh_projection || has_stale_rows)
}
async fn stale_projection_workspaces(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<String>> {
@@ -107,18 +171,31 @@ async fn metadata_backfill_is_complete(pool: &PgPool, workspace_id: &str) -> Run
async fn has_doc_ref(pool: &PgPool, workspace_id: &str, key: &str) -> RuntimeResult<bool> {
sqlx::query_scalar::<_, bool>(
"SELECT EXISTS(SELECT 1 FROM doc_blob_refs WHERE workspace_id = $1 AND blob_key = $2 AND status = 'fresh')",
r#"
SELECT EXISTS(
SELECT 1
FROM doc_blob_refs r
JOIN doc_blob_ref_projections p
ON p.workspace_id = r.workspace_id AND p.doc_id = r.doc_id
WHERE r.workspace_id = $1
AND r.blob_key = $2
AND r.status = 'fresh'
AND p.status = 'fresh'
AND p.parser_version = $3
AND r.parser_version = p.parser_version
AND r.snapshot_updated_at = p.source_revision
)
"#,
)
.bind(workspace_id)
.bind(key)
.bind(PARSER_VERSION)
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Blob cleanup doc ref check failed", err))
}
async fn has_other_ref(pool: &PgPool, workspace_id: &str, key: &str) -> RuntimeResult<bool> {
// Remove the ai_contexts branch after stable and beta no longer run binaries
// built with the 115-migration schema.
let required_ref = sqlx::query_scalar::<_, bool>(
r#"
SELECT EXISTS(SELECT 1 FROM workspaces WHERE id = $1 AND avatar_key = $2)
@@ -131,17 +208,6 @@ async fn has_other_ref(pool: &PgPool, workspace_id: &str, key: &str) -> RuntimeR
AND storage_key = concat($1, '/', $2)
AND status IN ('reserving', 'ready')
)
OR EXISTS(
SELECT 1
FROM ai_contexts c
JOIN ai_sessions_metadata s ON s.id = c.session_id
WHERE s.workspace_id = $1
AND jsonb_path_exists(
c.config::jsonb,
'$.** ? (@ == $blobKey)',
jsonb_build_object('blobKey', to_jsonb($2::text))
)
)
"#,
)
.bind(workspace_id)
packages/backend/native/src/runtime/storage_runtime/current_doc.rs
+39 -25
View File
@@ -6,8 +6,6 @@ use super::{RuntimeError, RuntimeResult};
#[derive(FromRow)]
pub(in crate::runtime) struct CurrentDoc {
pub(in crate::runtime) workspace_id: String,
pub(in crate::runtime) doc_id: String,
pub(in crate::runtime) blob: Vec<u8>,
pub(in crate::runtime) updated_at: DateTime<Utc>,
}
@@ -25,7 +23,7 @@ pub(in crate::runtime) async fn load_current_doc(
) -> RuntimeResult<Option<CurrentDoc>> {
let snapshot = sqlx::query_as::<_, CurrentDoc>(
r#"
SELECT workspace_id, guid AS doc_id, blob, updated_at
SELECT blob, updated_at
FROM snapshots
WHERE workspace_id = $1 AND guid = $2
"#,
@@ -48,12 +46,38 @@ pub(in crate::runtime) async fn load_current_doc(
.fetch_all(pool)
.await
.map_err(|err| RuntimeError::database("Current doc updates load failed", err))?;
merge_current_doc(workspace_id, doc_id, snapshot, updates)
merge_current_doc(snapshot, updates)
}
pub(super) async fn load_canonical_doc(
pool: &PgPool,
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<Option<CurrentDoc>> {
sqlx::query_as::<_, CurrentDoc>(
r#"
SELECT blob, updated_at
FROM snapshots
WHERE workspace_id = $1 AND guid = $2
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(pool)
.await
.map_err(|err| RuntimeError::database("Canonical doc snapshot load failed", err))
}
pub(super) async fn has_pending_updates(pool: &PgPool, workspace_id: &str, doc_id: &str) -> RuntimeResult<bool> {
sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM updates WHERE workspace_id = $1 AND guid = $2)")
.bind(workspace_id)
.bind(doc_id)
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Pending doc updates check failed", err))
}
pub(super) fn merge_current_doc(
workspace_id: &str,
doc_id: &str,
snapshot: Option<CurrentDoc>,
updates: Vec<CurrentDocUpdate>,
) -> RuntimeResult<Option<CurrentDoc>> {
@@ -84,16 +108,18 @@ pub(super) fn merge_current_doc(
.encode_update_v1()
.map_err(|err| RuntimeError::invalid_state(format!("Current doc encode failed: {err}")))?;
Ok(Some(CurrentDoc {
workspace_id: workspace_id.to_string(),
doc_id: doc_id.to_string(),
blob,
updated_at,
}))
Ok(Some(CurrentDoc { blob, updated_at }))
}
pub(super) async fn load_workspace_live_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<String>> {
workspace_live_doc_ids(load_current_doc(pool, workspace_id, workspace_id).await?)
load_workspace_canonical_doc_ids(pool, workspace_id).await
}
pub(super) async fn load_workspace_canonical_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<String>> {
if has_pending_updates(pool, workspace_id, workspace_id).await? {
return Err(RuntimeError::invalid_state("Workspace root doc has pending updates"));
}
workspace_live_doc_ids(load_canonical_doc(pool, workspace_id, workspace_id).await?)
}
fn workspace_live_doc_ids(root: Option<CurrentDoc>) -> RuntimeResult<Vec<String>> {
@@ -120,11 +146,7 @@ mod tests {
let snapshot = affine_doc_loader::add_doc_to_root_doc(Vec::new(), "live", None).unwrap();
let pending = affine_doc_loader::add_doc_to_root_doc(snapshot.clone(), "trash", None).unwrap();
let merged = merge_current_doc(
"workspace",
"workspace",
Some(CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id: "workspace".to_string(),
blob: snapshot,
updated_at: Utc::now(),
}),
@@ -149,8 +171,6 @@ mod tests {
trash.insert("trash".to_string(), Value::Any(Any::True)).unwrap();
let ids = workspace_live_doc_ids(Some(CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id: "workspace".to_string(),
blob: root.encode_update_v1().unwrap(),
updated_at: Utc::now(),
}))
@@ -165,8 +185,6 @@ mod tests {
.unwrap();
pages.remove(trash_index as u64, 1).unwrap();
let ids = workspace_live_doc_ids(Some(CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id: "workspace".to_string(),
blob: root.encode_update_v1().unwrap(),
updated_at: Utc::now(),
}))
@@ -179,8 +197,6 @@ mod tests {
assert!(workspace_live_doc_ids(None).is_err());
assert!(
workspace_live_doc_ids(Some(CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id: "workspace".to_string(),
blob: vec![0xff],
updated_at: Utc::now(),
}))
@@ -188,8 +204,6 @@ mod tests {
);
assert!(
workspace_live_doc_ids(Some(CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id: "workspace".to_string(),
blob: vec![
1, 1, 1, 1, 40, 0, 1, 0, 11, 115, 117, 98, 95, 109, 97, 112, 95, 107, 101, 121, 1, 119, 13, 115, 117, 98, 95,
109, 97, 112, 95, 118, 97, 108, 117, 101, 0,
packages/backend/native/src/runtime/storage_runtime/doc_blob_refs.rs
+501 -195
View File
@@ -1,24 +1,53 @@
use affine_doc_loader as doc_loader;
use chrono::{DateTime, Utc};
use sqlx::PgPool;
use sqlx::{Executor, FromRow, PgPool, Postgres};
use super::{
CurrentDoc, RuntimeDocBlobRefsResult, RuntimeError, RuntimeResult, StorageRuntime, load_current_doc,
load_workspace_live_doc_ids, napi_error,
CurrentDoc, RuntimeDocBlobRefsResult, RuntimeError, RuntimeResult, StorageRuntime, load_canonical_doc,
load_workspace_canonical_doc_ids, napi_error,
};
const PARSER_VERSION: i32 = 1;
pub(super) const PARSER_VERSION: i32 = 1;
const ERROR_SUMMARY_LIMIT: usize = 512;
type ExtractedRef = doc_loader::BlobRef;
#[derive(FromRow)]
struct DocSource {
updated_at: DateTime<Utc>,
has_pending_updates: bool,
}
#[derive(Default)]
struct ProjectionState {
cursor: Option<String>,
failed_docs: i64,
}
#[derive(Default)]
struct ProjectionStats {
result: RuntimeDocBlobRefsResult,
pending_docs: i64,
missing_docs: i64,
shadow_mismatches: i64,
}
#[derive(Default)]
struct ProjectionAttempt {
written: i64,
deleted: i64,
shadow_mismatch: bool,
}
enum ProjectionOutcome {
Fresh(ProjectionAttempt),
Pending,
Missing,
}
async fn load_workspace_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeResult<Vec<String>> {
let mut ids = load_workspace_live_doc_ids(pool, workspace_id).await?;
let mut ids = load_workspace_canonical_doc_ids(pool, workspace_id).await?;
ids.push(workspace_id.to_string());
let retained = sqlx::query_scalar::<_, String>(
"SELECT doc_id FROM document_cleanup_candidates WHERE workspace_id = $1 AND status IN ('marked', 'failed') ORDER \
BY doc_id",
@@ -33,15 +62,127 @@ async fn load_workspace_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeRes
Ok(ids)
}
async fn load_doc_source(pool: &PgPool, workspace_id: &str, doc_id: &str) -> RuntimeResult<Option<DocSource>> {
sqlx::query_as::<_, DocSource>(
r#"
SELECT s.updated_at,
EXISTS(
SELECT 1 FROM updates u
WHERE u.workspace_id = s.workspace_id AND u.guid = s.guid
) AS has_pending_updates
FROM snapshots s
WHERE s.workspace_id = $1 AND s.guid = $2
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(pool)
.await
.map_err(|err| RuntimeError::database("Doc blob refs source load failed", err))
}
async fn projection_is_fresh(
pool: &PgPool,
workspace_id: &str,
doc_id: &str,
source_revision: DateTime<Utc>,
) -> RuntimeResult<bool> {
sqlx::query_scalar::<_, bool>(
r#"
SELECT EXISTS(
SELECT 1 FROM doc_blob_ref_projections
WHERE workspace_id = $1
AND doc_id = $2
AND source_revision = $3
AND parser_version = $4
AND status = 'fresh'
)
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(source_revision)
.bind(PARSER_VERSION)
.fetch_one(pool)
.await
.map_err(|err| RuntimeError::database("Doc blob refs projection freshness load failed", err))
}
fn truncate_error_summary(error: &str) -> String {
let mut end = error.len().min(ERROR_SUMMARY_LIMIT);
while end > 0 && !error.is_char_boundary(end) {
end -= 1;
}
error[..end].to_string()
}
async fn upsert_projection_state<'e, E>(
executor: E,
workspace_id: &str,
doc_id: &str,
source_revision: Option<DateTime<Utc>>,
status: &str,
error_code: Option<&str>,
error_summary: Option<&str>,
) -> RuntimeResult<()>
where
E: Executor<'e, Database = Postgres>,
{
sqlx::query(
r#"
INSERT INTO doc_blob_ref_projections
(workspace_id, doc_id, source_revision, parser_version, status, indexed_at, error_code, error_summary, attempt_count)
VALUES ($1, $2, $3, $4, $5, CURRENT_TIMESTAMP, $6, $7, CASE WHEN $5 = 'fresh' THEN 0 ELSE 1 END)
ON CONFLICT (workspace_id, doc_id) DO UPDATE
SET source_revision = EXCLUDED.source_revision,
parser_version = EXCLUDED.parser_version,
status = EXCLUDED.status,
indexed_at = EXCLUDED.indexed_at,
error_code = EXCLUDED.error_code,
error_summary = EXCLUDED.error_summary,
attempt_count = CASE
WHEN EXCLUDED.status = 'fresh' THEN 0
ELSE doc_blob_ref_projections.attempt_count + 1
END,
updated_at = CURRENT_TIMESTAMP
WHERE (
EXCLUDED.source_revision IS NULL
AND doc_blob_ref_projections.source_revision IS NULL
AND doc_blob_ref_projections.parser_version <= EXCLUDED.parser_version
) OR (
EXCLUDED.source_revision IS NOT NULL
AND doc_blob_ref_projections.parser_version <= EXCLUDED.parser_version
AND (
doc_blob_ref_projections.source_revision IS NULL
OR EXCLUDED.source_revision >= doc_blob_ref_projections.source_revision
)
)
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(source_revision)
.bind(PARSER_VERSION)
.bind(status)
.bind(error_code)
.bind(error_summary.map(truncate_error_summary))
.execute(executor)
.await
.map_err(|err| RuntimeError::database("Doc blob refs projection state write failed", err))?;
Ok(())
}
async fn upsert_projection_checkpoint(
pool: &PgPool,
workspace_id: &str,
result: &RuntimeDocBlobRefsResult,
failed_docs: i64,
pending_docs: i64,
missing_docs: i64,
shadow_mismatches: i64,
) -> RuntimeResult<()> {
let status = if result.next_cursor.is_some() {
"running"
} else if failed_docs > 0 {
} else if result.failed_docs > 0 {
"failed"
} else {
"completed"
@@ -66,7 +207,10 @@ async fn upsert_projection_checkpoint(
.bind(completed)
.bind(serde_json::json!({
"parserVersion": PARSER_VERSION,
"failedDocs": failed_docs,
"failedDocs": result.failed_docs,
"pendingDocs": pending_docs,
"missingDocs": missing_docs,
"shadowMismatches": shadow_mismatches,
}))
.execute(pool)
.await
@@ -74,7 +218,7 @@ async fn upsert_projection_checkpoint(
Ok(())
}
async fn upsert_projection_failure_checkpoint(pool: &PgPool, workspace_id: &str, error: &str) -> RuntimeResult<()> {
async fn upsert_projection_failure_checkpoint(pool: &PgPool, workspace_id: &str) -> RuntimeResult<()> {
sqlx::query(
r#"
INSERT INTO storage_reconciliation_checkpoints
@@ -91,7 +235,7 @@ async fn upsert_projection_failure_checkpoint(pool: &PgPool, workspace_id: &str,
.bind(workspace_id)
.bind(serde_json::json!({
"parserVersion": PARSER_VERSION,
"error": error,
"errorCode": "root_projection_failed",
}))
.execute(pool)
.await
@@ -111,43 +255,46 @@ async fn load_projection_state(pool: &PgPool, workspace_id: &str) -> RuntimeResu
let Some((status, cursor, metadata)) = checkpoint else {
return Ok(ProjectionState::default());
};
if status != "running" && status != "failed" {
if status != "running" && status != "failed"
|| metadata.get("parserVersion").and_then(serde_json::Value::as_i64) != Some(i64::from(PARSER_VERSION))
{
return Ok(ProjectionState::default());
}
if metadata.get("parserVersion").and_then(serde_json::Value::as_i64) != Some(i64::from(PARSER_VERSION)) {
return Ok(ProjectionState::default());
}
let cursor = cursor
.get("lastDocId")
.and_then(|value| value.as_str())
.map(ToString::to_string);
let Some(cursor) = cursor else {
return Ok(ProjectionState::default());
};
let failed_docs = metadata
.get("failedDocs")
.and_then(serde_json::Value::as_i64)
.unwrap_or(i64::from(status == "failed"));
Ok(ProjectionState {
cursor: Some(cursor),
failed_docs,
cursor: cursor
.get("lastDocId")
.and_then(|value| value.as_str())
.map(ToString::to_string),
failed_docs: if status == "running" {
metadata
.get("failedDocs")
.and_then(serde_json::Value::as_i64)
.unwrap_or(0)
} else {
0
},
})
}
async fn purge_removed_doc_refs(pool: &PgPool, workspace_id: &str, current_doc_ids: &[String]) -> RuntimeResult<i64> {
let result = sqlx::query(
r#"
DELETE FROM doc_blob_refs
WHERE workspace_id = $1
AND NOT (doc_id = ANY($2))
"#,
)
.bind(workspace_id)
.bind(current_doc_ids)
.execute(pool)
.await
.map_err(|err| RuntimeError::database("Doc blob refs purge removed docs failed", err))?;
Ok(result.rows_affected() as i64)
async fn purge_removed_doc_projections(
pool: &PgPool,
workspace_id: &str,
current_doc_ids: &[String],
) -> RuntimeResult<i64> {
let refs = sqlx::query("DELETE FROM doc_blob_refs WHERE workspace_id = $1 AND NOT (doc_id = ANY($2))")
.bind(workspace_id)
.bind(current_doc_ids)
.execute(pool)
.await
.map_err(|err| RuntimeError::database("Doc blob refs purge removed docs failed", err))?
.rows_affected() as i64;
sqlx::query("DELETE FROM doc_blob_ref_projections WHERE workspace_id = $1 AND NOT (doc_id = ANY($2))")
.bind(workspace_id)
.bind(current_doc_ids)
.execute(pool)
.await
.map_err(|err| RuntimeError::database("Doc blob ref projections purge removed docs failed", err))?;
Ok(refs)
}
fn extract_refs(blob: Vec<u8>) -> RuntimeResult<Vec<ExtractedRef>> {
@@ -155,6 +302,283 @@ fn extract_refs(blob: Vec<u8>) -> RuntimeResult<Vec<ExtractedRef>> {
.map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs parse failed: {err}")))
}
async fn replace_doc_refs_if_current(
pool: &PgPool,
workspace_id: &str,
doc_id: &str,
source_revision: DateTime<Utc>,
refs: Vec<ExtractedRef>,
) -> RuntimeResult<ProjectionOutcome> {
let mut tx = pool
.begin()
.await
.map_err(|err| RuntimeError::database("Doc blob refs transaction failed", err))?;
let current = sqlx::query_as::<_, (DateTime<Utc>, bool)>(
r#"
SELECT s.updated_at,
EXISTS(
SELECT 1 FROM updates u
WHERE u.workspace_id = s.workspace_id AND u.guid = s.guid
) AS has_pending_updates
FROM snapshots s
WHERE s.workspace_id = $1 AND s.guid = $2
"#,
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs CAS source load failed", err))?;
let Some((current_revision, has_pending_updates)) = current else {
tx.rollback()
.await
.map_err(|err| RuntimeError::database("Doc blob refs CAS rollback failed", err))?;
upsert_projection_state(
pool,
workspace_id,
doc_id,
None,
"missing",
Some("snapshot_missing"),
None,
)
.await?;
return Ok(ProjectionOutcome::Missing);
};
if current_revision != source_revision || has_pending_updates {
tx.rollback()
.await
.map_err(|err| RuntimeError::database("Doc blob refs CAS rollback failed", err))?;
upsert_projection_state(
pool,
workspace_id,
doc_id,
Some(source_revision),
"pending",
Some("source_changed"),
None,
)
.await?;
return Ok(ProjectionOutcome::Pending);
}
let projection = sqlx::query_as::<_, (i32, Option<DateTime<Utc>>)>(
"SELECT parser_version, source_revision FROM doc_blob_ref_projections WHERE workspace_id = $1 AND doc_id = $2",
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs projection CAS load failed", err))?;
if projection.is_some_and(|(parser_version, projection_revision)| {
parser_version > PARSER_VERSION || projection_revision.is_some_and(|revision| revision > source_revision)
}) {
tx.rollback()
.await
.map_err(|err| RuntimeError::database("Doc blob refs projection CAS rollback failed", err))?;
upsert_projection_state(
pool,
workspace_id,
doc_id,
Some(source_revision),
"pending",
Some("projection_newer"),
None,
)
.await?;
return Ok(ProjectionOutcome::Pending);
}
let mut old_refs = sqlx::query_as::<_, (String, String, String)>(
"SELECT blob_key, block_id, flavour FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2",
)
.bind(workspace_id)
.bind(doc_id)
.fetch_all(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs shadow load failed", err))?;
old_refs.sort();
let mut new_refs = refs
.iter()
.map(|reference| {
(
reference.blob_key.clone(),
reference.block_id.clone(),
reference.flavour.clone(),
)
})
.collect::<Vec<_>>();
new_refs.sort();
let shadow_mismatch = old_refs != new_refs;
let deleted = sqlx::query("DELETE FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2")
.bind(workspace_id)
.bind(doc_id)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs delete failed", err))?
.rows_affected() as i64;
let mut written = 0;
for reference in refs {
written += sqlx::query(
r#"
INSERT INTO doc_blob_refs
(workspace_id, doc_id, blob_key, block_id, flavour, snapshot_updated_at, parser_version, status, error)
VALUES ($1, $2, $3, $4, $5, $6, $7, 'fresh', NULL)
ON CONFLICT (workspace_id, doc_id, blob_key, block_id) DO UPDATE
SET flavour = EXCLUDED.flavour,
snapshot_updated_at = EXCLUDED.snapshot_updated_at,
indexed_at = CURRENT_TIMESTAMP,
parser_version = EXCLUDED.parser_version,
status = 'fresh',
error = NULL
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(reference.blob_key)
.bind(reference.block_id)
.bind(reference.flavour)
.bind(source_revision)
.bind(PARSER_VERSION)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs insert failed", err))?
.rows_affected() as i64;
}
upsert_projection_state(
&mut *tx,
workspace_id,
doc_id,
Some(source_revision),
"fresh",
None,
None,
)
.await?;
tx.commit()
.await
.map_err(|err| RuntimeError::database("Doc blob refs transaction commit failed", err))?;
Ok(ProjectionOutcome::Fresh(ProjectionAttempt {
written,
deleted,
shadow_mismatch,
}))
}
async fn rebuild_doc_blob_refs_inner(
runtime: &StorageRuntime,
workspace_id: &str,
doc_id: &str,
expected_source_revision: Option<i64>,
) -> RuntimeResult<ProjectionStats> {
let pool = runtime.pool().await?;
let mut stats = ProjectionStats::default();
stats.result.scanned_docs = 1;
let Some(source) = load_doc_source(&pool, workspace_id, doc_id).await? else {
upsert_projection_state(
&pool,
workspace_id,
doc_id,
None,
"missing",
Some("snapshot_missing"),
None,
)
.await?;
stats.result.failed_docs = 1;
stats.missing_docs = 1;
return Ok(stats);
};
if expected_source_revision.is_some_and(|revision| source.updated_at.timestamp_millis() != revision) {
upsert_projection_state(
&pool,
workspace_id,
doc_id,
Some(source.updated_at),
"pending",
Some("source_changed"),
None,
)
.await?;
stats.pending_docs = 1;
return Ok(stats);
}
if source.has_pending_updates {
upsert_projection_state(
&pool,
workspace_id,
doc_id,
Some(source.updated_at),
"pending",
Some("pending_updates"),
None,
)
.await?;
stats.pending_docs = 1;
return Ok(stats);
}
if projection_is_fresh(&pool, workspace_id, doc_id, source.updated_at).await? {
return Ok(stats);
}
upsert_projection_state(
&pool,
workspace_id,
doc_id,
Some(source.updated_at),
"running",
None,
None,
)
.await?;
let Some(snapshot) = load_canonical_doc(&pool, workspace_id, doc_id).await? else {
upsert_projection_state(
&pool,
workspace_id,
doc_id,
None,
"missing",
Some("snapshot_missing"),
None,
)
.await?;
stats.result.failed_docs = 1;
stats.missing_docs = 1;
return Ok(stats);
};
let CurrentDoc { blob, updated_at, .. } = snapshot;
let refs = match extract_refs(blob) {
Ok(refs) => refs,
Err(_) => {
upsert_projection_state(
&pool,
workspace_id,
doc_id,
Some(updated_at),
"failed",
Some("parse_failed"),
Some("canonical snapshot parser rejected the document"),
)
.await?;
stats.result.failed_docs = 1;
return Ok(stats);
}
};
match replace_doc_refs_if_current(&pool, workspace_id, doc_id, updated_at, refs).await? {
ProjectionOutcome::Fresh(attempt) => {
stats.result.parsed_docs = 1;
stats.result.refs_written = attempt.written;
stats.result.refs_deleted = attempt.deleted;
stats.shadow_mismatches = i64::from(attempt.shadow_mismatch);
}
ProjectionOutcome::Pending => stats.pending_docs = 1,
ProjectionOutcome::Missing => {
stats.result.failed_docs = 1;
stats.missing_docs = 1;
}
}
Ok(stats)
}
#[cfg(test)]
mod tests {
use chrono::Utc;
@@ -168,14 +592,10 @@ mod tests {
let blob =
doc_loader::build_full_doc("Doc", "![Alt](blob://image-blob-key)", &doc_id).expect("doc fixture should build");
let snapshot = CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id,
blob,
updated_at: Utc::now(),
};
let refs = extract_refs(snapshot.blob).expect("refs should parse");
assert!(
refs
.iter()
@@ -201,144 +621,25 @@ mod tests {
meta
.insert("pages".to_string(), pages)
.expect("root pages should insert");
let root = root.encode_update_v1().expect("root doc should encode");
let ids = doc_loader::get_doc_ids_from_binary(root, true).expect("root doc ids should parse");
assert_eq!(ids, vec!["active-doc", "trashed-doc"]);
}
#[test]
fn doc_blob_refs_rejects_corrupt_docs() {
fn doc_blob_refs_rejects_corrupt_docs_without_a_failure_ref() {
let snapshot = CurrentDoc {
workspace_id: "workspace".to_string(),
doc_id: "corrupt".to_string(),
blob: vec![0xff],
updated_at: Utc::now(),
};
assert!(extract_refs(snapshot.blob).is_err());
}
}
async fn replace_doc_refs(
pool: &PgPool,
workspace_id: &str,
doc_id: &str,
updated_at: DateTime<Utc>,
refs: Vec<ExtractedRef>,
) -> RuntimeResult<(i64, i64)> {
let mut tx = pool
.begin()
.await
.map_err(|err| RuntimeError::database("Doc blob refs transaction failed", err))?;
let deleted = sqlx::query("DELETE FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2")
.bind(workspace_id)
.bind(doc_id)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs delete failed", err))?
.rows_affected() as i64;
let mut written = 0;
for reference in refs {
let affected = sqlx::query(
r#"
INSERT INTO doc_blob_refs
(workspace_id, doc_id, blob_key, block_id, flavour, snapshot_updated_at, parser_version, status)
VALUES ($1, $2, $3, $4, $5, $6, $7, 'fresh')
ON CONFLICT (workspace_id, doc_id, blob_key, block_id) DO UPDATE
SET flavour = EXCLUDED.flavour,
snapshot_updated_at = EXCLUDED.snapshot_updated_at,
indexed_at = CURRENT_TIMESTAMP,
parser_version = EXCLUDED.parser_version,
status = 'fresh',
error = NULL
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(reference.blob_key)
.bind(reference.block_id)
.bind(reference.flavour)
.bind(updated_at)
.bind(PARSER_VERSION)
.execute(&mut *tx)
.await
.map_err(|err| RuntimeError::database("Doc blob refs insert failed", err))?
.rows_affected() as i64;
written += affected;
#[test]
fn error_summary_is_bounded() {
let error = "x".repeat(ERROR_SUMMARY_LIMIT + 20);
assert_eq!(truncate_error_summary(&error).len(), ERROR_SUMMARY_LIMIT);
}
tx.commit()
.await
.map_err(|err| RuntimeError::database("Doc blob refs transaction commit failed", err))?;
Ok((written, deleted))
}
async fn mark_doc_failed(pool: &PgPool, workspace_id: &str, doc_id: &str, error: &str) -> RuntimeResult<()> {
sqlx::query(
r#"
INSERT INTO doc_blob_refs
(workspace_id, doc_id, blob_key, block_id, flavour, snapshot_updated_at, parser_version, status, error)
VALUES ($1, $2, '__parse_failed__', '__parse_failed__', '__parse_failed__', CURRENT_TIMESTAMP, $3, 'failed', $4)
ON CONFLICT (workspace_id, doc_id, blob_key, block_id) DO UPDATE
SET indexed_at = CURRENT_TIMESTAMP,
status = 'failed',
error = EXCLUDED.error
"#,
)
.bind(workspace_id)
.bind(doc_id)
.bind(PARSER_VERSION)
.bind(error)
.execute(pool)
.await
.map_err(|err| RuntimeError::database("Doc blob refs mark failure failed", err))?;
Ok(())
}
async fn rebuild_doc_blob_refs_inner(
runtime: &StorageRuntime,
workspace_id: String,
doc_id: String,
) -> RuntimeResult<RuntimeDocBlobRefsResult> {
let pool = runtime.pool().await?;
let mut result = RuntimeDocBlobRefsResult {
scanned_docs: 1,
parsed_docs: 0,
refs_written: 0,
refs_deleted: 0,
failed_docs: 0,
next_cursor: None,
};
let Some(snapshot) = load_current_doc(&pool, &workspace_id, &doc_id).await? else {
result.failed_docs = 1;
mark_doc_failed(&pool, &workspace_id, &doc_id, "snapshot_missing").await?;
return Ok(result);
};
let CurrentDoc {
workspace_id,
doc_id,
blob,
updated_at,
} = snapshot;
match extract_refs(blob) {
Ok(refs) => {
let (written, deleted) = replace_doc_refs(&pool, &workspace_id, &doc_id, updated_at, refs).await?;
result.parsed_docs = 1;
result.refs_written = written;
result.refs_deleted = deleted;
}
Err(err) => {
result.failed_docs = 1;
mark_doc_failed(&pool, &workspace_id, &doc_id, &err.to_string()).await?;
}
}
Ok(result)
}
#[napi_derive::napi]
@@ -348,8 +649,13 @@ impl StorageRuntime {
&self,
workspace_id: String,
doc_id: String,
source_revision: i64,
) -> napi::Result<RuntimeDocBlobRefsResult> {
Ok(rebuild_doc_blob_refs_inner(self, workspace_id, doc_id).await?)
Ok(
rebuild_doc_blob_refs_inner(self, &workspace_id, &doc_id, Some(source_revision))
.await?
.result,
)
}
#[napi]
@@ -361,12 +667,11 @@ impl StorageRuntime {
if limit <= 0 {
return Err(napi_error("doc blob refs rebuild limit must be positive"));
}
let pool = self.pool().await?;
let doc_ids = match load_workspace_doc_ids(&pool, &workspace_id).await {
Ok(doc_ids) => doc_ids,
Err(err) => {
upsert_projection_failure_checkpoint(&pool, &workspace_id, &err.to_string()).await?;
upsert_projection_failure_checkpoint(&pool, &workspace_id).await?;
return Err(err.into());
}
};
@@ -377,34 +682,35 @@ impl StorageRuntime {
.filter(|doc_id| state.cursor.as_ref().is_none_or(|cursor| doc_id > cursor))
.collect::<Vec<_>>();
let has_more = doc_ids.len() > limit as usize;
let mut total = RuntimeDocBlobRefsResult {
scanned_docs: 0,
parsed_docs: 0,
refs_written: 0,
refs_deleted: 0,
failed_docs: 0,
next_cursor: None,
};
let mut total = ProjectionStats::default();
total.result.failed_docs = state.failed_docs;
let mut last_doc_id = None;
for doc_id in doc_ids.into_iter().take(limit as usize) {
last_doc_id = Some(doc_id.clone());
let result = rebuild_doc_blob_refs_inner(self, workspace_id.clone(), doc_id).await?;
total.scanned_docs += result.scanned_docs;
total.parsed_docs += result.parsed_docs;
total.refs_written += result.refs_written;
total.refs_deleted += result.refs_deleted;
total.failed_docs += result.failed_docs;
let stats = rebuild_doc_blob_refs_inner(self, &workspace_id, &doc_id, None).await?;
total.result.scanned_docs += stats.result.scanned_docs;
total.result.parsed_docs += stats.result.parsed_docs;
total.result.refs_written += stats.result.refs_written;
total.result.refs_deleted += stats.result.refs_deleted;
total.result.failed_docs += stats.result.failed_docs;
total.pending_docs += stats.pending_docs;
total.missing_docs += stats.missing_docs;
total.shadow_mismatches += stats.shadow_mismatches;
}
let failed_docs = state.failed_docs + total.failed_docs;
if has_more {
total.next_cursor = last_doc_id;
} else if failed_docs == 0 {
total.refs_deleted += purge_removed_doc_refs(&pool, &workspace_id, &current_doc_ids).await?;
total.result.next_cursor = last_doc_id;
} else if total.result.failed_docs == 0 && total.pending_docs == 0 {
total.result.refs_deleted += purge_removed_doc_projections(&pool, &workspace_id, &current_doc_ids).await?;
}
upsert_projection_checkpoint(&pool, &workspace_id, &total, failed_docs).await?;
Ok(total)
upsert_projection_checkpoint(
&pool,
&workspace_id,
&total.result,
total.pending_docs,
total.missing_docs,
total.shadow_mismatches,
)
.await?;
Ok(total.result)
}
}
packages/backend/native/src/runtime/storage_runtime/document_cleanup.rs
+90 -18
View File
@@ -269,14 +269,13 @@ async fn load_current_doc_for_update(
workspace_id: &str,
doc_id: &str,
) -> RuntimeResult<Option<CurrentDoc>> {
let snapshot = sqlx::query_as::<_, CurrentDoc>(
"SELECT workspace_id, guid AS doc_id, blob, updated_at FROM snapshots WHERE workspace_id = $1 AND guid = $2",
)
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup current snapshot load failed", err))?;
let snapshot =
sqlx::query_as::<_, CurrentDoc>("SELECT blob, updated_at FROM snapshots WHERE workspace_id = $1 AND guid = $2")
.bind(workspace_id)
.bind(doc_id)
.fetch_optional(&mut **tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup current snapshot load failed", err))?;
let updates = sqlx::query_as::<_, CurrentDocUpdate>(
"SELECT blob, created_at FROM updates WHERE workspace_id = $1 AND guid = $2 ORDER BY created_at ASC",
)
@@ -285,7 +284,7 @@ async fn load_current_doc_for_update(
.fetch_all(&mut **tx)
.await
.map_err(|err| RuntimeError::database("Document cleanup current updates load failed", err))?;
merge_current_doc(workspace_id, doc_id, snapshot, updates)
merge_current_doc(snapshot, updates)
}
async fn current_activity(
@@ -358,6 +357,7 @@ async fn delete_doc_rows(tx: &mut Transaction<'_, Postgres>, candidate: &Candida
("doc_access_policies", "doc_id"),
("doc_grants", "doc_id"),
("doc_blob_refs", "doc_id"),
("doc_blob_ref_projections", "doc_id"),
("ai_workspace_ignored_docs", "doc_id"),
("comments", "doc_id"),
("comment_attachments", "doc_id"),
@@ -937,6 +937,7 @@ mod tests {
"blob_cleanup_candidates",
"document_cleanup_candidates",
"doc_blob_refs",
"doc_blob_ref_projections",
] {
sqlx::query(&format!("DELETE FROM {table} WHERE workspace_id = $1"))
.bind(workspace_id)
@@ -1001,6 +1002,59 @@ mod tests {
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(projection.failed_docs, 0);
assert_eq!(projection.parsed_docs, 3);
let projection_checkpoint = sqlx::query(
"SELECT metadata FROM storage_reconciliation_checkpoints WHERE kind = 'doc_blob_refs' AND scope = $1",
)
.bind(&workspace_id)
.fetch_one(&pool)
.await?;
assert_eq!(projection_checkpoint.get::<Value, _>("metadata")["shadowMismatches"], 1);
assert_eq!(
sqlx::query_scalar::<_, String>(
"SELECT status FROM doc_blob_ref_projections WHERE workspace_id = $1 AND doc_id = 'live-doc'",
)
.bind(&workspace_id)
.fetch_one(&pool)
.await?,
"fresh"
);
let unchanged = runtime
.rebuild_workspace_doc_blob_refs(workspace_id.clone(), 100)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(unchanged.parsed_docs, 0);
sqlx::query(
"INSERT INTO updates (workspace_id, guid, blob, created_at) VALUES ($1, 'live-doc', $2, CURRENT_TIMESTAMP)",
)
.bind(&workspace_id)
.bind(affine_doc_loader::build_full_doc("Live pending", "", "live-doc")?)
.execute(&pool)
.await?;
let pending = runtime
.rebuild_workspace_doc_blob_refs(workspace_id.clone(), 100)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(pending.failed_docs, 0);
assert_eq!(
sqlx::query_scalar::<_, String>(
"SELECT status FROM doc_blob_ref_projections WHERE workspace_id = $1 AND doc_id = 'live-doc'",
)
.bind(&workspace_id)
.fetch_one(&pool)
.await?,
"pending"
);
sqlx::query("DELETE FROM updates WHERE workspace_id = $1 AND guid = 'live-doc'")
.bind(&workspace_id)
.execute(&pool)
.await?;
let repaired = runtime
.rebuild_workspace_doc_blob_refs(workspace_id.clone(), 100)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(repaired.parsed_docs, 1);
assert_eq!(
sqlx::query_scalar::<_, i64>(
"SELECT COUNT(*) FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2 AND blob_key = 'candidate-blob'",
@@ -1023,7 +1077,7 @@ mod tests {
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!(
(partial.failed_docs, partial.next_cursor.as_deref()),
(1, Some("live-doc"))
(0, Some("live-doc"))
);
let partial_checkpoint = sqlx::query(
"SELECT status, metadata FROM storage_reconciliation_checkpoints WHERE kind = 'doc_blob_refs' AND scope = $1",
@@ -1032,11 +1086,11 @@ mod tests {
.fetch_one(&pool)
.await?;
assert_eq!(partial_checkpoint.get::<String, _>("status"), "running");
assert_eq!(partial_checkpoint.get::<Value, _>("metadata")["failedDocs"], 1);
assert_eq!(partial_checkpoint.get::<Value, _>("metadata")["failedDocs"], 0);
sqlx::query(
"UPDATE storage_reconciliation_checkpoints SET status = 'failed', metadata = '{\"parserVersion\":1}' WHERE kind \
= 'doc_blob_refs' AND scope = $1",
"UPDATE storage_reconciliation_checkpoints SET status = 'failed', metadata = \
'{\"parserVersion\":1,\"failedDocs\":99}' WHERE kind = 'doc_blob_refs' AND scope = $1",
)
.bind(&workspace_id)
.execute(&pool)
@@ -1045,15 +1099,18 @@ mod tests {
.rebuild_workspace_doc_blob_refs(workspace_id.clone(), 1)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!((resumed.failed_docs, resumed.next_cursor), (0, None));
assert_eq!(
(resumed.failed_docs, resumed.next_cursor.as_deref()),
(0, Some("missing-doc"))
);
let resumed_checkpoint = sqlx::query(
"SELECT status, metadata FROM storage_reconciliation_checkpoints WHERE kind = 'doc_blob_refs' AND scope = $1",
)
.bind(&workspace_id)
.fetch_one(&pool)
.await?;
assert_eq!(resumed_checkpoint.get::<String, _>("status"), "failed");
assert_eq!(resumed_checkpoint.get::<Value, _>("metadata")["failedDocs"], 1);
assert_eq!(resumed_checkpoint.get::<String, _>("status"), "running");
assert_eq!(resumed_checkpoint.get::<Value, _>("metadata")["failedDocs"], 0);
sqlx::query("UPDATE snapshots SET blob = $2 WHERE workspace_id = $1 AND guid = 'live-doc'")
.bind(&workspace_id)
@@ -1086,7 +1143,7 @@ mod tests {
.rebuild_workspace_doc_blob_refs(workspace_id.clone(), 100)
.await
.map_err(|err| anyhow::anyhow!(err.to_string()))?;
assert_eq!((parser_upgrade.scanned_docs, parser_upgrade.failed_docs), (2, 0));
assert_eq!((parser_upgrade.scanned_docs, parser_upgrade.failed_docs), (3, 0));
assert_eq!(
sqlx::query_scalar::<_, String>(
"SELECT status FROM storage_reconciliation_checkpoints WHERE kind = 'doc_blob_refs' AND scope = $1",
@@ -1159,6 +1216,10 @@ mod tests {
.bind(&workspace_id)
.execute(&pool)
.await?;
sqlx::query("DELETE FROM doc_blob_ref_projections WHERE workspace_id = $1")
.bind(&workspace_id)
.execute(&pool)
.await?;
sqlx::query("DELETE FROM document_cleanup_candidates WHERE workspace_id = $1")
.bind(&workspace_id)
.execute(&pool)
@@ -1259,7 +1320,17 @@ mod tests {
.bind(doc_id)
.fetch_one(&pool)
.await?;
assert_eq!(ref_count, 1);
assert_eq!(ref_count, 0);
assert_eq!(
sqlx::query_scalar::<_, String>(
"SELECT status FROM doc_blob_ref_projections WHERE workspace_id = $1 AND doc_id = $2",
)
.bind(&workspace_id)
.bind(doc_id)
.fetch_one(&pool)
.await?,
"pending"
);
let not_due = execute_one(&pool, Some(&workspace_id), 30).await?;
assert!(not_due.is_none());
@@ -1416,6 +1487,7 @@ mod tests {
("doc_access_policies", "doc_id"),
("doc_grants", "doc_id"),
("doc_blob_refs", "doc_id"),
("doc_blob_ref_projections", "doc_id"),
("ai_workspace_ignored_docs", "doc_id"),
("comments", "doc_id"),
("comment_attachments", "doc_id"),
packages/backend/native/src/runtime/storage_runtime/mod.rs
+4 -1
View File
@@ -17,7 +17,10 @@ pub use capabilities::StorageProviderCapabilities;
use capabilities::storage_provider_capabilities;
use config::StorageRuntimeConfig;
pub(super) use current_doc::load_current_doc;
use current_doc::{CurrentDoc, CurrentDocUpdate, load_workspace_live_doc_ids, merge_current_doc};
use current_doc::{
CurrentDoc, CurrentDocUpdate, load_canonical_doc, load_workspace_canonical_doc_ids, load_workspace_live_doc_ids,
merge_current_doc,
};
use super::object_storage::{
self, ObjectStorageService, StorageBackendConfig,
packages/backend/native/src/runtime/types.rs
+8
View File
@@ -444,6 +444,7 @@ pub struct RuntimeBlobMetadataBackfillResult {
pub workspace_ids: Vec<String>,
}
#[derive(Default)]
#[napi_derive::napi(object)]
pub struct RuntimeDocBlobRefsResult {
pub scanned_docs: i64,
@@ -552,3 +553,10 @@ pub struct RuntimeEmbeddingProgress {
pub total: i64,
pub embedded: i64,
}
#[napi_derive::napi(object)]
pub struct SearchOperationOutput {
pub ok: bool,
pub value: Option<serde_json::Value>,
pub error_code: Option<String>,
}
packages/backend/native/src/search_index/document.rs
+63
View File
@@ -0,0 +1,63 @@
use memory_indexer::{Document, FieldType, Value};
use serde_json::Value as JsonValue;
use super::{IndexError, Result, schema::TableSchema};
pub(super) fn compile_document(table: &TableSchema, value: JsonValue) -> Result<Document> {
let mut object = value
.as_object()
.cloned()
.ok_or_else(|| IndexError::InvalidInput("index document must be an object".into()))?;
let explicit_id = object
.remove("_id")
.map(|value| {
value
.as_str()
.map(str::to_owned)
.ok_or_else(|| IndexError::InvalidInput("index document _id must be a string".into()))
})
.transpose()?;
let id = match explicit_id {
Some(id) => id,
None => table.document_id(&object)?,
};
let mut document = Document::new(id);
for (name, value) in object {
if value.is_null() {
continue;
}
let field = table.field(&name)?;
let values = match value {
JsonValue::Array(values) => values.into_iter().filter(|value| !value.is_null()).collect(),
value => vec![value],
};
if values.is_empty() {
continue;
}
document.add_values(
field,
values
.into_iter()
.map(|value| compile_value(table.field_type(field), value))
.collect::<Result<Vec<_>>>()?,
);
}
Ok(document)
}
fn compile_value(field_type: &FieldType, value: JsonValue) -> Result<Value> {
match field_type {
FieldType::Text(_) | FieldType::Keyword => value
.as_str()
.map(|value| Value::String(value.into()))
.ok_or_else(|| IndexError::InvalidInput("string index value required".into())),
FieldType::I64 => value
.as_i64()
.map(Value::I64)
.ok_or_else(|| IndexError::InvalidInput("integer index value required".into())),
FieldType::Bool => value
.as_bool()
.map(Value::Bool)
.ok_or_else(|| IndexError::InvalidInput("boolean index value required".into())),
}
}
packages/backend/native/src/search_index/mod.rs
+371
View File
@@ -0,0 +1,371 @@
mod document;
mod query;
mod result;
mod schema;
use std::sync::Arc;
use memory_indexer::{MemoryIndex, Mutation, TermsAggregation};
use napi::{Status, bindgen_prelude::Buffer};
use serde_json::Value as JsonValue;
use tokio::sync::RwLock;
use self::{
document::compile_document,
query::{compile_options, compile_query},
result::{HighlightTags, aggregate_result, search_result},
schema::TableSchema,
};
type Result<T> = std::result::Result<T, IndexError>;
#[derive(Debug, thiserror::Error)]
enum IndexError {
#[error("Invalid index input: {0}")]
InvalidInput(String),
#[error(transparent)]
Memory(#[from] memory_indexer::Error),
#[error(transparent)]
Json(#[from] serde_json::Error),
}
impl From<IndexError> for napi::Error {
fn from(error: IndexError) -> Self {
napi::Error::new(Status::InvalidArg, error.to_string())
}
}
struct TableIndex {
schema: TableSchema,
index: RwLock<MemoryIndex>,
}
impl TableIndex {
fn new(schema: TableSchema) -> Self {
Self {
index: RwLock::new(MemoryIndex::new(schema.schema.clone())),
schema,
}
}
}
struct IndexManager {
doc: TableIndex,
block: TableIndex,
}
impl IndexManager {
fn new() -> Self {
Self {
doc: TableIndex::new(TableSchema::doc()),
block: TableIndex::new(TableSchema::block()),
}
}
fn table(&self, name: &str) -> Result<&TableIndex> {
match name {
"doc" => Ok(&self.doc),
"block" => Ok(&self.block),
_ => Err(IndexError::InvalidInput(format!("unknown index table {name}"))),
}
}
}
pub(crate) struct EmbeddedIndexCheckpoint {
pub sequence: i64,
pub data: Buffer,
}
pub(crate) struct EmbeddedSearchIndex {
manager: Arc<IndexManager>,
}
impl EmbeddedSearchIndex {
pub(crate) fn new() -> Self {
Self {
manager: Arc::new(IndexManager::new()),
}
}
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 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 write(&self, table: String, documents_json: String) -> napi::Result<()> {
let table = self.manager.table(&table)?;
let documents: Vec<JsonValue> = serde_json::from_str(&documents_json)?;
let documents = documents
.into_iter()
.map(|document| compile_document(&table.schema, document))
.collect::<Result<Vec<_>>>()?;
table
.index
.write()
.await
.apply_batch(documents.into_iter().map(Mutation::Upsert).collect())
.map_err(IndexError::from)?;
Ok(())
}
pub(crate) async fn delete(&self, table: String, id: String) -> napi::Result<()> {
self.manager.table(&table)?.index.write().await.delete(&id);
Ok(())
}
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)?;
let query = compile_query(
&table.schema,
dsl
.get("query")
.ok_or_else(|| IndexError::InvalidInput("search query is required".into()))?,
)?;
let result = table
.index
.read()
.await
.search(&query, compile_options(&table.schema, &dsl)?)
.map_err(IndexError::from)?;
Ok(serde_json::to_string(&search_result(
&table.schema,
result,
&highlight_tags(&dsl),
))?)
}
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)?;
let query = compile_query(
&table.schema,
dsl
.get("query")
.ok_or_else(|| IndexError::InvalidInput("aggregate query is required".into()))?,
)?;
let terms = dsl
.pointer("/aggs/result/terms")
.ok_or_else(|| IndexError::InvalidInput("terms aggregation is required".into()))?;
let top_hits = dsl
.pointer("/aggs/result/aggs/result/top_hits")
.map(|options| compile_options(&table.schema, options))
.transpose()?;
let limit = terms.get("size").and_then(JsonValue::as_u64).unwrap_or(10) as usize;
let result = table
.index
.read()
.await
.aggregate(
&query,
TermsAggregation {
field: table.schema.field(
terms
.get("field")
.and_then(JsonValue::as_str)
.ok_or_else(|| IndexError::InvalidInput("aggregation field is required".into()))?,
)?,
limit: limit.saturating_add(1),
offset: dsl.get("from").and_then(JsonValue::as_u64).unwrap_or(0) as usize,
top_hits,
},
)
.map_err(IndexError::from)?;
Ok(serde_json::to_string(&aggregate_result(
&table.schema,
result,
limit,
&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 highlight_tags(dsl: &JsonValue) -> HighlightTags {
dsl
.pointer("/highlight/fields")
.and_then(JsonValue::as_object)
.into_iter()
.flatten()
.filter_map(|(field, options)| {
Some((
field.clone(),
(
options.get("pre_tags")?.as_array()?.first()?.as_str()?.to_string(),
options.get("post_tags")?.as_array()?.first()?.as_str()?.to_string(),
),
))
})
.collect()
}
impl Default for EmbeddedSearchIndex {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use serde_json::{Value, json};
use super::EmbeddedSearchIndex;
fn doc(workspace: &str, id: &str, title: &str, updated_at: i64) -> Value {
json!({
"workspace_id": workspace,
"doc_id": id,
"title": title,
"summary": title,
"created_by_user_id": "user",
"updated_by_user_id": "user",
"created_at": updated_at,
"updated_at": updated_at
})
}
fn search(query: Value, cursor: Option<&str>) -> String {
json!({
"query": query,
"fields": ["doc_id", "title"],
"_source": ["doc_id"],
"sort": ["_score", { "updated_at": "desc" }, "doc_id"],
"size": 1,
"cursor": cursor
})
.to_string()
}
#[tokio::test]
async fn exact_search_cursor_and_checkpoint_roundtrip() {
let index = EmbeddedSearchIndex::new();
index
.write(
"doc".into(),
json!([
doc("workspace-1", "one", "设计文档", 1),
doc("workspace-1", "two", "设计方案", 2),
doc("workspace-2", "three", "设计文档", 3)
])
.to_string(),
)
.await
.unwrap();
let query = json!({ "term": { "workspace_id": { "value": "workspace-1" } } });
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]["fields"]["doc_id"], json!(["two"]));
let cursor = first["nextCursor"].as_str().unwrap();
let second: Value = serde_json::from_str(
&index
.search("doc".into(), search(query.clone(), Some(cursor)))
.await
.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);
let aggregate: Value = serde_json::from_str(
&index
.aggregate(
"doc".into(),
json!({
"query":{"match_all":{}},
"from":0,
"aggs":{"result":{"terms":{"field":"workspace_id","size":10},"aggs":{"result":{"top_hits":{
"size":1,"fields":["doc_id","title"],"sort":["updated_at","doc_id"]
}}}}}
})
.to_string(),
)
.await
.unwrap(),
)
.unwrap();
assert_eq!(aggregate["total"], 2);
assert_eq!(aggregate["buckets"][0]["key"], "workspace-1");
assert_eq!(aggregate["buckets"][0]["count"], 2);
assert!(aggregate["buckets"][0]["hits"][0]["fields"]["doc_id"].is_array());
}
#[tokio::test]
async fn write_is_atomic_and_corrupt_checkpoint_is_rejected() {
let index = EmbeddedSearchIndex::new();
index
.write(
"doc".into(),
json!([{ "workspace_id": "workspace", "doc_id": "null-values", "summary": [null] }]).to_string(),
)
.await
.unwrap();
let all = json!({ "match_all": {} });
let result: Value =
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 }
]);
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());
}
}
packages/backend/native/src/search_index/query.rs
+215
View File
@@ -0,0 +1,215 @@
use memory_indexer::{FieldType, Query, SearchMode, SearchOptions, Sort, SortOrder, SortValue, Value};
use serde_json::Value as JsonValue;
use super::{IndexError, Result, schema::TableSchema};
pub(super) fn compile_query(table: &TableSchema, value: &JsonValue) -> Result<Query> {
let query = if let Some(node) = value.get("match") {
let (field, options) = first_entry(node, "match")?;
Query::text(
table.field(field)?,
required_string(options, "query")?,
SearchMode::Auto,
)
} else if let Some(node) = value.get("term") {
let (field, options) = first_entry(node, "term")?;
let field_id = table.field(field)?;
Query::term(field_id, parse_term(table.field_type(field_id), options.get("value"))?)
} else if let Some(node) = value.get("exists") {
Query::Exists(table.field(required_string(node, "field")?)?)
} else if value.get("match_all").is_some() {
Query::All
} else if let Some(node) = value.get("bool") {
Query::boolean(
compile_clauses(table, node.get("must"))?,
compile_clauses(table, node.get("should"))?,
compile_clauses(table, node.get("must_not"))?,
)
} else {
return Err(IndexError::InvalidInput("unsupported search query".into()));
};
let boost = query_boost(value);
Ok(if boost == 1.0 {
query
} else {
Query::Boost {
query: Box::new(query),
factor: boost,
}
})
}
pub(super) fn compile_options(table: &TableSchema, dsl: &JsonValue) -> Result<SearchOptions> {
let limit = dsl.get("size").and_then(JsonValue::as_u64).unwrap_or(10) as usize;
let offset = dsl.get("from").and_then(JsonValue::as_u64).unwrap_or(0) as usize;
let mut stored_fields = string_array(dsl.get("fields"))
.into_iter()
.chain(string_array(dsl.get("_source")))
.map(|field| table.field(field))
.collect::<Result<Vec<_>>>()?;
let mut seen = std::collections::HashSet::new();
stored_fields.retain(|field| seen.insert(*field));
let highlight_fields = dsl
.pointer("/highlight/fields")
.and_then(JsonValue::as_object)
.map(|fields| {
fields
.keys()
.map(|field| table.field(field))
.collect::<Result<Vec<_>>>()
})
.transpose()?
.unwrap_or_default();
let sort = compile_sort(table, dsl.get("sort"))?;
let after = dsl
.get("cursor")
.and_then(JsonValue::as_str)
.map(|cursor| parse_cursor(cursor, &sort, table))
.transpose()?;
Ok(SearchOptions {
limit,
offset,
after,
sort,
stored_fields,
highlight_fields,
})
}
fn compile_sort(table: &TableSchema, value: Option<&JsonValue>) -> Result<Vec<Sort>> {
let mut sorts = Vec::new();
for item in value.and_then(JsonValue::as_array).into_iter().flatten() {
if let Some(field) = item.as_str() {
match field {
"_score" => sorts.push(Sort::ScoreDesc),
"id" | "_id" => sorts.push(Sort::DocumentId),
field => sorts.push(Sort::Field {
field: table.field(field)?,
order: SortOrder::Asc,
}),
}
} else if let Some((field, order)) = item.as_object().and_then(|value| value.iter().next()) {
sorts.push(Sort::Field {
field: table.field(field)?,
order: if order.as_str() == Some("desc") {
SortOrder::Desc
} else {
SortOrder::Asc
},
});
}
}
Ok(sorts)
}
fn parse_cursor(cursor: &str, sorts: &[Sort], table: &TableSchema) -> Result<Vec<SortValue>> {
let values: Vec<JsonValue> = serde_json::from_str(cursor)?;
let mut effective = sorts.to_vec();
if !effective.iter().any(|sort| matches!(sort, Sort::DocumentId)) {
effective.push(Sort::DocumentId);
}
if values.len() != effective.len() {
return Err(IndexError::InvalidInput("invalid search cursor".into()));
}
values
.into_iter()
.zip(effective)
.map(|(value, sort)| match sort {
_ if value.is_null() => Ok(SortValue::Missing),
Sort::ScoreDesc => value
.as_f64()
.map(|value| SortValue::Score(value as f32))
.ok_or_else(|| IndexError::InvalidInput("invalid score cursor".into())),
Sort::DocumentId => value
.as_str()
.map(|value| SortValue::String(value.into()))
.ok_or_else(|| IndexError::InvalidInput("invalid document cursor".into())),
Sort::Field { field, .. } => match table.field_type(field) {
FieldType::Keyword => value
.as_str()
.map(|value| SortValue::String(value.into()))
.ok_or_else(|| IndexError::InvalidInput("invalid keyword cursor".into())),
FieldType::I64 => value
.as_i64()
.map(SortValue::I64)
.ok_or_else(|| IndexError::InvalidInput("invalid integer cursor".into())),
FieldType::Bool => value
.as_bool()
.map(SortValue::Bool)
.ok_or_else(|| IndexError::InvalidInput("invalid boolean cursor".into())),
FieldType::Text(_) => Err(IndexError::InvalidInput("text fields are not sortable".into())),
},
})
.collect()
}
fn compile_clauses(table: &TableSchema, value: Option<&JsonValue>) -> Result<Vec<Query>> {
value
.and_then(JsonValue::as_array)
.into_iter()
.flatten()
.map(|query| compile_query(table, query))
.collect()
}
fn query_boost(value: &JsonValue) -> f32 {
for operator in ["match", "term", "exists", "match_all", "bool"] {
let Some(node) = value.get(operator) else {
continue;
};
if let Some(boost) = node.get("boost").and_then(JsonValue::as_f64) {
return boost as f32;
}
if let Some((_, options)) = node.as_object().and_then(|value| value.iter().next())
&& let Some(boost) = options.get("boost").and_then(JsonValue::as_f64)
{
return boost as f32;
}
}
1.0
}
fn first_entry<'a>(value: &'a JsonValue, operator: &str) -> Result<(&'a str, &'a JsonValue)> {
value
.as_object()
.and_then(|value| value.iter().next())
.map(|(field, value)| (field.as_str(), value))
.ok_or_else(|| IndexError::InvalidInput(format!("invalid {operator} query")))
}
fn required_string<'a>(value: &'a JsonValue, field: &str) -> Result<&'a str> {
value
.get(field)
.and_then(JsonValue::as_str)
.ok_or_else(|| IndexError::InvalidInput(format!("{field} must be a string")))
}
fn parse_term(field_type: &FieldType, value: Option<&JsonValue>) -> Result<Value> {
let value = value.ok_or_else(|| IndexError::InvalidInput("term value is required".into()))?;
match field_type {
FieldType::Keyword => value
.as_str()
.map(|value| Value::String(value.into()))
.ok_or_else(|| IndexError::InvalidInput("keyword term must be a string".into())),
FieldType::I64 => value
.as_i64()
.map(Value::I64)
.ok_or_else(|| IndexError::InvalidInput("integer term must be an integer".into())),
FieldType::Bool => value
.as_bool()
.map(Value::Bool)
.ok_or_else(|| IndexError::InvalidInput("boolean term must be a boolean".into())),
FieldType::Text(_) => Err(IndexError::InvalidInput(
"term query does not accept text fields".into(),
)),
}
}
fn string_array(value: Option<&JsonValue>) -> Vec<&str> {
value
.and_then(JsonValue::as_array)
.into_iter()
.flatten()
.filter_map(JsonValue::as_str)
.collect()
}
packages/backend/native/src/search_index/result.rs
+173
View File
@@ -0,0 +1,173 @@
use std::collections::HashMap;
use memory_indexer::{AggregationResult, SearchHit, SearchResult, SortValue, Value};
use serde::Serialize;
use super::schema::TableSchema;
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct NativeSearchResult {
pub total: usize,
pub nodes: Vec<NativeHit>,
pub next_cursor: Option<String>,
}
#[derive(Serialize)]
pub(super) struct NativeAggregateResult {
pub total: usize,
#[serde(rename = "hasMore")]
pub has_more: bool,
pub buckets: Vec<NativeBucket>,
}
#[derive(Serialize)]
pub(super) struct NativeBucket {
pub key: serde_json::Value,
pub count: u64,
pub hits: Vec<NativeHit>,
}
#[derive(Serialize)]
pub(super) struct NativeHit {
pub id: String,
pub score: f32,
pub fields: serde_json::Map<String, serde_json::Value>,
pub highlights: serde_json::Map<String, serde_json::Value>,
}
pub(super) type HighlightTags = HashMap<String, (String, String)>;
pub(super) fn search_result(
table: &TableSchema,
result: SearchResult,
highlight_tags: &HighlightTags,
) -> NativeSearchResult {
let next_cursor = result.hits.last().map(|hit| cursor(&hit.sort_values));
NativeSearchResult {
total: result.total,
nodes: result
.hits
.into_iter()
.map(|hit| native_hit(table, hit, highlight_tags))
.collect(),
next_cursor,
}
}
pub(super) fn aggregate_result(
table: &TableSchema,
mut result: AggregationResult,
limit: usize,
highlight_tags: &HighlightTags,
) -> NativeAggregateResult {
let total = result.buckets.len();
let has_more = result.buckets.len() > limit;
result.buckets.truncate(limit);
NativeAggregateResult {
total,
has_more,
buckets: result
.buckets
.into_iter()
.map(|bucket| NativeBucket {
key: json_value(bucket.key),
count: bucket.count,
hits: bucket
.hits
.into_iter()
.map(|hit| native_hit(table, hit, highlight_tags))
.collect(),
})
.collect(),
}
}
fn native_hit(table: &TableSchema, hit: SearchHit, highlight_tags: &HighlightTags) -> NativeHit {
let mut fields = serde_json::Map::new();
for (field, values) in hit.fields {
let name = table.field_name(field).to_string();
let values = values.into_iter().map(json_value).collect::<Vec<_>>();
fields.insert(name, serde_json::Value::Array(values));
}
let mut highlights: serde_json::Map<String, serde_json::Value> = serde_json::Map::new();
for highlight in hit.highlights {
let name = table.field_name(highlight.field).to_string();
let Some((before, after)) = highlight_tags.get(&name) else {
continue;
};
let Some(text) = fields
.get(&name)
.and_then(serde_json::Value::as_array)
.and_then(|values| values.get(highlight.value_index as usize))
.and_then(serde_json::Value::as_str)
else {
continue;
};
let value = render_highlight(text, &highlight.spans, before, after);
highlights
.entry(name)
.or_insert_with(|| serde_json::Value::Array(Vec::new()))
.as_array_mut()
.expect("highlight value is an array")
.push(serde_json::Value::String(value));
}
NativeHit {
id: hit.id,
score: hit.score,
fields,
highlights,
}
}
fn render_highlight(text: &str, spans: &[(u32, u32)], before: &str, after: &str) -> String {
let mut output = String::new();
let mut cursor = 0;
for &(start, end) in spans {
let start = utf16_to_byte(text, start as usize);
let end = utf16_to_byte(text, end as usize);
if start < cursor || end < start || end > text.len() {
continue;
}
output.push_str(&text[cursor..start]);
output.push_str(before);
output.push_str(&text[start..end]);
output.push_str(after);
cursor = end;
}
output.push_str(&text[cursor..]);
output
}
fn utf16_to_byte(text: &str, offset: usize) -> usize {
let mut units = 0;
for (byte, character) in text.char_indices() {
if units >= offset {
return byte;
}
units += character.len_utf16();
}
text.len()
}
fn cursor(values: &[SortValue]) -> String {
serde_json::to_string(&values.iter().map(sort_value).collect::<Vec<_>>()).expect("cursor values serialize")
}
fn sort_value(value: &SortValue) -> serde_json::Value {
match value {
SortValue::Score(value) => serde_json::json!(value),
SortValue::String(value) => serde_json::json!(value),
SortValue::I64(value) => serde_json::json!(value),
SortValue::Bool(value) => serde_json::json!(value),
SortValue::Missing => serde_json::Value::Null,
}
}
fn json_value(value: Value) -> serde_json::Value {
match value {
Value::String(value) => serde_json::Value::String(value),
Value::I64(value) => serde_json::json!(value),
Value::Bool(value) => serde_json::json!(value),
}
}
packages/backend/native/src/search_index/schema.rs
+175
View File
@@ -0,0 +1,175 @@
use std::collections::HashMap;
use memory_indexer::{FieldId, FieldOptions, FieldType, PositionEncoding, Schema, TextOptions};
use super::{IndexError, Result};
pub(super) struct TableSchema {
pub schema: Schema,
fields: HashMap<String, FieldId>,
id_fields: &'static [&'static str],
}
impl TableSchema {
pub fn doc() -> Self {
let mut builder = Schema::builder().position_encoding(PositionEncoding::Utf16);
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, "doc_id", true, true);
keyword(&mut builder, &mut fields, "doc_token", true, false);
fields.insert(
"title".into(),
builder.text("title", text_options(), FieldOptions::indexed_stored()),
);
keyword(&mut builder, &mut fields, "summary", false, false);
keyword(&mut builder, &mut fields, "journal", false, false);
keyword(&mut builder, &mut fields, "created_by_user_id", true, false);
keyword(&mut builder, &mut fields, "updated_by_user_id", true, false);
keyword(&mut builder, &mut fields, "acl_read_tokens", true, false);
boolean(&mut builder, &mut fields, "acl_public_readable");
boolean(&mut builder, &mut fields, "acl_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"])
}
pub fn block() -> Self {
let mut builder = Schema::builder().position_encoding(PositionEncoding::Utf16);
let mut fields = HashMap::new();
for field in [
"workspace_id",
"unit_id",
"source_hash",
"visibility",
"element_id",
"frame_id",
"source_block_id",
"flavour",
"blob",
"ref_doc_id",
"parent_flavour",
"parent_block_id",
"created_by_user_id",
"updated_by_user_id",
] {
keyword(&mut builder, &mut fields, field, true, false);
}
keyword(&mut builder, &mut fields, "doc_id", true, true);
keyword(&mut builder, &mut fields, "workspace_token", true, false);
keyword(&mut builder, &mut fields, "doc_token", true, false);
keyword(&mut builder, &mut fields, "block_id", true, true);
keyword(&mut builder, &mut fields, "block_token", true, false);
fields.insert(
"content".into(),
builder.text("content", text_options(), FieldOptions::indexed_stored().multi_value()),
);
for field in ["ref", "additional", "markdown_preview"] {
keyword(&mut builder, &mut fields, field, false, false);
}
integer(&mut builder, &mut fields, "projection_version", true, false);
integer(&mut builder, &mut fields, "created_at", true, true);
integer(&mut builder, &mut fields, "updated_at", true, true);
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"])
}
fn finish(
builder: memory_indexer::SchemaBuilder,
fields: HashMap<String, FieldId>,
id_fields: &'static [&'static str],
) -> Self {
Self {
schema: builder.build().expect("static server index schema must be valid"),
fields,
id_fields,
}
}
pub fn document_id(&self, document: &serde_json::Map<String, serde_json::Value>) -> Result<String> {
self
.id_fields
.iter()
.map(|field| {
document
.get(*field)
.and_then(serde_json::Value::as_str)
.ok_or_else(|| IndexError::InvalidInput(format!("index document {field} is required")))
})
.collect::<Result<Vec<_>>>()
.map(|parts| parts.join("/"))
}
pub fn field(&self, name: &str) -> Result<FieldId> {
self
.fields
.get(name)
.copied()
.ok_or_else(|| IndexError::InvalidInput(format!("unknown index field {name}")))
}
pub fn field_name(&self, field: FieldId) -> &str {
&self.schema.field(field).expect("field belongs to table schema").name
}
pub fn field_type(&self, field: FieldId) -> &FieldType {
&self
.schema
.field(field)
.expect("field belongs to table schema")
.field_type
}
}
fn keyword(
builder: &mut memory_indexer::SchemaBuilder,
fields: &mut HashMap<String, FieldId>,
name: &str,
indexed: bool,
sortable: bool,
) {
let mut options = FieldOptions::new().stored();
if !sortable {
options = options.multi_value();
}
if indexed {
options = options.indexed();
}
if sortable {
options = options.sortable();
}
fields.insert(name.into(), builder.keyword(name, options));
}
fn integer(
builder: &mut memory_indexer::SchemaBuilder,
fields: &mut HashMap<String, FieldId>,
name: &str,
indexed: bool,
sortable: bool,
) {
let mut options = FieldOptions::new().stored();
if indexed {
options = options.indexed();
}
if sortable {
options = options.sortable();
}
fields.insert(name.into(), builder.i64(name, options));
}
fn boolean(builder: &mut memory_indexer::SchemaBuilder, fields: &mut HashMap<String, FieldId>, name: &str) {
fields.insert(name.into(), builder.bool(name, FieldOptions::indexed_stored()));
}
fn text_options() -> TextOptions {
TextOptions::multilingual()
.with_pinyin()
.with_prefix()
.with_fuzzy()
.with_positions()
}