lancedb/nodejs/src/merge.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

use std::time::Duration;

use lancedb::{ipc::ipc_file_to_batches, table::merge::MergeInsertBuilder};
use napi::bindgen_prelude::*;
use napi_derive::napi;

use crate::{error::convert_error, table::MergeResult};

#[napi]
#[derive(Clone)]
/// A builder used to create and run a merge insert operation
pub struct NativeMergeInsertBuilder {
    pub(crate) inner: MergeInsertBuilder,
}

#[napi]
impl NativeMergeInsertBuilder {
    #[napi]
    pub fn when_matched_update_all(&self, condition: Option<String>) -> Self {
        let mut this = self.clone();
        this.inner.when_matched_update_all(condition);
        this
    }

    #[napi]
    pub fn when_not_matched_insert_all(&self) -> Self {
        let mut this = self.clone();
        this.inner.when_not_matched_insert_all();
        this
    }
    #[napi]
    pub fn when_not_matched_by_source_delete(&self, filter: Option<String>) -> Self {
        let mut this = self.clone();
        this.inner.when_not_matched_by_source_delete(filter);
        this
    }

    #[napi]
    pub fn set_timeout(&mut self, timeout: u32) {
        self.inner.timeout(Duration::from_millis(timeout as u64));
    }

    #[napi]
    pub fn use_index(&self, use_index: bool) -> Self {
        let mut this = self.clone();
        this.inner.use_index(use_index);
        this
    }

    #[napi(js_name = "useRTE")]
    pub fn use_rte(&self, enable: bool) -> Self {
        let mut this = self.clone();
        this.inner.use_rte(enable);
        this
    }

    #[napi]
    pub fn validate_single_shard(&self, validate_single_shard: bool) -> Self {
        let mut this = self.clone();
        this.inner.validate_single_shard(validate_single_shard);
        this
    }

    #[napi(catch_unwind)]
    pub async fn execute(&self, buf: Buffer) -> napi::Result<MergeResult> {
        let data = ipc_file_to_batches(buf.to_vec()).map_err(|e| {
            napi::Error::from_reason(format!("Failed to read IPC file: {}", convert_error(&e)))
        })?;

        let this = self.clone();

        let res = this.inner.execute(data).await.map_err(|e| {
            napi::Error::from_reason(format!(
                "Failed to execute merge insert: {}",
                convert_error(&e)
            ))
        })?;
        Ok(res.into())
    }
}

impl From<MergeInsertBuilder> for NativeMergeInsertBuilder {
    fn from(inner: MergeInsertBuilder) -> Self {
        Self { inner }
    }
}