nodedb 0.4.0

// SPDX-License-Identifier: BUSL-1.1

//! Surrogate → user-PK translation for vector search responses.
//!
//! The Data Plane emits each hit's `id` as the bound `Surrogate.as_u32()`
//! (or the local node id for headless rows) and leaves `doc_id` as
//! `None`. The Control Plane runs this translator at the response
//! boundary so pgwire / HTTP / native clients still see human-readable
//! document identifiers without the engine ever consulting the catalog.
//!
//! Behaviour:
//!  - non-msgpack payloads (already JSON, empty, or non-array) round-
//!    trip unchanged.
//!  - decode failures are non-fatal — the original payload is returned
//!    so the client still sees the raw search hits.
//!
//! [`resolve_surrogate_pk`] is the single catalog lookup shared by this
//! translator and the full-text / hybrid translators in `text_hybrid.rs` —
//! see `dispatch.rs` for the plan-to-translator routing.

use nodedb_types::DatabaseId;
use nodedb_types::Surrogate;
use nodedb_types::TenantId;
use serde::{Deserialize, Serialize};

use crate::bridge::scan_filter::ScanFilter;
use crate::control::state::SharedState;

#[derive(Serialize, Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack)]
#[msgpack(map)]
struct Hit {
    id: u32,
    distance: f32,
    #[serde(skip_serializing_if = "Option::is_none")]
    doc_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    body: Option<Vec<u8>>,
}

/// Apply RLS post-filter at the Control-Plane security boundary, then
/// truncate to `top_k` and strip the body bytes that the Data Plane
/// attached purely for the predicate evaluation. A no-op when
/// `rls_filters` is empty.
fn apply_rls_filter(hits: &mut Vec<Hit>, rls_filters: &[u8], top_k: usize) {
    // When filters are empty the Data Plane never attaches a body, so there
    // is nothing to evaluate or strip — return immediately.
    if rls_filters.is_empty() {
        return;
    }
    let filters: Vec<ScanFilter> = match zerompk::from_msgpack(rls_filters) {
        Ok(f) => f,
        Err(_) => {
            // fail-closed: drop everything if filters are corrupt.
            tracing::warn!("RLS filter decode failed at CP boundary — denying all hits");
            hits.clear();
            return;
        }
    };
    hits.retain(|h| match h.body.as_deref() {
        Some(body) => filters.iter().all(|f| f.matches_binary(body)),
        None => false,
    });
    if hits.len() > top_k {
        hits.truncate(top_k);
    }
    for h in hits.iter_mut() {
        h.body = None;
    }
}

/// Resolve one surrogate to its user-facing primary key via the catalog.
///
/// Shared by the vector, full-text, and hybrid (RRF) response translators —
/// every search-hit path resolves the internal `u32` surrogate back to the
/// user's PK through this single catalog call. Returns `None` when the
/// catalog has no PK mapping for the surrogate (headless row, or a document
/// that was never written) — callers must leave the row's identifier
/// untouched in that case rather than fabricate a value.
pub(super) fn resolve_surrogate_pk(
    state: &SharedState,
    database_id: DatabaseId,
    tenant_id: TenantId,
    collection: &str,
    surrogate: Surrogate,
) -> Option<String> {
    let catalog = state.credentials.catalog();
    let pk_bytes = catalog
        .get_pk_for_surrogate(database_id, tenant_id, collection, surrogate)
        .ok()??;
    String::from_utf8(pk_bytes).ok()
}

/// Decode the DP-side msgpack array of `VectorSearchHit`, fill each
/// row's `doc_id` from the catalog using `id` as the surrogate, and
/// re-encode. On any decode failure the payload is returned unchanged.
pub fn translate_vector_search_payload(
    payload: &[u8],
    state: &SharedState,
    database_id: DatabaseId,
    tenant_id: TenantId,
    collection: &str,
    rls_filters: &[u8],
    top_k: usize,
) -> Vec<u8> {
    if payload.is_empty() {
        return payload.to_vec();
    }
    let first = payload[0];
    if first == b'[' || first == b'{' || first == b'"' {
        return payload.to_vec();
    }

    let mut hits: Vec<Hit> = match zerompk::from_msgpack(payload) {
        Ok(h) => h,
        Err(_) => return payload.to_vec(),
    };

    apply_rls_filter(&mut hits, rls_filters, top_k);

    for hit in &mut hits {
        if hit.doc_id.is_some() {
            continue;
        }
        if let Some(pk) = resolve_surrogate_pk(
            state,
            database_id,
            tenant_id,
            collection,
            Surrogate::new(hit.id),
        ) {
            hit.doc_id = Some(pk);
        }
    }

    // Slow-path column projection: when `body` is a msgpack-encoded payload
    // map, decode it and surface fields alongside `id` / `distance` so client
    // SQL projections like `SELECT id, label, vector_distance(...)` see the
    // payload columns. The base hit fields stay top-level; payload fields are
    // serialized as JSON-style siblings.
    use std::collections::BTreeMap;
    let flattened: Vec<BTreeMap<String, serde_json::Value>> = hits
        .iter()
        .map(|h| {
            let mut obj: BTreeMap<String, serde_json::Value> = BTreeMap::new();
            obj.insert("distance".into(), serde_json::json!(h.distance));
            // Flatten body fields first so the document's own `id` field wins
            // over the internal surrogate integer. Body fields are only present
            // when the Data Plane performed the slow-path fetch.
            if let Some(ref body) = h.body
                && let Ok(map) = zerompk::from_msgpack::<
                    std::collections::HashMap<String, nodedb_types::Value>,
                >(body)
            {
                for (k, v) in map {
                    if let Ok(j) = serde_json::to_value(&v) {
                        obj.insert(k, j);
                    }
                }
            }
            // Fall back to catalog-resolved doc_id or raw surrogate when
            // the body didn't provide an "id" field (e.g. skip_payload_fetch).
            if !obj.contains_key("id") {
                if let Some(ref doc) = h.doc_id {
                    obj.insert("id".into(), serde_json::json!(doc));
                } else {
                    obj.insert("id".into(), serde_json::json!(h.id));
                }
            }
            // Always expose internal surrogate for debugging / join use.
            obj.insert("_surrogate".into(), serde_json::json!(h.id));
            obj
        })
        .collect();
    if let Ok(s) = sonic_rs::to_string(&flattened) {
        return s.into_bytes();
    }
    zerompk::to_msgpack_vec(&hits).unwrap_or_else(|_| payload.to_vec())
}