reifydb_engine/bulk_insert/

builder.rs

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (c) 2025 ReifyDB

use std::{any, marker::PhantomData};

use any::TypeId;
use reifydb_catalog::{
	catalog::Catalog,
	error::{CatalogError, CatalogObjectKind},
};
use reifydb_core::{
	encoded::schema::Schema,
	error::CoreError,
	interface::catalog::id::IndexId,
	internal_error,
	key::{EncodableKey, index_entry::IndexEntryKey},
};
use reifydb_transaction::transaction::{Transaction, command::CommandTransaction};
use reifydb_type::{
	fragment::Fragment,
	value::{Value, identity::IdentityId, row_number::RowNumber, r#type::Type},
};

use super::{
	BulkInsertResult, RingBufferInsertResult, TableInsertResult,
	validation::{
		reorder_rows_trusted, reorder_rows_trusted_rb, validate_and_coerce_rows, validate_and_coerce_rows_rb,
	},
};
use crate::{
	Result,
	bulk_insert::primitive::{
		ringbuffer::{PendingRingBufferInsert, RingBufferInsertBuilder},
		table::{PendingTableInsert, TableInsertBuilder},
	},
	engine::StandardEngine,
	transaction::operation::{
		dictionary::DictionaryOperations, ringbuffer::RingBufferOperations, table::TableOperations,
	},
	vm::instruction::dml::{
		primary_key,
		schema::{get_or_create_ringbuffer_schema, get_or_create_table_schema},
	},
};

/// Marker trait for validation mode (sealed)
pub trait ValidationMode: sealed::Sealed + 'static {}

/// Validated mode - performs full type checking and constraint validation
pub struct Validated;
impl ValidationMode for Validated {}

/// Trusted mode - skips validation for pre-validated internal data
pub struct Trusted;
impl ValidationMode for Trusted {}

pub mod sealed {

	use super::{Trusted, Validated};
	pub trait Sealed {}
	impl Sealed for Validated {}
	impl Sealed for Trusted {}
}

/// Main builder for bulk insert operations.
///
/// Type parameter `V` tracks the validation mode at compile time.
pub struct BulkInsertBuilder<'e, V: ValidationMode = Validated> {
	engine: &'e StandardEngine,
	_identity: IdentityId,
	pending_tables: Vec<PendingTableInsert>,
	pending_ringbuffers: Vec<PendingRingBufferInsert>,
	_validation: PhantomData<V>,
}

impl<'e> BulkInsertBuilder<'e, Validated> {
	/// Create a new bulk insert builder with full validation enabled.
	pub(crate) fn new(engine: &'e StandardEngine, identity: IdentityId) -> Self {
		Self {
			engine,
			_identity: identity,
			pending_tables: Vec::new(),
			pending_ringbuffers: Vec::new(),
			_validation: PhantomData,
		}
	}
}

impl<'e> BulkInsertBuilder<'e, Trusted> {
	/// Create a new bulk insert builder with validation disabled (trusted mode).
	pub(crate) fn new_trusted(engine: &'e StandardEngine, identity: IdentityId) -> Self {
		Self {
			engine,
			_identity: identity,
			pending_tables: Vec::new(),
			pending_ringbuffers: Vec::new(),
			_validation: PhantomData,
		}
	}
}

impl<'e, V: ValidationMode> BulkInsertBuilder<'e, V> {
	/// Begin inserting into a table.
	///
	/// The qualified name can be either "namespace::table" or just "table"
	/// (which uses the default namespace).
	pub fn table<'a>(&'a mut self, qualified_name: &str) -> TableInsertBuilder<'a, 'e, V> {
		let (namespace, table) = parse_qualified_name(qualified_name);
		TableInsertBuilder::new(self, namespace, table)
	}

	/// Begin inserting into a ring buffer.
	///
	/// The qualified name can be either "namespace::ringbuffer" or just "ringbuffer"
	/// (which uses the default namespace).
	pub fn ringbuffer<'a>(&'a mut self, qualified_name: &str) -> RingBufferInsertBuilder<'a, 'e, V> {
		let (namespace, ringbuffer) = parse_qualified_name(qualified_name);
		RingBufferInsertBuilder::new(self, namespace, ringbuffer)
	}

	/// Add a pending table insert (called by TableInsertBuilder::done)
	pub(super) fn add_table_insert(&mut self, pending: PendingTableInsert) {
		self.pending_tables.push(pending);
	}

	/// Add a pending ring buffer insert (called by RingBufferInsertBuilder::done)
	pub(super) fn add_ringbuffer_insert(&mut self, pending: PendingRingBufferInsert) {
		self.pending_ringbuffers.push(pending);
	}

	/// Execute all pending inserts in a single transaction.
	///
	/// Returns a summary of what was inserted. On error, the entire
	/// transaction is rolled back (no partial inserts).
	pub fn execute(self) -> Result<BulkInsertResult> {
		let mut txn = self.engine.begin_command()?;
		let catalog = self.engine.catalog();
		let mut result = BulkInsertResult::default();

		// Process all pending table inserts
		for pending in self.pending_tables {
			let table_result = execute_table_insert::<V>(&catalog, &mut txn, &pending, TypeId::of::<V>())?;
			result.tables.push(table_result);
		}

		// Process all pending ring buffer inserts
		for pending in self.pending_ringbuffers {
			let rb_result =
				execute_ringbuffer_insert::<V>(&catalog, &mut txn, &pending, TypeId::of::<V>())?;
			result.ringbuffers.push(rb_result);
		}

		// Commit the transaction
		txn.commit()?;

		Ok(result)
	}
}

