qql-cli 0.4.0

//! First-class collection dump → `.qql` script exporter.
//!
//! Emits:
//! 1. `CREATE COLLECTION` from typed vector schema (sizes, distances, sparse)
//! 2. `CREATE INDEX` from typed payload index specs
//! 3. `CREATE SHARD KEY` for every custom shard key (sharded collections only)
//! 4. Batched `UPSERT` with real `vector:` values (not re-embed stubs),
//!    each batch carrying its `SHARD` key on custom-sharded collections
//!
//! Scroll uses cursor pagination and `with_vector: true` so every page
//! includes stored vectors. Output is streamed to disk to bound memory.

mod cursor;
mod escape;
mod indexes;
mod point;
mod quant;
mod schema;

use std::error::Error;
use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::Path;

use qql::client::QdrantOps;
use qql::executor::Executor;
use qql_core::ast::ShardKey;
use qql_plan::PlannedOperation;
use qql_plan::semantic::{PlanPointId, PlanShardKey};
use qql_plan::types::{FilterExpression, PayloadSelectorReq, ScrollRequest, VectorSelectorReq};

use point::write_upsert_batch;

pub(crate) use cursor::{extract_scroll_page, json_to_plan_point_id};
pub(crate) use escape::format_ident;
pub(crate) use indexes::generate_index_statements;
pub(crate) use point::point_to_upsert_object;
pub(crate) use schema::generate_create_statement;
// ── Public API ──────────────────────────────────────────────────

/// Result summary of a dump run.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DumpStats {
    pub written: usize,
    pub skipped: usize,
    pub batches: usize,
}

/// Progress details emitted after each batch during dump execution.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DumpProgress {
    pub collection: String,
    pub written: usize,
    pub skipped: usize,
    pub batches: usize,
}

/// Dump a collection to a `.qql` script file atomically.
///
/// Output is streamed to a temporary file (`.tmp`) and renamed to `output_path`
/// on successful completion. If an error occurs, any pre-existing file at `output_path`
/// is preserved untouched, and the temporary file is removed.
pub async fn dump_collection(
    executor: &Executor,
    collection: &str,
    output_path: &str,
    batch_size: u32,
    progress: Option<&(dyn Fn(DumpProgress) + Sync)>,
) -> Result<DumpStats, Box<dyn Error>> {
    let tmp_path = format!("{}.tmp", output_path);
    let res = dump_collection_inner(executor, collection, &tmp_path, batch_size, progress).await;
    match res {
        Ok(stats) => {
            // Atomic replace when possible; fall back to copy on cross-device rename.
            if let Err(rename_err) = fs::rename(&tmp_path, output_path) {
                if let Err(copy_err) = fs::copy(&tmp_path, output_path) {
                    // Keep .tmp for recovery; surface both failures.
                    return Err(format!(
                        "failed to finalize dump at '{}': rename error: {}; copy error: {}",
                        output_path, rename_err, copy_err
                    )
                    .into());
                }
                let _ = fs::remove_file(&tmp_path);
            }
            Ok(stats)
        }
        Err(err) => {
            if Path::new(&tmp_path).exists() {
                let _ = fs::remove_file(&tmp_path);
            }
            Err(err)
        }
    }
}

