use serde_json::Value;
use super::super::types::Context;
use super::super::QueryExecutor;
use crate::error::DbResult;
use crate::sdbql::ast::*;
use crate::sync::log::LogEntry;
use crate::sync::protocol::Operation;
impl<'a> QueryExecutor<'a> {
pub(super) fn try_streaming_bulk_insert(
&self,
query: &Query,
initial_bindings: &Context,
) -> DbResult<Option<(Vec<Value>, usize)>> {
if query.body_clauses.len() != 2
|| query.sort_clause.is_some()
|| query.limit_clause.is_some()
{
return Ok(None);
}
let for_clause = match &query.body_clauses[0] {
BodyClause::For(fc) => fc,
_ => return Ok(None),
};
let insert_clause = match &query.body_clauses[1] {
BodyClause::Insert(ic) => ic,
_ => return Ok(None),
};
let range_expr = match &for_clause.source_expression {
Some(Expression::Range(start, end)) => (start, end),
_ => return Ok(None),
};
let start_val = self.evaluate_expr_with_context(range_expr.0, initial_bindings)?;
let end_val = self.evaluate_expr_with_context(range_expr.1, initial_bindings)?;
let start = match &start_val {
Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)),
_ => None,
};
let end = match &end_val {
Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)),
_ => None,
};
let (start, end) = match (start, end) {
(Some(s), Some(e)) => (s, e),
_ => return Ok(None),
};
const STREAMING_THRESHOLD: i64 = 5_000;
const BATCH_SIZE: i64 = 5_000;
let total_count = (end - start + 1).max(0);
if total_count < STREAMING_THRESHOLD {
return Ok(None); }
tracing::info!(
"STREAMING INSERT: Processing {} documents in batches of {}",
total_count,
BATCH_SIZE
);
let collection = self.get_collection(&insert_clause.collection)?;
if let Some(config) = collection.get_shard_config() {
if config.num_shards > 0 {
tracing::debug!(
"Streaming insert disabled for sharded collection: {}",
insert_clause.collection
);
return Ok(None);
}
}
let has_indexes = !collection.list_indexes().is_empty();
let var_name = &for_clause.variable;
let mut all_results: Vec<Value> = Vec::new();
let mut current = start;
let total_start = std::time::Instant::now();
while current <= end {
let batch_end = (current + BATCH_SIZE - 1).min(end);
let batch_size = (batch_end - current + 1) as usize;
let mut documents = Vec::with_capacity(batch_size);
for i in current..=batch_end {
let mut ctx = initial_bindings.clone();
ctx.insert(var_name.clone(), Value::Number(serde_json::Number::from(i)));
let doc_value = self.evaluate_expr_with_context(&insert_clause.document, &ctx)?;
documents.push(doc_value);
}
let inserted_docs = collection.insert_batch(documents)?;
if query.return_clause.is_some() {
for i in current..=batch_end {
all_results.push(Value::Number(serde_json::Number::from(i)));
}
}
self.log_mutations_async(&insert_clause.collection, Operation::Insert, &inserted_docs);
if has_indexes && !inserted_docs.is_empty() {
let _ = collection.index_documents(&inserted_docs);
}
current = batch_end + 1;
collection.flush_stats_throttled();
if total_count > 100_000 && (current - start) % 100_000 == 0 {
tracing::info!(
"STREAMING INSERT: Processed {}/{} documents",
current - start,
total_count
);
}
}
let elapsed = total_start.elapsed();
tracing::info!(
"STREAMING INSERT: Completed {} documents in {:?} ({:.0} docs/sec)",
total_count,
elapsed,
total_count as f64 / elapsed.as_secs_f64()
);
collection.flush_stats();
Ok(Some((all_results, total_count as usize)))
}
pub(super) fn log_mutation(
&self,
collection: &str,
operation: Operation,
key: &str,
data: Option<&Value>,
) {
if let (Some(repl), Some(ref db)) = (&self.replication, &self.database) {
let entry = LogEntry {
sequence: 0,
node_id: "".to_string(),
database: db.clone(),
collection: collection.to_string(),
operation,
key: key.to_string(),
data: data.and_then(|v| serde_json::to_vec(v).ok()),
timestamp: chrono::Utc::now().timestamp_millis() as u64,
origin_sequence: None,
};
let _ = repl.append(entry);
}
}
pub(super) fn log_mutations_async(
&self,
collection: &str,
operation: Operation,
docs: &[crate::storage::Document],
) {
let repl_clone = self.replication.cloned();
let db_clone = self.database.clone();
if let (Some(repl), Some(db)) = (repl_clone, db_clone) {
let collection = collection.to_string();
let entries: Vec<LogEntry> = docs
.iter()
.map(|doc| LogEntry {
sequence: 0,
node_id: "".to_string(),
database: db.clone(),
collection: collection.clone(),
operation, key: doc.key.clone(),
data: serde_json::to_vec(&doc.to_value()).ok(),
timestamp: chrono::Utc::now().timestamp_millis() as u64,
origin_sequence: None,
})
.collect();
let count = entries.len();
tracing::debug!(
"INSERT: Starting async replication logging for {} docs",
count
);
let start = std::time::Instant::now();
let _ = repl.append_batch(entries);
let elapsed = start.elapsed();
tracing::debug!(
"INSERT: Replication logging of {} docs completed in {:?}",
count,
elapsed
);
}
}
}