herodb/src/search_cmd.rs

use crate::{
    error::DBError,
    protocol::Protocol,
    server::Server,
    tantivy_search::{
        FieldDef, Filter, FilterType, IndexConfig, NumericType, SearchOptions, TantivySearch,
    },
};
use std::collections::HashMap;
use std::sync::Arc;

pub async fn ft_create_cmd(
    server: &Server,
    index_name: String,
    schema: Vec<(String, String, Vec<String>)>,
) -> Result<Protocol, DBError> {
    if server.selected_db == 0 {
        return Ok(Protocol::err("FT commands are not allowed on DB 0"));
    }
    // Enforce Tantivy backend for selected DB
    let is_tantivy = crate::admin_meta::get_database_backend(
        &server.option.dir,
        server.option.backend.clone(),
        &server.option.admin_secret,
        server.selected_db,
    )
    .ok()
    .flatten()
    .map(|b| matches!(b, crate::options::BackendType::Tantivy))
    .unwrap_or(false);
    if !is_tantivy {
        return Ok(Protocol::err("ERR DB backend is not Tantivy; FT.* commands are not allowed"));
    }

    // Parse schema into field definitions
    let mut field_definitions = Vec::new();
    for (field_name, field_type, options) in schema {
        let field_def = match field_type.to_uppercase().as_str() {
            "TEXT" => {
                let mut sortable = false;
                let mut no_index = false;
                // Weight is not used in current implementation
                let mut _weight = 1.0f32;
                let mut i = 0;
                while i < options.len() {
                    match options[i].to_uppercase().as_str() {
                        "WEIGHT" => {
                            if i + 1 < options.len() {
                                _weight = options[i + 1].parse::<f32>().unwrap_or(1.0);
                                i += 2;
                                continue;
                            }
                        }
                        "SORTABLE" => {
                            sortable = true;
                        }
                        "NOINDEX" => {
                            no_index = true;
                        }
                        _ => {}
                    }
                    i += 1;
                }
                FieldDef::Text {
                    stored: true,
                    indexed: !no_index,
                    tokenized: true,
                    fast: sortable,
                }
            }
            "NUMERIC" => {
                // default to F64
                let mut sortable = false;
                for opt in &options {
                    if opt.to_uppercase() == "SORTABLE" {
                        sortable = true;
                    }
                }
                FieldDef::Numeric {
                    stored: true,
                    indexed: true,
                    fast: sortable,
                    precision: NumericType::F64,
                }
            }
            "TAG" => {
                let mut separator = ",".to_string();
                let mut case_sensitive = false;
                let mut i = 0;
                while i < options.len() {
                    match options[i].to_uppercase().as_str() {
                        "SEPARATOR" => {
                            if i + 1 < options.len() {
                                separator = options[i + 1].clone();
                                i += 2;
                                continue;
                            }
                        }
                        "CASESENSITIVE" => {
                            case_sensitive = true;
                        }
                        _ => {}
                    }
                    i += 1;
                }
                FieldDef::Tag {
                    stored: true,
                    separator,
                    case_sensitive,
                }
            }
            "GEO" => FieldDef::Geo { stored: true },
            _ => {
                return Err(DBError(format!("Unknown field type: {}", field_type)));
            }
        };
        field_definitions.push((field_name, field_def));
    }

    // Create the search index
    let search_path = server.search_index_path();
    let config = IndexConfig::default();
    let search_index = TantivySearch::new_with_schema(
        search_path,
        index_name.clone(),
        field_definitions,
        Some(config),
    )?;

    // Store in registry
    let mut indexes = server.search_indexes.write().unwrap();
    indexes.insert(index_name, Arc::new(search_index));

    Ok(Protocol::SimpleString("OK".to_string()))
}

pub async fn ft_add_cmd(
    server: &Server,
    index_name: String,
    doc_id: String,
    _score: f64,
    fields: HashMap<String, String>,
) -> Result<Protocol, DBError> {
    if server.selected_db == 0 {
        return Ok(Protocol::err("FT commands are not allowed on DB 0"));
    }
    // Enforce Tantivy backend for selected DB
    let is_tantivy = crate::admin_meta::get_database_backend(
        &server.option.dir,
        server.option.backend.clone(),
        &server.option.admin_secret,
        server.selected_db,
    )
    .ok()
    .flatten()
    .map(|b| matches!(b, crate::options::BackendType::Tantivy))
    .unwrap_or(false);
    if !is_tantivy {
        return Ok(Protocol::err("ERR DB backend is not Tantivy; FT.* commands are not allowed"));
    }
    let indexes = server.search_indexes.read().unwrap();
    let search_index = indexes
        .get(&index_name)
        .ok_or_else(|| DBError(format!("Index '{}' not found", index_name)))?;
    search_index.add_document_with_fields(&doc_id, fields)?;
    Ok(Protocol::SimpleString("OK".to_string()))
}

