use crate::{EntityType, FalkorDBError, FalkorResult, GraphSchema, IndexType};
use std::{collections::HashMap, fmt::Display};
pub(crate) mod blocking;
pub(crate) mod query_builder;
pub(crate) mod batch;
pub(crate) mod ops;
#[cfg(feature = "tokio")]
pub(crate) mod asynchronous;
pub trait HasGraphSchema {
fn get_graph_schema_mut(&mut self) -> &mut GraphSchema;
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Default, strum::Display)]
#[strum(serialize_all = "lowercase")]
#[non_exhaustive]
pub enum VectorSimilarity {
#[default]
Euclidean,
Cosine,
}
pub(crate) fn generate_create_index_query<P: Display>(
index_field_type: IndexType,
entity_type: EntityType,
label: &str,
properties: &[P],
options: Option<&HashMap<String, String>>,
) -> FalkorResult<String> {
let properties_string = properties
.iter()
.map(|element| format!("l.{}", element))
.collect::<Vec<_>>()
.join(", ");
let pattern = match entity_type {
EntityType::Node => format!("(l:{})", label),
EntityType::Edge => format!("()-[l:{}]->()", label),
};
let idx_type = match index_field_type {
IndexType::Range => "",
IndexType::Vector => "VECTOR ",
IndexType::Fulltext => "FULLTEXT ",
};
let options_string = match options {
Some(hashmap) => {
let mut entries = hashmap
.iter()
.map(|(key, val)| {
if !crate::value::param::is_bare_identifier(key) {
return Err(FalkorDBError::InvalidIndexOption {
key: key.clone(),
message: "must be a Cypher identifier ([A-Za-z_][A-Za-z0-9_]*)"
.to_string(),
});
}
Ok(format!("{key}: {}", format_index_option_value(val)?))
})
.collect::<FalkorResult<Vec<_>>>()?;
entries.sort();
format!(" OPTIONS {{ {} }}", entries.join(", "))
}
None => String::new(),
};
Ok(format!(
"CREATE {idx_type}INDEX FOR {pattern} ON ({}){}",
properties_string, options_string
))
}
fn format_index_option_value(value: &str) -> FalkorResult<String> {
if value.parse::<i64>().is_ok() {
Ok(value.to_string())
} else {
let mut out = String::new();
crate::value::param::encode_str(value, &mut out)?;
Ok(out)
}
}
pub(crate) fn vector_index_options(
dimension: u32,
similarity_function: VectorSimilarity,
) -> HashMap<String, String> {
HashMap::from([
("dimension".to_string(), dimension.to_string()),
(
"similarityFunction".to_string(),
similarity_function.to_string(),
),
])
}
pub(crate) fn generate_drop_index_query<P: Display>(
index_field_type: IndexType,
entity_type: EntityType,
label: &str,
properties: &[P],
) -> String {
let properties_string = properties
.iter()
.map(|element| format!("e.{}", element))
.collect::<Vec<_>>()
.join(", ");
let pattern = match entity_type {
EntityType::Node => format!("(e:{})", label),
EntityType::Edge => format!("()-[e:{}]->()", label),
};
let idx_type = match index_field_type {
IndexType::Range => "",
IndexType::Vector => "VECTOR",
IndexType::Fulltext => "FULLTEXT",
}
.to_string();
format!(
"DROP {idx_type} INDEX for {pattern} ON ({})",
properties_string
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_generate_create_index_query_range_node() {
let query = generate_create_index_query(
IndexType::Range,
EntityType::Node,
"Person",
&["name"],
None,
)
.unwrap();
assert!(query.contains("CREATE INDEX"));
assert!(query.contains("(l:Person)"));
assert!(query.contains("l.name"));
}
#[test]
fn test_generate_create_index_query_range_edge() {
let query = generate_create_index_query(
IndexType::Range,
EntityType::Edge,
"KNOWS",
&["since"],
None,
)
.unwrap();
assert!(query.contains("CREATE INDEX"));
assert!(query.contains("()-[l:KNOWS]->()"));
assert!(query.contains("l.since"));
}
#[test]
fn test_generate_create_index_query_vector() {
let query = generate_create_index_query(
IndexType::Vector,
EntityType::Node,
"Item",
&["embedding"],
None,
)
.unwrap();
assert!(query.contains("CREATE VECTOR INDEX"));
assert!(query.contains("(l:Item)"));
assert!(query.contains("l.embedding"));
}
#[test]
fn test_generate_create_index_query_fulltext() {
let query = generate_create_index_query(
IndexType::Fulltext,
EntityType::Node,
"Document",
&["content"],
None,
)
.unwrap();
assert!(query.contains("CREATE FULLTEXT INDEX"));
assert!(query.contains("(l:Document)"));
assert!(query.contains("l.content"));
}
#[test]
fn test_generate_create_index_query_with_options() {
let mut options = HashMap::new();
options.insert("option1".to_string(), "value1".to_string());
let query = generate_create_index_query(
IndexType::Range,
EntityType::Node,
"Test",
&["field"],
Some(&options),
)
.unwrap();
assert!(query.contains("OPTIONS"));
assert!(query.contains("option1"));
assert!(query.contains("value1"));
}
#[test]
fn test_vector_index_options_are_valid_cypher() {
let options = vector_index_options(128, VectorSimilarity::Euclidean);
let query = generate_create_index_query(
IndexType::Vector,
EntityType::Node,
"Item",
&["embedding"],
Some(&options),
)
.unwrap();
assert!(query.contains("CREATE VECTOR INDEX"));
assert!(
query.contains("OPTIONS { dimension: 128, similarityFunction: 'euclidean' }"),
"{query}"
);
assert!(!query.contains("'dimension'"));
}
#[test]
fn test_format_index_option_value() {
assert_eq!(format_index_option_value("128").unwrap(), "128");
assert_eq!(format_index_option_value("-3").unwrap(), "-3");
assert_eq!(
format_index_option_value("euclidean").unwrap(),
"'euclidean'"
);
assert_eq!(format_index_option_value("1.5").unwrap(), "'1.5'");
assert_eq!(format_index_option_value("a'b\\c").unwrap(), "'a\\'b\\\\c'");
assert_eq!(format_index_option_value("a\nb").unwrap(), "'a\\nb'");
assert!(format_index_option_value("a\0b").is_err());
}
#[test]
fn test_generate_create_index_query_rejects_unencodable_option() {
let options = HashMap::from([("similarityFunction".to_string(), "a\0b".to_string())]);
let err = generate_create_index_query(
IndexType::Vector,
EntityType::Node,
"Item",
&["embedding"],
Some(&options),
)
.unwrap_err();
assert!(matches!(err, FalkorDBError::ParamEncoding { .. }));
assert!(err.to_string().contains("NUL byte"));
}
#[test]
fn test_generate_create_index_query_rejects_non_identifier_key() {
let options = HashMap::from([("bad-key".to_string(), "4".to_string())]);
let err = generate_create_index_query(
IndexType::Vector,
EntityType::Node,
"Item",
&["embedding"],
Some(&options),
)
.unwrap_err();
assert!(matches!(
err,
FalkorDBError::InvalidIndexOption { ref key, .. } if key == "bad-key"
));
assert!(err.to_string().contains("'bad-key'"));
}
#[test]
fn test_vector_similarity_display() {
assert_eq!(VectorSimilarity::Euclidean.to_string(), "euclidean");
assert_eq!(VectorSimilarity::Cosine.to_string(), "cosine");
assert_eq!(VectorSimilarity::default(), VectorSimilarity::Euclidean);
}
#[test]
fn test_generate_create_index_query_multiple_properties() {
let query = generate_create_index_query(
IndexType::Range,
EntityType::Node,
"Person",
&["firstName", "lastName"],
None,
)
.unwrap();
assert!(query.contains("l.firstName"));
assert!(query.contains("l.lastName"));
}
#[test]
fn test_generate_drop_index_query_range_node() {
let query =
generate_drop_index_query(IndexType::Range, EntityType::Node, "Person", &["name"]);
assert!(query.contains("DROP"));
assert!(query.contains("INDEX"));
assert!(query.contains("(e:Person)"));
assert!(query.contains("e.name"));
}
#[test]
fn test_generate_drop_index_query_range_edge() {
let query =
generate_drop_index_query(IndexType::Range, EntityType::Edge, "KNOWS", &["since"]);
assert!(query.contains("DROP"));
assert!(query.contains("INDEX"));
assert!(query.contains("()-[e:KNOWS]->()"));
assert!(query.contains("e.since"));
}
#[test]
fn test_generate_drop_index_query_vector() {
let query =
generate_drop_index_query(IndexType::Vector, EntityType::Node, "Item", &["embedding"]);
assert!(query.contains("DROP VECTOR INDEX"));
assert!(query.contains("(e:Item)"));
assert!(query.contains("e.embedding"));
}
#[test]
fn test_generate_drop_index_query_fulltext() {
let query = generate_drop_index_query(
IndexType::Fulltext,
EntityType::Node,
"Document",
&["content"],
);
assert!(query.contains("DROP FULLTEXT INDEX"));
assert!(query.contains("(e:Document)"));
assert!(query.contains("e.content"));
}
#[test]
fn test_generate_drop_index_query_multiple_properties() {
let query = generate_drop_index_query(
IndexType::Range,
EntityType::Node,
"Person",
&["firstName", "lastName"],
);
assert!(query.contains("e.firstName"));
assert!(query.contains("e.lastName"));
}
}