feat: add SQL-style escaping for LanceDB predicate values

- Introduced `escape_predicate_value` function to safely escape strings for use in LanceDB predicate expressions, preventing SQL injection.
- Updated delete operations in the LanceDB vector store to utilize the new escaping function for `chunk_id` and `document_id`.
- Enhanced filtering logic to apply escaping for `user_id` and `agent_id` in query conditions.

This change improves security by ensuring that user inputs are properly sanitized before being used in database queries.
This commit is contained in:
ILGIN KANAT
2026-02-18 16:57:46 +04:00
parent 327e009622
commit 6489c433f6
+23 -4
View File
@@ -82,6 +82,12 @@ mod impl_lancedb {
const TABLE_NAME: &str = "memory_chunks";
/// Escapes a string for safe use in LanceDB predicate expressions.
/// Uses SQL-style escaping: single quotes are doubled to prevent injection.
fn escape_predicate_value(s: &str) -> String {
s.replace('\'', "''")
}
/// LanceDB-backed vector store.
pub struct LanceDbVectorStore {
db: Arc<lancedb::Connection>,
@@ -258,7 +264,10 @@ mod impl_lancedb {
})?;
table
.delete(&format!("chunk_id = '{}'", chunk_id))
.delete(&format!(
"chunk_id = '{}'",
escape_predicate_value(&chunk_id.to_string())
))
.await
.map_err(|e| WorkspaceError::EmbeddingFailed {
reason: format!("Failed to delete chunk for update: {}", e),
@@ -276,7 +285,10 @@ mod impl_lancedb {
})?;
table
.delete(&format!("document_id = '{}'", document_id))
.delete(&format!(
"document_id = '{}'",
escape_predicate_value(&document_id.to_string())
))
.await
.map_err(|e| WorkspaceError::ChunkingFailed {
reason: format!("Failed to delete chunks: {}", e),
@@ -299,9 +311,16 @@ mod impl_lancedb {
})?;
let filter = if let Some(aid) = agent_id {
format!("user_id = '{}' AND agent_id = '{}'", user_id, aid)
format!(
"user_id = '{}' AND agent_id = '{}'",
escape_predicate_value(user_id),
escape_predicate_value(&aid.to_string())
)
} else {
format!("user_id = '{}' AND agent_id IS NULL", user_id)
format!(
"user_id = '{}' AND agent_id IS NULL",
escape_predicate_value(user_id)
)
};
let mut stream = table