mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user