1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum QueryMode {
8 Sql,
10 Gremlin,
12 Cypher,
14 Sparql,
16 Path,
18 Natural,
20 Unknown,
22}
23
24pub fn detect_mode(input: &str) -> QueryMode {
26 let trimmed = input.trim();
27 let lower = trimmed.to_lowercase();
28
29 if trimmed.starts_with('"') || trimmed.starts_with('\'') {
31 return QueryMode::Natural;
32 }
33
34 if lower.starts_with("g.") || lower.starts_with("__.") {
36 return QueryMode::Gremlin;
37 }
38
39 if lower.starts_with("path ") || lower.starts_with("paths ") {
41 return QueryMode::Path;
42 }
43
44 if lower.starts_with("prefix ") || has_sparql_pattern(&lower) {
46 return QueryMode::Sparql;
47 }
48
49 if lower.starts_with("match ") || lower.starts_with("match(") {
51 return QueryMode::Cypher;
52 }
53
54 let first_token = lower.split_whitespace().next().unwrap_or("");
59 if matches!(
60 first_token,
61 "begin"
62 | "start"
63 | "commit"
64 | "rollback"
65 | "savepoint"
66 | "release"
67 | "end"
68 | "vacuum"
69 | "analyze"
70 | "reset"
71 | "checkpoint"
72 | "checkout"
73 | "merge"
74 | "cherry"
75 | "revert"
76 | "resolve"
77 | "fork"
78 | "promote"
79 | "copy"
80 | "refresh"
81 | "explain"
82 | "grant"
83 | "revoke"
84 | "attach"
85 | "detach"
86 | "simulate"
87 | "lint"
88 | "migrate"
89 | "apply"
90 | "events"
91 | "describe"
92 | "desc"
93 ) {
94 return QueryMode::Sql;
95 }
96 if lower.starts_with("select ")
97 || lower.starts_with("from ")
98 || lower.starts_with("insert ")
99 || lower.starts_with("update ")
100 || lower.starts_with("delete ")
101 || lower.starts_with("truncate ")
102 || lower.starts_with("create ")
103 || lower.starts_with("drop ")
104 || lower.starts_with("alter ")
105 || lower.starts_with("vector ")
106 || lower.starts_with("hybrid ")
107 || lower.starts_with("graph ")
108 || lower.starts_with("queue ")
109 || lower.starts_with("events ")
110 || lower.starts_with("tree ")
111 || lower.starts_with("hll ")
112 || lower.starts_with("sketch ")
113 || lower.starts_with("filter ")
114 || lower.starts_with("vault ")
115 || lower.starts_with("unseal vault ")
116 || lower.starts_with("rotate vault ")
117 || lower.starts_with("history vault ")
118 || lower.starts_with("list vault ")
119 || lower.starts_with("list kv ")
120 || lower.starts_with("watch vault ")
121 || lower.starts_with("delete vault ")
122 || lower.starts_with("purge vault ")
123 || lower.starts_with("search ")
124 || lower.starts_with("ask ")
125 || lower.starts_with("put config ")
126 || lower.starts_with("get config ")
127 || lower.starts_with("resolve config ")
128 || lower.starts_with("rotate config ")
129 || lower.starts_with("delete config ")
130 || lower.starts_with("history config ")
131 || lower.starts_with("list config ")
132 || lower.starts_with("watch config ")
133 || lower.starts_with("incr config ")
134 || lower.starts_with("decr config ")
135 || lower.starts_with("add config ")
136 || lower.starts_with("invalidate config ")
137 || lower.starts_with("invalidate tags ")
138 || lower.starts_with("set config ")
139 || lower.starts_with("set secret ")
140 || lower.starts_with("set kv ")
141 || lower.starts_with("set tenant")
142 || lower.starts_with("show create ")
143 || lower.starts_with("show config")
144 || lower.starts_with("show collections")
145 || lower.starts_with("show tables")
146 || lower.starts_with("show queues")
147 || lower.starts_with("show vectors")
148 || lower.starts_with("show documents")
149 || lower.starts_with("show timeseries")
150 || lower.starts_with("show graphs")
151 || lower.starts_with("kv ")
152 || lower.starts_with("show kv")
153 || lower.starts_with("show configs")
154 || lower.starts_with("show vaults")
155 || lower.starts_with("show schema")
156 || lower.starts_with("show indices")
157 || lower.starts_with("show indexes")
158 || lower.starts_with("show sample ")
159 || lower.starts_with("show secret")
160 || lower.starts_with("show stats")
161 || lower.starts_with("show tenant")
162 || lower.starts_with("show policies")
163 || lower.starts_with("show effective ")
164 || lower.starts_with("rank of ")
165 || lower.starts_with("rank range ")
166 || lower.starts_with("approx rank of ")
167 || lower.starts_with("approximate rank of ")
168 || lower.starts_with("zrank ")
169 || lower.starts_with("zrange ")
170 || lower.starts_with("describe ")
171 || lower.starts_with("desc ")
172 {
173 if lower.starts_with("select ") && has_sparql_variable(&lower) {
177 return QueryMode::Sparql;
178 }
179 return QueryMode::Sql;
180 }
181
182 if is_natural_language(&lower) {
184 return QueryMode::Natural;
185 }
186
187 QueryMode::Unknown
188}
189
190fn has_sparql_pattern(lower: &str) -> bool {
192 let has_var = has_sparql_variable(lower);
198
199 let has_triple_pattern = lower.contains(" where {") || lower.contains(" where{");
201
202 let has_prefix_pattern = lower.contains(":")
204 && (lower.contains(":<")
205 || lower.contains("> :")
206 || lower.contains(" :") && lower.contains("?"));
207
208 has_var || has_triple_pattern || has_prefix_pattern
209}
210
211fn has_sparql_variable(input: &str) -> bool {
212 let bytes = input.as_bytes();
213 bytes
214 .windows(2)
215 .any(|pair| pair[0] == b'?' && is_sparql_variable_start(pair[1]))
216}
217
218fn is_sparql_variable_start(byte: u8) -> bool {
219 byte.is_ascii_alphabetic() || byte == b'_'
220}
221
222fn is_natural_language(lower: &str) -> bool {
224 let question_starters = [
226 "find ", "show ", "list ", "what ", "which ", "where ", "how ", "who ", "get ", "give ",
227 "tell ", "display ", "search ", "look ",
228 ];
229
230 let nl_patterns = [
232 " with ",
233 " for ",
234 " that ",
235 " have ",
236 " has ",
237 " can ",
238 " are ",
239 " is ",
240 " all ",
241 " me ",
242 " the ",
243 " from ",
244 " to ",
245 " on ",
246 " in ",
247 "vulnerable",
248 "credential",
249 "password",
250 "user",
251 "host",
252 "service",
253 "connected",
254 "reachable",
255 "exposed",
256 "critical",
257 ];
258
259 for starter in question_starters.iter() {
261 if lower.starts_with(starter) {
262 return true;
263 }
264 }
265
266 let pattern_count = nl_patterns.iter().filter(|p| lower.contains(*p)).count();
268
269 pattern_count >= 2
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275
276 #[test]
277 fn test_sql_detection() {
278 assert_eq!(
279 detect_mode("SELECT * FROM users WHERE id = 1"),
280 QueryMode::Sql
281 );
282 assert_eq!(detect_mode("select name, age from hosts"), QueryMode::Sql);
283 assert_eq!(
284 detect_mode("FROM hosts h WHERE h.os = 'Linux'"),
285 QueryMode::Sql
286 );
287 assert_eq!(
288 detect_mode("INSERT INTO users VALUES (1, 'alice')"),
289 QueryMode::Sql
290 );
291 assert_eq!(
292 detect_mode("UPDATE hosts SET status = 'active'"),
293 QueryMode::Sql
294 );
295 assert_eq!(
296 detect_mode("DELETE FROM logs WHERE age > 30"),
297 QueryMode::Sql
298 );
299 assert_eq!(
300 detect_mode("QUEUE GROUP CREATE tasks workers"),
301 QueryMode::Sql
302 );
303 assert_eq!(
304 detect_mode("EVENTS BACKFILL users TO audit"),
305 QueryMode::Sql
306 );
307 assert_eq!(detect_mode("TREE VALIDATE forest.org"), QueryMode::Sql);
308 assert_eq!(
309 detect_mode("VECTOR SEARCH embeddings SIMILAR TO [1.0, 0.0] LIMIT 5"),
310 QueryMode::Sql
311 );
312 assert_eq!(
313 detect_mode("HYBRID FROM hosts VECTOR SEARCH embeddings SIMILAR TO [1.0, 0.0] LIMIT 5"),
314 QueryMode::Sql
315 );
316 assert_eq!(
317 detect_mode("ASK 'what happened on host 10.0.0.1?' USING groq"),
318 QueryMode::Sql
319 );
320 assert_eq!(
321 detect_mode("SELECT name FROM t WHERE id = ?"),
322 QueryMode::Sql
323 );
324 assert_eq!(
325 detect_mode("SELECT name FROM t WHERE id = ?1"),
326 QueryMode::Sql
327 );
328 assert_eq!(
329 detect_mode("INSERT INTO t (id, name) VALUES (?, ?)"),
330 QueryMode::Sql
331 );
332 assert_eq!(
333 detect_mode("SET SECRET red.secret.api = 'x'"),
334 QueryMode::Sql
335 );
336 assert_eq!(detect_mode("SHOW SECRET red.secret"), QueryMode::Sql);
337 assert_eq!(detect_mode("SHOW SECRETS"), QueryMode::Sql);
338 assert_eq!(detect_mode("VAULT PUT secrets.api = 'x'"), QueryMode::Sql);
339 assert_eq!(
340 detect_mode("LIST KV settings PREFIX feature LIMIT 10"),
341 QueryMode::Sql
342 );
343 assert_eq!(detect_mode("SHOW SAMPLE users"), QueryMode::Sql);
344 assert_eq!(detect_mode("SHOW TABLES"), QueryMode::Sql);
345 assert_eq!(detect_mode("SHOW QUEUES"), QueryMode::Sql);
346 assert_eq!(detect_mode("SHOW VECTORS"), QueryMode::Sql);
347 assert_eq!(detect_mode("SHOW DOCUMENTS"), QueryMode::Sql);
348 assert_eq!(detect_mode("SHOW TIMESERIES"), QueryMode::Sql);
349 assert_eq!(detect_mode("SHOW GRAPHS"), QueryMode::Sql);
350 assert_eq!(detect_mode("SHOW KV"), QueryMode::Sql);
351 assert_eq!(detect_mode("SHOW KVS"), QueryMode::Sql);
352 assert_eq!(detect_mode("SHOW CONFIGS"), QueryMode::Sql);
353 assert_eq!(detect_mode("SHOW VAULTS"), QueryMode::Sql);
354 assert_eq!(detect_mode("SHOW SCHEMA users"), QueryMode::Sql);
355 assert_eq!(detect_mode("SHOW CREATE TABLE users"), QueryMode::Sql);
356 assert_eq!(detect_mode("DESCRIBE users"), QueryMode::Sql);
357 assert_eq!(detect_mode("DESC users"), QueryMode::Sql);
358 assert_eq!(detect_mode("SHOW INDICES"), QueryMode::Sql);
359 assert_eq!(detect_mode("SHOW INDEXES"), QueryMode::Sql);
360 assert_eq!(detect_mode("SHOW STATS users"), QueryMode::Sql);
361 assert_eq!(detect_mode("EXPLAIN MIGRATION *"), QueryMode::Sql);
362 }
363
364 #[test]
365 fn test_gremlin_detection() {
366 assert_eq!(detect_mode("g.V()"), QueryMode::Gremlin);
367 assert_eq!(detect_mode("g.V().hasLabel('host')"), QueryMode::Gremlin);
368 assert_eq!(
369 detect_mode("g.V().out('connects').in('has_service')"),
370 QueryMode::Gremlin
371 );
372 assert_eq!(
373 detect_mode("g.E().hasLabel('auth_access')"),
374 QueryMode::Gremlin
375 );
376 assert_eq!(
377 detect_mode("__.out('knows').has('name', 'bob')"),
378 QueryMode::Gremlin
379 );
380 assert_eq!(
381 detect_mode("g.V('host:10.0.0.1').repeat(out()).times(3)"),
382 QueryMode::Gremlin
383 );
384 }
385
386 #[test]
387 fn test_cypher_detection() {
388 assert_eq!(
389 detect_mode("MATCH (a)-[r]->(b) RETURN a, b"),
390 QueryMode::Cypher
391 );
392 assert_eq!(
393 detect_mode("MATCH (h:Host)-[:HAS_SERVICE]->(s:Service)"),
394 QueryMode::Cypher
395 );
396 assert_eq!(
397 detect_mode("match (n) where n.ip = '10.0.0.1' return n"),
398 QueryMode::Cypher
399 );
400 assert_eq!(
401 detect_mode("MATCH(a:User) RETURN a.name"),
402 QueryMode::Cypher
403 );
404 }
405
406 #[test]
407 fn test_sparql_detection() {
408 assert_eq!(
409 detect_mode("SELECT ?name WHERE { ?s :name ?name }"),
410 QueryMode::Sparql
411 );
412 assert_eq!(
413 detect_mode("PREFIX ex: <http://example.org/> SELECT ?x WHERE { ?x ex:type ?t }"),
414 QueryMode::Sparql
415 );
416 assert_eq!(
417 detect_mode("SELECT ?host ?ip WHERE { ?host :hasIP ?ip }"),
418 QueryMode::Sparql
419 );
420 assert_eq!(
421 detect_mode("SELECT ?x WHERE { ?x rdf:type :Foo }"),
422 QueryMode::Sparql
423 );
424 }
425
426 #[test]
427 fn test_path_detection() {
428 assert_eq!(
429 detect_mode("PATH FROM host('10.0.0.1') TO host('10.0.0.2')"),
430 QueryMode::Path
431 );
432 assert_eq!(
433 detect_mode("PATHS ALL FROM credential('admin') TO host('db')"),
434 QueryMode::Path
435 );
436 assert_eq!(
437 detect_mode("path from user('root') to service('ssh') via auth_access"),
438 QueryMode::Path
439 );
440 }
441
442 #[test]
443 fn test_natural_detection() {
444 assert_eq!(
445 detect_mode("find all hosts with ssh open"),
446 QueryMode::Natural
447 );
448 assert_eq!(
449 detect_mode("show me vulnerable services"),
450 QueryMode::Natural
451 );
452 assert_eq!(
453 detect_mode("what credentials can reach the database?"),
454 QueryMode::Natural
455 );
456 assert_eq!(
457 detect_mode("list users with weak passwords"),
458 QueryMode::Natural
459 );
460 assert_eq!(
461 detect_mode("\"find hosts connected to 10.0.0.1\""),
462 QueryMode::Natural
463 );
464 assert_eq!(
465 detect_mode("which hosts have critical vulnerabilities?"),
466 QueryMode::Natural
467 );
468 }
469
470 #[test]
471 fn test_edge_cases() {
472 assert_eq!(detect_mode(""), QueryMode::Unknown);
474
475 assert_eq!(detect_mode(" "), QueryMode::Unknown);
477
478 assert_eq!(detect_mode("SELECT"), QueryMode::Unknown); assert_eq!(detect_mode("G.V()"), QueryMode::Gremlin);
481 assert_eq!(detect_mode("Match (a) RETURN a"), QueryMode::Cypher);
482 }
483}