/// Execute a table insert within a transaction
fn execute_table_insert<V: ValidationMode>(
	catalog: &Catalog,
	txn: &mut CommandTransaction,
	pending: &PendingTableInsert,
	type_id: TypeId,
) -> Result<TableInsertResult> {
	// 1. Look up namespace and table from catalog
	let namespace = catalog
		.find_namespace_by_name(&mut Transaction::Command(txn), &pending.namespace)?
		.ok_or_else(|| CatalogError::NotFound {
			kind: CatalogObjectKind::Namespace,
			namespace: pending.namespace.to_string(),
			name: String::new(),
			fragment: Fragment::None,
		})?;

	let table = catalog
		.find_table_by_name(&mut Transaction::Command(txn), namespace.id, &pending.table)?
		.ok_or_else(|| CatalogError::NotFound {
			kind: CatalogObjectKind::Table,
			namespace: pending.namespace.to_string(),
			name: pending.table.to_string(),
			fragment: Fragment::None,
		})?;

	// 2. Get or create schema with proper field names and constraints
	let schema = get_or_create_table_schema(catalog, &table, &mut Transaction::Command(txn))?;

	// 3. Validate and coerce all rows in batch (fail-fast)
	let is_validated = type_id == TypeId::of::<Validated>();
	let coerced_rows = if is_validated {
		validate_and_coerce_rows(&pending.rows, &table)?
	} else {
		reorder_rows_trusted(&pending.rows, &table)?
	};

	let mut encoded_rows = Vec::with_capacity(coerced_rows.len());

	for mut values in coerced_rows {
		// Handle auto-increment columns
		for (idx, col) in table.columns.iter().enumerate() {
			if col.auto_increment && matches!(values[idx], Value::None { .. }) {
				values[idx] = catalog.column_sequence_next_value(txn, table.id, col.id)?;
			}
		}

		// Handle dictionary encoding
		for (idx, col) in table.columns.iter().enumerate() {
			if let Some(dict_id) = col.dictionary_id {
				let dictionary = catalog
					.find_dictionary(&mut Transaction::Command(txn), dict_id)?
					.ok_or_else(|| {
						internal_error!(
							"Dictionary {:?} not found for column {}",
							dict_id,
							col.name
						)
					})?;
				let entry_id = txn.insert_into_dictionary(&dictionary, &values[idx])?;
				values[idx] = entry_id.to_value();
			}
		}

		// Validate constraints (coercion is done in batch, but final constraint check still needed)
		if is_validated {
			for (idx, col) in table.columns.iter().enumerate() {
				col.constraint.validate(&values[idx])?;
			}
		}

		// Encode the row
		let mut row = schema.allocate();
		for (idx, value) in values.iter().enumerate() {
			schema.set_value(&mut row, idx, value);
		}
		encoded_rows.push(row);
	}

	// 4. Batch allocate row numbers
	let total_rows = encoded_rows.len();
	if total_rows == 0 {
		return Ok(TableInsertResult {
			namespace: pending.namespace.clone(),
			table: pending.table.clone(),
			inserted: 0,
		});
	}

	let row_numbers = catalog.next_row_number_batch(txn, table.id, total_rows as u64)?;

	// Hoist loop-invariant computations out of the insertion loop
	let pk_def = primary_key::get_primary_key(catalog, &mut Transaction::Command(txn), &table)?;
	let row_number_schema = if pk_def.is_some() {
		Some(Schema::testing(&[Type::Uint8]))
	} else {
		None
	};

	// 5. Insert all rows with their row numbers
	for (row, &row_number) in encoded_rows.iter().zip(row_numbers.iter()) {
		txn.insert_table(&table, &schema, row.clone(), row_number)?;

		// Handle primary key index if table has one
		if let Some(ref pk_def) = pk_def {
			let index_key = primary_key::encode_primary_key(pk_def, row, &table, &schema)?;
			let index_entry_key =
				IndexEntryKey::new(table.id, IndexId::primary(pk_def.id), index_key.clone());

			// Check for primary key violation
			if txn.contains_key(&index_entry_key.encode())? {
				let key_columns = pk_def.columns.iter().map(|c| c.name.clone()).collect();
				return Err(CoreError::PrimaryKeyViolation {
					fragment: Fragment::None,
					table_name: table.name.clone(),
					key_columns,
				}
				.into());
			}

			// Store the index entry
			let rns = row_number_schema.as_ref().unwrap();
			let mut row_number_encoded = rns.allocate();
			rns.set_u64(&mut row_number_encoded, 0, u64::from(row_number));
			txn.set(&index_entry_key.encode(), row_number_encoded)?;
		}
	}

	Ok(TableInsertResult {
		namespace: pending.namespace.clone(),
		table: pending.table.clone(),
		inserted: total_rows as u64,
	})
}

