use std::{collections::HashMap, sync::Arc};
use postcard::to_stdvec;
use reifydb_codec::row::{
bytes::{EncodedBytes, RowBuilder},
queue::EncodedQueueRow,
queue_deduplication::EncodedQueueDeduplicationRow,
shape::RowShape,
};
use reifydb_core::{
error::diagnostic::catalog::{
namespace_not_found, queue_deduplication_key_not_utf8, queue_not_before_not_datetime, queue_not_found,
},
interface::{
catalog::{
config::{ConfigKey, GetConfig},
namespace::Namespace,
policy::{DataOp, PolicyTargetType},
queue::{Queue, decode_queue_deduplication, encode_queue_deduplication},
},
resolved::{ResolvedColumn, ResolvedNamespace, ResolvedObject, ResolvedQueue},
},
internal_error,
key::{queue::QueueDeduplicationKey, row::RowKey},
return_internal_error,
value::column::{buffer::ColumnBuffer, columns::Columns},
};
use reifydb_evaluate::stack::SymbolTable;
use reifydb_rql::{
expression::Expression,
nodes::{InsertQueueNode, QUEUE_CREATED_COLUMN, QUEUE_DEDUPLICATION_KEY_FIELD, QUEUE_NOT_BEFORE_FIELD},
};
use reifydb_transaction::transaction::Transaction;
use reifydb_value::{
fragment::Fragment,
params::Params,
return_error,
value::{Value, datetime::DateTime, duration::Duration, identity::IdentityId, row_number::RowNumber},
};
use tracing::instrument;
use super::{
returning::{decode_returning_dictionaries, decode_rows_to_columns, evaluate_returning},
shape::get_or_create_queue_shape,
};
use crate::{
Result,
policy::PolicyEvaluator,
queue::partition::{ordered_by_index, placement_of},
transaction::operation::{
dictionary::DictionaryOperations,
queue::{QueueInsertRow, QueueOperations},
},
vm::{
instruction::dml::{coerce::coerce_value_to_column_type, time::resolve_time},
services::Services,
volcano::{
compile::compile,
query::{QueryContext, QueryNode, query_budget},
},
},
};
struct QueueTarget<'a> {
namespace: &'a Namespace,
queue: &'a Queue,
}
#[instrument(name = "mutate::queue::insert", level = "trace", skip_all)]
pub(crate) fn insert_queue(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
plan: InsertQueueNode,
symbols: &mut SymbolTable,
) -> Result<Columns> {
let InsertQueueNode {
input,
target,
has_deduplication,
has_not_before,
returning,
} = plan;
let (namespace, queue) = resolve_insert_queue_target(services, txn, &target)?;
let shape = get_or_create_queue_shape(&services.catalog, &queue, txn)?;
let target_data = QueueTarget {
namespace: &namespace,
queue: &queue,
};
let context = build_insert_queue_query_context(services, &target_data, symbols, txn.identity());
let mut input_node = compile(*input, txn, context.clone());
input_node.initialize(txn, &context)?;
let pending = validate_and_encode_input_rows(
services,
txn,
&target_data,
&shape,
&context,
symbols,
&mut input_node,
has_deduplication,
has_not_before,
)?;
if pending.is_empty() {
return Ok(insert_queue_result(namespace.name(), &queue.name, 0, 0));
}
let now = services.runtime_context.clock.now();
let outcomes = resolve_duplicates(txn, &queue, &shape, &pending, now)?;
let fresh_count = outcomes.iter().filter(|outcome| matches!(outcome, Outcome::Fresh)).count();
let duplicates = outcomes.len() - fresh_count;
let row_numbers = if fresh_count == 0 {
Vec::new()
} else {
services.catalog.next_row_number_batch_for_queue(txn, queue.id, fresh_count as u64)?
};
let ordered_by_index = ordered_by_index(&queue)?;
let mut assigned = row_numbers.into_iter();
let mut rows: Vec<QueueInsertRow> = Vec::with_capacity(fresh_count);
let mut returned: Vec<ReturnedRow> = Vec::with_capacity(outcomes.len());
for (item, outcome) in pending.iter().zip(outcomes.into_iter()) {
match outcome {
Outcome::Fresh => {
let row_number = assigned.next().expect("a row number per fresh item");
if let Some(key) = &item.deduplication_key {
write_deduplication_record(txn, &queue, key, row_number, now)?;
}
let placement = placement_of(
&queue,
&shape,
EncodedQueueRow::view(&item.encoded),
ordered_by_index,
row_number,
);
rows.push(QueueInsertRow {
row_number,
partition: placement.partition,
key_hash: placement.key_hash,
not_before: item.not_before,
encoded: item.encoded.clone(),
});
returned.push(ReturnedRow {
created: true,
row_number,
encoded: item.encoded.clone(),
});
}
Outcome::Duplicate {
row_number,
encoded,
} => returned.push(ReturnedRow {
created: false,
row_number,
encoded: encoded.unwrap_or_else(|| shape.allocate_queue().freeze_bytes()),
}),
Outcome::DuplicateInBatch {
origin,
} => {
let row_number = returned[origin].row_number;
let encoded = returned[origin].encoded.clone();
returned.push(ReturnedRow {
created: false,
row_number,
encoded,
});
}
}
}
txn.insert_queue(&queue, &rows)?;
if let Some(returning_exprs) = &returning {
return project_returning(services, txn, symbols, &queue, &shape, returning_exprs, &returned);
}
Ok(insert_queue_result(namespace.name(), &queue.name, fresh_count as u64, duplicates as u64))
}
struct PendingItem {
encoded: EncodedBytes,
deduplication_key: Option<Vec<u8>>,
not_before: Option<DateTime>,
}
enum Outcome {
Fresh,
Duplicate {
row_number: RowNumber,
encoded: Option<EncodedBytes>,
},
DuplicateInBatch {
origin: usize,
},
}
struct ReturnedRow {
created: bool,
row_number: RowNumber,
encoded: EncodedBytes,
}
fn write_deduplication_record(
txn: &mut Transaction<'_>,
queue: &Queue,
key: &[u8],
row_number: RowNumber,
now: DateTime,
) -> Result<()> {
let ttl = queue.deduplicate.as_ref().map(|d| d.ttl).unwrap_or(Duration::MAX);
let record = encode_queue_deduplication(row_number, now.saturating_add(ttl));
txn.set(&QueueDeduplicationKey::new(queue.id, key), record.into_bytes())?;
Ok(())
}
fn resolve_duplicates(
txn: &mut Transaction<'_>,
queue: &Queue,
shape: &RowShape,
pending: &[PendingItem],
now: DateTime,
) -> Result<Vec<Outcome>> {
let mut outcomes = Vec::with_capacity(pending.len());
let mut seen: HashMap<Vec<u8>, usize> = HashMap::new();
for (index, item) in pending.iter().enumerate() {
let Some(key) = &item.deduplication_key else {
outcomes.push(Outcome::Fresh);
continue;
};
if let Some(&origin) = seen.get(key) {
outcomes.push(Outcome::DuplicateInBatch {
origin,
});
continue;
}
let stored = txn.get(&QueueDeduplicationKey::new(queue.id, key.clone()))?;
if let Some(stored) = stored {
let Some((row_number, expires_at)) =
decode_queue_deduplication(EncodedQueueDeduplicationRow::view(&stored.bytes))
else {
return_internal_error!(
"Queue {} deduplication record is {} bytes wide, too short for its header. This indicates a corrupt record.",
queue.name,
stored.bytes.len()
)
};
if expires_at > now {
let encoded = txn.get(&RowKey::new(queue.id, row_number))?.map(|item| item.bytes);
outcomes.push(Outcome::Duplicate {
row_number,
encoded,
});
continue;
}
}
seen.insert(key.clone(), index);
outcomes.push(Outcome::Fresh);
}
let _ = shape;
Ok(outcomes)
}
fn project_returning(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
symbols: &mut SymbolTable,
queue: &Queue,
shape: &RowShape,
returning_exprs: &[Expression],
returned: &[ReturnedRow],
) -> Result<Columns> {
let rows: Vec<(RowNumber, EncodedBytes)> =
returned.iter().map(|row| (row.row_number, row.encoded.clone())).collect();
let mut columns = decode_rows_to_columns(shape, &rows);
truncate_to_declared(&mut columns, queue.columns.len());
decode_returning_dictionaries(services, txn, &queue.columns, &mut columns)?;
let mut created = ColumnBuffer::bool_with_capacity(returned.len());
for row in returned {
created.push_value(Value::Boolean(row.created));
}
columns.columns.push(created);
columns.names.push(Fragment::internal(QUEUE_CREATED_COLUMN));
evaluate_returning(services, symbols, returning_exprs, columns, txn.identity())
}
fn declared_key_indices(queue: &Queue) -> Result<Option<Vec<usize>>> {
let Some(deduplicate) = &queue.deduplicate else {
return Ok(None);
};
let mut indices = Vec::with_capacity(deduplicate.by.len());
for column in &deduplicate.by {
let index = queue.columns.iter().position(|c| c.name == *column).ok_or_else(|| {
internal_error!("queue {} deduplicates by {} which is not a column", queue.name, column)
})?;
indices.push(index);
}
Ok(Some(indices))
}
fn declared_key_bytes(shape: &RowShape, bytes: &EncodedBytes, indices: &[usize]) -> Vec<u8> {
let values: Vec<Value> = indices.iter().map(|&index| shape.get_value(bytes, index)).collect();
to_stdvec(&values).expect("postcard serialization of a Value list is total")
}
#[inline]
fn truncate_to_declared(columns: &mut Columns, declared: usize) {
columns.columns.truncate(declared);
columns.names.truncate(declared);
}
#[inline]
fn resolve_insert_queue_target(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
target: &ResolvedQueue,
) -> Result<(Namespace, Queue)> {
let namespace_name = target.namespace().name();
let Some(namespace) = services.catalog.find_namespace_by_name(txn, namespace_name)? else {
return_error!(namespace_not_found(Fragment::internal(namespace_name), namespace_name));
};
let queue_name = target.name();
let Some(queue) = services.catalog.find_queue_by_name(txn, namespace.id(), queue_name)? else {
return_error!(queue_not_found(target.identifier().clone(), namespace_name, queue_name));
};
Ok((namespace, queue))
}
#[inline]
fn build_insert_queue_query_context(
services: &Arc<Services>,
target: &QueueTarget<'_>,
symbols: &SymbolTable,
identity: IdentityId,
) -> Arc<QueryContext> {
let namespace_ident = Fragment::internal(target.namespace.name());
let resolved_namespace = ResolvedNamespace::new(namespace_ident, target.namespace.clone());
let queue_ident = Fragment::internal(target.queue.name.clone());
let resolved_queue = ResolvedQueue::new(queue_ident, resolved_namespace, target.queue.clone());
Arc::new(QueryContext {
services: services.clone(),
source: Some(ResolvedObject::Queue(resolved_queue)),
batch_size: services.catalog.get_config_uint2(ConfigKey::QueryRowBatchSize) as u64,
params: Params::None,
symbols: symbols.clone(),
identity,
memory: query_budget(services),
})
}
#[allow(clippy::too_many_arguments)]
fn validate_and_encode_input_rows(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
target: &QueueTarget<'_>,
shape: &RowShape,
context: &Arc<QueryContext>,
symbols: &SymbolTable,
input_node: &mut Box<dyn QueryNode>,
has_deduplication: bool,
has_not_before: bool,
) -> Result<Vec<PendingItem>> {
let mut pending: Vec<PendingItem> = Vec::new();
let mut mutable_context = (**context).clone();
let declared_key_indices = declared_key_indices(target.queue)?;
while let Some(columns) = input_node.next(txn, &mut mutable_context)? {
PolicyEvaluator::new(services, symbols).enforce_write_policies(
txn,
target.namespace.name(),
&target.queue.name,
DataOp::Insert,
&columns,
PolicyTargetType::Queue,
)?;
let mut column_map: HashMap<&str, usize> = HashMap::new();
for (idx, col) in columns.iter().enumerate() {
column_map.insert(col.name().text(), idx);
}
for row_idx in 0..columns.row_count() {
let declared_key_indices = declared_key_indices.as_deref();
let not_before = if has_not_before {
read_not_before(target, &columns, &column_map, row_idx)?
} else {
None
};
let encoded = build_insert_queue_row(
services,
txn,
target,
shape,
&columns,
&column_map,
context,
row_idx,
not_before,
)?;
let deduplication_key = match declared_key_indices {
Some(indices) => Some(declared_key_bytes(shape, &encoded, indices)),
None if has_deduplication => {
read_deduplication_key(target, &columns, &column_map, row_idx)?
}
None => None,
};
pending.push(PendingItem {
encoded,
deduplication_key,
not_before,
});
}
}
Ok(pending)
}
#[inline]
fn read_deduplication_key(
target: &QueueTarget<'_>,
columns: &Columns,
column_map: &HashMap<&str, usize>,
row_idx: usize,
) -> Result<Option<Vec<u8>>> {
let Some(&idx) = column_map.get(QUEUE_DEDUPLICATION_KEY_FIELD) else {
return Ok(None);
};
let value = columns[idx].get_value(row_idx);
match value {
Value::None {
..
}
| Value::Utf8(_) => Ok(Some(statement_key_bytes(&value))),
other => return_error!(queue_deduplication_key_not_utf8(
Fragment::internal(target.queue.name.clone()),
other.get_type().to_string().as_str()
)),
}
}
fn statement_key_bytes(value: &Value) -> Vec<u8> {
to_stdvec(value).expect("postcard serialization of a Value is total")
}
#[inline]
fn read_not_before(
target: &QueueTarget<'_>,
columns: &Columns,
column_map: &HashMap<&str, usize>,
row_idx: usize,
) -> Result<Option<DateTime>> {
let Some(&idx) = column_map.get(QUEUE_NOT_BEFORE_FIELD) else {
return Ok(None);
};
match columns[idx].get_value(row_idx) {
Value::None {
..
} => Ok(None),
Value::DateTime(instant) => Ok(Some(instant)),
other => return_error!(queue_not_before_not_datetime(
Fragment::internal(target.queue.name.clone()),
other.get_type().to_string().as_str()
)),
}
}
#[allow(clippy::too_many_arguments)]
#[inline]
fn build_insert_queue_row(
services: &Arc<Services>,
txn: &mut Transaction<'_>,
target: &QueueTarget<'_>,
shape: &RowShape,
columns: &Columns,
column_map: &HashMap<&str, usize>,
context: &Arc<QueryContext>,
row_idx: usize,
not_before: Option<DateTime>,
) -> Result<EncodedBytes> {
let mut row = shape.allocate_queue();
for (queue_idx, queue_column) in target.queue.columns.iter().enumerate() {
let mut value = if let Some(&input_idx) = column_map.get(queue_column.name.as_str()) {
columns[input_idx].get_value(row_idx)
} else {
Value::none()
};
if queue_column.auto_increment && matches!(value, Value::None { .. }) {
value = services.catalog.column_sequence_next_value(txn, target.queue.id, queue_column.id)?;
}
let column_ident = column_map
.get(queue_column.name.as_str())
.map(|&idx| columns.name_at(idx).clone())
.unwrap_or_else(|| Fragment::internal(queue_column.name.clone()));
let resolved_column = ResolvedColumn::new(
column_ident.clone(),
context.source.clone().unwrap(),
queue_column.clone(),
);
value = coerce_value_to_column_type(
value,
queue_column.constraint.get_type(),
resolved_column,
context,
)?;
if let Err(mut e) = queue_column.constraint.validate(&value) {
e.0.fragment = column_ident.clone();
return Err(e);
}
let value = if let Some(dict_id) = queue_column.dictionary_id {
let dictionary = services.catalog.find_dictionary(txn, dict_id)?.ok_or_else(|| {
internal_error!("Dictionary {:?} not found for column {}", dict_id, queue_column.name)
})?;
let entry_id = if matches!(value, Value::None { .. }) {
dictionary.id_type.none()
} else {
txn.insert_into_dictionary(&dictionary, &value)?
};
entry_id.to_value()
} else {
value
};
shape.set_value(&mut row, queue_idx, &value);
}
if let Some(instant) = not_before {
row.set_not_before(instant);
}
let now = services.runtime_context.clock.now();
row.set_timestamps(now, now);
if let Some(time) =
resolve_time(&target.queue.name, &target.queue.columns, &target.queue.time, shape, &row, now)?
{
row.set_time(time);
}
Ok(row.freeze_bytes())
}
#[inline]
fn insert_queue_result(namespace: &str, queue: &str, inserted: u64, duplicates: u64) -> Columns {
Columns::single_row([
("namespace", Value::Utf8(namespace.to_string())),
("queue", Value::Utf8(queue.to_string())),
("inserted", Value::Uint8(inserted)),
("duplicates", Value::Uint8(duplicates)),
])
}