async fn dump_collection_inner(
    executor: &Executor,
    collection: &str,
    output_path: &str,
    batch_size: u32,
    progress: Option<&(dyn Fn(DumpProgress) + Sync)>,
) -> Result<DumpStats, Box<dyn Error>> {
    if batch_size == 0 {
        return Err("batch_size must be >= 1".into());
    }

    let ops = executor.ops();

    let exists = ops.collection_exists(collection).await?;
    if !exists {
        return Err(format!("collection '{}' does not exist", collection).into());
    }

    let info = ops.get_collection_info(collection).await?;

    if let Some(parent) = Path::new(output_path).parent()
        && !parent.as_os_str().is_empty()
    {
        fs::create_dir_all(parent)?;
    }

    let file = File::create(output_path)?;
    let mut out = BufWriter::new(file);

    writeln!(out, "-- QQL dump for {}", collection)?;
    writeln!(out, "-- Generated by qql dump")?;
    writeln!(out)?;

    let create = generate_create_statement(collection, &info);
    writeln!(out, "{};", create)?;
    writeln!(out)?;

    let indexes = generate_index_statements(collection, &info.schema.payload_indexes);
    if !indexes.is_empty() {
        for idx in &indexes {
            writeln!(out, "{};", idx)?;
        }
        writeln!(out)?;
    }

    // Every batch on a custom-sharded collection carries its `SHARD` key:
    // an unrouted replay against a `custom` collection fails with
    // "Shard key not specified". Auto-sharded collections dump as one stream.
    let shard_keys = list_shard_keys(ops, collection).await?;
    for key in &shard_keys {
        writeln!(
            out,
            "CREATE SHARD KEY {} ON COLLECTION {};",
            key,
            format_ident(collection)
        )?;
    }
    if !shard_keys.is_empty() {
        writeln!(out)?;
    }

    let shards: Vec<Option<&ShardKey>> = if shard_keys.is_empty() {
        vec![None]
    } else {
        shard_keys.iter().map(Some).collect()
    };

    let mut written = 0usize;
    let mut skipped = 0usize;
    let mut batches = 0usize;

    for shard in shards {
        let mut pages =
            ScrollPages::new(ops, collection, batch_size).shard(shard.map(PlanShardKey::from));
        while let Some(points) = pages.next().await? {
            let mut records = Vec::with_capacity(points.len());
            for point in &points {
                match point_to_upsert_object(point) {
                    Some(rec) => records.push(rec),
                    None => skipped += 1,
                }
            }

            if !records.is_empty() {
                write_upsert_batch(&mut out, collection, &records, shard)?;
                written += records.len();
                batches += 1;

                if let Some(cb) = progress {
                    cb(DumpProgress {
                        collection: collection.to_string(),
                        written,
                        skipped,
                        batches,
                    });
                }
            }
        }
    }

    writeln!(out)?;
    writeln!(out, "-- Written: {}", written)?;
    writeln!(out, "-- Skipped: {}", skipped)?;
    writeln!(out, "-- Batches: {}", batches)?;
    out.flush()?;

    Ok(DumpStats {
        written,
        skipped,
        batches,
    })
}

/// Custom shard keys on the collection, sorted and deduplicated.
///
/// Empty for auto-sharded collections. Reads the typed
/// [`ExecData::ShardKeys`](qql::executor::ExecData::ShardKeys) payload — both
/// transports produce the same `PlanShardKey` values.
pub(crate) async fn list_shard_keys(
    ops: &dyn QdrantOps,
    collection: &str,
) -> Result<Vec<ShardKey>, Box<dyn Error>> {
    let op = PlannedOperation::ListShardKeys {
        collection: collection.to_string(),
    };
    let response = ops.execute_planned(&op).await?;
    Ok(response
        .data
        .shard_keys()
        .map(parse_shard_key_list)
        .unwrap_or_default())
}

/// Sorted, deduplicated shard keys from a typed `ListShardKeys` payload.
pub(crate) fn parse_shard_key_list(keys: &[PlanShardKey]) -> Vec<ShardKey> {
    let mut deduped: std::collections::HashMap<String, ShardKey> = std::collections::HashMap::new();
    for key in keys {
        let key = match key {
            PlanShardKey::Keyword(text) => ShardKey::Keyword(text.clone()),
            PlanShardKey::Number(number) => ShardKey::Number(*number),
        };
        deduped.insert(key.to_string(), key);
    }
    let mut out: Vec<(String, ShardKey)> = deduped.into_iter().collect();
    out.sort_by(|a, b| a.0.cmp(&b.0));
    out.into_iter().map(|(_, k)| k).collect()
}

/// Sequential scroll stream shared by dump and migrate.
///
/// Qdrant scroll `offset` is inclusive while the server `next_page_offset`
/// already points past the page: resuming from a server cursor never repeats,
/// but resuming from the last-id fallback (used when the server omits the
/// cursor) repeats that point first, so it is dropped. Point ids are unique
/// per page, making the drop a no-op for exclusive-offset backends.
pub(crate) struct ScrollPages<'a> {
    ops: &'a dyn QdrantOps,
    collection: String,
    batch_size: u32,
    filter: Option<FilterExpression>,
    shard_key: Option<PlanShardKey>,
    with_payload: PayloadSelectorReq,
    with_vector: VectorSelectorReq,
    after: Option<PlanPointId>,
    fell_back: bool,
    done: bool,
}