/// Execute a ring buffer insert within a transaction
fn execute_ringbuffer_insert<V: ValidationMode>(
	catalog: &Catalog,
	txn: &mut CommandTransaction,
	pending: &PendingRingBufferInsert,
	type_id: TypeId,
) -> Result<RingBufferInsertResult> {
	let namespace = catalog
		.find_namespace_by_name(&mut Transaction::Command(txn), &pending.namespace)?
		.ok_or_else(|| CatalogError::NotFound {
			kind: CatalogObjectKind::Namespace,
			namespace: pending.namespace.to_string(),
			name: String::new(),
			fragment: Fragment::None,
		})?;

	let ringbuffer = catalog
		.find_ringbuffer_by_name(&mut Transaction::Command(txn), namespace.id, &pending.ringbuffer)?
		.ok_or_else(|| CatalogError::NotFound {
			kind: CatalogObjectKind::RingBuffer,
			namespace: pending.namespace.to_string(),
			name: pending.ringbuffer.to_string(),
			fragment: Fragment::None,
		})?;

	let mut metadata = catalog
		.find_ringbuffer_metadata(&mut Transaction::Command(txn), ringbuffer.id)?
		.ok_or_else(|| CatalogError::NotFound {
			kind: CatalogObjectKind::RingBuffer,
			namespace: pending.namespace.to_string(),
			name: pending.ringbuffer.to_string(),
			fragment: Fragment::None,
		})?;

	// Get or create schema with proper field names and constraints
	let schema = get_or_create_ringbuffer_schema(catalog, &ringbuffer, &mut Transaction::Command(txn))?;

	// 3. Validate and coerce all rows in batch (fail-fast)
	let is_validated = type_id == TypeId::of::<Validated>();
	let coerced_rows = if is_validated {
		validate_and_coerce_rows_rb(&pending.rows, &ringbuffer)?
	} else {
		reorder_rows_trusted_rb(&pending.rows, &ringbuffer)?
	};

	let mut inserted_count = 0u64;

	// 4. Process each coerced row
	for mut values in coerced_rows {
		// Handle dictionary encoding
		for (idx, col) in ringbuffer.columns.iter().enumerate() {
			if let Some(dict_id) = col.dictionary_id {
				let dictionary = catalog
					.find_dictionary(&mut Transaction::Command(txn), dict_id)?
					.ok_or_else(|| {
						internal_error!(
							"Dictionary {:?} not found for column {}",
							dict_id,
							col.name
						)
					})?;
				let entry_id = txn.insert_into_dictionary(&dictionary, &values[idx])?;
				values[idx] = entry_id.to_value();
			}
		}

		// Validate constraints (coercion is done in batch, but final constraint check still needed)
		if is_validated {
			for (idx, col) in ringbuffer.columns.iter().enumerate() {
				col.constraint.validate(&values[idx])?;
			}
		}

		// Encode the row
		let mut row = schema.allocate();
		for (idx, value) in values.iter().enumerate() {
			schema.set_value(&mut row, idx, value);
		}

		// Handle ring buffer overflow - delete oldest entry if full
		if metadata.is_full() {
			let oldest_row = RowNumber(metadata.head);
			txn.remove_from_ringbuffer(&ringbuffer, oldest_row)?;
			metadata.head += 1;
			metadata.count -= 1;
		}

		// Allocate row number
		let row_number = catalog.next_row_number_for_ringbuffer(txn, ringbuffer.id)?;

		// Store the row
		txn.insert_ringbuffer_at(&ringbuffer, &schema, row_number, row)?;

		// Update metadata
		if metadata.is_empty() {
			metadata.head = row_number.0;
		}
		metadata.count += 1;
		metadata.tail = row_number.0 + 1;

		inserted_count += 1;
	}

	// Save updated metadata
	catalog.update_ringbuffer_metadata(txn, metadata)?;

	Ok(RingBufferInsertResult {
		namespace: pending.namespace.clone(),
		ringbuffer: pending.ringbuffer.clone(),
		inserted: inserted_count,
	})
}

/// Parse a qualified name like "namespace::table" into (namespace, name).
/// If no namespace is provided, uses "default".
fn parse_qualified_name(qualified_name: &str) -> (String, String) {
	if let Some((ns, name)) = qualified_name.split_once("::") {
		(ns.to_string(), name.to_string())
	} else {
		("default".to_string(), qualified_name.to_string())
	}
}