mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-08-22 04:21:40 +08:00
refactor(server): indexer & worker & sync perf (#15504)
This commit is contained in:
@@ -34,6 +34,7 @@ little_exif = { workspace = true }
|
||||
llm_adapter = { workspace = true, features = ["schema", "ureq-client"] }
|
||||
llm_runtime = { workspace = true, features = ["schema", "ureq-client"] }
|
||||
matroska = { workspace = true }
|
||||
memory-indexer = { workspace = true }
|
||||
mp4parse = { workspace = true }
|
||||
napi = { workspace = true, features = ["async", "serde-json"] }
|
||||
napi-derive = { workspace = true }
|
||||
|
||||
Vendored
+72
-1
@@ -65,6 +65,14 @@ export declare class BackendRuntime {
|
||||
reloadConfig(privateKey?: string | undefined | null): Promise<void>
|
||||
health(): Promise<BackendRuntimeHealth>
|
||||
runMigrations(): Promise<void>
|
||||
searchAuthorized(actorUserId: string, workspaceId: string, request: RuntimeSearchRequest): Promise<SearchOperationOutput>
|
||||
aggregateAuthorized(actorUserId: string, workspaceId: string, request: RuntimeAggregateRequest): Promise<SearchOperationOutput>
|
||||
indexSearchDocument(workspaceId: string, docId: string): Promise<void>
|
||||
deleteSearchDocument(workspaceId: string, docId: string): Promise<void>
|
||||
reconcileSearchWorkspace(workspaceId: string): Promise<void>
|
||||
deleteSearchWorkspace(workspaceId: string): Promise<void>
|
||||
filterReadableDocs(actorUserId: string, workspaceId: string, docIds: Array<string>): Promise<Array<string>>
|
||||
searchStatus(): Promise<any>
|
||||
embeddingHealth(): Promise<EmbeddingHealth>
|
||||
syncEmbeddingState(input: SyncEmbeddingStateInput): Promise<RuntimeEmbeddingWorkspaceState>
|
||||
embeddingQueueCounts(): Promise<RuntimeEmbeddingQueueCounts>
|
||||
@@ -100,7 +108,7 @@ export declare class StorageRuntime {
|
||||
cleanupExpiredPendingBlobs(cutoffMs: number, limit: number): Promise<RuntimeBlobCleanupResult>
|
||||
releaseDeletedBlobs(workspaceId: string, limit: number): Promise<RuntimeBlobCleanupResult>
|
||||
backfillMissingBlobMetadata(workspaceId: string | undefined | null, limit: number): Promise<RuntimeBlobMetadataBackfillResult>
|
||||
rebuildDocBlobRefs(workspaceId: string, docId: string): Promise<RuntimeDocBlobRefsResult>
|
||||
rebuildDocBlobRefs(workspaceId: string, docId: string, sourceRevision: number): Promise<RuntimeDocBlobRefsResult>
|
||||
rebuildWorkspaceDocBlobRefs(workspaceId: string, limit: number): Promise<RuntimeDocBlobRefsResult>
|
||||
reconcileWorkspaceDocuments(workspaceId: string): Promise<RuntimeDocumentCleanupReconcileResult>
|
||||
executeDocumentCleanupCandidates(workspaceId: string | undefined | null, gracePeriodDays: number, limit: number): Promise<RuntimeDocumentCleanupExecuteResult>
|
||||
@@ -152,6 +160,17 @@ export const AFFINE_PRO_LICENSE_AES_KEY: string | undefined | null
|
||||
|
||||
export const AFFINE_PRO_PUBLIC_KEY: string | undefined | null
|
||||
|
||||
export interface AggregateHitsOptions {
|
||||
fields: Array<string>
|
||||
highlights: Array<SearchHighlight>
|
||||
pagination: SearchPagination
|
||||
}
|
||||
|
||||
export interface AggregateOptions {
|
||||
hits: AggregateHitsOptions
|
||||
pagination: SearchPagination
|
||||
}
|
||||
|
||||
export interface AppConfigDescriptor {
|
||||
key: string
|
||||
description: string
|
||||
@@ -1159,6 +1178,14 @@ export interface RotateByokCredentialInput {
|
||||
actorUserId: string
|
||||
}
|
||||
|
||||
export interface RuntimeAggregateRequest {
|
||||
table: SearchTable
|
||||
queries: Array<RuntimeSearchQuery>
|
||||
rootQuery: number
|
||||
field: string
|
||||
options: AggregateOptions
|
||||
}
|
||||
|
||||
export interface RuntimeBlobCleanupExecuteResult {
|
||||
scannedCandidates: number
|
||||
deletedObjects: number
|
||||
@@ -1435,6 +1462,23 @@ export interface RuntimeRetrievalScope {
|
||||
preferredSourceIds: Array<string>
|
||||
}
|
||||
|
||||
export interface RuntimeSearchQuery {
|
||||
queryType: string
|
||||
field?: string
|
||||
matchValue?: string
|
||||
query?: number
|
||||
queries?: Array<number>
|
||||
occur?: string
|
||||
boost?: number
|
||||
}
|
||||
|
||||
export interface RuntimeSearchRequest {
|
||||
table: SearchTable
|
||||
queries: Array<RuntimeSearchQuery>
|
||||
rootQuery: number
|
||||
options: SearchOptions
|
||||
}
|
||||
|
||||
export interface RuntimeTurnScopeSnapshot {
|
||||
version: number
|
||||
resolvedAt: string
|
||||
@@ -1565,6 +1609,33 @@ export interface ScopeSelectorInput {
|
||||
source: string
|
||||
}
|
||||
|
||||
export interface SearchHighlight {
|
||||
field: string
|
||||
before: string
|
||||
end: string
|
||||
}
|
||||
|
||||
export interface SearchOperationOutput {
|
||||
ok: boolean
|
||||
value?: any
|
||||
errorCode?: string
|
||||
}
|
||||
|
||||
export interface SearchOptions {
|
||||
fields: Array<string>
|
||||
highlights: Array<SearchHighlight>
|
||||
pagination: SearchPagination
|
||||
}
|
||||
|
||||
export interface SearchPagination {
|
||||
limit?: number
|
||||
skip?: number
|
||||
cursor?: string
|
||||
}
|
||||
|
||||
export type SearchTable = 'doc'|
|
||||
'block';
|
||||
|
||||
export declare function signAuthSessionAccessToken(userId: string, authSessionId: string, keyId: string, secret: Buffer, issuedAt: number, expiresAt: number): string
|
||||
|
||||
export interface StorageProviderCapabilities {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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}");
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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))?,
|
||||
})
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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}}
|
||||
]}}})
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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'"
|
||||
)
|
||||
}
|
||||
@@ -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!({})),
|
||||
})
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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>,
|
||||
}
|
||||
@@ -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",
|
||||
&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");
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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"]));
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
|
||||
@@ -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());
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
@@ -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()
|
||||
);
|
||||
@@ -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(¤t_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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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", "", &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, ¤t_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, ¤t_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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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())),
|
||||
}
|
||||
}
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
-- This migration is intentionally fail-closed. The data migration with the
|
||||
-- same release must have admitted every live legacy context blob through the
|
||||
-- artifact runtime before these product-owned tables are removed.
|
||||
DO $$
|
||||
BEGIN
|
||||
IF to_regclass('public.ai_contexts') IS NOT NULL AND EXISTS (
|
||||
SELECT 1
|
||||
FROM ai_contexts context
|
||||
JOIN ai_sessions_metadata session ON session.id = context.session_id
|
||||
JOIN blobs blob
|
||||
ON blob.workspace_id = session.workspace_id
|
||||
AND blob.deleted_at IS NULL
|
||||
AND blob.status = 'completed'
|
||||
WHERE jsonb_path_exists(
|
||||
context.config::jsonb,
|
||||
'$.** ? (@ == $blobKey)',
|
||||
jsonb_build_object('blobKey', to_jsonb(blob.key::text))
|
||||
)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM workspace_artifacts artifact
|
||||
WHERE artifact.workspace_id = session.workspace_id
|
||||
AND artifact.status = 'ready'
|
||||
AND artifact.storage_scope = 'blob'
|
||||
AND artifact.storage_key = concat(session.workspace_id, '/', blob.key)
|
||||
)
|
||||
) THEN
|
||||
RAISE EXCEPTION
|
||||
'legacy context blob artifact admission is incomplete; run the data migration before cleanup';
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DELETE FROM app_configs WHERE id = 'copilot.providers.defaults';
|
||||
|
||||
ALTER TABLE ai_workspace_byok_configs
|
||||
ALTER COLUMN definition DROP DEFAULT,
|
||||
DROP COLUMN IF EXISTS endpoint,
|
||||
DROP COLUMN IF EXISTS disabled_reason,
|
||||
DROP COLUMN IF EXISTS last_validated_at,
|
||||
DROP COLUMN IF EXISTS last_validation_error;
|
||||
|
||||
DELETE FROM ai_workspace_byok_configs WHERE definition = '{}'::jsonb;
|
||||
|
||||
DROP TABLE IF EXISTS ai_context_embeddings;
|
||||
DROP TABLE IF EXISTS ai_workspace_embeddings;
|
||||
DROP TABLE IF EXISTS ai_contexts;
|
||||
@@ -19,7 +19,7 @@
|
||||
"seed": "r ./src/seed/index.ts",
|
||||
"genconfig": "r ./scripts/genconfig.ts",
|
||||
"cli": "cross-env SERVER_FLAVOR=script node ./dist/main.js",
|
||||
"predeploy": "yarn prisma migrate deploy && yarn cli run",
|
||||
"predeploy": "yarn cli admit-legacy-context-blobs && yarn prisma migrate deploy && yarn cli run",
|
||||
"postinstall": "prisma generate"
|
||||
},
|
||||
"dependencies": {
|
||||
|
||||
@@ -16,6 +16,7 @@ const test = ava as TestFn<{
|
||||
app: TestingApp;
|
||||
db: PrismaClient;
|
||||
}>;
|
||||
let originalDeploymentType: typeof env.DEPLOYMENT_TYPE;
|
||||
|
||||
const mobileUAString =
|
||||
'Mozilla/5.0 (Linux; Android 6.0; Nexus 5 Build/MRA58N) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Mobile Safari/537.36';
|
||||
@@ -47,6 +48,7 @@ export class TestResolver {
|
||||
}
|
||||
|
||||
test.before('init selfhost server', async t => {
|
||||
originalDeploymentType = globalThis.env.DEPLOYMENT_TYPE;
|
||||
// @ts-expect-error override
|
||||
globalThis.env.DEPLOYMENT_TYPE = 'selfhosted';
|
||||
const app = await createTestingApp({
|
||||
@@ -69,7 +71,12 @@ test.beforeEach(async t => {
|
||||
});
|
||||
|
||||
test.after.always(async t => {
|
||||
await t.context.app.close();
|
||||
try {
|
||||
await t.context.app.close();
|
||||
} finally {
|
||||
// @ts-expect-error restore mutable test env singleton
|
||||
globalThis.env.DEPLOYMENT_TYPE = originalDeploymentType;
|
||||
}
|
||||
});
|
||||
|
||||
test('do not allow visit index.html directly', async t => {
|
||||
|
||||
@@ -10,6 +10,7 @@ import {
|
||||
type Config,
|
||||
type EventBus,
|
||||
type JobQueue,
|
||||
SearchProviderNotFound,
|
||||
} from '../../base';
|
||||
import { ServerFeature, type ServerService } from '../../core';
|
||||
import type { DocReader } from '../../core/doc';
|
||||
@@ -417,7 +418,6 @@ test('document tools enforce the user-selected hard scope', async t => {
|
||||
) => candidates,
|
||||
};
|
||||
const hybrid = new DocumentRetrievalService(
|
||||
{ indexer: { enabled: true } } as Config,
|
||||
readableAc,
|
||||
lexicalIndexer,
|
||||
vectorSearch,
|
||||
@@ -432,7 +432,6 @@ test('document tools enforce the user-selected hard scope', async t => {
|
||||
t.true(hybridResult.hits[0].score > 1 / 61);
|
||||
|
||||
const lexicalOnly = new DocumentRetrievalService(
|
||||
{ indexer: { enabled: true } } as Config,
|
||||
readableAc,
|
||||
lexicalIndexer,
|
||||
{ ...vectorSearch, canEmbedding: false },
|
||||
@@ -448,9 +447,12 @@ test('document tools enforce the user-selected hard scope', async t => {
|
||||
t.is(lexicalResult.degradedReason, 'VECTOR_UNAVAILABLE');
|
||||
|
||||
const vectorOnly = new DocumentRetrievalService(
|
||||
{ indexer: { enabled: false } } as Config,
|
||||
readableAc,
|
||||
lexicalIndexer,
|
||||
{
|
||||
searchDocsByKeyword: async () => {
|
||||
throw new SearchProviderNotFound();
|
||||
},
|
||||
} as unknown as IndexerService,
|
||||
vectorSearch,
|
||||
documentModels
|
||||
);
|
||||
|
||||
@@ -7,7 +7,18 @@ import {
|
||||
import { PrismaClient } from '@prisma/client';
|
||||
|
||||
import { FunctionalityModules } from '../app.module';
|
||||
import { AFFiNELogger, EventBus, JobModule, JobQueue } from '../base';
|
||||
import {
|
||||
AFFiNELogger,
|
||||
ConfigFactory,
|
||||
EventBus,
|
||||
JobModule,
|
||||
JobQueue,
|
||||
} from '../base';
|
||||
import {
|
||||
BACKEND_RUNTIME_CONFIG_PATHS,
|
||||
BackendRuntimeProvider,
|
||||
} from '../core/backend-runtime';
|
||||
import { StorageRuntimeProvider } from '../core/storage-runtime';
|
||||
import {
|
||||
createFactory,
|
||||
MockEventBus,
|
||||
@@ -15,6 +26,7 @@ import {
|
||||
MockJobQueue,
|
||||
} from './mocks';
|
||||
import { TEST_LOG_LEVEL } from './utils';
|
||||
import { createTestRuntimeConfig } from './utils/runtime-config';
|
||||
|
||||
interface TestingModuleMetadata extends ModuleMetadata {
|
||||
tapModule?(m: TestingModuleBuilder): void;
|
||||
@@ -30,6 +42,11 @@ export interface TestingModule extends NestjsTestingModule {
|
||||
export async function createModule(
|
||||
metadata: TestingModuleMetadata = {}
|
||||
): Promise<TestingModule> {
|
||||
const config = new ConfigFactory().config;
|
||||
const runtimeConfig = await createTestRuntimeConfig(
|
||||
config.db.datasourceUrl,
|
||||
config.indexer
|
||||
);
|
||||
const { tapModule, ...meta } = metadata;
|
||||
const functionalityModules = [
|
||||
...FunctionalityModules.filter(module => {
|
||||
@@ -48,14 +65,22 @@ export async function createModule(
|
||||
.overrideProvider(JobQueue)
|
||||
.useValue(new MockJobQueue())
|
||||
.overrideProvider(EventBus)
|
||||
.useValue(new MockEventBus());
|
||||
.useValue(new MockEventBus())
|
||||
.overrideProvider(BACKEND_RUNTIME_CONFIG_PATHS)
|
||||
.useValue([runtimeConfig.configPath]);
|
||||
|
||||
// when custom override happens
|
||||
if (tapModule) {
|
||||
tapModule(builder);
|
||||
}
|
||||
|
||||
const module = (await builder.compile()) as TestingModule;
|
||||
let module: TestingModule;
|
||||
try {
|
||||
module = (await builder.compile()) as TestingModule;
|
||||
} catch (error) {
|
||||
await runtimeConfig.cleanup();
|
||||
throw error;
|
||||
}
|
||||
|
||||
const logger = new AFFiNELogger();
|
||||
// we got a lot smoking tests try to break nestjs
|
||||
@@ -63,7 +88,33 @@ export async function createModule(
|
||||
logger.setLogLevels([TEST_LOG_LEVEL]);
|
||||
module.useLogger(logger);
|
||||
|
||||
await module.init();
|
||||
const close = module.close.bind(module);
|
||||
let closePromise: Promise<void> | undefined;
|
||||
module.close = () => {
|
||||
return (closePromise ??= (async () => {
|
||||
try {
|
||||
await close();
|
||||
} finally {
|
||||
await runtimeConfig.cleanup();
|
||||
}
|
||||
})());
|
||||
};
|
||||
|
||||
try {
|
||||
await module.init();
|
||||
} catch (error) {
|
||||
await module.close();
|
||||
throw error;
|
||||
}
|
||||
const backendRuntime = module.get(BackendRuntimeProvider);
|
||||
if (backendRuntime instanceof BackendRuntimeProvider) {
|
||||
await backendRuntime.runMigrations();
|
||||
await backendRuntime.onConfigChanged({ updates: { indexer: {} } });
|
||||
}
|
||||
const storageRuntime = module.get(StorageRuntimeProvider);
|
||||
if (storageRuntime instanceof StorageRuntimeProvider) {
|
||||
await storageRuntime.runMigrations();
|
||||
}
|
||||
module[Symbol.asyncDispose] = async () => {
|
||||
await module.close();
|
||||
};
|
||||
|
||||
@@ -4,7 +4,7 @@ import ava, { TestFn } from 'ava';
|
||||
import Sinon from 'sinon';
|
||||
|
||||
import { BackendRuntimeProvider } from '../../core/backend-runtime';
|
||||
import { DocStorageModule } from '../../core/doc';
|
||||
import { DocStorageModule, DocStorageWorkerModule } from '../../core/doc';
|
||||
import { DocStorageCronJob } from '../../core/doc/job';
|
||||
import { createTestingModule, type TestingModule } from '../utils';
|
||||
|
||||
@@ -23,7 +23,11 @@ test.before(async t => {
|
||||
cleanupExpiredSnapshotHistories: Sinon.stub(),
|
||||
};
|
||||
t.context.module = await createTestingModule({
|
||||
imports: [ScheduleModule.forRoot(), DocStorageModule],
|
||||
imports: [
|
||||
ScheduleModule.forRoot(),
|
||||
DocStorageModule,
|
||||
DocStorageWorkerModule,
|
||||
],
|
||||
tapModule: builder => {
|
||||
builder
|
||||
.overrideProvider(BackendRuntimeProvider)
|
||||
|
||||
@@ -1,66 +1,93 @@
|
||||
import { getCurrentUserQuery } from '@affine/graphql';
|
||||
|
||||
import { JobExecutor } from '../../../base/job/queue/executor';
|
||||
import { JobHandlerScanner } from '../../../base/job/queue/scanner';
|
||||
import { DatabaseDocReader, DocReader } from '../../../core/doc';
|
||||
import { createApp } from '../create-app';
|
||||
import { e2e } from '../test';
|
||||
|
||||
type TestFlavor = 'doc' | 'graphql' | 'sync' | 'renderer' | 'front';
|
||||
type TestFlavor =
|
||||
| 'allinone'
|
||||
| 'worker'
|
||||
| 'graphql'
|
||||
| 'sync'
|
||||
| 'renderer'
|
||||
| 'front';
|
||||
|
||||
const createFlavorApp = async (flavor: TestFlavor) => {
|
||||
const withFlavor = async <T>(
|
||||
flavor: TestFlavor,
|
||||
run: (app: Awaited<ReturnType<typeof createApp>>) => Promise<T>
|
||||
) => {
|
||||
const mutableEnv = globalThis.env as unknown as { FLAVOR: string };
|
||||
const previousFlavor = mutableEnv.FLAVOR;
|
||||
// @ts-expect-error override
|
||||
globalThis.env.FLAVOR = flavor;
|
||||
return await createApp({
|
||||
tapModule(module) {
|
||||
module.overrideProvider(JobExecutor).useValue({
|
||||
onConfigInit: async () => {},
|
||||
onConfigChanged: async () => {},
|
||||
onModuleDestroy: async () => {},
|
||||
});
|
||||
},
|
||||
});
|
||||
try {
|
||||
await using app = await createApp({
|
||||
tapModule(module) {
|
||||
module.overrideProvider(JobExecutor).useValue({
|
||||
onConfigInit: async () => {},
|
||||
onConfigChanged: async () => {},
|
||||
onModuleDestroy: async () => {},
|
||||
});
|
||||
},
|
||||
});
|
||||
return await run(app);
|
||||
} finally {
|
||||
mutableEnv.FLAVOR = previousFlavor;
|
||||
}
|
||||
};
|
||||
|
||||
e2e('should init doc service', async t => {
|
||||
await using app = await createFlavorApp('doc');
|
||||
e2e('should init worker service', async t => {
|
||||
await withFlavor('worker', async app => {
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'worker');
|
||||
t.truthy(app.get(JobHandlerScanner).getHandler('indexer.indexDoc'));
|
||||
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'doc');
|
||||
await t.throwsAsync(app.gql({ query: getCurrentUserQuery }));
|
||||
await app.PUT('/api/storage/upload').expect(404);
|
||||
});
|
||||
});
|
||||
|
||||
await t.throwsAsync(app.gql({ query: getCurrentUserQuery }));
|
||||
e2e('should init allinone service with worker handlers', async t => {
|
||||
await withFlavor('allinone', async app => {
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'allinone');
|
||||
t.truthy(app.get(JobHandlerScanner).getHandler('indexer.indexDoc'));
|
||||
});
|
||||
});
|
||||
|
||||
e2e('should init graphql service', async t => {
|
||||
await using app = await createFlavorApp('graphql');
|
||||
await withFlavor('graphql', async app => {
|
||||
const res = await app.GET('/info').expect(200);
|
||||
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'graphql');
|
||||
|
||||
t.is(res.body.flavor, 'graphql');
|
||||
|
||||
const user = await app.gql({ query: getCurrentUserQuery });
|
||||
t.is(user.currentUser, null);
|
||||
const user = await app.gql({ query: getCurrentUserQuery });
|
||||
t.is(user.currentUser, null);
|
||||
});
|
||||
});
|
||||
|
||||
e2e('should init sync service', async t => {
|
||||
await using app = await createFlavorApp('sync');
|
||||
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'sync');
|
||||
await withFlavor('sync', async app => {
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'sync');
|
||||
});
|
||||
});
|
||||
|
||||
e2e('should init renderer service', async t => {
|
||||
await using app = await createFlavorApp('renderer');
|
||||
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'renderer');
|
||||
await withFlavor('renderer', async app => {
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'renderer');
|
||||
});
|
||||
});
|
||||
|
||||
e2e('should init front service', async t => {
|
||||
await using app = await createFlavorApp('front');
|
||||
await withFlavor('front', async app => {
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'front');
|
||||
|
||||
const res = await app.GET('/info').expect(200);
|
||||
t.is(res.body.flavor, 'front');
|
||||
|
||||
const docReader = app.get(DocReader);
|
||||
t.true(docReader instanceof DatabaseDocReader);
|
||||
const docReader = app.get(DocReader);
|
||||
t.true(docReader instanceof DatabaseDocReader);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import assert from 'node:assert';
|
||||
|
||||
import { gqlFetcherFactory } from '@affine/graphql';
|
||||
import { INestApplication, ModuleMetadata } from '@nestjs/common';
|
||||
import { INestApplication, ModuleMetadata, Type } from '@nestjs/common';
|
||||
import { NestApplication } from '@nestjs/core';
|
||||
import {
|
||||
Test,
|
||||
@@ -26,9 +26,15 @@ import {
|
||||
import { ThrottlerStorage } from '../../base/throttler';
|
||||
import { SocketIoAdapter } from '../../base/websocket';
|
||||
import { AuthGuard, AuthService } from '../../core/auth';
|
||||
import { BACKEND_RUNTIME_CONFIG_PATHS } from '../../core/backend-runtime';
|
||||
import {
|
||||
BACKEND_RUNTIME_CONFIG_PATHS,
|
||||
BackendRuntimeProvider,
|
||||
} from '../../core/backend-runtime';
|
||||
import { Mailer } from '../../core/mail';
|
||||
import { StorageRuntimeProvider } from '../../core/storage-runtime';
|
||||
import { ServerRole } from '../../env';
|
||||
import { Models } from '../../models';
|
||||
import { IndexerService } from '../../plugins/indexer/service';
|
||||
import {
|
||||
createFactory,
|
||||
MockedUser,
|
||||
@@ -52,8 +58,16 @@ export class TestingApp extends NestApplication {
|
||||
private csrfCookie: string | null = null;
|
||||
private readonly userCookies: Set<string> = new Set();
|
||||
|
||||
private getOptional<T>(token: Type<T>) {
|
||||
try {
|
||||
return this.get(token, { strict: false });
|
||||
} catch {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
create = createFactory(this.get(PrismaClient, { strict: false }));
|
||||
mails = this.get(Mailer, { strict: false }) as MockMailer;
|
||||
mails = this.getOptional(Mailer) as unknown as MockMailer;
|
||||
queue = this.get(JobQueue, { strict: false }) as MockJobQueue;
|
||||
eventBus = this.get(EventBus, { strict: false });
|
||||
models = this.get(Models, { strict: false });
|
||||
@@ -241,8 +255,10 @@ export class TestingApp extends NestApplication {
|
||||
export async function createApp(
|
||||
metadata: TestingAppMetadata = {}
|
||||
): Promise<TestingApp> {
|
||||
const config = new ConfigFactory().config;
|
||||
const runtimeConfig = await createTestRuntimeConfig(
|
||||
new ConfigFactory().config.db.datasourceUrl
|
||||
config.db.datasourceUrl,
|
||||
config.indexer
|
||||
);
|
||||
const { buildAppModule } = await import('../../app.module');
|
||||
const { tapModule, tapApp } = metadata;
|
||||
@@ -326,7 +342,11 @@ export async function createApp(
|
||||
})
|
||||
);
|
||||
|
||||
app.useGlobalGuards(app.get(AuthGuard), app.get(CloudThrottlerGuard));
|
||||
if (globalThis.env.role === ServerRole.Worker) {
|
||||
app.useGlobalGuards(app.get(CloudThrottlerGuard));
|
||||
} else {
|
||||
app.useGlobalGuards(app.get(AuthGuard), app.get(CloudThrottlerGuard));
|
||||
}
|
||||
app.useGlobalInterceptors(app.get(CacheInterceptor));
|
||||
app.useGlobalFilters(new GlobalExceptionFilter(app.getHttpAdapter()));
|
||||
|
||||
@@ -340,6 +360,9 @@ export async function createApp(
|
||||
|
||||
try {
|
||||
await app.init();
|
||||
await app.get(BackendRuntimeProvider, { strict: false }).runMigrations();
|
||||
await app.get(StorageRuntimeProvider, { strict: false }).runMigrations();
|
||||
await app.get(IndexerService, { strict: false }).onApplicationBootstrap();
|
||||
} catch (error) {
|
||||
await app.close();
|
||||
throw error;
|
||||
|
||||
-98
@@ -1,98 +0,0 @@
|
||||
# Snapshot report for `src/__tests__/e2e/doc-service/controller.spec.ts`
|
||||
|
||||
The actual snapshot is saved in `controller.spec.ts.snap`.
|
||||
|
||||
Generated by [AVA](https://avajs.dev).
|
||||
|
||||
## should get doc markdown success
|
||||
|
||||
> Snapshot 1
|
||||
|
||||
{
|
||||
knownUnsupportedBlocks: [
|
||||
'RX4CG2zsBk:affine:note',
|
||||
'S1mkc8zUoU:affine:note',
|
||||
'yGlBdshAqN:affine:note',
|
||||
'6lDiuDqZGL:affine:note',
|
||||
'cauvaHOQmh:affine:note',
|
||||
'2jwCeO8Yot:affine:note',
|
||||
'c9MF_JiRgx:affine:note',
|
||||
'6x7ALjUDjj:affine:surface',
|
||||
],
|
||||
markdown: `AFFiNE is an open source all in one workspace, an operating system for all the building blocks of your team wiki, knowledge management and digital assets and a better alternative to Notion and Miro.␊
|
||||
␊
|
||||
␊
|
||||
␊
|
||||
# You own your data, with no compromises␊
|
||||
␊
|
||||
## Local-first & Real-time collaborative␊
|
||||
␊
|
||||
We love the idea proposed by Ink & Switch in the famous article about you owning your data, despite the cloud. Furthermore, AFFiNE is the first all-in-one workspace that keeps your data ownership with no compromises on real-time collaboration and editing experience.␊
|
||||
␊
|
||||
AFFiNE is a local-first application upon CRDTs with real-time collaboration support. Your data is always stored locally while multiple nodes remain synced in real-time.␊
|
||||
␊
|
||||
␊
|
||||
␊
|
||||
### Blocks that assemble your next docs, tasks kanban or whiteboard␊
|
||||
␊
|
||||
There is a large overlap of their atomic "building blocks" between these apps. They are neither open source nor have a plugin system like VS Code for contributors to customize. We want to have something that contains all the features we love and goes one step further.␊
|
||||
␊
|
||||
We are building AFFiNE to be a fundamental open source platform that contains all the building blocks for docs, task management and visual collaboration, hoping you can shape your next workflow with us that can make your life better and also connect others, too.␊
|
||||
␊
|
||||
If you want to learn more about the product design of AFFiNE, here goes the concepts:␊
|
||||
␊
|
||||
To Shape, not to adapt. AFFiNE is built for individuals & teams who care about their data, who refuse vendor lock-in, and who want to have control over their essential tools.␊
|
||||
␊
|
||||
## A true canvas for blocks in any form␊
|
||||
␊
|
||||
[Many editor apps](http://notion.so) claimed to be a canvas for productivity. Since _the Mother of All Demos,_ Douglas Engelbart, a creative and programable digital workspace has been a pursuit and an ultimate mission for generations of tool makers.␊
|
||||
␊
|
||||
␊
|
||||
␊
|
||||
"We shape our tools and thereafter our tools shape us”. A lot of pioneers have inspired us a long the way, e.g.:␊
|
||||
␊
|
||||
* Quip & Notion with their great concept of "everything is a block"␊
|
||||
* Trello with their Kanban␊
|
||||
* Airtable & Miro with their no-code programable datasheets␊
|
||||
* Miro & Whimiscal with their edgeless visual whiteboard␊
|
||||
* Remnote & Capacities with their object-based tag system␊
|
||||
For more details, please refer to our [RoadMap](https://docs.affine.pro/docs/core-concepts/roadmap)␊
|
||||
␊
|
||||
## Self Host␊
|
||||
␊
|
||||
Self host AFFiNE␊
|
||||
␊
|
||||
␊
|
||||
### Learning From␊
|
||||
||Title|Tag|␊
|
||||
|---|---|---|␊
|
||||
|Affine Development|Affine Development|<span data-affine-option data-value="AxSe-53xjX" data-option-color="var(--affine-tag-pink)">AFFiNE</span>|␊
|
||||
|For developers or installations guides, please go to AFFiNE Doc|For developers or installations guides, please go to AFFiNE Doc|<span data-affine-option data-value="0jh9gNw4Yl" data-option-color="var(--affine-tag-orange)">Developers</span>|␊
|
||||
|Quip & Notion with their great concept of "everything is a block"|Quip & Notion with their great concept of "everything is a block"|<span data-affine-option data-value="HgHsKOUINZ" data-option-color="var(--affine-tag-blue)">Reference</span>|␊
|
||||
|Trello with their Kanban|Trello with their Kanban|<span data-affine-option data-value="HgHsKOUINZ" data-option-color="var(--affine-tag-blue)">Reference</span>|␊
|
||||
|Airtable & Miro with their no-code programable datasheets|Airtable & Miro with their no-code programable datasheets|<span data-affine-option data-value="HgHsKOUINZ" data-option-color="var(--affine-tag-blue)">Reference</span>|␊
|
||||
|Miro & Whimiscal with their edgeless visual whiteboard|Miro & Whimiscal with their edgeless visual whiteboard|<span data-affine-option data-value="HgHsKOUINZ" data-option-color="var(--affine-tag-blue)">Reference</span>|␊
|
||||
|Remnote & Capacities with their object-based tag system|Remnote & Capacities with their object-based tag system||␊
|
||||
␊
|
||||
## Affine Development␊
|
||||
␊
|
||||
For developer or installation guides, please go to [AFFiNE Development](https://docs.affine.pro/docs/development/quick-start)␊
|
||||
␊
|
||||
␊
|
||||
␊
|
||||
`,
|
||||
title: 'Write, Draw, Plan all at Once.',
|
||||
unknownBlocks: [],
|
||||
}
|
||||
|
||||
## should get doc markdown return null when doc not exists
|
||||
|
||||
> Snapshot 1
|
||||
|
||||
{
|
||||
code: 'Not Found',
|
||||
message: 'Doc not found',
|
||||
name: 'NOT_FOUND',
|
||||
status: 404,
|
||||
type: 'RESOURCE_NOT_FOUND',
|
||||
}
|
||||
BIN
Binary file not shown.
@@ -1,52 +0,0 @@
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
import { CryptoHelper } from '../../../base';
|
||||
import { app, e2e, Mockers } from '../test';
|
||||
|
||||
const crypto = app.get(CryptoHelper);
|
||||
|
||||
e2e('should get doc markdown success', async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const docSnapshot = await app.create(Mockers.DocSnapshot, {
|
||||
workspaceId: workspace.id,
|
||||
user: owner,
|
||||
});
|
||||
|
||||
const path = `/rpc/workspaces/${workspace.id}/docs/${docSnapshot.id}/markdown`;
|
||||
const res = await app
|
||||
.GET(path)
|
||||
.set(
|
||||
'x-access-token',
|
||||
crypto.signInternalAccessToken({ method: 'GET', path })
|
||||
)
|
||||
.expect(200)
|
||||
.expect('Content-Type', 'application/json; charset=utf-8');
|
||||
|
||||
const { revision, ...body } = res.body;
|
||||
t.regex(revision, /^\d+$/);
|
||||
t.snapshot(body);
|
||||
});
|
||||
|
||||
e2e('should get doc markdown return null when doc not exists', async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const docId = randomUUID();
|
||||
const path = `/rpc/workspaces/${workspace.id}/docs/${docId}/markdown`;
|
||||
const res = await app
|
||||
.GET(path)
|
||||
.set(
|
||||
'x-access-token',
|
||||
crypto.signInternalAccessToken({ method: 'GET', path })
|
||||
)
|
||||
.expect(404)
|
||||
.expect('Content-Type', 'application/json; charset=utf-8');
|
||||
|
||||
t.snapshot(res.body);
|
||||
});
|
||||
@@ -1,85 +1,30 @@
|
||||
import { indexerAggregateQuery, SearchTable } from '@affine/graphql';
|
||||
import {
|
||||
indexerAggregateQuery,
|
||||
SearchQueryType,
|
||||
SearchTable,
|
||||
} from '@affine/graphql';
|
||||
|
||||
import { createDocWithMarkdown } from '../../../native';
|
||||
import { IndexerService } from '../../../plugins/indexer/service';
|
||||
import { Mockers } from '../../mocks';
|
||||
import { app, e2e } from '../test';
|
||||
|
||||
e2e('should aggregate by docId', async t => {
|
||||
const owner = await app.signup();
|
||||
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner: { id: owner.id },
|
||||
});
|
||||
|
||||
const indexerService = app.get(IndexerService);
|
||||
|
||||
await indexerService.write(
|
||||
SearchTable.block,
|
||||
[
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test1 hello world top2',
|
||||
flavour: 'affine:text',
|
||||
blockId: 'block-0',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
},
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test2 hello hello top3',
|
||||
flavour: 'affine:text',
|
||||
blockId: 'block-1',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
},
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test3 hello title top1',
|
||||
flavour: 'affine:page',
|
||||
blockId: 'block-2',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
},
|
||||
{
|
||||
docId: 'doc-1',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test4 hello world',
|
||||
flavour: 'affine:text',
|
||||
blockId: 'block-3',
|
||||
refDocId: 'doc-0',
|
||||
ref: ['{"foo": "bar1"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
},
|
||||
{
|
||||
docId: 'doc-2',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test5 hello',
|
||||
flavour: 'affine:text',
|
||||
blockId: 'block-4',
|
||||
refDocId: 'doc-0',
|
||||
ref: ['{"foo": "bar2"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
},
|
||||
],
|
||||
{
|
||||
refresh: true,
|
||||
}
|
||||
);
|
||||
const workspace = await app.create(Mockers.Workspace, { owner });
|
||||
for (const [docId, markdown] of [
|
||||
['doc-0', 'hello world\n\nhello again'],
|
||||
['doc-1', 'hello world'],
|
||||
] as const) {
|
||||
await app.create(Mockers.DocMeta, { workspaceId: workspace.id, docId });
|
||||
await app.create(Mockers.DocSnapshot, {
|
||||
workspaceId: workspace.id,
|
||||
docId,
|
||||
user: owner,
|
||||
blob: createDocWithMarkdown(docId, markdown, docId),
|
||||
});
|
||||
await app.get(IndexerService).indexDoc(workspace.id, docId);
|
||||
}
|
||||
|
||||
const result = await app.gql({
|
||||
query: indexerAggregateQuery,
|
||||
@@ -88,72 +33,25 @@ e2e('should aggregate by docId', async t => {
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
// @ts-expect-error allow to use string as enum
|
||||
type: 'boolean',
|
||||
// @ts-expect-error allow to use string as enum
|
||||
occur: 'must',
|
||||
queries: [
|
||||
{
|
||||
// @ts-expect-error allow to use string as enum
|
||||
type: 'match',
|
||||
field: 'content',
|
||||
match: 'hello world',
|
||||
},
|
||||
{
|
||||
// @ts-expect-error allow to use string as enum
|
||||
type: 'boolean',
|
||||
// @ts-expect-error allow to use string as enum
|
||||
occur: 'should',
|
||||
queries: [
|
||||
{
|
||||
// @ts-expect-error allow to use string as enum
|
||||
type: 'match',
|
||||
field: 'content',
|
||||
match: 'hello world',
|
||||
},
|
||||
{
|
||||
// @ts-expect-error allow to use string as enum
|
||||
type: 'boost',
|
||||
boost: 1.5,
|
||||
query: {
|
||||
// @ts-expect-error allow to use string as enum
|
||||
type: 'match',
|
||||
field: 'flavour',
|
||||
match: 'affine:page',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
type: SearchQueryType.match,
|
||||
field: 'content',
|
||||
match: 'hello',
|
||||
},
|
||||
field: 'docId',
|
||||
options: {
|
||||
pagination: {
|
||||
limit: 50,
|
||||
skip: 0,
|
||||
},
|
||||
pagination: { limit: 50, skip: 0 },
|
||||
hits: {
|
||||
pagination: {
|
||||
limit: 2,
|
||||
skip: 0,
|
||||
},
|
||||
fields: ['blockId', 'flavour'],
|
||||
highlights: [
|
||||
{
|
||||
field: 'content',
|
||||
before: '<b>',
|
||||
end: '</b>',
|
||||
},
|
||||
],
|
||||
pagination: { limit: 2, skip: 0 },
|
||||
fields: ['docId', 'blockId', 'content'],
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.truthy(result.workspace.aggregate, 'failed to aggregate');
|
||||
t.is(result.workspace.aggregate.pagination.count, 5);
|
||||
t.is(result.workspace.aggregate.pagination.hasMore, true);
|
||||
t.truthy(result.workspace.aggregate.pagination.nextCursor);
|
||||
t.snapshot(result.workspace.aggregate.buckets);
|
||||
t.is(result.workspace.aggregate.pagination.count, 2);
|
||||
t.deepEqual(
|
||||
result.workspace.aggregate.buckets.map(bucket => bucket.key).sort(),
|
||||
['doc-0', 'doc-1']
|
||||
);
|
||||
});
|
||||
|
||||
@@ -1,182 +1,57 @@
|
||||
import { indexerSearchDocsQuery, SearchTable } from '@affine/graphql';
|
||||
import { omit } from 'lodash-es';
|
||||
import { indexerSearchDocsQuery } from '@affine/graphql';
|
||||
|
||||
import { ConfigFactory } from '../../../base';
|
||||
import { createDocWithMarkdown } from '../../../native';
|
||||
import { SearchProviderType } from '../../../plugins/indexer/config';
|
||||
import { IndexerService } from '../../../plugins/indexer/service';
|
||||
import { Mockers } from '../../mocks';
|
||||
import { app, e2e } from '../test';
|
||||
|
||||
e2e('should search docs by keyword', async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, { owner });
|
||||
for (const docId of ['doc-0', 'doc-1', 'doc-2']) {
|
||||
await app.create(Mockers.DocMeta, { workspaceId: workspace.id, docId });
|
||||
await app.create(Mockers.DocSnapshot, {
|
||||
workspaceId: workspace.id,
|
||||
docId,
|
||||
user: owner,
|
||||
blob: createDocWithMarkdown(docId, `${docId} hello`, docId),
|
||||
});
|
||||
await app.get(IndexerService).indexDoc(workspace.id, docId);
|
||||
}
|
||||
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const indexerService = app.get(IndexerService);
|
||||
|
||||
await indexerService.write(
|
||||
SearchTable.block,
|
||||
[
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test1 hello',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-0',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-1',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test2 hello',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-1',
|
||||
refDocId: ['doc-0'],
|
||||
ref: ['{"foo": "bar1"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-2',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test3 hello',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-2',
|
||||
refDocId: ['doc-0', 'doc-2'],
|
||||
ref: ['{"foo": "bar1"}', '{"foo": "bar3"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-03-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-03-22T03:00:01.000Z'),
|
||||
},
|
||||
],
|
||||
{
|
||||
refresh: true,
|
||||
}
|
||||
);
|
||||
|
||||
const result = await app.gql({
|
||||
const search = app.gql({
|
||||
query: indexerSearchDocsQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
keyword: 'hello',
|
||||
},
|
||||
},
|
||||
variables: { id: workspace.id, input: { keyword: 'hello', limit: 2 } },
|
||||
});
|
||||
if (
|
||||
app.get(ConfigFactory).config.indexer.provider.type ===
|
||||
SearchProviderType.Manticoresearch
|
||||
) {
|
||||
await t.throwsAsync(search, {
|
||||
message: /Invalid indexer input: unsupported_query/,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
t.is(result.workspace.searchDocs.length, 3);
|
||||
t.snapshot(
|
||||
result.workspace.searchDocs.map(doc =>
|
||||
omit(doc, 'createdByUser', 'updatedByUser')
|
||||
)
|
||||
);
|
||||
});
|
||||
|
||||
e2e('should search docs by keyword with limit 1', async t => {
|
||||
const owner = await app.signup();
|
||||
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const indexerService = app.get(IndexerService);
|
||||
|
||||
await indexerService.write(
|
||||
SearchTable.block,
|
||||
[
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test1 hello',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-0',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-1',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test2 hello',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-1',
|
||||
refDocId: ['doc-0'],
|
||||
ref: ['{"foo": "bar1"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-2',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test3 hello',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-2',
|
||||
refDocId: ['doc-0', 'doc-2'],
|
||||
ref: ['{"foo": "bar1"}', '{"foo": "bar3"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-03-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-03-22T03:00:01.000Z'),
|
||||
},
|
||||
],
|
||||
{
|
||||
refresh: true,
|
||||
}
|
||||
);
|
||||
|
||||
const result = await app.gql({
|
||||
query: indexerSearchDocsQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
keyword: 'hello',
|
||||
limit: 1,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.is(result.workspace.searchDocs.length, 1);
|
||||
t.snapshot(
|
||||
result.workspace.searchDocs.map(doc =>
|
||||
omit(doc, 'createdByUser', 'updatedByUser')
|
||||
)
|
||||
);
|
||||
const result = await search;
|
||||
t.is(result.workspace.searchDocs.length, 2);
|
||||
t.true(result.workspace.searchDocs.every(doc => doc.highlight.length > 0));
|
||||
});
|
||||
|
||||
e2e(
|
||||
'should search docs by keyword failed when workspace is no permission',
|
||||
async t => {
|
||||
const owner = await app.signup();
|
||||
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
// signup another user
|
||||
const workspace = await app.create(Mockers.Workspace, { owner });
|
||||
await app.signup();
|
||||
|
||||
await t.throwsAsync(
|
||||
app.gql({
|
||||
query: indexerSearchDocsQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
keyword: 'hello',
|
||||
},
|
||||
},
|
||||
variables: { id: workspace.id, input: { keyword: 'hello' } },
|
||||
}),
|
||||
{
|
||||
message: /You do not have permission to access Space/,
|
||||
}
|
||||
{ message: /You do not have permission to access Space/ }
|
||||
);
|
||||
}
|
||||
);
|
||||
|
||||
@@ -1,68 +1,40 @@
|
||||
import {
|
||||
indexerSearchQuery,
|
||||
SearchQueryOccur,
|
||||
SearchQueryType,
|
||||
SearchTable,
|
||||
} from '@affine/graphql';
|
||||
|
||||
import { DocRole } from '../../../models';
|
||||
import { createDocWithMarkdown } from '../../../native';
|
||||
import { IndexerService } from '../../../plugins/indexer/service';
|
||||
import { Mockers } from '../../mocks';
|
||||
import { app, e2e } from '../test';
|
||||
|
||||
async function indexDoc(
|
||||
workspaceId: string,
|
||||
user: { id: string },
|
||||
docId: string,
|
||||
markdown: string,
|
||||
defaultRole = DocRole.Manager
|
||||
) {
|
||||
await app.create(Mockers.DocMeta, { workspaceId, docId, defaultRole });
|
||||
await app.create(Mockers.DocSnapshot, {
|
||||
workspaceId,
|
||||
docId,
|
||||
user,
|
||||
blob: createDocWithMarkdown(docId, markdown, docId),
|
||||
});
|
||||
await app.get(IndexerService).indexDoc(workspaceId, docId);
|
||||
}
|
||||
|
||||
e2e('should search with query', async t => {
|
||||
const owner = await app.signup();
|
||||
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner: { id: owner.id },
|
||||
});
|
||||
|
||||
const indexerService = app.get(IndexerService);
|
||||
|
||||
await indexerService.write(
|
||||
SearchTable.block,
|
||||
[
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test1',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-0',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-1',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test2',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-1',
|
||||
refDocId: ['doc-0'],
|
||||
ref: ['{"foo": "bar1"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-2',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test3',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-2',
|
||||
refDocId: ['doc-0', 'doc-2'],
|
||||
ref: ['{"foo": "bar1"}', '{"foo": "bar3"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-03-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-03-22T00:00:00.000Z'),
|
||||
},
|
||||
],
|
||||
{
|
||||
refresh: true,
|
||||
}
|
||||
const workspace = await app.create(Mockers.Workspace, { owner });
|
||||
await indexDoc(
|
||||
workspace.id,
|
||||
owner,
|
||||
'doc-0',
|
||||
'searchable first\n\nsearchable second'
|
||||
);
|
||||
|
||||
const result = await app.gql({
|
||||
@@ -72,158 +44,95 @@ e2e('should search with query', async t => {
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.boolean,
|
||||
occur: SearchQueryOccur.must,
|
||||
queries: [
|
||||
{
|
||||
type: SearchQueryType.boolean,
|
||||
occur: SearchQueryOccur.should,
|
||||
queries: ['doc-0', 'doc-1', 'doc-2'].map(id => ({
|
||||
type: SearchQueryType.match,
|
||||
field: 'docId',
|
||||
match: id,
|
||||
})),
|
||||
},
|
||||
{
|
||||
type: SearchQueryType.exists,
|
||||
field: 'refDocId',
|
||||
},
|
||||
],
|
||||
type: SearchQueryType.match,
|
||||
field: 'content',
|
||||
match: 'searchable',
|
||||
},
|
||||
options: {
|
||||
fields: ['refDocId', 'ref'],
|
||||
pagination: {
|
||||
limit: 100,
|
||||
},
|
||||
fields: ['docId', 'blockId', 'content'],
|
||||
highlights: [{ field: 'content', before: '<b>', end: '</b>' }],
|
||||
pagination: { limit: 100 },
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.truthy(result.workspace.search, 'failed to search');
|
||||
t.is(result.workspace.search.pagination.count, 2);
|
||||
t.is(result.workspace.search.pagination.hasMore, true);
|
||||
t.truthy(result.workspace.search.pagination.nextCursor);
|
||||
t.is(result.workspace.search.nodes.length, 2);
|
||||
t.snapshot(result.workspace.search.nodes);
|
||||
t.true(result.workspace.search.pagination.count > 0);
|
||||
t.true(
|
||||
result.workspace.search.nodes.every(node =>
|
||||
node.fields.docId.includes('doc-0')
|
||||
)
|
||||
);
|
||||
t.true(
|
||||
result.workspace.search.nodes.some(node =>
|
||||
node.highlights?.content?.some((value: string) => value.includes('<b>'))
|
||||
)
|
||||
);
|
||||
|
||||
const firstPage = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'content',
|
||||
match: 'searchable',
|
||||
},
|
||||
options: {
|
||||
fields: ['docId', 'blockId'],
|
||||
pagination: { limit: 1 },
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
const secondPage = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'content',
|
||||
match: 'searchable',
|
||||
},
|
||||
options: {
|
||||
fields: ['docId', 'blockId'],
|
||||
pagination: {
|
||||
limit: 1,
|
||||
cursor: firstPage.workspace.search.pagination.nextCursor,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
t.not(
|
||||
firstPage.workspace.search.nodes[0].fields.blockId[0],
|
||||
secondPage.workspace.search.nodes[0].fields.blockId[0]
|
||||
);
|
||||
});
|
||||
|
||||
e2e('should filter no read permission docs on team workspace', async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
const workspace = await app.create(Mockers.Workspace, { owner });
|
||||
await app.create(Mockers.TeamWorkspace, { id: workspace.id });
|
||||
await indexDoc(
|
||||
workspace.id,
|
||||
owner,
|
||||
});
|
||||
await app.create(Mockers.TeamWorkspace, {
|
||||
id: workspace.id,
|
||||
});
|
||||
|
||||
const indexerService = app.get(IndexerService);
|
||||
await indexerService.write(
|
||||
SearchTable.block,
|
||||
[
|
||||
{
|
||||
docId: 'doc-0',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test1',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-0',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-1',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test2',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-1',
|
||||
refDocId: ['doc-0'],
|
||||
ref: ['{"foo": "bar1"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2021-04-22T00:00:00.000Z'),
|
||||
},
|
||||
{
|
||||
docId: 'doc-2',
|
||||
workspaceId: workspace.id,
|
||||
content: 'test3',
|
||||
flavour: 'markdown',
|
||||
blockId: 'block-2',
|
||||
refDocId: ['doc-0', 'doc-2'],
|
||||
ref: ['{"foo": "bar1"}', '{"foo": "bar3"}'],
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-03-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-03-22T00:00:00.000Z'),
|
||||
},
|
||||
],
|
||||
{
|
||||
refresh: true,
|
||||
}
|
||||
'private-doc',
|
||||
'team secret searchable',
|
||||
DocRole.None
|
||||
);
|
||||
// set all docs to no access
|
||||
await app.create(Mockers.DocMeta, {
|
||||
workspaceId: workspace.id,
|
||||
docId: 'doc-0',
|
||||
defaultRole: DocRole.None,
|
||||
});
|
||||
await app.create(Mockers.DocMeta, {
|
||||
workspaceId: workspace.id,
|
||||
docId: 'doc-1',
|
||||
defaultRole: DocRole.None,
|
||||
});
|
||||
await app.create(Mockers.DocMeta, {
|
||||
workspaceId: workspace.id,
|
||||
docId: 'doc-2',
|
||||
defaultRole: DocRole.None,
|
||||
});
|
||||
|
||||
// owner can read all docs
|
||||
const result = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'workspaceId',
|
||||
match: workspace.id,
|
||||
},
|
||||
options: {
|
||||
fields: ['docId', 'blockId', 'refDocId', 'ref'],
|
||||
pagination: {
|
||||
limit: 100,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.snapshot(result.workspace.search.nodes);
|
||||
|
||||
// other user can only read docs that they have read permission
|
||||
const other = await app.signup();
|
||||
const member = await app.signup();
|
||||
await app.create(Mockers.WorkspaceUser, {
|
||||
workspaceId: workspace.id,
|
||||
userId: other.id,
|
||||
userId: member.id,
|
||||
});
|
||||
await app.create(Mockers.DocUser, {
|
||||
workspaceId: workspace.id,
|
||||
docId: 'doc-0',
|
||||
userId: other.id,
|
||||
type: DocRole.Reader,
|
||||
});
|
||||
await app.create(Mockers.DocUser, {
|
||||
workspaceId: workspace.id,
|
||||
docId: 'doc-1',
|
||||
userId: other.id,
|
||||
type: DocRole.Manager,
|
||||
});
|
||||
|
||||
const otherResult = await app.gql({
|
||||
await app.get(IndexerService).reconcileWorkspace(workspace.id);
|
||||
const denied = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
@@ -231,132 +140,74 @@ e2e('should filter no read permission docs on team workspace', async t => {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'workspaceId',
|
||||
match: workspace.id,
|
||||
},
|
||||
options: {
|
||||
fields: ['docId', 'blockId', 'refDocId', 'ref'],
|
||||
pagination: {
|
||||
limit: 100,
|
||||
},
|
||||
field: 'content',
|
||||
match: 'secret',
|
||||
},
|
||||
options: { fields: ['docId'], pagination: { limit: 10 } },
|
||||
},
|
||||
},
|
||||
});
|
||||
t.is(denied.workspace.search.pagination.count, 0);
|
||||
|
||||
t.snapshot(otherResult.workspace.search.nodes);
|
||||
await app.create(Mockers.DocUser, {
|
||||
workspaceId: workspace.id,
|
||||
docId: 'private-doc',
|
||||
userId: member.id,
|
||||
type: DocRole.Reader,
|
||||
});
|
||||
await app.get(IndexerService).reconcileWorkspace(workspace.id);
|
||||
const allowed = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'content',
|
||||
match: 'secret',
|
||||
},
|
||||
options: { fields: ['docId'], pagination: { limit: 10 } },
|
||||
},
|
||||
},
|
||||
});
|
||||
t.true(allowed.workspace.search.pagination.count > 0);
|
||||
|
||||
await app.models.docUser.delete(workspace.id, 'private-doc', member.id);
|
||||
await app.get(IndexerService).reconcileWorkspace(workspace.id);
|
||||
const revoked = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'content',
|
||||
match: 'secret',
|
||||
},
|
||||
options: { fields: ['docId'], pagination: { limit: 10 } },
|
||||
},
|
||||
},
|
||||
});
|
||||
t.is(revoked.workspace.search.pagination.count, 0);
|
||||
});
|
||||
|
||||
e2e('should return empty results when search not match any docs', async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const result = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.block,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'workspaceId',
|
||||
match: workspace.id,
|
||||
},
|
||||
options: {
|
||||
fields: ['docId', 'blockId', 'refDocId', 'ref'],
|
||||
pagination: {
|
||||
limit: 100,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.snapshot(result);
|
||||
});
|
||||
|
||||
e2e('should return empty nodes when docId not exists', async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const workspace = await app.create(Mockers.Workspace, { owner });
|
||||
await app.get(IndexerService).reconcileWorkspace(workspace.id);
|
||||
const result = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.doc,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'docId',
|
||||
match: 'not-exists-doc-id',
|
||||
},
|
||||
options: {
|
||||
fields: ['summary'],
|
||||
pagination: {
|
||||
limit: 1,
|
||||
},
|
||||
},
|
||||
query: { type: SearchQueryType.match, field: 'title', match: 'absent' },
|
||||
options: { fields: ['docId'], pagination: { limit: 10 } },
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.snapshot(result);
|
||||
t.is(result.workspace.search.pagination.count, 0);
|
||||
t.deepEqual(result.workspace.search.nodes, []);
|
||||
});
|
||||
|
||||
e2e(
|
||||
'should empty doc summary string when doc exists but no summary',
|
||||
async t => {
|
||||
const owner = await app.signup();
|
||||
const workspace = await app.create(Mockers.Workspace, {
|
||||
owner,
|
||||
});
|
||||
|
||||
const indexerService = app.get(IndexerService);
|
||||
|
||||
await indexerService.write(
|
||||
SearchTable.doc,
|
||||
[
|
||||
{
|
||||
docId: 'doc-1-without-summary',
|
||||
workspaceId: workspace.id,
|
||||
title: 'test1',
|
||||
summary: '',
|
||||
createdByUserId: owner.id,
|
||||
updatedByUserId: owner.id,
|
||||
createdAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
updatedAt: new Date('2025-04-22T00:00:00.000Z'),
|
||||
},
|
||||
],
|
||||
{
|
||||
refresh: true,
|
||||
}
|
||||
);
|
||||
|
||||
const result = await app.gql({
|
||||
query: indexerSearchQuery,
|
||||
variables: {
|
||||
id: workspace.id,
|
||||
input: {
|
||||
table: SearchTable.doc,
|
||||
query: {
|
||||
type: SearchQueryType.match,
|
||||
field: 'docId',
|
||||
match: 'doc-1-without-summary',
|
||||
},
|
||||
options: {
|
||||
fields: ['summary'],
|
||||
pagination: {
|
||||
limit: 1,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
t.snapshot(result.workspace.search.nodes);
|
||||
}
|
||||
);
|
||||
|
||||
@@ -121,6 +121,51 @@ e2e('should mention user in a doc', async t => {
|
||||
t.falsy(body2.workspace!.avatarUrl);
|
||||
});
|
||||
|
||||
e2e(
|
||||
'notification totalCount selection does not load the notification list',
|
||||
async t => {
|
||||
const { member, owner, workspace } = await init();
|
||||
|
||||
await app.login(owner);
|
||||
await app.gql({
|
||||
query: mentionUserMutation,
|
||||
variables: {
|
||||
input: {
|
||||
userId: member.id,
|
||||
workspaceId: workspace.id,
|
||||
doc: {
|
||||
id: 'count-only-doc',
|
||||
title: 'count-only-doc',
|
||||
mode: DocMode.page,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await app.login(member);
|
||||
const result = (await app.gql({
|
||||
query: {
|
||||
...listNotificationsQuery,
|
||||
op: 'CountOnlyNotifications',
|
||||
query: `
|
||||
query CountOnlyNotifications($pagination: PaginationInput!) {
|
||||
currentUser {
|
||||
notifications(pagination: $pagination) {
|
||||
totalCount
|
||||
}
|
||||
}
|
||||
}
|
||||
`,
|
||||
},
|
||||
variables: { pagination: { first: 10, offset: 0 } },
|
||||
})) as unknown as {
|
||||
currentUser: { notifications: { totalCount: number } };
|
||||
};
|
||||
|
||||
t.is(result.currentUser.notifications.totalCount, 1);
|
||||
}
|
||||
);
|
||||
|
||||
e2e('should mention doc mode support string value', async t => {
|
||||
const { member, owner, workspace } = await init();
|
||||
|
||||
|
||||
@@ -68,14 +68,20 @@ test('should read DEPLOYMENT_TYPE', t => {
|
||||
|
||||
test('should read FLAVOR', t => {
|
||||
t.deepEqual(
|
||||
['allinone', 'graphql', 'sync', 'renderer', 'front', 'doc', 'script'].map(
|
||||
envVal => {
|
||||
process.env.SERVER_FLAVOR = envVal;
|
||||
const env = new Env();
|
||||
return env.FLAVOR;
|
||||
}
|
||||
),
|
||||
['allinone', 'graphql', 'sync', 'renderer', 'front', 'doc', 'script']
|
||||
[
|
||||
'allinone',
|
||||
'graphql',
|
||||
'sync',
|
||||
'renderer',
|
||||
'front',
|
||||
'worker',
|
||||
'script',
|
||||
].map(envVal => {
|
||||
process.env.SERVER_FLAVOR = envVal;
|
||||
const env = new Env();
|
||||
return env.FLAVOR;
|
||||
}),
|
||||
['allinone', 'graphql', 'sync', 'renderer', 'front', 'worker', 'script']
|
||||
);
|
||||
|
||||
t.throws(
|
||||
@@ -85,7 +91,7 @@ test('should read FLAVOR', t => {
|
||||
},
|
||||
{
|
||||
message:
|
||||
'Invalid value "unknown" for environment variable SERVER_FLAVOR, expected one of ["allinone","graphql","sync","renderer","front","doc","script"]',
|
||||
'Invalid value "unknown" for environment variable SERVER_FLAVOR, expected one of ["allinone","graphql","sync","renderer","front","worker","script"]',
|
||||
}
|
||||
);
|
||||
});
|
||||
@@ -113,7 +119,7 @@ test('should tell flavors correctly', t => {
|
||||
sync: true,
|
||||
renderer: true,
|
||||
front: false,
|
||||
doc: true,
|
||||
worker: true,
|
||||
script: false,
|
||||
});
|
||||
|
||||
@@ -123,7 +129,7 @@ test('should tell flavors correctly', t => {
|
||||
sync: false,
|
||||
renderer: false,
|
||||
front: false,
|
||||
doc: false,
|
||||
worker: false,
|
||||
script: false,
|
||||
});
|
||||
|
||||
@@ -133,7 +139,7 @@ test('should tell flavors correctly', t => {
|
||||
sync: false,
|
||||
renderer: false,
|
||||
front: true,
|
||||
doc: false,
|
||||
worker: false,
|
||||
script: false,
|
||||
});
|
||||
|
||||
@@ -143,7 +149,7 @@ test('should tell flavors correctly', t => {
|
||||
sync: false,
|
||||
renderer: false,
|
||||
front: false,
|
||||
doc: false,
|
||||
worker: false,
|
||||
script: true,
|
||||
});
|
||||
});
|
||||
|
||||
@@ -63,6 +63,21 @@ test('should broadcast event to cluster instances', async t => {
|
||||
off();
|
||||
});
|
||||
|
||||
test('should preserve encoded binary updates across cluster instances', async t => {
|
||||
const { app1, app2 } = t.context;
|
||||
const eventbus1 = app1.get(EventBus);
|
||||
const eventbus2 = app2.get(EventBus);
|
||||
const listener = Sinon.spy(app1.get(Listeners), 'onEncodedBinaryEvent');
|
||||
const payload = {
|
||||
updates: [Buffer.from(new Uint8Array([1, 2, 3])).toString('base64')],
|
||||
};
|
||||
|
||||
eventbus2.broadcast('__test__.encodedBinary', payload);
|
||||
await eventbus1.waitFor('__test__.encodedBinary');
|
||||
|
||||
t.true(listener.calledOnceWith(payload));
|
||||
});
|
||||
|
||||
test('should continuously use the same request id', async t => {
|
||||
const { app1, app2 } = t.context;
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ declare global {
|
||||
interface Events {
|
||||
'__test__.event': { count: number };
|
||||
'__test__.event2': { count: number };
|
||||
'__test__.encodedBinary': { updates: string[] };
|
||||
'__test__.throw': { count: number };
|
||||
'__test__.suppressThrow': {};
|
||||
'__test__.requestId': {};
|
||||
@@ -28,6 +29,11 @@ export class Listeners {
|
||||
return payload;
|
||||
}
|
||||
|
||||
@OnEvent('__test__.encodedBinary')
|
||||
onEncodedBinaryEvent(payload: Events['__test__.encodedBinary']) {
|
||||
return payload;
|
||||
}
|
||||
|
||||
@OnEvent('__test__.throw')
|
||||
onThrow() {
|
||||
throw new Error('Error in event handler');
|
||||
|
||||
@@ -6,10 +6,10 @@ import { EventName } from '../../base/event/def';
|
||||
export class MockEventBus {
|
||||
private readonly stub = Sinon.createStubInstance(EventBus);
|
||||
|
||||
emit = this.stub.emitAsync;
|
||||
emitAsync = this.stub.emitAsync;
|
||||
emitDetached = this.stub.emitAsync;
|
||||
broadcast = this.stub.broadcast;
|
||||
emit: Sinon.SinonStub = this.stub.emitAsync;
|
||||
emitAsync: Sinon.SinonStub = this.stub.emitAsync;
|
||||
emitDetached: Sinon.SinonStub = this.stub.emitAsync;
|
||||
broadcast: Sinon.SinonStub = this.stub.broadcast;
|
||||
|
||||
last<Event extends EventName>(
|
||||
name: Event
|
||||
@@ -22,7 +22,6 @@ export class MockEventBus {
|
||||
throw new Error(`Event ${name} never called`);
|
||||
}
|
||||
|
||||
// @ts-expect-error allow
|
||||
return {
|
||||
name,
|
||||
payload: call.args[1],
|
||||
|
||||
@@ -136,127 +136,3 @@ test('should claim job', async t => {
|
||||
'should update job status to claimed'
|
||||
);
|
||||
});
|
||||
|
||||
test('should fence transcript dispatch generations atomically', async t => {
|
||||
const task = await t.context.transcriptTask.create({
|
||||
userId: user.id,
|
||||
workspaceId: workspace.id,
|
||||
blobId: 'transcript-blob',
|
||||
recipeId: 'transcript.audio',
|
||||
recipeVersion: 'v1',
|
||||
inputSnapshot: { normalizedTranscript: 'source' },
|
||||
});
|
||||
const adoptions = await Promise.all([
|
||||
t.context.transcriptTask.adoptLegacyDispatch(
|
||||
task.id,
|
||||
null,
|
||||
'legacy-generation-a'
|
||||
),
|
||||
t.context.transcriptTask.adoptLegacyDispatch(
|
||||
task.id,
|
||||
null,
|
||||
'legacy-generation-b'
|
||||
),
|
||||
]);
|
||||
t.is(adoptions.filter(Boolean).length, 1);
|
||||
const adopted = await t.context.transcriptTask.get(task.id);
|
||||
const adoptedGeneration = adopted?.dispatchGeneration;
|
||||
if (!adoptedGeneration) {
|
||||
t.fail('legacy dispatch should have a generation');
|
||||
return;
|
||||
}
|
||||
t.true(
|
||||
await t.context.transcriptTask.claimDispatch(
|
||||
task.id,
|
||||
adoptedGeneration,
|
||||
null
|
||||
)
|
||||
);
|
||||
t.true(
|
||||
await t.context.transcriptTask.completeDispatch(
|
||||
task.id,
|
||||
adoptedGeneration,
|
||||
null,
|
||||
{
|
||||
status: 'failed',
|
||||
protectedResult: { normalizedTranscript: 'source' },
|
||||
errorCode: 'provider_failed',
|
||||
}
|
||||
)
|
||||
);
|
||||
|
||||
const claims = await Promise.all([
|
||||
t.context.transcriptTask.claimRetry(
|
||||
task.id,
|
||||
user.id,
|
||||
workspace.id,
|
||||
null,
|
||||
'generation-a'
|
||||
),
|
||||
t.context.transcriptTask.claimRetry(
|
||||
task.id,
|
||||
user.id,
|
||||
workspace.id,
|
||||
null,
|
||||
'generation-b'
|
||||
),
|
||||
]);
|
||||
t.is(claims.filter(Boolean).length, 1);
|
||||
|
||||
const claimed = await t.context.transcriptTask.get(task.id);
|
||||
const generation = claimed?.dispatchGeneration;
|
||||
if (!generation) {
|
||||
t.fail('retry should have a dispatch generation');
|
||||
return;
|
||||
}
|
||||
t.false(
|
||||
await t.context.transcriptTask.claimDispatch(
|
||||
task.id,
|
||||
generation === 'generation-a' ? 'generation-b' : 'generation-a',
|
||||
null
|
||||
)
|
||||
);
|
||||
t.true(
|
||||
await t.context.transcriptTask.claimDispatch(task.id, generation, null)
|
||||
);
|
||||
t.true(
|
||||
await t.context.transcriptTask.attachActionRun(
|
||||
task.id,
|
||||
generation,
|
||||
null,
|
||||
'run-next'
|
||||
)
|
||||
);
|
||||
t.false(
|
||||
await t.context.transcriptTask.attachActionRun(
|
||||
task.id,
|
||||
generation,
|
||||
null,
|
||||
'run-duplicate'
|
||||
)
|
||||
);
|
||||
t.false(
|
||||
await t.context.transcriptTask.completeDispatch(
|
||||
task.id,
|
||||
generation,
|
||||
'run-duplicate',
|
||||
{ status: 'ready' }
|
||||
)
|
||||
);
|
||||
t.true(
|
||||
await t.context.transcriptTask.completeDispatch(
|
||||
task.id,
|
||||
generation,
|
||||
'run-next',
|
||||
{
|
||||
status: 'ready',
|
||||
protectedResult: { normalizedTranscript: 'result' },
|
||||
}
|
||||
)
|
||||
);
|
||||
t.like(await t.context.transcriptTask.get(task.id), {
|
||||
status: 'ready',
|
||||
dispatchGeneration: null,
|
||||
actionRunId: 'run-next',
|
||||
});
|
||||
});
|
||||
|
||||
@@ -4,7 +4,7 @@ import ava, { TestFn } from 'ava';
|
||||
import { Config } from '../../base/config';
|
||||
import { SessionModel } from '../../models/session';
|
||||
import { UserModel } from '../../models/user';
|
||||
import { createTestingModule, type TestingModule } from '../utils';
|
||||
import { createTestingModule, sleep, type TestingModule } from '../utils';
|
||||
|
||||
interface Context {
|
||||
config: Config;
|
||||
@@ -109,6 +109,7 @@ test('should refresh exists userSession', async t => {
|
||||
t.is(userSession.userId, user.id);
|
||||
t.not(userSession.expiresAt, null);
|
||||
|
||||
await sleep(1);
|
||||
const existsUserSession = await t.context.session.createOrRefreshUserSession(
|
||||
user.id,
|
||||
session.id
|
||||
|
||||
@@ -4,7 +4,11 @@ import ava, { TestFn } from 'ava';
|
||||
import Sinon from 'sinon';
|
||||
|
||||
import { OneDay } from '../../base';
|
||||
import { StorageModule, WorkspaceBlobStorage } from '../../core/storage';
|
||||
import {
|
||||
StorageModule,
|
||||
StorageWorkerModule,
|
||||
WorkspaceBlobStorage,
|
||||
} from '../../core/storage';
|
||||
import { BlobUploadCleanupJob } from '../../core/storage/job';
|
||||
import { StorageRuntimeProvider } from '../../core/storage-runtime';
|
||||
import { MockUser, MockWorkspace } from '../mocks';
|
||||
@@ -25,7 +29,7 @@ test.before(async t => {
|
||||
cleanupExpiredPendingBlobs: Sinon.stub(),
|
||||
};
|
||||
t.context.module = await createTestingModule({
|
||||
imports: [ScheduleModule.forRoot(), StorageModule],
|
||||
imports: [ScheduleModule.forRoot(), StorageModule, StorageWorkerModule],
|
||||
tapModule: builder => {
|
||||
builder
|
||||
.overrideProvider(StorageRuntimeProvider)
|
||||
|
||||
@@ -3,7 +3,7 @@ import test, { type ExecutionContext } from 'ava';
|
||||
import { io, type Socket as SocketIOClient } from 'socket.io-client';
|
||||
import { Doc, encodeStateAsUpdate } from 'yjs';
|
||||
|
||||
import { CANARY_CLIENT_VERSION_MAX_AGE_DAYS } from '../../base';
|
||||
import { CANARY_CLIENT_VERSION_MAX_AGE_DAYS, EventBus } from '../../base';
|
||||
import {
|
||||
DocRole,
|
||||
Models,
|
||||
@@ -312,71 +312,7 @@ test('should reject websocket jwt auth after session deletion', async t => {
|
||||
}
|
||||
});
|
||||
|
||||
test('clientVersion=0.25.0 should only receive space:broadcast-doc-update', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const spaceId = user.id;
|
||||
const update = createYjsUpdateBase64();
|
||||
|
||||
const sender = createClient(url, cookieHeader);
|
||||
const receiver = createClient(url, cookieHeader);
|
||||
|
||||
try {
|
||||
await Promise.all([waitForConnect(sender), waitForConnect(receiver)]);
|
||||
|
||||
const receiverJoin = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
receiver,
|
||||
'space:join',
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.25.0' }
|
||||
)
|
||||
);
|
||||
t.true(receiverJoin.success);
|
||||
|
||||
const senderJoin = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
sender,
|
||||
'space:join',
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.26.0' }
|
||||
)
|
||||
);
|
||||
t.true(senderJoin.success);
|
||||
|
||||
const onUpdate = waitForEvent<{
|
||||
spaceType: string;
|
||||
spaceId: string;
|
||||
docId: string;
|
||||
update: string;
|
||||
}>(receiver, 'space:broadcast-doc-update');
|
||||
const noUpdates = expectNoEvent(receiver, 'space:broadcast-doc-updates');
|
||||
|
||||
const pushRes = await emitWithAck<{ accepted: true; timestamp?: number }>(
|
||||
sender,
|
||||
'space:push-doc-update',
|
||||
{
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
docId: 'doc-1',
|
||||
update,
|
||||
}
|
||||
);
|
||||
unwrapResponse(t, pushRes);
|
||||
|
||||
const message = await onUpdate;
|
||||
t.is(message.spaceType, 'userspace');
|
||||
t.is(message.spaceId, spaceId);
|
||||
t.is(message.docId, 'doc-1');
|
||||
t.is(message.update, update);
|
||||
|
||||
await noUpdates;
|
||||
} finally {
|
||||
sender.disconnect();
|
||||
receiver.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('clientVersion>=0.26.0 should only receive space:broadcast-doc-updates', async t => {
|
||||
test('clientVersion>=0.26.0 should receive legacy space:broadcast-doc-updates', async t => {
|
||||
const { user, cookieHeader } = await loginWithCookie(app);
|
||||
const spaceId = user.id;
|
||||
const update = createYjsUpdateBase64();
|
||||
@@ -402,7 +338,7 @@ test('clientVersion>=0.26.0 should only receive space:broadcast-doc-updates', as
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
sender,
|
||||
'space:join',
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.25.0' }
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.26.0' }
|
||||
)
|
||||
);
|
||||
t.true(senderJoin.success);
|
||||
@@ -413,7 +349,6 @@ test('clientVersion>=0.26.0 should only receive space:broadcast-doc-updates', as
|
||||
docId: string;
|
||||
updates: string[];
|
||||
}>(receiver, 'space:broadcast-doc-updates');
|
||||
const noUpdate = expectNoEvent(receiver, 'space:broadcast-doc-update');
|
||||
|
||||
const pushRes = await emitWithAck<{ accepted: true; timestamp?: number }>(
|
||||
sender,
|
||||
@@ -432,15 +367,13 @@ test('clientVersion>=0.26.0 should only receive space:broadcast-doc-updates', as
|
||||
t.is(message.spaceId, spaceId);
|
||||
t.is(message.docId, 'doc-2');
|
||||
t.deepEqual(message.updates, [update]);
|
||||
|
||||
await noUpdate;
|
||||
} finally {
|
||||
sender.disconnect();
|
||||
receiver.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('canary date clientVersion should use sync-026 in canary namespace', async t => {
|
||||
test('canary date clientVersion should use sync-027 in canary namespace', async t => {
|
||||
const prevNamespace = env.NAMESPACE;
|
||||
// @ts-expect-error test
|
||||
env.NAMESPACE = 'dev';
|
||||
@@ -456,15 +389,18 @@ test('canary date clientVersion should use sync-026 in canary namespace', async
|
||||
try {
|
||||
await Promise.all([waitForConnect(sender), waitForConnect(receiver)]);
|
||||
|
||||
const canaryVersion = makeCanaryDateVersion(new Date(), '015');
|
||||
const receiverJoin = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
receiver,
|
||||
'space:join',
|
||||
'space:join-batch',
|
||||
{
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
clientVersion: makeCanaryDateVersion(new Date(), '015'),
|
||||
spaces: [
|
||||
{ spaceType: 'userspace', spaceId },
|
||||
{ spaceType: 'userspace', spaceId, docId: 'doc-canary' },
|
||||
],
|
||||
clientVersion: canaryVersion,
|
||||
}
|
||||
)
|
||||
);
|
||||
@@ -474,8 +410,14 @@ test('canary date clientVersion should use sync-026 in canary namespace', async
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
sender,
|
||||
'space:join',
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.25.0' }
|
||||
'space:join-batch',
|
||||
{
|
||||
spaces: [
|
||||
{ spaceType: 'userspace', spaceId },
|
||||
{ spaceType: 'userspace', spaceId, docId: 'doc-canary' },
|
||||
],
|
||||
clientVersion: canaryVersion,
|
||||
}
|
||||
)
|
||||
);
|
||||
t.true(senderJoin.success);
|
||||
@@ -486,7 +428,6 @@ test('canary date clientVersion should use sync-026 in canary namespace', async
|
||||
docId: string;
|
||||
updates: string[];
|
||||
}>(receiver, 'space:broadcast-doc-updates');
|
||||
const noUpdate = expectNoEvent(receiver, 'space:broadcast-doc-update');
|
||||
|
||||
const pushRes = await emitWithAck<{ accepted: true; timestamp?: number }>(
|
||||
sender,
|
||||
@@ -505,8 +446,6 @@ test('canary date clientVersion should use sync-026 in canary namespace', async
|
||||
t.is(message.spaceId, spaceId);
|
||||
t.is(message.docId, 'doc-canary');
|
||||
t.deepEqual(message.updates, [update]);
|
||||
|
||||
await noUpdate;
|
||||
} finally {
|
||||
sender.disconnect();
|
||||
receiver.disconnect();
|
||||
@@ -517,7 +456,7 @@ test('canary date clientVersion should use sync-026 in canary namespace', async
|
||||
}
|
||||
});
|
||||
|
||||
test('clientVersion<0.25.0 should be rejected and disconnected', async t => {
|
||||
test('clientVersion<0.26.0 should be rejected and disconnected', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const spaceId = user.id;
|
||||
|
||||
@@ -530,7 +469,7 @@ test('clientVersion<0.25.0 should be rejected and disconnected', async t => {
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
'space:join',
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.24.4' }
|
||||
{ spaceType: 'userspace', spaceId, clientVersion: '0.25.0' }
|
||||
)
|
||||
);
|
||||
t.false(res.success);
|
||||
@@ -620,7 +559,7 @@ test('canary date clientVersion should be rejected outside canary namespace', as
|
||||
}
|
||||
});
|
||||
|
||||
test('space:join-awareness should reject clientVersion<0.25.0', async t => {
|
||||
test('space:join-awareness should reject clientVersion<0.26.0', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const spaceId = user.id;
|
||||
|
||||
@@ -637,7 +576,7 @@ test('space:join-awareness should reject clientVersion<0.25.0', async t => {
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
docId: 'doc-awareness',
|
||||
clientVersion: '0.24.4',
|
||||
clientVersion: '0.25.0',
|
||||
}
|
||||
)
|
||||
);
|
||||
@@ -649,6 +588,611 @@ test('space:join-awareness should reject clientVersion<0.25.0', async t => {
|
||||
}
|
||||
});
|
||||
|
||||
test('new clients must use batch join endpoints on new servers', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const requests = [
|
||||
{
|
||||
event: 'space:join',
|
||||
payload: {
|
||||
spaceType: 'userspace',
|
||||
spaceId: user.id,
|
||||
clientVersion: '0.27.5',
|
||||
},
|
||||
},
|
||||
{
|
||||
event: 'space:join-awareness',
|
||||
payload: {
|
||||
spaceType: 'userspace',
|
||||
spaceId: user.id,
|
||||
docId: 'doc-awareness',
|
||||
clientVersion: '0.27.5',
|
||||
},
|
||||
},
|
||||
] as const;
|
||||
|
||||
for (const request of requests) {
|
||||
const socket = createClient(url, cookieHeader);
|
||||
try {
|
||||
await waitForConnect(socket);
|
||||
const result = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
request.event,
|
||||
request.payload
|
||||
)
|
||||
);
|
||||
t.false(result.success);
|
||||
await waitForDisconnect(socket);
|
||||
} finally {
|
||||
socket.disconnect();
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
test('space:join-batch should validate entries before joining', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const socket = createClient(url, cookieHeader);
|
||||
const spaceId = user.id;
|
||||
|
||||
try {
|
||||
await waitForConnect(socket);
|
||||
|
||||
const invalidBatches = [
|
||||
{
|
||||
label: 'empty',
|
||||
payload: { spaces: [], clientVersion: '0.27.5' },
|
||||
},
|
||||
{
|
||||
label: 'missing client version',
|
||||
payload: { spaces: [{ spaceType: 'userspace', spaceId }] },
|
||||
},
|
||||
{
|
||||
label: 'cross workspace',
|
||||
payload: {
|
||||
spaces: [
|
||||
{ spaceType: 'userspace', spaceId },
|
||||
{ spaceType: 'userspace', spaceId: `${spaceId}-other` },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
},
|
||||
},
|
||||
{
|
||||
label: 'duplicate',
|
||||
payload: {
|
||||
spaces: [
|
||||
{ spaceType: 'userspace', spaceId, docId: 'doc-1' },
|
||||
{ spaceType: 'userspace', spaceId, docId: 'doc-1' },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
},
|
||||
},
|
||||
{
|
||||
label: 'invalid entry',
|
||||
payload: {
|
||||
spaces: [{ spaceType: 'invalid', spaceId }],
|
||||
clientVersion: '0.27.5',
|
||||
},
|
||||
},
|
||||
{
|
||||
label: 'over limit',
|
||||
payload: {
|
||||
spaces: Array.from({ length: 101 }, (_, index) => ({
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
docId: `doc-${index}`,
|
||||
})),
|
||||
clientVersion: '0.27.5',
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
for (const { label, payload } of invalidBatches) {
|
||||
const error = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(socket, 'space:join-batch', payload)
|
||||
);
|
||||
t.is(error.name, 'BAD_REQUEST', label);
|
||||
}
|
||||
} finally {
|
||||
socket.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('space:join-batch should reject clients before 0.27.5', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const socket = createClient(url, cookieHeader);
|
||||
|
||||
try {
|
||||
await waitForConnect(socket);
|
||||
const result = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
'space:join-batch',
|
||||
{
|
||||
spaces: [{ spaceType: 'userspace', spaceId: user.id }],
|
||||
clientVersion: '0.27.4',
|
||||
}
|
||||
)
|
||||
);
|
||||
t.false(result.success);
|
||||
await waitForDisconnect(socket);
|
||||
} finally {
|
||||
socket.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('space:join-batch should authorize once and join all requested rooms', async t => {
|
||||
const models = app.get(Models);
|
||||
const { user: owner, cookieHeader: ownerCookieHeader } = await login(app);
|
||||
const { cookieHeader: deniedCookieHeader } = await login(app);
|
||||
const workspace = await models.workspace.create(owner.id);
|
||||
|
||||
const ownerSocket = createClient(url, ownerCookieHeader);
|
||||
const receiverSocket = createClient(url, ownerCookieHeader);
|
||||
const deniedSocket = createClient(url, deniedCookieHeader);
|
||||
|
||||
try {
|
||||
await Promise.all([
|
||||
waitForConnect(ownerSocket),
|
||||
waitForConnect(receiverSocket),
|
||||
waitForConnect(deniedSocket),
|
||||
]);
|
||||
|
||||
const batch = {
|
||||
spaces: [
|
||||
{ spaceType: 'workspace', spaceId: workspace.id },
|
||||
{ spaceType: 'workspace', spaceId: workspace.id, docId: 'doc-a' },
|
||||
{ spaceType: 'workspace', spaceId: workspace.id, docId: 'doc-b' },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
};
|
||||
|
||||
for (const socket of [ownerSocket, receiverSocket]) {
|
||||
const result = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
'space:join-batch',
|
||||
batch
|
||||
)
|
||||
);
|
||||
t.true(result.success);
|
||||
}
|
||||
|
||||
const awarenessOnlyResult = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
ownerSocket,
|
||||
'space:join-batch',
|
||||
{
|
||||
spaces: [
|
||||
{
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId: 'doc-c',
|
||||
},
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
}
|
||||
)
|
||||
);
|
||||
t.true(awarenessOnlyResult.success);
|
||||
|
||||
const timestamps = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<Record<string, number>>(
|
||||
ownerSocket,
|
||||
'space:load-doc-timestamps',
|
||||
{
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
}
|
||||
)
|
||||
);
|
||||
t.deepEqual(timestamps, {});
|
||||
|
||||
const deniedError = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(deniedSocket, 'space:join-batch', batch)
|
||||
);
|
||||
t.is(deniedError.name, 'SPACE_ACCESS_DENIED');
|
||||
|
||||
const deniedSyncRoomError = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(deniedSocket, 'space:load-doc-timestamps', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
})
|
||||
);
|
||||
t.is(deniedSyncRoomError.name, 'NOT_IN_SPACE');
|
||||
|
||||
const deniedAwarenessRoomError = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(deniedSocket, 'space:load-awarenesses', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId: 'doc-a',
|
||||
})
|
||||
);
|
||||
t.is(deniedAwarenessRoomError.name, 'NOT_IN_SPACE');
|
||||
|
||||
const receivedA = waitForEvent<{
|
||||
spaceType: string;
|
||||
spaceId: string;
|
||||
docId: string;
|
||||
awarenessUpdate: string;
|
||||
}>(receiverSocket, 'space:broadcast-awareness-update');
|
||||
const noDeniedEvent = expectNoEvent(
|
||||
deniedSocket,
|
||||
'space:broadcast-awareness-update'
|
||||
);
|
||||
|
||||
ownerSocket.emit('space:update-awareness', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId: 'doc-a',
|
||||
awarenessUpdate: 'AQID',
|
||||
});
|
||||
const messageA = await receivedA;
|
||||
|
||||
const receivedB = waitForEvent<{
|
||||
spaceType: string;
|
||||
spaceId: string;
|
||||
docId: string;
|
||||
awarenessUpdate: string;
|
||||
}>(receiverSocket, 'space:broadcast-awareness-update');
|
||||
ownerSocket.emit('space:update-awareness', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId: 'doc-b',
|
||||
awarenessUpdate: 'BAUG',
|
||||
});
|
||||
const messageB = await receivedB;
|
||||
|
||||
t.deepEqual(
|
||||
new Set([messageA.docId, messageB.docId]),
|
||||
new Set(['doc-a', 'doc-b'])
|
||||
);
|
||||
await noDeniedEvent;
|
||||
} finally {
|
||||
ownerSocket.disconnect();
|
||||
receiverSocket.disconnect();
|
||||
deniedSocket.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('batch doc entries require Doc.Read atomically', async t => {
|
||||
const db = app.get(PrismaClient);
|
||||
const models = app.get(Models);
|
||||
const { user: owner } = await login(app);
|
||||
const { user: collaborator, cookieHeader } = await login(app);
|
||||
const workspace = await models.workspace.create(owner.id);
|
||||
const docId = 'batch-private-doc';
|
||||
|
||||
await models.workspaceUser.set(
|
||||
workspace.id,
|
||||
collaborator.id,
|
||||
WorkspaceRole.Collaborator,
|
||||
{ status: WorkspaceMemberStatus.Accepted }
|
||||
);
|
||||
await models.doc.setDefaultRole(workspace.id, docId, DocRole.None);
|
||||
await createSnapshot(db, {
|
||||
workspaceId: workspace.id,
|
||||
docId,
|
||||
userId: owner.id,
|
||||
});
|
||||
|
||||
const socket = createClient(url, cookieHeader);
|
||||
try {
|
||||
await waitForConnect(socket);
|
||||
|
||||
const error = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(socket, 'space:join-batch', {
|
||||
spaces: [
|
||||
{ spaceType: 'workspace', spaceId: workspace.id },
|
||||
{ spaceType: 'workspace', spaceId: workspace.id, docId },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
})
|
||||
);
|
||||
t.true(error.message.includes('Doc.Read'));
|
||||
|
||||
const timestampsError = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(socket, 'space:load-doc-timestamps', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
})
|
||||
);
|
||||
t.is(timestampsError.name, 'NOT_IN_SPACE');
|
||||
} finally {
|
||||
socket.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('batch sync routes active updates and only broadcasts invalidation to control room', async t => {
|
||||
const { user, cookieHeader } = await login(app);
|
||||
const spaceId = user.id;
|
||||
const sender = createClient(url, cookieHeader);
|
||||
const receiver = createClient(url, cookieHeader);
|
||||
const passive = createClient(url, cookieHeader);
|
||||
|
||||
try {
|
||||
await Promise.all([
|
||||
waitForConnect(sender),
|
||||
waitForConnect(receiver),
|
||||
waitForConnect(passive),
|
||||
]);
|
||||
|
||||
const activeBatch = {
|
||||
spaces: [
|
||||
{ spaceType: 'userspace', spaceId },
|
||||
{ spaceType: 'userspace', spaceId, docId: 'some-doc' },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
};
|
||||
for (const socket of [sender, receiver]) {
|
||||
const result = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
'space:join-batch',
|
||||
activeBatch
|
||||
)
|
||||
);
|
||||
t.true(result.success);
|
||||
}
|
||||
const passiveJoin = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
passive,
|
||||
'space:join-batch',
|
||||
{
|
||||
spaces: [{ spaceType: 'userspace', spaceId }],
|
||||
clientVersion: '0.27.5',
|
||||
}
|
||||
)
|
||||
);
|
||||
t.true(passiveJoin.success);
|
||||
|
||||
const receivedUpdate = waitForEvent<{
|
||||
docId: string;
|
||||
updates: string[];
|
||||
}>(receiver, 'space:broadcast-doc-updates');
|
||||
const receivedInvalidation = waitForEvent<{
|
||||
spaceType: string;
|
||||
spaceId: string;
|
||||
timestamp: number;
|
||||
docId?: string;
|
||||
updates?: string[];
|
||||
}>(receiver, 'space:broadcast-doc-invalidation');
|
||||
const noPassiveUpdate = expectNoEvent(
|
||||
passive,
|
||||
'space:broadcast-doc-updates'
|
||||
);
|
||||
|
||||
unwrapResponse(
|
||||
t,
|
||||
await emitWithAck(sender, 'space:push-doc-update', {
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
docId: 'some-doc',
|
||||
update: createYjsUpdateBase64(),
|
||||
})
|
||||
);
|
||||
|
||||
const [update, invalidation] = await Promise.all([
|
||||
receivedUpdate,
|
||||
receivedInvalidation,
|
||||
]);
|
||||
t.is(update.docId, 'some-doc');
|
||||
t.deepEqual(Object.keys(invalidation).sort(), [
|
||||
'spaceId',
|
||||
'spaceType',
|
||||
'timestamp',
|
||||
]);
|
||||
await noPassiveUpdate;
|
||||
|
||||
const leave = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
receiver,
|
||||
'space:leave-batch',
|
||||
{
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
docIds: ['some-doc'],
|
||||
}
|
||||
)
|
||||
);
|
||||
t.true(leave.success);
|
||||
|
||||
const noLeftUpdate = expectNoEvent(receiver, 'space:broadcast-doc-updates');
|
||||
const receivedAfterLeave = waitForEvent(
|
||||
receiver,
|
||||
'space:broadcast-doc-invalidation'
|
||||
);
|
||||
unwrapResponse(
|
||||
t,
|
||||
await emitWithAck(sender, 'space:push-doc-update', {
|
||||
spaceType: 'userspace',
|
||||
spaceId,
|
||||
docId: 'some-doc',
|
||||
update: createYjsUpdateBase64(),
|
||||
})
|
||||
);
|
||||
await Promise.all([noLeftUpdate, receivedAfterLeave]);
|
||||
} finally {
|
||||
sender.disconnect();
|
||||
receiver.disconnect();
|
||||
passive.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('permission revocation removes a active document subscription', async t => {
|
||||
const db = app.get(PrismaClient);
|
||||
const models = app.get(Models);
|
||||
const { user: owner, cookieHeader: ownerCookie } = await login(app);
|
||||
const { user: collaborator, cookieHeader: collaboratorCookie } =
|
||||
await login(app);
|
||||
const workspace = await models.workspace.create(owner.id);
|
||||
const docId = 'revoked-doc';
|
||||
|
||||
await models.workspaceUser.set(
|
||||
workspace.id,
|
||||
collaborator.id,
|
||||
WorkspaceRole.Collaborator,
|
||||
{ status: WorkspaceMemberStatus.Accepted }
|
||||
);
|
||||
await models.doc.setDefaultRole(workspace.id, docId, DocRole.None);
|
||||
await models.docUser.set(
|
||||
workspace.id,
|
||||
docId,
|
||||
collaborator.id,
|
||||
DocRole.Reader
|
||||
);
|
||||
await createSnapshot(db, {
|
||||
workspaceId: workspace.id,
|
||||
docId,
|
||||
userId: owner.id,
|
||||
});
|
||||
|
||||
const ownerSocket = createClient(url, ownerCookie);
|
||||
const collaboratorSocket = createClient(url, collaboratorCookie);
|
||||
try {
|
||||
await Promise.all([
|
||||
waitForConnect(ownerSocket),
|
||||
waitForConnect(collaboratorSocket),
|
||||
]);
|
||||
|
||||
for (const socket of [ownerSocket, collaboratorSocket]) {
|
||||
const response = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
'space:join-batch',
|
||||
{
|
||||
spaces: [
|
||||
{ spaceType: 'workspace', spaceId: workspace.id },
|
||||
{ spaceType: 'workspace', spaceId: workspace.id, docId },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
}
|
||||
)
|
||||
);
|
||||
t.true(response.success);
|
||||
}
|
||||
|
||||
await models.docUser.delete(workspace.id, docId, collaborator.id);
|
||||
await app.get(EventBus).emitAsync('doc.grants.changed', {
|
||||
workspaceId: workspace.id,
|
||||
docId,
|
||||
});
|
||||
|
||||
const noRevokedUpdate = expectNoEvent(
|
||||
collaboratorSocket,
|
||||
'space:broadcast-doc-updates'
|
||||
);
|
||||
unwrapResponse(
|
||||
t,
|
||||
await emitWithAck(ownerSocket, 'space:push-doc-update', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId,
|
||||
update: createYjsUpdateBase64(),
|
||||
})
|
||||
);
|
||||
await noRevokedUpdate;
|
||||
} finally {
|
||||
ownerSocket.disconnect();
|
||||
collaboratorSocket.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('awareness requires Doc.Read but not Doc.Update', async t => {
|
||||
const db = app.get(PrismaClient);
|
||||
const models = app.get(Models);
|
||||
const { user: owner, cookieHeader: ownerCookie } = await login(app);
|
||||
const { user: reader, cookieHeader: readerCookie } = await login(app);
|
||||
const workspace = await models.workspace.create(owner.id);
|
||||
const docId = 'awareness-reader-doc';
|
||||
|
||||
await models.workspaceUser.set(
|
||||
workspace.id,
|
||||
reader.id,
|
||||
WorkspaceRole.Collaborator,
|
||||
{ status: WorkspaceMemberStatus.Accepted }
|
||||
);
|
||||
await models.doc.setDefaultRole(workspace.id, docId, DocRole.None);
|
||||
await models.docUser.set(workspace.id, docId, reader.id, DocRole.Reader);
|
||||
await createSnapshot(db, {
|
||||
workspaceId: workspace.id,
|
||||
docId,
|
||||
userId: owner.id,
|
||||
});
|
||||
|
||||
const ownerSocket = createClient(url, ownerCookie);
|
||||
const readerSocket = createClient(url, readerCookie);
|
||||
try {
|
||||
await Promise.all([
|
||||
waitForConnect(ownerSocket),
|
||||
waitForConnect(readerSocket),
|
||||
]);
|
||||
|
||||
for (const socket of [ownerSocket, readerSocket]) {
|
||||
const response = unwrapResponse(
|
||||
t,
|
||||
await emitWithAck<{ clientId: string; success: boolean }>(
|
||||
socket,
|
||||
'space:join-batch',
|
||||
{
|
||||
spaces: [
|
||||
{ spaceType: 'workspace', spaceId: workspace.id },
|
||||
{ spaceType: 'workspace', spaceId: workspace.id, docId },
|
||||
],
|
||||
clientVersion: '0.27.5',
|
||||
}
|
||||
)
|
||||
);
|
||||
t.true(response.success);
|
||||
}
|
||||
|
||||
const receivedAwareness = waitForEvent<{
|
||||
docId: string;
|
||||
awarenessUpdate: string;
|
||||
}>(readerSocket, 'space:broadcast-awareness-update');
|
||||
ownerSocket.emit('space:update-awareness', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId,
|
||||
awarenessUpdate: 'AQID',
|
||||
});
|
||||
t.deepEqual(await receivedAwareness, {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId,
|
||||
awarenessUpdate: 'AQID',
|
||||
});
|
||||
|
||||
const updateError = getErrorResponse(
|
||||
t,
|
||||
await emitWithAck(readerSocket, 'space:push-doc-update', {
|
||||
spaceType: 'workspace',
|
||||
spaceId: workspace.id,
|
||||
docId,
|
||||
update: createYjsUpdateBase64(),
|
||||
})
|
||||
);
|
||||
t.is(updateError.name, 'DOC_ACTION_DENIED');
|
||||
} finally {
|
||||
ownerSocket.disconnect();
|
||||
readerSocket.disconnect();
|
||||
}
|
||||
});
|
||||
|
||||
test('active users metric should dedupe multiple sockets for one user', async t => {
|
||||
const db = app.get(PrismaClient);
|
||||
await ensureSyncActiveUsersTable(db);
|
||||
|
||||
@@ -8,7 +8,10 @@ const testPrivateKey = privateKey
|
||||
.export({ format: 'pem', type: 'pkcs8' })
|
||||
.toString();
|
||||
|
||||
export async function createTestRuntimeConfig(databaseUrl: string) {
|
||||
export async function createTestRuntimeConfig(
|
||||
databaseUrl: string,
|
||||
indexer: AppConfig['indexer']
|
||||
) {
|
||||
const directory = await mkdtemp(join(tmpdir(), 'affine-server-test-'));
|
||||
const storagePath = join(directory, 'storage');
|
||||
const storage = (bucket: string) => ({
|
||||
@@ -30,6 +33,10 @@ export async function createTestRuntimeConfig(databaseUrl: string) {
|
||||
enabled: true,
|
||||
storage: storage('copilot'),
|
||||
},
|
||||
indexer: {
|
||||
enabled: indexer.enabled,
|
||||
provider: indexer.provider,
|
||||
},
|
||||
})
|
||||
);
|
||||
return {
|
||||
|
||||
@@ -75,8 +75,10 @@ export async function createTestingModule(
|
||||
moduleDef: TestingModuleMetadata = {},
|
||||
autoInitialize = true
|
||||
): Promise<TestingModule> {
|
||||
const config = new ConfigFactory().config;
|
||||
const runtimeConfig = await createTestRuntimeConfig(
|
||||
new ConfigFactory().config.db.datasourceUrl
|
||||
config.db.datasourceUrl,
|
||||
config.indexer
|
||||
);
|
||||
// setting up
|
||||
let imports = moduleDef.imports ?? [buildAppModule(globalThis.env)];
|
||||
|
||||
@@ -38,12 +38,17 @@ test.before(async t => {
|
||||
|
||||
const db = app.get(PrismaClient);
|
||||
|
||||
t.context.u1 = await app.signupV1('u1@affine.pro');
|
||||
t.context.db = db;
|
||||
t.context.app = app;
|
||||
t.context.storage = app.get(WorkspaceBlobStorage);
|
||||
t.context.workspace = app.get(PgWorkspaceDocStorageAdapter);
|
||||
t.context.models = app.get(Models);
|
||||
});
|
||||
|
||||
test.beforeEach(async t => {
|
||||
const { app, db } = t.context;
|
||||
await app.initTestingDB();
|
||||
t.context.u1 = await app.signupV1('u1@affine.pro');
|
||||
|
||||
await db.workspaceDoc.create({
|
||||
data: {
|
||||
|
||||
@@ -28,12 +28,16 @@ import { RedisModule } from './base/redis';
|
||||
import { RateLimiterModule } from './base/throttler';
|
||||
import { WebSocketModule } from './base/websocket';
|
||||
import { AuthModule } from './core/auth';
|
||||
import { BackendRuntimeModule } from './core/backend-runtime';
|
||||
import {
|
||||
BackendRuntimeModule,
|
||||
BackendRuntimeProducerModule,
|
||||
BackendRuntimeWorkerModule,
|
||||
} from './core/backend-runtime';
|
||||
import { CommentModule } from './core/comment';
|
||||
import { ServerConfigModule, ServerConfigResolverModule } from './core/config';
|
||||
import { DocStorageModule } from './core/doc';
|
||||
import { DocJobsModule } from './core/doc-jobs';
|
||||
import { DocRendererModule } from './core/doc-renderer';
|
||||
import { DocServiceModule } from './core/doc-service';
|
||||
import { FeatureModule } from './core/features';
|
||||
import { MailModule } from './core/mail';
|
||||
import { MonitorModule } from './core/monitor';
|
||||
@@ -44,20 +48,20 @@ import { QuotaModule } from './core/quota';
|
||||
import { RealtimeModule } from './core/realtime';
|
||||
import { SelfhostModule } from './core/selfhost';
|
||||
import { StaticFileModule } from './core/static-files';
|
||||
import { StorageModule } from './core/storage';
|
||||
import { StorageApiModule, StorageWorkerModule } from './core/storage';
|
||||
import { StorageRuntimeModule } from './core/storage-runtime';
|
||||
import { SyncModule } from './core/sync';
|
||||
import { TelemetryModule } from './core/telemetry';
|
||||
import { UserModule } from './core/user';
|
||||
import { VersionModule } from './core/version';
|
||||
import { WorkspaceModule } from './core/workspaces';
|
||||
import { Env } from './env';
|
||||
import { Env, ServerRole } from './env';
|
||||
import { ModelsModule } from './models';
|
||||
import { CalendarModule } from './plugins/calendar';
|
||||
import { CaptchaModule } from './plugins/captcha';
|
||||
import { CopilotModule } from './plugins/copilot';
|
||||
import { GCloudModule } from './plugins/gcloud';
|
||||
import { IndexerModule } from './plugins/indexer';
|
||||
import { IndexerModule, IndexerWorkerModule } from './plugins/indexer';
|
||||
import { LicenseModule } from './plugins/license';
|
||||
import { OAuthModule } from './plugins/oauth';
|
||||
import { PaymentModule } from './plugins/payment';
|
||||
@@ -120,6 +124,7 @@ export const FunctionalityModules = [
|
||||
RealtimeModule,
|
||||
ModelsModule,
|
||||
BackendRuntimeModule,
|
||||
BackendRuntimeProducerModule,
|
||||
StorageRuntimeModule,
|
||||
ScheduleModule.forRoot(),
|
||||
MonitorModule,
|
||||
@@ -157,29 +162,27 @@ export class AppModuleBuilder {
|
||||
|
||||
export function buildAppModule(env: Env) {
|
||||
const factor = new AppModuleBuilder();
|
||||
const workerOnly = env.role === ServerRole.Worker;
|
||||
|
||||
factor
|
||||
// basic
|
||||
.use(...FunctionalityModules)
|
||||
|
||||
// enable indexer module on graphql, doc and front service
|
||||
.useIf(
|
||||
() => env.flavors.graphql || env.flavors.doc || env.flavors.front,
|
||||
IndexerModule
|
||||
)
|
||||
// online roles publish indexer events; only the worker registers consumers
|
||||
.useIf(() => env.isApi || env.isFrontend, IndexerModule)
|
||||
.useIf(() => env.isWorker, IndexerWorkerModule)
|
||||
|
||||
// auth
|
||||
.use(UserModule, AuthModule, PermissionModule)
|
||||
// the worker owns doc consumers and schedulers
|
||||
.useIf(() => env.isWorker, DocJobsModule)
|
||||
.useIf(() => env.isWorker, BackendRuntimeWorkerModule)
|
||||
|
||||
// auth and business APIs are not part of the queue worker application
|
||||
.useIf(() => !workerOnly, UserModule, AuthModule, PermissionModule)
|
||||
|
||||
// business modules
|
||||
.use(
|
||||
ServerConfigModule,
|
||||
FeatureModule,
|
||||
QuotaModule,
|
||||
DocStorageModule,
|
||||
NotificationModule,
|
||||
MailModule
|
||||
)
|
||||
.use(ServerConfigModule, QuotaModule, DocStorageModule)
|
||||
.useIf(() => env.isWorker, StorageWorkerModule)
|
||||
.useIf(() => !workerOnly, FeatureModule, NotificationModule, MailModule)
|
||||
// renderer server and front server
|
||||
.useIf(() => env.flavors.renderer || env.flavors.front, DocRendererModule)
|
||||
// sync server and front server
|
||||
@@ -197,7 +200,7 @@ export function buildAppModule(env: Env) {
|
||||
() => env.flavors.graphql,
|
||||
GqlModule,
|
||||
VersionModule,
|
||||
StorageModule,
|
||||
StorageApiModule,
|
||||
ServerConfigResolverModule,
|
||||
WorkspaceModule,
|
||||
LicenseModule,
|
||||
@@ -210,10 +213,12 @@ export function buildAppModule(env: Env) {
|
||||
CommentModule,
|
||||
QueueDashboardModule
|
||||
)
|
||||
// doc service and front service
|
||||
.useIf(() => env.flavors.doc || env.flavors.front, DocServiceModule)
|
||||
// worker for and self-hosted API only for self-host and local development only
|
||||
.useIf(() => env.dev || env.selfhosted, WorkerModule, SelfhostModule)
|
||||
.useIf(
|
||||
() => !workerOnly && (env.dev || env.selfhosted),
|
||||
WorkerModule,
|
||||
SelfhostModule
|
||||
)
|
||||
// static frontend routes for front flavor
|
||||
.useIf(() => env.flavors.front, StaticFileModule)
|
||||
|
||||
|
||||
@@ -102,24 +102,6 @@ test('should be able to safe compare', t => {
|
||||
t.false(t.context.crypto.compare('abc', 'def'));
|
||||
});
|
||||
|
||||
test('should sign and parse internal access token', t => {
|
||||
const token = t.context.crypto.signInternalAccessToken({
|
||||
method: 'GET',
|
||||
path: '/rpc/workspaces/123/docs/456',
|
||||
now: 1700000000000,
|
||||
nonce: 'nonce-123',
|
||||
});
|
||||
|
||||
const payload = t.context.crypto.parseInternalAccessToken(token);
|
||||
t.deepEqual(payload, {
|
||||
v: 1,
|
||||
ts: 1700000000000,
|
||||
nonce: 'nonce-123',
|
||||
m: 'GET',
|
||||
p: '/rpc/workspaces/123/docs/456',
|
||||
});
|
||||
});
|
||||
|
||||
test('should be able to hash and verify password', async t => {
|
||||
const password = 'mySecurePassword';
|
||||
const hash = await t.context.crypto.encryptPassword(password);
|
||||
|
||||
@@ -173,67 +173,6 @@ export class CryptoHelper implements OnModuleInit {
|
||||
});
|
||||
}
|
||||
|
||||
signInternalAccessToken(input: {
|
||||
method: string;
|
||||
path: string;
|
||||
now?: number;
|
||||
nonce?: string;
|
||||
}) {
|
||||
const payload = {
|
||||
v: 1 as const,
|
||||
ts: input.now ?? Date.now(),
|
||||
nonce: input.nonce ?? this.randomBytes(16).toString('base64url'),
|
||||
m: input.method.toUpperCase(),
|
||||
p: input.path,
|
||||
};
|
||||
const data = Buffer.from(JSON.stringify(payload), 'utf8').toString(
|
||||
'base64url'
|
||||
);
|
||||
return this.sign(data);
|
||||
}
|
||||
|
||||
parseInternalAccessToken(signatureWithData: string): {
|
||||
v: 1;
|
||||
ts: number;
|
||||
nonce: string;
|
||||
m: string;
|
||||
p: string;
|
||||
} | null {
|
||||
const [data, signature] = signatureWithData.split(',');
|
||||
if (!signature) {
|
||||
return null;
|
||||
}
|
||||
if (!this.verify(signatureWithData)) {
|
||||
return null;
|
||||
}
|
||||
try {
|
||||
const json = Buffer.from(data, 'base64url').toString('utf8');
|
||||
const payload = JSON.parse(json) as unknown;
|
||||
if (!payload || typeof payload !== 'object') {
|
||||
return null;
|
||||
}
|
||||
const val = payload as {
|
||||
v?: unknown;
|
||||
ts?: unknown;
|
||||
nonce?: unknown;
|
||||
m?: unknown;
|
||||
p?: unknown;
|
||||
};
|
||||
if (
|
||||
val.v !== 1 ||
|
||||
typeof val.ts !== 'number' ||
|
||||
typeof val.nonce !== 'string' ||
|
||||
typeof val.m !== 'string' ||
|
||||
typeof val.p !== 'string'
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
return { v: 1, ts: val.ts, nonce: val.nonce, m: val.m, p: val.p };
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
encrypt(data: string) {
|
||||
const iv = this.randomBytes();
|
||||
const cipher = createCipheriv(
|
||||
|
||||
@@ -2,7 +2,7 @@ import { getQueueToken, getSharedConfigToken } from '@nestjs/bullmq';
|
||||
import { Injectable, Logger, OnModuleDestroy } from '@nestjs/common';
|
||||
import { ModuleRef } from '@nestjs/core';
|
||||
import { Job, Queue as Bullmq, Worker, WorkerOptions } from 'bullmq';
|
||||
import { difference, merge } from 'lodash-es';
|
||||
import { merge } from 'lodash-es';
|
||||
import { CLS_ID, ClsServiceManager } from 'nestjs-cls';
|
||||
|
||||
import { Config } from '../../config';
|
||||
@@ -10,7 +10,8 @@ import { OnEvent } from '../../event';
|
||||
import { metrics, wrapCallMetric } from '../../metrics';
|
||||
import { QueueRedis } from '../../redis';
|
||||
import { genRequestId } from '../../utils';
|
||||
import { JOB_SIGNAL, namespace, Queue, QUEUES } from './def';
|
||||
import { JOB_SIGNAL, namespace, Queue } from './def';
|
||||
import { queuesForRole } from './owner';
|
||||
import { JobHandlerScanner } from './scanner';
|
||||
|
||||
@Injectable()
|
||||
@@ -27,18 +28,7 @@ export class JobExecutor implements OnModuleDestroy {
|
||||
|
||||
@OnEvent('config.init')
|
||||
async onConfigInit() {
|
||||
const queues = env.flavors.graphql
|
||||
? difference(QUEUES, [Queue.DOC, Queue.INDEXER])
|
||||
: [];
|
||||
|
||||
// Enable doc/indexer queues in both doc and front service.
|
||||
if (env.flavors.doc || env.flavors.front) {
|
||||
queues.push(Queue.DOC);
|
||||
// NOTE(@fengmk2): Once the index task cannot be processed in time, it needs to be separated from the doc service and deployed independently.
|
||||
queues.push(Queue.INDEXER);
|
||||
}
|
||||
|
||||
await this.startWorkers(queues);
|
||||
await this.startWorkers(queuesForRole(env.role));
|
||||
}
|
||||
|
||||
@OnEvent('config.changed')
|
||||
|
||||
@@ -55,3 +55,4 @@ export class JobModule {
|
||||
|
||||
export { JobQueue };
|
||||
export { JOB_SIGNAL, OnJob } from './def';
|
||||
export { queuesForRole, WORKER_QUEUES } from './owner';
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
import { ServerRole } from '../../../env';
|
||||
import { Queue, QUEUES } from './def';
|
||||
|
||||
export const WORKER_QUEUES = [
|
||||
Queue.DOC,
|
||||
Queue.INDEXER,
|
||||
Queue.BACKENDRUNTIME,
|
||||
] as const;
|
||||
|
||||
export function queuesForRole(role: ServerRole | undefined): Queue[] {
|
||||
switch (role) {
|
||||
case ServerRole.AllInOne:
|
||||
return [...QUEUES];
|
||||
case ServerRole.Api:
|
||||
return QUEUES.filter(
|
||||
queue => !(WORKER_QUEUES as readonly Queue[]).includes(queue)
|
||||
);
|
||||
case ServerRole.Worker:
|
||||
return [...WORKER_QUEUES];
|
||||
case ServerRole.Frontend:
|
||||
case undefined:
|
||||
return [];
|
||||
}
|
||||
}
|
||||
@@ -15,13 +15,98 @@ import {
|
||||
import type { Request, Response } from 'express';
|
||||
|
||||
import { Config } from '../config';
|
||||
import { CacheRedis } from '../redis';
|
||||
import { getRequestResponseFromContext } from '../utils/request';
|
||||
import { getRequestTrackerId } from '../utils/request-tracker';
|
||||
import type { ThrottlerType } from './config';
|
||||
import { THROTTLER_PROTECTED, Throttlers } from './decorators';
|
||||
|
||||
const REDIS_THROTTLE_SCRIPT = `
|
||||
local now = redis.call("TIME")
|
||||
local nowMs = now[1] * 1000 + math.floor(now[2] / 1000)
|
||||
local blockedUntil = tonumber(redis.call("HGET", KEYS[1], "blockedUntil")) or 0
|
||||
|
||||
if blockedUntil > nowMs then
|
||||
return {
|
||||
tonumber(redis.call("HGET", KEYS[1], "hits")) or 0,
|
||||
redis.call("PTTL", KEYS[1]),
|
||||
blockedUntil - nowMs
|
||||
}
|
||||
end
|
||||
|
||||
if blockedUntil > 0 then
|
||||
redis.call("HDEL", KEYS[1], "blockedUntil")
|
||||
redis.call("HSET", KEYS[1], "hits", 0)
|
||||
end
|
||||
|
||||
local hits = redis.call("HINCRBY", KEYS[1], "hits", 1)
|
||||
if hits == 1 then
|
||||
redis.call("PEXPIRE", KEYS[1], ARGV[1])
|
||||
end
|
||||
|
||||
local blockTtl = 0
|
||||
if hits > tonumber(ARGV[2]) then
|
||||
blockedUntil = nowMs + tonumber(ARGV[3])
|
||||
redis.call("HSET", KEYS[1], "blockedUntil", blockedUntil)
|
||||
if redis.call("PTTL", KEYS[1]) < tonumber(ARGV[3]) then
|
||||
redis.call("PEXPIRE", KEYS[1], ARGV[3])
|
||||
end
|
||||
blockTtl = tonumber(ARGV[3])
|
||||
end
|
||||
|
||||
return { hits, redis.call("PTTL", KEYS[1]), blockTtl }
|
||||
`;
|
||||
|
||||
@Injectable()
|
||||
export class ThrottlerStorage extends ThrottlerStorageService {}
|
||||
export class ThrottlerStorage extends ThrottlerStorageService {
|
||||
constructor(private readonly redis: CacheRedis) {
|
||||
super();
|
||||
}
|
||||
|
||||
override async increment(
|
||||
key: string,
|
||||
ttl: number,
|
||||
limit: number,
|
||||
blockDuration: number,
|
||||
throttlerName: string
|
||||
) {
|
||||
if (env.testing) {
|
||||
return super.increment(key, ttl, limit, blockDuration, throttlerName);
|
||||
}
|
||||
|
||||
try {
|
||||
const result = await this.redis.eval(
|
||||
REDIS_THROTTLE_SCRIPT,
|
||||
1,
|
||||
key,
|
||||
ttl,
|
||||
limit,
|
||||
Math.max(blockDuration, 1)
|
||||
);
|
||||
if (!Array.isArray(result) || result.length !== 3) {
|
||||
throw new Error('Unexpected Redis throttler response');
|
||||
}
|
||||
|
||||
const totalHits = Number(result[0]);
|
||||
const timeToExpire = Math.max(0, Math.ceil(Number(result[1]) / 1000));
|
||||
const timeToBlockExpire = Math.max(
|
||||
0,
|
||||
Math.ceil(Number(result[2]) / 1000)
|
||||
);
|
||||
|
||||
return {
|
||||
totalHits,
|
||||
timeToExpire,
|
||||
isBlocked: timeToBlockExpire > 0,
|
||||
timeToBlockExpire,
|
||||
};
|
||||
} catch {
|
||||
// Preserve availability if Redis is unavailable. The inherited local
|
||||
// storage still protects each process while the shared limiter recovers.
|
||||
return super.increment(key, ttl, limit, blockDuration, throttlerName);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
class CustomOptionsFactory implements ThrottlerOptionsFactory {
|
||||
|
||||
@@ -55,6 +55,17 @@ function buildProgram(logger: Logger) {
|
||||
});
|
||||
});
|
||||
|
||||
program
|
||||
.command('admit-legacy-context-blobs')
|
||||
.description(
|
||||
'Admit legacy context blobs before the cleanup schema migration'
|
||||
)
|
||||
.action(async () => {
|
||||
await withCliApp(logger, async app => {
|
||||
await app.get(RunCommand).admitLegacyContextBlobs();
|
||||
});
|
||||
});
|
||||
|
||||
program
|
||||
.command('revert [name]')
|
||||
.description('Revert one data migration with given name')
|
||||
|
||||
@@ -20,6 +20,11 @@ export interface AuthConfig {
|
||||
requireEmailVerification: boolean;
|
||||
newAccountShareActionDelay: number;
|
||||
trustedCloudflareHeaders: boolean;
|
||||
signInRateLimit: ConfigItem<{
|
||||
ttl: number;
|
||||
ipLimit: number;
|
||||
emailLimit: number;
|
||||
}>;
|
||||
inviteQuotaShadowMode: boolean;
|
||||
inviteQuotaFailOpenOnRuntimeError: boolean;
|
||||
passwordRequirements: ConfigItem<{
|
||||
@@ -61,6 +66,21 @@ defineModuleConfig('auth', {
|
||||
default: false,
|
||||
shape: z.boolean(),
|
||||
},
|
||||
signInRateLimit: {
|
||||
desc: 'Limits for sign-in attempts shared through Redis by source IP and email. ttl is measured in milliseconds.',
|
||||
default: {
|
||||
ttl: 60_000,
|
||||
ipLimit: 20,
|
||||
emailLimit: 5,
|
||||
},
|
||||
shape: z
|
||||
.object({
|
||||
ttl: z.number().int().positive(),
|
||||
ipLimit: z.number().int().positive(),
|
||||
emailLimit: z.number().int().positive(),
|
||||
})
|
||||
.strict(),
|
||||
},
|
||||
inviteQuotaShadowMode: {
|
||||
desc: 'Whether workspace invite quota should record would-block decisions without rejecting requests or executing abuse actions.',
|
||||
default: false,
|
||||
|
||||
@@ -118,7 +118,7 @@ export class AuthController {
|
||||
) {
|
||||
const credential = SignInBodySchema.parse(body);
|
||||
validators.assertValidEmail(credential.email);
|
||||
const canSignIn = await this.auth.canSignIn(credential.email);
|
||||
const canSignIn = await this.auth.canSignIn(credential.email, req);
|
||||
if (!canSignIn) {
|
||||
throw new ActionForbidden();
|
||||
}
|
||||
|
||||
@@ -11,12 +11,9 @@ import semver from 'semver';
|
||||
import { Socket } from 'socket.io';
|
||||
|
||||
import {
|
||||
AccessDenied,
|
||||
AuthenticationRequired,
|
||||
Cache,
|
||||
checkCanaryDateClientVersion,
|
||||
Config,
|
||||
CryptoHelper,
|
||||
getClientVersionFromRequest,
|
||||
getRequestResponseFromContext,
|
||||
parseCookies,
|
||||
@@ -32,9 +29,6 @@ import { AuthSessionHttpError } from './session-exchange';
|
||||
import { isLikelyJwt } from './token';
|
||||
|
||||
const PUBLIC_ENTRYPOINT_SYMBOL = Symbol('public');
|
||||
const INTERNAL_ENTRYPOINT_SYMBOL = Symbol('internal');
|
||||
const INTERNAL_ACCESS_TOKEN_TTL_MS = 5 * 60 * 1000;
|
||||
const INTERNAL_ACCESS_TOKEN_CLOCK_SKEW_MS = 30 * 1000;
|
||||
|
||||
type AuthenticatedRequestSession =
|
||||
| { type: 'jwt'; session: Session }
|
||||
@@ -50,8 +44,6 @@ export class AuthGuard implements CanActivate, OnModuleInit {
|
||||
private static readonly CANARY_REQUIRED_VERSION = 'canary (within 2 months)';
|
||||
|
||||
constructor(
|
||||
private readonly crypto: CryptoHelper,
|
||||
private readonly cache: Cache,
|
||||
private readonly config: Config,
|
||||
private readonly ref: ModuleRef,
|
||||
private readonly reflector: Reflector
|
||||
@@ -67,38 +59,6 @@ export class AuthGuard implements CanActivate, OnModuleInit {
|
||||
const { req, res } = getRequestResponseFromContext(context);
|
||||
const clazz = context.getClass();
|
||||
const handler = context.getHandler();
|
||||
// rpc request is internal
|
||||
const isInternal = this.reflector.getAllAndOverride<boolean>(
|
||||
INTERNAL_ENTRYPOINT_SYMBOL,
|
||||
[clazz, handler]
|
||||
);
|
||||
if (isInternal) {
|
||||
const accessToken = req.get('x-access-token');
|
||||
if (accessToken) {
|
||||
const payload = this.crypto.parseInternalAccessToken(accessToken);
|
||||
if (payload) {
|
||||
const now = Date.now();
|
||||
const method = req.method.toUpperCase();
|
||||
const path = req.path;
|
||||
|
||||
const timestampInRange =
|
||||
payload.ts <= now + INTERNAL_ACCESS_TOKEN_CLOCK_SKEW_MS &&
|
||||
now - payload.ts <= INTERNAL_ACCESS_TOKEN_TTL_MS;
|
||||
|
||||
if (timestampInRange && payload.m === method && payload.p === path) {
|
||||
const nonceKey = `rpc:nonce:${payload.nonce}`;
|
||||
const ok = await this.cache.setnx(nonceKey, 1, {
|
||||
ttl: INTERNAL_ACCESS_TOKEN_TTL_MS,
|
||||
});
|
||||
if (ok) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
throw new AccessDenied('Invalid internal request');
|
||||
}
|
||||
|
||||
// api is public
|
||||
const isPublic = this.reflector.getAllAndOverride<boolean>(
|
||||
PUBLIC_ENTRYPOINT_SYMBOL,
|
||||
@@ -327,11 +287,6 @@ export class AuthGuard implements CanActivate, OnModuleInit {
|
||||
*/
|
||||
export const Public = () => SetMetadata(PUBLIC_ENTRYPOINT_SYMBOL, true);
|
||||
|
||||
/**
|
||||
* Mark rpc api to be internal accessible
|
||||
*/
|
||||
export const Internal = () => SetMetadata(INTERNAL_ENTRYPOINT_SYMBOL, true);
|
||||
|
||||
export const AuthWebsocketOptionsProvider: FactoryProvider = {
|
||||
provide: WEBSOCKET_OPTIONS,
|
||||
useFactory: (config: Config, guard: AuthGuard) => {
|
||||
|
||||
@@ -1,11 +1,18 @@
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { createHash, randomUUID } from 'node:crypto';
|
||||
|
||||
import { Injectable, OnApplicationBootstrap } from '@nestjs/common';
|
||||
import { Transactional } from '@nestjs-cls/transactional';
|
||||
import type { CookieOptions, Request, Response } from 'express';
|
||||
import { assign, pick } from 'lodash-es';
|
||||
|
||||
import { Config, OnEvent, SignUpForbidden } from '../../base';
|
||||
import {
|
||||
Cache,
|
||||
Config,
|
||||
getRequestClientIp,
|
||||
OnEvent,
|
||||
SignUpForbidden,
|
||||
TooManyRequest,
|
||||
} from '../../base';
|
||||
import { Models, type User, type UserSession } from '../../models';
|
||||
import { EntitlementService } from '../entitlement';
|
||||
import { Mailer } from '../mail/mailer';
|
||||
@@ -46,7 +53,8 @@ export class AuthService implements OnApplicationBootstrap {
|
||||
private readonly models: Models,
|
||||
private readonly mailer: Mailer,
|
||||
private readonly authSessions: AuthSessionService,
|
||||
private readonly entitlement: EntitlementService
|
||||
private readonly entitlement: EntitlementService,
|
||||
private readonly cache: Cache
|
||||
) {
|
||||
this.cookieOptions = {
|
||||
sameSite: 'lax',
|
||||
@@ -69,11 +77,38 @@ export class AuthService implements OnApplicationBootstrap {
|
||||
}
|
||||
}
|
||||
|
||||
async canSignIn(_email: string) {
|
||||
async canSignIn(email: string, req: Request) {
|
||||
if (!env.testing) {
|
||||
const { ttl, ipLimit, emailLimit } = this.config.auth.signInRateLimit;
|
||||
const normalizedEmail = email.toLowerCase();
|
||||
const ip = getRequestClientIp(req);
|
||||
|
||||
const emailAttempts = this.cache.increaseWithTtl(
|
||||
this.signInRateLimitKey('email', normalizedEmail),
|
||||
ttl
|
||||
);
|
||||
const ipAttempts = ip
|
||||
? this.cache.increaseWithTtl(this.signInRateLimitKey('ip', ip), ttl)
|
||||
: Promise.resolve(0);
|
||||
const [emailCount, ipCount] = await Promise.all([
|
||||
emailAttempts,
|
||||
ipAttempts,
|
||||
]);
|
||||
|
||||
if (emailCount > emailLimit || ipCount > ipLimit) {
|
||||
throw new TooManyRequest();
|
||||
}
|
||||
}
|
||||
|
||||
// may add more sign-in check later
|
||||
return true;
|
||||
}
|
||||
|
||||
private signInRateLimitKey(scope: 'email' | 'ip', value: string) {
|
||||
const digest = createHash('sha256').update(value).digest('hex');
|
||||
return `auth:sign-in-rate:${scope}:${digest}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated
|
||||
*
|
||||
|
||||
@@ -16,15 +16,24 @@ import {
|
||||
CopilotSelectedSourcesUnavailable,
|
||||
} from '../../../base';
|
||||
import { Models } from '../../../models';
|
||||
import { BackendRuntimeModule, BackendRuntimeProvider } from '../index';
|
||||
import {
|
||||
BackendRuntimeModule,
|
||||
BackendRuntimeProducerModule,
|
||||
BackendRuntimeProvider,
|
||||
BackendRuntimeWorkerModule,
|
||||
} from '../index';
|
||||
import {
|
||||
BackendRuntimeEmbeddingJob,
|
||||
BackendRuntimeEmbeddingProducer,
|
||||
BackendRuntimeEmbeddingService,
|
||||
BackendRuntimeHousekeepingJob,
|
||||
} from '../job';
|
||||
|
||||
interface Context {
|
||||
module: TestingModule;
|
||||
embeddingJob: BackendRuntimeEmbeddingJob;
|
||||
embeddingProducer: BackendRuntimeEmbeddingProducer;
|
||||
embeddingService: BackendRuntimeEmbeddingService;
|
||||
job: BackendRuntimeHousekeepingJob;
|
||||
getSnapshot: Sinon.SinonStub;
|
||||
allowEmbedding: Sinon.SinonStub;
|
||||
@@ -55,7 +64,12 @@ test.before(async t => {
|
||||
syncEmbeddingState: Sinon.stub(),
|
||||
};
|
||||
t.context.module = await createTestingModule({
|
||||
imports: [ScheduleModule.forRoot(), BackendRuntimeModule],
|
||||
imports: [
|
||||
ScheduleModule.forRoot(),
|
||||
BackendRuntimeModule,
|
||||
BackendRuntimeProducerModule,
|
||||
BackendRuntimeWorkerModule,
|
||||
],
|
||||
tapModule: builder => {
|
||||
builder
|
||||
.overrideProvider(BackendRuntimeProvider)
|
||||
@@ -81,6 +95,12 @@ test.before(async t => {
|
||||
'allowEmbedding'
|
||||
).resolves(true);
|
||||
t.context.embeddingJob = t.context.module.get(BackendRuntimeEmbeddingJob);
|
||||
t.context.embeddingProducer = t.context.module.get(
|
||||
BackendRuntimeEmbeddingProducer
|
||||
);
|
||||
t.context.embeddingService = t.context.module.get(
|
||||
BackendRuntimeEmbeddingService
|
||||
);
|
||||
t.context.job = t.context.module.get(BackendRuntimeHousekeepingJob);
|
||||
});
|
||||
|
||||
@@ -102,7 +122,7 @@ test.after.always(async t => {
|
||||
});
|
||||
|
||||
test('backend-runtime jobs ingest documents and clean runtime state', async t => {
|
||||
await t.context.embeddingJob.onDocSnapshotUpdated({
|
||||
await t.context.embeddingProducer.onDocSnapshotUpdated({
|
||||
workspaceId: 'workspace-1',
|
||||
docId: 'doc-1',
|
||||
blob: Buffer.alloc(0),
|
||||
@@ -130,7 +150,7 @@ test('backend-runtime jobs ingest documents and clean runtime state', async t =>
|
||||
const documentJobCount = t.context.module.queue.count(
|
||||
'backendRuntime.syncDocumentEmbedding'
|
||||
);
|
||||
await t.context.embeddingJob.onDocSnapshotUpdated({
|
||||
await t.context.embeddingProducer.onDocSnapshotUpdated({
|
||||
workspaceId: 'workspace-1',
|
||||
docId: 'db$docProperties',
|
||||
blob: Buffer.alloc(0),
|
||||
@@ -140,7 +160,7 @@ test('backend-runtime jobs ingest documents and clean runtime state', async t =>
|
||||
documentJobCount
|
||||
);
|
||||
|
||||
await t.context.embeddingJob.onDocSnapshotUpdated({
|
||||
await t.context.embeddingProducer.onDocSnapshotUpdated({
|
||||
workspaceId: 'workspace-1',
|
||||
docId: 'workspace-1',
|
||||
blob: Buffer.alloc(0),
|
||||
@@ -155,7 +175,7 @@ test('backend-runtime jobs ingest documents and clean runtime state', async t =>
|
||||
reconcileDocuments: true,
|
||||
});
|
||||
|
||||
await t.context.embeddingJob.prepareSelectedDocuments('workspace-1', [
|
||||
await t.context.embeddingService.prepareSelectedDocuments('workspace-1', [
|
||||
'doc-1',
|
||||
'doc-1',
|
||||
]);
|
||||
@@ -181,7 +201,9 @@ test('backend-runtime jobs ingest documents and clean runtime state', async t =>
|
||||
] as const) {
|
||||
t.context.runtime.syncEmbeddingState.rejects(new Error(nativeError));
|
||||
const error = await t.throwsAsync(() =>
|
||||
t.context.embeddingJob.prepareSelectedDocuments('workspace-1', ['doc-1'])
|
||||
t.context.embeddingService.prepareSelectedDocuments('workspace-1', [
|
||||
'doc-1',
|
||||
])
|
||||
);
|
||||
t.true(error instanceof expectedError);
|
||||
}
|
||||
@@ -189,7 +211,7 @@ test('backend-runtime jobs ingest documents and clean runtime state', async t =>
|
||||
|
||||
await t.throwsAsync(
|
||||
() =>
|
||||
t.context.embeddingJob.prepareSelectedDocuments(
|
||||
t.context.embeddingService.prepareSelectedDocuments(
|
||||
'workspace-1',
|
||||
Array.from({ length: 65 }, (_, index) => `doc-${index}`)
|
||||
),
|
||||
@@ -198,7 +220,7 @@ test('backend-runtime jobs ingest documents and clean runtime state', async t =>
|
||||
t.context.getSnapshot.resolves(null);
|
||||
await t.throwsAsync(
|
||||
() =>
|
||||
t.context.embeddingJob.prepareSelectedDocuments('workspace-1', [
|
||||
t.context.embeddingService.prepareSelectedDocuments('workspace-1', [
|
||||
'missing-doc',
|
||||
]),
|
||||
{ instanceOf: CopilotSelectedSourcesUnavailable }
|
||||
|
||||
@@ -11,7 +11,7 @@ const privateKey = generateKeyPairSync('ec', {
|
||||
}).privateKey.export({ format: 'pem', type: 'pkcs8' }) as string;
|
||||
const config = { crypto: { privateKey } } as Config;
|
||||
|
||||
test('backend-runtime provider starts once, runs migrations once, and reports health', async t => {
|
||||
test('backend-runtime provider starts without migrations and exposes explicit migration', async t => {
|
||||
const provider = new BackendRuntimeProvider(config);
|
||||
const runtime = {
|
||||
start: Sinon.stub().resolves(),
|
||||
@@ -27,6 +27,7 @@ test('backend-runtime provider starts once, runs migrations once, and reports he
|
||||
|
||||
await provider.start();
|
||||
await provider.start();
|
||||
await provider.runMigrations();
|
||||
await provider.onConfigChanged({ updates: { mailer: {} } });
|
||||
await provider.onConfigChanged({ updates: { copilot: {} } });
|
||||
await provider.onConfigChanged({ updates: { storages: {} } });
|
||||
@@ -66,6 +67,106 @@ test('backend-runtime provider measures explicit typed methods', async t => {
|
||||
t.true(runtime.assertCopilotRoute.calledOnceWithExactly(routeInput));
|
||||
});
|
||||
|
||||
test('backend-runtime provider encodes recursive search contracts at the native boundary', async t => {
|
||||
const provider = new BackendRuntimeProvider(config);
|
||||
const runtime = {
|
||||
searchAuthorized: Sinon.stub().resolves({
|
||||
ok: true,
|
||||
value: { total: 0, nodes: [] },
|
||||
}),
|
||||
aggregateAuthorized: Sinon.stub().resolves({
|
||||
ok: true,
|
||||
value: { total: 0, buckets: [] },
|
||||
}),
|
||||
};
|
||||
(provider as unknown as { runtime: typeof runtime }).runtime = runtime;
|
||||
const query = {
|
||||
type: 'boolean',
|
||||
occur: 'must',
|
||||
queries: [
|
||||
{ type: 'exists', field: 'refDocId' },
|
||||
{
|
||||
type: 'boost',
|
||||
boost: 1.5,
|
||||
query: { type: 'match', field: 'content', match: 'hello' },
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
await provider.searchAuthorized('actor', 'workspace', {
|
||||
table: 'block',
|
||||
query,
|
||||
options: {
|
||||
fields: ['docId'],
|
||||
highlights: [{ field: 'content', before: '<b>', end: '</b>' }],
|
||||
pagination: { limit: 10, cursor: 'cursor' },
|
||||
},
|
||||
});
|
||||
await provider.aggregateAuthorized('actor', 'workspace', {
|
||||
table: 'block',
|
||||
query,
|
||||
field: 'docId',
|
||||
options: {
|
||||
hits: { fields: ['content'] },
|
||||
pagination: { limit: 5, skip: 2 },
|
||||
},
|
||||
});
|
||||
|
||||
const search = runtime.searchAuthorized.firstCall.args[2];
|
||||
t.is(search.rootQuery, 0);
|
||||
t.deepEqual(search.queries, [
|
||||
{
|
||||
queryType: 'boolean',
|
||||
field: undefined,
|
||||
matchValue: undefined,
|
||||
query: undefined,
|
||||
queries: [1, 2],
|
||||
occur: 'must',
|
||||
boost: undefined,
|
||||
},
|
||||
{
|
||||
queryType: 'exists',
|
||||
field: 'refDocId',
|
||||
matchValue: undefined,
|
||||
query: undefined,
|
||||
queries: undefined,
|
||||
occur: undefined,
|
||||
boost: undefined,
|
||||
},
|
||||
{
|
||||
queryType: 'boost',
|
||||
field: undefined,
|
||||
matchValue: undefined,
|
||||
query: 3,
|
||||
queries: undefined,
|
||||
occur: undefined,
|
||||
boost: 1.5,
|
||||
},
|
||||
{
|
||||
queryType: 'match',
|
||||
field: 'content',
|
||||
matchValue: 'hello',
|
||||
query: undefined,
|
||||
queries: undefined,
|
||||
occur: undefined,
|
||||
boost: undefined,
|
||||
},
|
||||
]);
|
||||
t.deepEqual(search.options, {
|
||||
fields: ['docId'],
|
||||
highlights: [{ field: 'content', before: '<b>', end: '</b>' }],
|
||||
pagination: { limit: 10, cursor: 'cursor' },
|
||||
});
|
||||
t.deepEqual(runtime.aggregateAuthorized.firstCall.args[2].options, {
|
||||
hits: { fields: ['content'], highlights: [], pagination: {} },
|
||||
pagination: { limit: 5, skip: 2 },
|
||||
});
|
||||
t.true(runtime.searchAuthorized.calledOnce);
|
||||
t.true(runtime.searchAuthorized.calledWithMatch('actor', 'workspace'));
|
||||
t.true(runtime.aggregateAuthorized.calledOnce);
|
||||
t.true(runtime.aggregateAuthorized.calledWithMatch('actor', 'workspace'));
|
||||
});
|
||||
|
||||
test('backend-runtime provider aborts a stream handle that resolves after iterator cancellation', async t => {
|
||||
const provider = new BackendRuntimeProvider(config);
|
||||
const abort = Sinon.stub();
|
||||
|
||||
@@ -2,6 +2,8 @@ import { Global, Module } from '@nestjs/common';
|
||||
|
||||
import {
|
||||
BackendRuntimeEmbeddingJob,
|
||||
BackendRuntimeEmbeddingProducer,
|
||||
BackendRuntimeEmbeddingService,
|
||||
BackendRuntimeHousekeepingJob,
|
||||
} from './job';
|
||||
import {
|
||||
@@ -17,14 +19,30 @@ import {
|
||||
useValue: undefined,
|
||||
},
|
||||
BackendRuntimeProvider,
|
||||
BackendRuntimeEmbeddingJob,
|
||||
BackendRuntimeHousekeepingJob,
|
||||
BackendRuntimeEmbeddingService,
|
||||
],
|
||||
exports: [BackendRuntimeProvider, BackendRuntimeEmbeddingJob],
|
||||
exports: [BackendRuntimeProvider, BackendRuntimeEmbeddingService],
|
||||
})
|
||||
export class BackendRuntimeModule {}
|
||||
|
||||
export { BackendRuntimeEmbeddingJob } from './job';
|
||||
@Module({
|
||||
imports: [BackendRuntimeModule],
|
||||
providers: [BackendRuntimeEmbeddingProducer],
|
||||
})
|
||||
export class BackendRuntimeProducerModule {}
|
||||
|
||||
@Module({
|
||||
imports: [BackendRuntimeModule],
|
||||
providers: [BackendRuntimeEmbeddingJob, BackendRuntimeHousekeepingJob],
|
||||
})
|
||||
export class BackendRuntimeWorkerModule {}
|
||||
|
||||
export {
|
||||
BackendRuntimeEmbeddingJob,
|
||||
BackendRuntimeEmbeddingProducer,
|
||||
BackendRuntimeEmbeddingService,
|
||||
BackendRuntimeHousekeepingJob,
|
||||
} from './job';
|
||||
export {
|
||||
BACKEND_RUNTIME_CONFIG_PATHS,
|
||||
BackendRuntimeProvider,
|
||||
|
||||
@@ -22,7 +22,7 @@ const SELECTED_DOCUMENT_WAIT_MS = 90_000;
|
||||
|
||||
declare global {
|
||||
interface Jobs {
|
||||
'nightly.cleanExpiredBackendRuntimeHousekeeping': {};
|
||||
'backendRuntime.cleanExpiredHousekeeping': {};
|
||||
'backendRuntime.syncDocumentEmbedding': {
|
||||
workspaceId: string;
|
||||
docId: string;
|
||||
@@ -34,54 +34,13 @@ declare global {
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class BackendRuntimeEmbeddingJob {
|
||||
export class BackendRuntimeEmbeddingService {
|
||||
constructor(
|
||||
private readonly rt: BackendRuntimeProvider,
|
||||
private readonly queue: JobQueue,
|
||||
private readonly models: Models
|
||||
) {}
|
||||
|
||||
@OnEvent('doc.updated')
|
||||
async onDocUpdated({ workspaceId, docId }: Events['doc.updated']) {
|
||||
await this.queueDocument(workspaceId, docId);
|
||||
}
|
||||
|
||||
@OnEvent('doc.snapshot.updated')
|
||||
async onDocSnapshotUpdated({
|
||||
workspaceId,
|
||||
docId,
|
||||
}: Events['doc.snapshot.updated']) {
|
||||
if (workspaceId === docId) {
|
||||
await this.queue.add(
|
||||
'backendRuntime.reconcileDocumentEmbeddings',
|
||||
{ workspaceId },
|
||||
{ jobId: `reconcileDocumentEmbeddings/${workspaceId}` }
|
||||
);
|
||||
return;
|
||||
}
|
||||
await this.queueDocument(workspaceId, docId);
|
||||
}
|
||||
|
||||
private async queueDocument(workspaceId: string, docId: string) {
|
||||
if (
|
||||
workspaceId === docId ||
|
||||
docId.startsWith('db$') ||
|
||||
docId.startsWith('userdata$')
|
||||
) {
|
||||
return;
|
||||
}
|
||||
await this.queue.add(
|
||||
'backendRuntime.syncDocumentEmbedding',
|
||||
{ workspaceId, docId },
|
||||
{ jobId: `syncDocumentEmbedding/${workspaceId}/${docId}` }
|
||||
);
|
||||
}
|
||||
|
||||
@OnJob('backendRuntime.syncDocumentEmbedding')
|
||||
async syncDocument({
|
||||
workspaceId,
|
||||
docId,
|
||||
}: Jobs['backendRuntime.syncDocumentEmbedding']) {
|
||||
async syncDocument(workspaceId: string, docId: string) {
|
||||
await this.syncDocuments(workspaceId, [docId], true);
|
||||
}
|
||||
|
||||
@@ -174,10 +133,7 @@ export class BackendRuntimeEmbeddingJob {
|
||||
});
|
||||
}
|
||||
|
||||
@OnJob('backendRuntime.reconcileDocumentEmbeddings')
|
||||
async reconcileDocuments({
|
||||
workspaceId,
|
||||
}: Jobs['backendRuntime.reconcileDocumentEmbeddings']) {
|
||||
async reconcileDocuments(workspaceId: string) {
|
||||
if (!(await this.rt.embeddingHealth()).enabled) return;
|
||||
await this.rt.syncEmbeddingState({
|
||||
workspaceId,
|
||||
@@ -187,6 +143,67 @@ export class BackendRuntimeEmbeddingJob {
|
||||
}
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class BackendRuntimeEmbeddingProducer {
|
||||
constructor(private readonly queue: JobQueue) {}
|
||||
|
||||
@OnEvent('doc.updated')
|
||||
async onDocUpdated({ workspaceId, docId }: Events['doc.updated']) {
|
||||
await this.queueDocument(workspaceId, docId);
|
||||
}
|
||||
|
||||
@OnEvent('doc.snapshot.updated')
|
||||
async onDocSnapshotUpdated({
|
||||
workspaceId,
|
||||
docId,
|
||||
}: Events['doc.snapshot.updated']) {
|
||||
if (workspaceId === docId) {
|
||||
await this.queue.add(
|
||||
'backendRuntime.reconcileDocumentEmbeddings',
|
||||
{ workspaceId },
|
||||
{ jobId: `reconcileDocumentEmbeddings/${workspaceId}` }
|
||||
);
|
||||
return;
|
||||
}
|
||||
await this.queueDocument(workspaceId, docId);
|
||||
}
|
||||
|
||||
private async queueDocument(workspaceId: string, docId: string) {
|
||||
if (
|
||||
workspaceId === docId ||
|
||||
docId.startsWith('db$') ||
|
||||
docId.startsWith('userdata$')
|
||||
) {
|
||||
return;
|
||||
}
|
||||
await this.queue.add(
|
||||
'backendRuntime.syncDocumentEmbedding',
|
||||
{ workspaceId, docId },
|
||||
{ jobId: `syncDocumentEmbedding/${workspaceId}/${docId}` }
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class BackendRuntimeEmbeddingJob {
|
||||
constructor(private readonly service: BackendRuntimeEmbeddingService) {}
|
||||
|
||||
@OnJob('backendRuntime.syncDocumentEmbedding')
|
||||
async syncDocument({
|
||||
workspaceId,
|
||||
docId,
|
||||
}: Jobs['backendRuntime.syncDocumentEmbedding']) {
|
||||
await this.service.syncDocument(workspaceId, docId);
|
||||
}
|
||||
|
||||
@OnJob('backendRuntime.reconcileDocumentEmbeddings')
|
||||
async reconcileDocuments({
|
||||
workspaceId,
|
||||
}: Jobs['backendRuntime.reconcileDocumentEmbeddings']) {
|
||||
await this.service.reconcileDocuments(workspaceId);
|
||||
}
|
||||
}
|
||||
|
||||
@Injectable()
|
||||
export class BackendRuntimeHousekeepingJob {
|
||||
private readonly logger = new Logger(BackendRuntimeHousekeepingJob.name);
|
||||
@@ -199,7 +216,7 @@ export class BackendRuntimeHousekeepingJob {
|
||||
@Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT)
|
||||
async nightlyJob() {
|
||||
await this.queue.add(
|
||||
'nightly.cleanExpiredBackendRuntimeHousekeeping',
|
||||
'backendRuntime.cleanExpiredHousekeeping',
|
||||
{},
|
||||
{
|
||||
jobId: 'nightly-backend-runtime-housekeeping',
|
||||
@@ -207,7 +224,7 @@ export class BackendRuntimeHousekeepingJob {
|
||||
);
|
||||
}
|
||||
|
||||
@OnJob('nightly.cleanExpiredBackendRuntimeHousekeeping')
|
||||
@OnJob('backendRuntime.cleanExpiredHousekeeping')
|
||||
async cleanExpiredRuntimeHousekeeping() {
|
||||
const states = await this.cleanBatches(() =>
|
||||
this.rt.cleanupExpiredRuntimeStates(1000)
|
||||
|
||||
@@ -35,6 +35,12 @@ import {
|
||||
type RuntimeWorkspaceArtifact,
|
||||
type SyncEmbeddingStateInput,
|
||||
} from '../../native';
|
||||
import {
|
||||
type AggregateRequestInput,
|
||||
encodeAggregateRequest,
|
||||
encodeSearchRequest,
|
||||
type SearchRequestInput,
|
||||
} from './search';
|
||||
|
||||
type RuntimeInstance = InstanceType<typeof BackendRuntime>;
|
||||
|
||||
@@ -299,11 +305,18 @@ export class BackendRuntimeProvider
|
||||
|
||||
async start() {
|
||||
await this.runtime.start();
|
||||
await this.runMigrationsOnce();
|
||||
const health = await this.runtime.health();
|
||||
this.logger.log(`backend runtime started: db=${health.databaseConnected}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Schema changes belong to the explicit predeploy path. Runtime startup only
|
||||
* connects services and must not mutate the database schema.
|
||||
*/
|
||||
async runMigrations() {
|
||||
await this.runMigrationsOnce();
|
||||
}
|
||||
|
||||
async stop() {
|
||||
await this.runtime.stop();
|
||||
this.logger.log('backend runtime stopped');
|
||||
@@ -315,6 +328,7 @@ export class BackendRuntimeProvider
|
||||
!updates.copilot &&
|
||||
!updates.crypto &&
|
||||
!updates.db &&
|
||||
!updates.indexer &&
|
||||
!updates.storages
|
||||
) {
|
||||
return;
|
||||
@@ -332,6 +346,74 @@ export class BackendRuntimeProvider
|
||||
);
|
||||
}
|
||||
|
||||
async searchAuthorized(
|
||||
actorUserId: string,
|
||||
workspaceId: string,
|
||||
request: SearchRequestInput
|
||||
) {
|
||||
return await this.measured('searchAuthorized', runtime =>
|
||||
runtime.searchAuthorized(
|
||||
actorUserId,
|
||||
workspaceId,
|
||||
encodeSearchRequest(request)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
async aggregateAuthorized(
|
||||
actorUserId: string,
|
||||
workspaceId: string,
|
||||
request: AggregateRequestInput
|
||||
) {
|
||||
return await this.measured('aggregateAuthorized', runtime =>
|
||||
runtime.aggregateAuthorized(
|
||||
actorUserId,
|
||||
workspaceId,
|
||||
encodeAggregateRequest(request)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
async indexSearchDocument(workspaceId: string, docId: string) {
|
||||
await this.measured('indexSearchDocument', runtime =>
|
||||
runtime.indexSearchDocument(workspaceId, docId)
|
||||
);
|
||||
}
|
||||
|
||||
async deleteSearchDocument(workspaceId: string, docId: string) {
|
||||
await this.measured('deleteSearchDocument', runtime =>
|
||||
runtime.deleteSearchDocument(workspaceId, docId)
|
||||
);
|
||||
}
|
||||
|
||||
async reconcileSearchWorkspace(workspaceId: string) {
|
||||
await this.measured('reconcileSearchWorkspace', runtime =>
|
||||
runtime.reconcileSearchWorkspace(workspaceId)
|
||||
);
|
||||
}
|
||||
|
||||
async deleteSearchWorkspace(workspaceId: string) {
|
||||
await this.measured('deleteSearchWorkspace', runtime =>
|
||||
runtime.deleteSearchWorkspace(workspaceId)
|
||||
);
|
||||
}
|
||||
|
||||
async filterReadableDocs(
|
||||
actorUserId: string,
|
||||
workspaceId: string,
|
||||
docIds: string[]
|
||||
) {
|
||||
return await this.measured('filterReadableDocs', runtime =>
|
||||
runtime.filterReadableDocs(actorUserId, workspaceId, docIds)
|
||||
);
|
||||
}
|
||||
|
||||
async searchStatus() {
|
||||
return await this.measured('searchStatus', runtime =>
|
||||
runtime.searchStatus()
|
||||
);
|
||||
}
|
||||
|
||||
async embeddingQueueCounts() {
|
||||
return await this.measured('embeddingQueueCounts', runtime =>
|
||||
runtime.embeddingQueueCounts()
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
import type {
|
||||
RuntimeAggregateRequest,
|
||||
RuntimeSearchQuery,
|
||||
RuntimeSearchRequest,
|
||||
} from '../../native';
|
||||
|
||||
type SearchQueryInput = {
|
||||
type: string;
|
||||
field?: string;
|
||||
match?: string;
|
||||
query?: SearchQueryInput;
|
||||
queries?: SearchQueryInput[];
|
||||
occur?: string;
|
||||
boost?: number;
|
||||
};
|
||||
|
||||
type SearchPaginationInput = {
|
||||
limit?: number;
|
||||
skip?: number;
|
||||
cursor?: string;
|
||||
};
|
||||
|
||||
type SearchHighlightInput = {
|
||||
field: string;
|
||||
before: string;
|
||||
end: string;
|
||||
};
|
||||
|
||||
type SearchOptionsInput = {
|
||||
fields: string[];
|
||||
highlights?: SearchHighlightInput[];
|
||||
pagination?: SearchPaginationInput;
|
||||
};
|
||||
|
||||
export type SearchRequestInput = {
|
||||
table: 'doc' | 'block';
|
||||
query: SearchQueryInput;
|
||||
options: SearchOptionsInput;
|
||||
};
|
||||
|
||||
export type AggregateRequestInput = {
|
||||
table: 'doc' | 'block';
|
||||
query: SearchQueryInput;
|
||||
field: string;
|
||||
options: {
|
||||
hits: SearchOptionsInput;
|
||||
pagination?: SearchPaginationInput;
|
||||
};
|
||||
};
|
||||
|
||||
export function encodeSearchRequest(
|
||||
request: SearchRequestInput
|
||||
): RuntimeSearchRequest {
|
||||
const { queries, rootQuery } = encodeQuery(request.query);
|
||||
return {
|
||||
table: request.table,
|
||||
queries,
|
||||
rootQuery,
|
||||
options: encodeOptions(request.options),
|
||||
};
|
||||
}
|
||||
|
||||
export function encodeAggregateRequest(
|
||||
request: AggregateRequestInput
|
||||
): RuntimeAggregateRequest {
|
||||
const { queries, rootQuery } = encodeQuery(request.query);
|
||||
return {
|
||||
table: request.table,
|
||||
queries,
|
||||
rootQuery,
|
||||
field: request.field,
|
||||
options: {
|
||||
hits: encodeOptions(request.options.hits),
|
||||
pagination: request.options.pagination ?? {},
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function encodeOptions(options: SearchOptionsInput) {
|
||||
return {
|
||||
fields: options.fields,
|
||||
highlights: options.highlights ?? [],
|
||||
pagination: options.pagination ?? {},
|
||||
};
|
||||
}
|
||||
|
||||
function encodeQuery(root: SearchQueryInput) {
|
||||
const nodes: RuntimeSearchQuery[] = [];
|
||||
const visit = (query: SearchQueryInput): number => {
|
||||
const index = nodes.length;
|
||||
nodes.push({ queryType: query.type });
|
||||
nodes[index] = {
|
||||
queryType: query.type,
|
||||
field: query.field,
|
||||
matchValue: query.match,
|
||||
query: query.query ? visit(query.query) : undefined,
|
||||
queries: query.queries?.map(visit),
|
||||
occur: query.occur,
|
||||
boost: query.boost,
|
||||
};
|
||||
return index;
|
||||
};
|
||||
return { queries: nodes, rootQuery: visit(root) };
|
||||
}
|
||||
@@ -73,6 +73,25 @@ export class ServerService implements OnApplicationBootstrap {
|
||||
user: string,
|
||||
updates: Array<{ module: string; key: string; value: any }>
|
||||
): Promise<DeepPartial<AppConfig>> {
|
||||
const providerType = updates.find(
|
||||
update => update.module === 'indexer' && update.key === 'provider.type'
|
||||
);
|
||||
if (providerType?.value === 'embedded') {
|
||||
updates = updates.filter(update => update !== providerType);
|
||||
updates = [
|
||||
...updates.filter(
|
||||
update => !(update.module === 'indexer' && update.key === 'enabled')
|
||||
),
|
||||
{ module: 'indexer', key: 'enabled', value: false },
|
||||
];
|
||||
} else if (providerType) {
|
||||
updates = [
|
||||
...updates.filter(
|
||||
update => !(update.module === 'indexer' && update.key === 'enabled')
|
||||
),
|
||||
{ module: 'indexer', key: 'enabled', value: true },
|
||||
];
|
||||
}
|
||||
const errors = this.validateConfig(updates);
|
||||
|
||||
if (errors?.length) {
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
import { Module } from '@nestjs/common';
|
||||
|
||||
import { DocStorageModule, DocStorageWorkerModule } from '../doc';
|
||||
import { DocJobConsumer, DocJobScheduler } from './job';
|
||||
|
||||
@Module({
|
||||
imports: [DocStorageModule, DocStorageWorkerModule],
|
||||
providers: [DocJobConsumer, DocJobScheduler],
|
||||
})
|
||||
export class DocJobsModule {}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user