pub async fn ft_search_cmd(
    server: &Server,
    index_name: String,
    query: String,
    filters: Vec<(String, String)>,
    limit: Option<usize>,
    offset: Option<usize>,
    return_fields: Option<Vec<String>>,
) -> Result<Protocol, DBError> {
    if server.selected_db == 0 {
        return Ok(Protocol::err("FT commands are not allowed on DB 0"));
    }
    // Enforce Tantivy backend for selected DB
    let is_tantivy = crate::admin_meta::get_database_backend(
        &server.option.dir,
        server.option.backend.clone(),
        &server.option.admin_secret,
        server.selected_db,
    )
    .ok()
    .flatten()
    .map(|b| matches!(b, crate::options::BackendType::Tantivy))
    .unwrap_or(false);
    if !is_tantivy {
        return Ok(Protocol::err("ERR DB backend is not Tantivy; FT.* commands are not allowed"));
    }
    let indexes = server.search_indexes.read().unwrap();
    let search_index = indexes
        .get(&index_name)
        .ok_or_else(|| DBError(format!("Index '{}' not found", index_name)))?;

    let search_filters = filters
        .into_iter()
        .map(|(field, value)| Filter {
            field,
            filter_type: FilterType::Equals(value),
        })
        .collect();

    let options = SearchOptions {
        limit: limit.unwrap_or(10),
        offset: offset.unwrap_or(0),
        filters: search_filters,
        sort_by: None,
        return_fields,
        highlight: false,
    };

    let results = search_index.search_with_options(&query, options)?;

    // Format results as Redis protocol
    let mut response = Vec::new();
    // First element is the total count
    response.push(Protocol::SimpleString(results.total.to_string()));
    // Then each document
    for doc in results.documents {
        let mut doc_array = Vec::new();
        // Add document ID if it exists
        if let Some(id) = doc.fields.get("_id") {
            doc_array.push(Protocol::BulkString(id.clone()));
        }
        // Add score
        doc_array.push(Protocol::BulkString(doc.score.to_string()));
        // Add fields as key-value pairs
        for (field_name, field_value) in doc.fields {
            if field_name != "_id" {
                doc_array.push(Protocol::BulkString(field_name));
                doc_array.push(Protocol::BulkString(field_value));
            }
        }
        response.push(Protocol::Array(doc_array));
    }

    Ok(Protocol::Array(response))
}

pub async fn ft_del_cmd(
    server: &Server,
    index_name: String,
    doc_id: String,
) -> Result<Protocol, DBError> {
    if server.selected_db == 0 {
        return Ok(Protocol::err("FT commands are not allowed on DB 0"));
    }
    // Enforce Tantivy backend for selected DB
    let is_tantivy = crate::admin_meta::get_database_backend(
        &server.option.dir,
        server.option.backend.clone(),
        &server.option.admin_secret,
        server.selected_db,
    )
    .ok()
    .flatten()
    .map(|b| matches!(b, crate::options::BackendType::Tantivy))
    .unwrap_or(false);
    if !is_tantivy {
        return Ok(Protocol::err("ERR DB backend is not Tantivy; FT.* commands are not allowed"));
    }
    let indexes = server.search_indexes.read().unwrap();
    let _search_index = indexes
        .get(&index_name)
        .ok_or_else(|| DBError(format!("Index '{}' not found", index_name)))?;
    // Not fully implemented yet: Tantivy delete by term would require a writer session and commit coordination.
    println!("Deleting document '{}' from index '{}'", doc_id, index_name);
    Ok(Protocol::SimpleString("1".to_string()))
}

pub async fn ft_info_cmd(server: &Server, index_name: String) -> Result<Protocol, DBError> {
    if server.selected_db == 0 {
        return Ok(Protocol::err("FT commands are not allowed on DB 0"));
    }
    // Enforce Tantivy backend for selected DB
    let is_tantivy = crate::admin_meta::get_database_backend(
        &server.option.dir,
        server.option.backend.clone(),
        &server.option.admin_secret,
        server.selected_db,
    )
    .ok()
    .flatten()
    .map(|b| matches!(b, crate::options::BackendType::Tantivy))
    .unwrap_or(false);
    if !is_tantivy {
        return Ok(Protocol::err("ERR DB backend is not Tantivy; FT.* commands are not allowed"));
    }
    let indexes = server.search_indexes.read().unwrap();
    let search_index = indexes
        .get(&index_name)
        .ok_or_else(|| DBError(format!("Index '{}' not found", index_name)))?;
    let info = search_index.get_info()?;

    // Format info as Redis protocol
    let mut response = Vec::new();
    response.push(Protocol::BulkString("index_name".to_string()));
    response.push(Protocol::BulkString(info.name));
    response.push(Protocol::BulkString("num_docs".to_string()));
    response.push(Protocol::BulkString(info.num_docs.to_string()));
    response.push(Protocol::BulkString("num_fields".to_string()));
    response.push(Protocol::BulkString(info.fields.len().to_string()));
    response.push(Protocol::BulkString("fields".to_string()));
    let fields_str = info
        .fields
        .iter()
        .map(|f| format!("{}:{}", f.name, f.field_type))
        .collect::<Vec<_>>()
        .join(", ");
    response.push(Protocol::BulkString(fields_str));
    Ok(Protocol::Array(response))
}

pub async fn ft_drop_cmd(server: &Server, index_name: String) -> Result<Protocol, DBError> {
    if server.selected_db == 0 {
        return Ok(Protocol::err("FT commands are not allowed on DB 0"));
    }
    // Enforce Tantivy backend for selected DB
    let is_tantivy = crate::admin_meta::get_database_backend(
        &server.option.dir,
        server.option.backend.clone(),
        &server.option.admin_secret,
        server.selected_db,
    )
    .ok()
    .flatten()
    .map(|b| matches!(b, crate::options::BackendType::Tantivy))
    .unwrap_or(false);
    if !is_tantivy {
        return Ok(Protocol::err("ERR DB backend is not Tantivy; FT.* commands are not allowed"));
    }

    // Remove from registry
    {
        let mut indexes = server.search_indexes.write().unwrap();
        indexes.remove(&index_name);
    }

    // Remove the index files from disk
    let index_path = server.search_index_path().join(&index_name);
    if index_path.exists() {
        std::fs::remove_dir_all(&index_path)
            .map_err(|e| DBError(format!("Failed to remove index files: {}", e)))?;
    }

    Ok(Protocol::SimpleString("OK".to_string()))
}