sqlite_vector_rs/vtab/
shadow.rs

1use crate::vtab::config::VectorTableConfig;
2
3/// SQL statements for shadow table management.
4pub struct ShadowOps;
5
6impl ShadowOps {
7    pub fn create_data_table_sql(config: &VectorTableConfig) -> String {
8        let mut cols = vec![
9            "id INTEGER PRIMARY KEY AUTOINCREMENT".to_string(),
10            "vector BLOB NOT NULL".to_string(),
11        ];
12        for (name, sql_type) in &config.metadata_columns {
13            cols.push(format!("{name} {sql_type}"));
14        }
15        format!(
16            "CREATE TABLE IF NOT EXISTS \"{}_data\"({})",
17            config.table_name,
18            cols.join(", ")
19        )
20    }
21
22    pub fn create_index_table_sql(config: &VectorTableConfig) -> String {
23        format!(
24            "CREATE TABLE IF NOT EXISTS \"{}_index\"(key TEXT PRIMARY KEY, value BLOB)",
25            config.table_name
26        )
27    }
28
29    pub fn drop_shadow_tables_sql(table_name: &str) -> Vec<String> {
30        vec![
31            format!("DROP TABLE IF EXISTS \"{table_name}_data\""),
32            format!("DROP TABLE IF EXISTS \"{table_name}_index\""),
33        ]
34    }
35
36    pub fn insert_data_sql(config: &VectorTableConfig) -> String {
37        let mut col_names = vec!["vector".to_string()];
38        let mut placeholders = vec!["?".to_string()];
39        for (name, _) in &config.metadata_columns {
40            col_names.push(name.clone());
41            placeholders.push("?".to_string());
42        }
43        format!(
44            "INSERT INTO \"{}_data\"({}) VALUES({})",
45            config.table_name,
46            col_names.join(", "),
47            placeholders.join(", ")
48        )
49    }
50
51    pub fn insert_vector_only_sql(table_name: &str) -> String {
52        format!("INSERT INTO \"{table_name}_data\"(vector) VALUES(?)")
53    }
54
55    pub fn delete_data_sql(table_name: &str) -> String {
56        format!("DELETE FROM \"{table_name}_data\" WHERE id = ?")
57    }
58
59    pub fn select_data_sql(table_name: &str) -> String {
60        format!("SELECT * FROM \"{table_name}_data\" WHERE id = ?")
61    }
62
63    pub fn select_all_data_sql(table_name: &str) -> String {
64        format!("SELECT * FROM \"{table_name}_data\"")
65    }
66
67    pub fn upsert_index_sql(table_name: &str) -> String {
68        format!("INSERT OR REPLACE INTO \"{table_name}_index\"(key, value) VALUES(?, ?)")
69    }
70
71    pub fn select_index_sql(table_name: &str) -> String {
72        format!("SELECT value FROM \"{table_name}_index\" WHERE key = ?")
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use crate::vtab::config::VectorTableConfig;
80
81    fn base_config() -> VectorTableConfig {
82        VectorTableConfig::parse(&["vector", "main", "emb", "dim=3"]).unwrap()
83    }
84
85    fn config_with_metadata() -> VectorTableConfig {
86        VectorTableConfig::parse(&[
87            "vector",
88            "main",
89            "emb",
90            "dim=3",
91            "metadata=label TEXT,score REAL",
92        ])
93        .unwrap()
94    }
95
96    // --- create_data_table_sql ---
97
98    #[test]
99    fn create_data_table_sql_no_metadata_contains_table_name() {
100        let sql = ShadowOps::create_data_table_sql(&base_config());
101        assert!(sql.contains("emb_data"), "expected 'emb_data' in: {sql}");
102    }
103
104    #[test]
105    fn create_data_table_sql_no_metadata_has_id_column() {
106        let sql = ShadowOps::create_data_table_sql(&base_config());
107        assert!(
108            sql.contains("id INTEGER PRIMARY KEY AUTOINCREMENT"),
109            "expected id column in: {sql}"
110        );
111    }
112
113    #[test]
114    fn create_data_table_sql_no_metadata_has_vector_column() {
115        let sql = ShadowOps::create_data_table_sql(&base_config());
116        assert!(
117            sql.contains("vector BLOB NOT NULL"),
118            "expected vector column in: {sql}"
119        );
120    }
121
122    #[test]
123    fn create_data_table_sql_no_metadata_exact() {
124        let sql = ShadowOps::create_data_table_sql(&base_config());
125        assert_eq!(
126            sql,
127            "CREATE TABLE IF NOT EXISTS \"emb_data\"(id INTEGER PRIMARY KEY AUTOINCREMENT, vector BLOB NOT NULL)"
128        );
129    }
130
131    #[test]
132    fn create_data_table_sql_with_metadata_contains_label_column() {
133        let sql = ShadowOps::create_data_table_sql(&config_with_metadata());
134        assert!(
135            sql.contains("label TEXT"),
136            "expected 'label TEXT' in: {sql}"
137        );
138    }
139
140    #[test]
141    fn create_data_table_sql_with_metadata_contains_score_column() {
142        let sql = ShadowOps::create_data_table_sql(&config_with_metadata());
143        assert!(
144            sql.contains("score REAL"),
145            "expected 'score REAL' in: {sql}"
146        );
147    }
148
149    #[test]
150    fn create_data_table_sql_with_metadata_exact() {
151        let sql = ShadowOps::create_data_table_sql(&config_with_metadata());
152        assert_eq!(
153            sql,
154            "CREATE TABLE IF NOT EXISTS \"emb_data\"(id INTEGER PRIMARY KEY AUTOINCREMENT, vector BLOB NOT NULL, label TEXT, score REAL)"
155        );
156    }
157
158    // --- create_index_table_sql ---
159
160    #[test]
161    fn create_index_table_sql_contains_table_name() {
162        let sql = ShadowOps::create_index_table_sql(&base_config());
163        assert!(sql.contains("emb_index"), "expected 'emb_index' in: {sql}");
164    }
165
166    #[test]
167    fn create_index_table_sql_has_key_column() {
168        let sql = ShadowOps::create_index_table_sql(&base_config());
169        assert!(
170            sql.contains("key TEXT PRIMARY KEY"),
171            "expected 'key TEXT PRIMARY KEY' in: {sql}"
172        );
173    }
174
175    #[test]
176    fn create_index_table_sql_has_value_column() {
177        let sql = ShadowOps::create_index_table_sql(&base_config());
178        assert!(
179            sql.contains("value BLOB"),
180            "expected 'value BLOB' in: {sql}"
181        );
182    }
183
184    #[test]
185    fn create_index_table_sql_exact() {
186        let sql = ShadowOps::create_index_table_sql(&base_config());
187        assert_eq!(
188            sql,
189            "CREATE TABLE IF NOT EXISTS \"emb_index\"(key TEXT PRIMARY KEY, value BLOB)"
190        );
191    }
192
193    // --- drop_shadow_tables_sql ---
194
195    #[test]
196    fn drop_shadow_tables_sql_returns_two_statements() {
197        let stmts = ShadowOps::drop_shadow_tables_sql("emb");
198        assert_eq!(stmts.len(), 2);
199    }
200
201    #[test]
202    fn drop_shadow_tables_sql_drops_data_table() {
203        let stmts = ShadowOps::drop_shadow_tables_sql("emb");
204        assert!(
205            stmts[0].contains("emb_data"),
206            "expected 'emb_data' in: {}",
207            stmts[0]
208        );
209        assert_eq!(stmts[0], "DROP TABLE IF EXISTS \"emb_data\"");
210    }
211
212    #[test]
213    fn drop_shadow_tables_sql_drops_index_table() {
214        let stmts = ShadowOps::drop_shadow_tables_sql("emb");
215        assert!(
216            stmts[1].contains("emb_index"),
217            "expected 'emb_index' in: {}",
218            stmts[1]
219        );
220        assert_eq!(stmts[1], "DROP TABLE IF EXISTS \"emb_index\"");
221    }
222
223    // --- insert_data_sql ---
224
225    #[test]
226    fn insert_data_sql_no_metadata_exact() {
227        let sql = ShadowOps::insert_data_sql(&base_config());
228        assert_eq!(sql, "INSERT INTO \"emb_data\"(vector) VALUES(?)");
229    }
230
231    #[test]
232    fn insert_data_sql_with_metadata_contains_label() {
233        let sql = ShadowOps::insert_data_sql(&config_with_metadata());
234        assert!(sql.contains("label"), "expected 'label' in: {sql}");
235    }
236
237    #[test]
238    fn insert_data_sql_with_metadata_contains_score() {
239        let sql = ShadowOps::insert_data_sql(&config_with_metadata());
240        assert!(sql.contains("score"), "expected 'score' in: {sql}");
241    }
242
243    #[test]
244    fn insert_data_sql_with_metadata_has_correct_placeholder_count() {
245        let sql = ShadowOps::insert_data_sql(&config_with_metadata());
246        // vector + label + score = 3 placeholders
247        let placeholder_count = sql.matches('?').count();
248        assert_eq!(placeholder_count, 3, "expected 3 placeholders in: {sql}");
249    }
250
251    #[test]
252    fn insert_data_sql_with_metadata_exact() {
253        let sql = ShadowOps::insert_data_sql(&config_with_metadata());
254        assert_eq!(
255            sql,
256            "INSERT INTO \"emb_data\"(vector, label, score) VALUES(?, ?, ?)"
257        );
258    }
259
260    // --- insert_vector_only_sql ---
261
262    #[test]
263    fn insert_vector_only_sql_exact() {
264        let sql = ShadowOps::insert_vector_only_sql("emb");
265        assert_eq!(sql, "INSERT INTO \"emb_data\"(vector) VALUES(?)");
266    }
267
268    // --- delete_data_sql ---
269
270    #[test]
271    fn delete_data_sql_exact() {
272        let sql = ShadowOps::delete_data_sql("emb");
273        assert_eq!(sql, "DELETE FROM \"emb_data\" WHERE id = ?");
274    }
275
276    // --- select_data_sql ---
277
278    #[test]
279    fn select_data_sql_exact() {
280        let sql = ShadowOps::select_data_sql("emb");
281        assert_eq!(sql, "SELECT * FROM \"emb_data\" WHERE id = ?");
282    }
283
284    // --- select_all_data_sql ---
285
286    #[test]
287    fn select_all_data_sql_exact() {
288        let sql = ShadowOps::select_all_data_sql("emb");
289        assert_eq!(sql, "SELECT * FROM \"emb_data\"");
290    }
291
292    #[test]
293    fn select_all_data_sql_no_where_clause() {
294        let sql = ShadowOps::select_all_data_sql("emb");
295        assert!(!sql.contains("WHERE"), "unexpected WHERE clause in: {sql}");
296    }
297
298    // --- upsert_index_sql ---
299
300    #[test]
301    fn upsert_index_sql_exact() {
302        let sql = ShadowOps::upsert_index_sql("emb");
303        assert_eq!(
304            sql,
305            "INSERT OR REPLACE INTO \"emb_index\"(key, value) VALUES(?, ?)"
306        );
307    }
308
309    #[test]
310    fn upsert_index_sql_contains_insert_or_replace() {
311        let sql = ShadowOps::upsert_index_sql("emb");
312        assert!(
313            sql.contains("INSERT OR REPLACE INTO \"emb_index\""),
314            "expected INSERT OR REPLACE into emb_index in: {sql}"
315        );
316    }
317
318    // --- select_index_sql ---
319
320    #[test]
321    fn select_index_sql_exact() {
322        let sql = ShadowOps::select_index_sql("emb");
323        assert_eq!(sql, "SELECT value FROM \"emb_index\" WHERE key = ?");
324    }
325
326    // --- special characters in table name ---
327
328    #[test]
329    fn special_table_name_data_table() {
330        let sql = ShadowOps::create_data_table_sql(
331            &VectorTableConfig::parse(&["vector", "main", "my_table", "dim=3"]).unwrap(),
332        );
333        assert!(
334            sql.contains("my_table_data"),
335            "expected 'my_table_data' in: {sql}"
336        );
337    }
338
339    #[test]
340    fn special_table_name_index_table() {
341        let sql = ShadowOps::create_index_table_sql(
342            &VectorTableConfig::parse(&["vector", "main", "my_table", "dim=3"]).unwrap(),
343        );
344        assert!(
345            sql.contains("my_table_index"),
346            "expected 'my_table_index' in: {sql}"
347        );
348    }
349
350    #[test]
351    fn special_table_name_drop_shadow_tables() {
352        let stmts = ShadowOps::drop_shadow_tables_sql("my_table");
353        assert_eq!(stmts[0], "DROP TABLE IF EXISTS \"my_table_data\"");
354        assert_eq!(stmts[1], "DROP TABLE IF EXISTS \"my_table_index\"");
355    }
356
357    #[test]
358    fn special_table_name_insert_vector_only() {
359        let sql = ShadowOps::insert_vector_only_sql("my_table");
360        assert_eq!(sql, "INSERT INTO \"my_table_data\"(vector) VALUES(?)");
361    }
362
363    #[test]
364    fn special_table_name_delete_data() {
365        let sql = ShadowOps::delete_data_sql("my_table");
366        assert_eq!(sql, "DELETE FROM \"my_table_data\" WHERE id = ?");
367    }
368
369    #[test]
370    fn special_table_name_select_data() {
371        let sql = ShadowOps::select_data_sql("my_table");
372        assert_eq!(sql, "SELECT * FROM \"my_table_data\" WHERE id = ?");
373    }
374
375    #[test]
376    fn special_table_name_select_all_data() {
377        let sql = ShadowOps::select_all_data_sql("my_table");
378        assert_eq!(sql, "SELECT * FROM \"my_table_data\"");
379    }
380
381    #[test]
382    fn special_table_name_upsert_index() {
383        let sql = ShadowOps::upsert_index_sql("my_table");
384        assert_eq!(
385            sql,
386            "INSERT OR REPLACE INTO \"my_table_index\"(key, value) VALUES(?, ?)"
387        );
388    }
389
390    #[test]
391    fn special_table_name_select_index() {
392        let sql = ShadowOps::select_index_sql("my_table");
393        assert_eq!(sql, "SELECT value FROM \"my_table_index\" WHERE key = ?");
394    }
395}
sqlite_vector_rs/vtab/shadow.rs

sqlite_vector_rs/vtab/
shadow.rs