impl<'a> ScrollPages<'a> {
    pub(crate) fn new(ops: &'a dyn QdrantOps, collection: &str, batch_size: u32) -> Self {
        Self {
            ops,
            collection: collection.to_string(),
            batch_size,
            filter: None,
            shard_key: None,
            with_payload: PayloadSelectorReq::All(true),
            with_vector: VectorSelectorReq::All(true),
            after: None,
            fell_back: false,
            done: false,
        }
    }

    pub(crate) fn filter(mut self, filter: Option<FilterExpression>) -> Self {
        self.filter = filter;
        self
    }

    pub(crate) fn shard(mut self, shard_key: Option<PlanShardKey>) -> Self {
        self.shard_key = shard_key;
        self
    }

    /// Start from a stored checkpoint cursor. Checkpoint cursors resume
    /// exactly like the fetch that stored them, so they are never treated as
    /// a last-id fallback.
    pub(crate) fn resume(mut self, after: Option<PlanPointId>) -> Self {
        self.after = after;
        self
    }

    pub(crate) fn select(
        mut self,
        with_payload: PayloadSelectorReq,
        with_vector: VectorSelectorReq,
    ) -> Self {
        self.with_payload = with_payload;
        self.with_vector = with_vector;
        self
    }

    /// Cursor the next page resumes from.
    pub(crate) fn cursor(&self) -> Option<&PlanPointId> {
        self.after.as_ref()
    }

    /// Next page, or `None` when the stream is exhausted.
    pub(crate) async fn next(&mut self) -> Result<Option<Vec<serde_json::Value>>, Box<dyn Error>> {
        if self.done {
            return Ok(None);
        }
        let op = PlannedOperation::Scroll {
            collection: self.collection.clone(),
            request: ScrollRequest {
                filter: self.filter.clone(),
                offset: self.after.clone(),
                limit: Some(self.batch_size as u64),
                with_payload: Some(self.with_payload.clone()),
                with_vector: Some(self.with_vector.clone()),
                order_by: None,
                shard_key: self.shard_key.clone(),
            },
        };
        let response = self.ops.execute_planned(&op).await?;
        let (points, next) = extract_scroll_page(&response);
        let mut points = points;
        if self.fell_back {
            points = drop_resumed_point(points, self.after.as_ref());
        }
        if points.is_empty() {
            self.done = true;
            return Ok(None);
        }
        self.fell_back = next.is_none();
        self.after = next_scroll_cursor(next, &points);
        Ok(Some(points))
    }
}

/// Drop the resume-cursor point Qdrant repeats at page boundaries.
///
/// Scroll `offset` is inclusive, so resuming from the previous page's last id
/// re-emits that point first. The server `next_page_offset` already points
/// past it — this only triggers on the last-id fallback — and point ids are
/// unique per page, so the drop is a no-op for exclusive-offset backends.
pub(crate) fn drop_resumed_point(
    mut points: Vec<serde_json::Value>,
    after: Option<&PlanPointId>,
) -> Vec<serde_json::Value> {
    let duplicate = match (after, points.first()) {
        (Some(cursor), Some(first)) => {
            first.get("id").and_then(json_to_plan_point_id).as_ref() == Some(cursor)
        }
        _ => false,
    };
    if duplicate {
        points.remove(0);
    }
    points
}

/// Prefer Qdrant's `next_page_offset`; fall back to the last point id on the page.
///
/// [`ScrollPages`] uses the fallback to keep cursor-less backends streaming;
/// the inclusive-offset repeat it causes is removed by [`drop_resumed_point`].
pub(crate) fn next_scroll_cursor(
    next: Option<PlanPointId>,
    points: &[serde_json::Value],
) -> Option<PlanPointId> {
    next.or_else(|| {
        points
            .last()
            .and_then(|p| p.get("id"))
            .and_then(json_to_plan_point_id)
    })
}

#[cfg(test)]
mod tests;