1use crate::data::csv_fixes::quote_if_needed;
2use crate::parser::{ParseState, Schema};
3use crate::recursive_parser::{detect_cursor_context, CursorContext, LogicalOp};
4
5#[derive(Debug, Clone)]
6pub struct CursorAwareParser {
7 schema: Schema,
8}
9
10#[derive(Debug)]
11pub struct ParseResult {
12 pub suggestions: Vec<String>,
13 pub context: String,
14 pub partial_word: Option<String>,
15}
16
17impl Default for CursorAwareParser {
18 fn default() -> Self {
19 Self::new()
20 }
21}
22
23impl CursorAwareParser {
24 #[must_use]
25 pub fn new() -> Self {
26 Self {
27 schema: Schema::new(),
28 }
29 }
30
31 pub fn set_schema(&mut self, schema: Schema) {
32 self.schema = schema;
33 }
34
35 pub fn update_single_table(&mut self, table_name: String, columns: Vec<String>) {
36 self.schema.set_single_table(&table_name, columns);
37 }
38
39 #[must_use]
40 pub fn get_table_columns(&self, table_name: &str) -> Vec<String> {
41 self.schema.get_columns(table_name)
42 }
43
44 #[must_use]
45 pub fn get_completions(&self, query: &str, cursor_pos: usize) -> ParseResult {
46 let (cursor_context, partial_word) = detect_cursor_context(query, cursor_pos);
48
49 let partial_word = partial_word.or_else(|| self.extract_word_at_cursor(query, cursor_pos));
51
52 let default_table = self
53 .schema
54 .get_first_table_name()
55 .unwrap_or("trade_deal".to_string());
56
57 let (suggestions, context_str) = match &cursor_context {
58 CursorContext::SelectClause => {
59 let mut cols = self
61 .schema
62 .get_columns(&default_table)
63 .into_iter()
64 .map(|col| quote_if_needed(&col))
65 .collect::<Vec<_>>();
66 cols.push("*".to_string());
67
68 cols.extend(vec![
70 "ROUND(".to_string(),
71 "ABS(".to_string(),
72 "FLOOR(".to_string(),
73 "CEILING(".to_string(),
74 "CEIL(".to_string(),
75 "MOD(".to_string(),
76 "QUOTIENT(".to_string(),
77 "POWER(".to_string(),
78 "POW(".to_string(),
79 "SQRT(".to_string(),
80 "EXP(".to_string(),
81 "LN(".to_string(),
82 "LOG(".to_string(),
83 "LOG10(".to_string(),
84 "PI(".to_string(),
85 "TEXTJOIN(".to_string(),
86 "DATEDIFF(".to_string(),
87 "DATEADD(".to_string(),
88 "NOW(".to_string(),
89 "TODAY(".to_string(),
90 ]);
91
92 (cols, "SelectClause".to_string())
98 }
99 CursorContext::FromClause => {
100 let tables = self.schema.get_table_names();
101 (tables, "FromClause".to_string())
102 }
103 CursorContext::WhereClause | CursorContext::AfterLogicalOp(_) => {
104 let mut suggestions = self
106 .schema
107 .get_columns(&default_table)
108 .into_iter()
109 .map(|col| quote_if_needed(&col))
110 .collect::<Vec<_>>();
111
112 suggestions.extend(vec![
114 "ROUND(".to_string(),
115 "ABS(".to_string(),
116 "FLOOR(".to_string(),
117 "CEILING(".to_string(),
118 "CEIL(".to_string(),
119 "MOD(".to_string(),
120 "QUOTIENT(".to_string(),
121 "POWER(".to_string(),
122 "POW(".to_string(),
123 "SQRT(".to_string(),
124 "EXP(".to_string(),
125 "LN(".to_string(),
126 "LOG(".to_string(),
127 "LOG10(".to_string(),
128 "PI(".to_string(),
129 "TEXTJOIN(".to_string(),
130 "DATEDIFF(".to_string(),
131 "DATEADD(".to_string(),
132 "NOW(".to_string(),
133 "TODAY(".to_string(),
134 ]);
135
136 let add_keywords = if let Some(ref partial) = partial_word {
138 let partial_lower = partial.to_lowercase();
139 !suggestions
140 .iter()
141 .any(|col| col.to_lowercase().starts_with(&partial_lower))
142 } else {
143 true
144 };
145
146 if add_keywords {
147 suggestions.extend(vec![
148 "AND".to_string(),
149 "OR".to_string(),
150 "IN".to_string(),
151 "ORDER BY".to_string(),
152 ]);
153 }
154
155 let ctx = match &cursor_context {
156 CursorContext::AfterLogicalOp(LogicalOp::And) => "AfterAND",
157 CursorContext::AfterLogicalOp(LogicalOp::Or) => "AfterOR",
158 _ => "WhereClause",
159 };
160 (suggestions, ctx.to_string())
161 }
162 CursorContext::AfterColumn(col_name) => {
163 let property_type = self
165 .get_property_type(col_name)
166 .unwrap_or("string".to_string());
167 let suggestions = self.get_string_method_suggestions(&property_type, &partial_word);
168 (suggestions, "AfterColumn".to_string())
169 }
170 CursorContext::AfterComparisonOp(col_name, op) => {
171 let property_type = self
173 .get_property_type(col_name)
174 .unwrap_or("string".to_string());
175 let suggestions = match property_type.as_str() {
176 "datetime" => {
177 let mut suggestions = vec!["DateTime(".to_string()];
179 suggestions.extend(vec![
181 "DateTime.Today".to_string(),
182 "DateTime.Now".to_string(),
183 ]);
184 suggestions
185 }
186 "string" => {
187 vec!["''".to_string()]
189 }
190 "numeric" => {
191 vec![]
193 }
194 _ => vec![],
195 };
196 (suggestions, format!("AfterComparison({col_name} {op})"))
197 }
198 CursorContext::InMethodCall(obj, method) => {
199 let property_type = self.get_property_type(obj).unwrap_or("string".to_string());
200 let suggestions = self.get_string_method_suggestions(&property_type, &partial_word);
201 (suggestions, format!("InMethodCall({obj}.{method})"))
202 }
203 CursorContext::InExpression => {
204 let mut suggestions = self
206 .schema
207 .get_columns(&default_table)
208 .into_iter()
209 .map(|col| quote_if_needed(&col))
210 .collect::<Vec<_>>();
211
212 suggestions.extend(vec![
214 "ROUND(".to_string(),
215 "ABS(".to_string(),
216 "FLOOR(".to_string(),
217 "CEILING(".to_string(),
218 "CEIL(".to_string(),
219 "MOD(".to_string(),
220 "QUOTIENT(".to_string(),
221 "POWER(".to_string(),
222 "POW(".to_string(),
223 "SQRT(".to_string(),
224 "EXP(".to_string(),
225 "LN(".to_string(),
226 "LOG(".to_string(),
227 "LOG10(".to_string(),
228 "PI(".to_string(),
229 "AND".to_string(),
230 "OR".to_string(),
231 ]);
232 (suggestions, "InExpression".to_string())
233 }
234 CursorContext::OrderByClause => {
235 let mut suggestions = Vec::new();
237
238 let selected_columns = self.extract_selected_columns(query, query.len());
240
241 if !selected_columns.is_empty() && !selected_columns.contains(&"*".to_string()) {
243 suggestions.extend(selected_columns);
244 } else {
245 suggestions.extend(
248 self.schema
249 .get_columns(&default_table)
250 .into_iter()
251 .map(|col| quote_if_needed(&col)),
252 );
253 }
254
255 suggestions.extend(vec!["ASC".to_string(), "DESC".to_string()]);
257 (suggestions, "OrderByClause".to_string())
258 }
259 CursorContext::Unknown => {
260 let safe_cursor_pos = self.find_safe_boundary(query, cursor_pos.min(query.len()));
263 let query_before_cursor = &query[..safe_cursor_pos];
264 let context = self.determine_context(query_before_cursor);
265 let suggestions = self.get_suggestions_for_context(&context, &partial_word, query);
266 return ParseResult {
267 suggestions,
268 context: format!("{context:?} (partial: {partial_word:?})"),
269 partial_word,
270 };
271 }
272 };
273
274 let mut final_suggestions = suggestions;
276 let is_method_context = matches!(
277 cursor_context,
278 CursorContext::AfterColumn(_)
279 | CursorContext::InMethodCall(_, _)
280 | CursorContext::AfterComparisonOp(_, _)
281 );
282
283 if let Some(ref partial) = partial_word {
284 if !is_method_context {
285 final_suggestions.retain(|suggestion| {
287 if let Some(partial_without_quote) = partial.strip_prefix('"') {
289 if suggestion.starts_with('"')
294 && suggestion.ends_with('"')
295 && suggestion.len() > 2
296 {
297 let suggestion_without_quotes = &suggestion[1..suggestion.len() - 1];
299 suggestion_without_quotes
300 .to_lowercase()
301 .starts_with(&partial_without_quote.to_lowercase())
302 } else if suggestion.starts_with('"') && suggestion.len() > 1 {
303 let suggestion_without_quote = &suggestion[1..];
305 suggestion_without_quote
306 .to_lowercase()
307 .starts_with(&partial_without_quote.to_lowercase())
308 } else {
309 suggestion
311 .to_lowercase()
312 .starts_with(&partial_without_quote.to_lowercase())
313 }
314 } else {
315 let suggestion_to_check = if suggestion.starts_with('"')
318 && suggestion.ends_with('"')
319 && suggestion.len() > 2
320 {
321 &suggestion[1..suggestion.len() - 1]
323 } else if suggestion.starts_with('"') && suggestion.len() > 1 {
324 &suggestion[1..]
326 } else {
327 suggestion
328 };
329
330 suggestion_to_check
332 .to_lowercase()
333 .starts_with(&partial.to_lowercase())
334 }
335 });
336 }
337 }
338
339 ParseResult {
340 suggestions: final_suggestions,
341 context: format!("{context_str} (partial: {partial_word:?})"),
342 partial_word,
343 }
344 }
345
346 fn extract_word_at_cursor(&self, query: &str, cursor_pos: usize) -> Option<String> {
347 if cursor_pos == 0 || cursor_pos > query.len() {
348 return None;
349 }
350
351 let chars: Vec<char> = query.chars().collect();
352
353 let mut start = cursor_pos;
355 let mut end = cursor_pos;
356
357 while start > 0 && Self::is_word_char(chars.get(start - 1).copied().unwrap_or(' ')) {
359 start -= 1;
360 }
361
362 while end < chars.len() && Self::is_word_char(chars.get(end).copied().unwrap_or(' ')) {
364 end += 1;
365 }
366
367 if start < end {
369 let partial: String = chars[start..cursor_pos.min(end)].iter().collect();
371 if partial.is_empty() {
372 None
373 } else {
374 Some(partial)
375 }
376 } else {
377 None
378 }
379 }
380
381 fn is_word_char(ch: char) -> bool {
382 ch.is_alphanumeric() || ch == '_'
383 }
384
385 fn determine_context(&self, query_before_cursor: &str) -> ParseState {
386 let query_upper = query_before_cursor.to_uppercase();
387
388 let trimmed = query_before_cursor.trim();
391 let upper_trimmed = trimmed.to_uppercase();
395 let ends_with_and_or = upper_trimmed.ends_with(" AND") ||
396 upper_trimmed.ends_with(" OR") ||
397 upper_trimmed.ends_with(" AND ") || upper_trimmed.ends_with(" OR "); let words_check: Vec<&str> = query_upper.split_whitespace().collect();
402 let last_word_is_and_or = words_check
403 .last()
404 .is_some_and(|w| *w == "AND" || *w == "OR");
405
406 if ends_with_and_or || last_word_is_and_or {
407 if query_upper.contains("WHERE") {
409 return ParseState::InWhere;
411 }
412 }
413
414 let words: Vec<&str> = query_upper.split_whitespace().collect();
415
416 if words.is_empty() {
417 return ParseState::Start;
418 }
419
420 let mut last_keyword_idx = None;
422 let mut last_keyword = "";
423
424 for (i, word) in words.iter().enumerate() {
425 match *word {
426 "SELECT" => {
427 last_keyword_idx = Some(i);
428 last_keyword = "SELECT";
429 }
430 "FROM" => {
431 last_keyword_idx = Some(i);
432 last_keyword = "FROM";
433 }
434 "WHERE" => {
435 last_keyword_idx = Some(i);
436 last_keyword = "WHERE";
437 }
438 "AND" | "OR" => {
439 if last_keyword == "WHERE" {
441 last_keyword_idx = Some(i);
442 last_keyword = "WHERE"; }
444 }
445 "IN" => {
446 if last_keyword == "WHERE" {
448 last_keyword_idx = Some(i);
449 last_keyword = "WHERE";
450 }
451 }
452 "ORDER" => {
453 if i + 1 < words.len() && words[i + 1] == "BY" {
455 last_keyword_idx = Some(i);
456 last_keyword = "ORDER BY";
457 }
458 }
459 _ => {}
460 }
461 }
462
463 match last_keyword {
464 "SELECT" => {
465 if let Some(idx) = last_keyword_idx {
466 let tokens_after_select = words.len() - idx - 1;
468 if tokens_after_select == 0 {
469 ParseState::AfterSelect
470 } else {
471 if words[(idx + 1)..].contains(&"FROM") {
473 ParseState::AfterTable } else {
475 ParseState::InColumnList
476 }
477 }
478 } else {
479 ParseState::AfterSelect
480 }
481 }
482 "FROM" => {
483 if let Some(idx) = last_keyword_idx {
484 let tokens_after_from = words.len() - idx - 1;
485 if tokens_after_from == 0 {
486 ParseState::AfterFrom
487 } else {
488 ParseState::AfterTable
489 }
490 } else {
491 ParseState::AfterFrom
492 }
493 }
494 "WHERE" => ParseState::InWhere,
495 "ORDER BY" => ParseState::InOrderBy,
496 _ => {
497 if query_upper.contains("SELECT")
499 && query_upper.contains("FROM")
500 && query_upper.contains("WHERE")
501 {
502 ParseState::InWhere
503 } else if query_upper.contains("SELECT") && query_upper.contains("FROM") {
504 ParseState::AfterTable
505 } else if query_upper.contains("SELECT") {
506 ParseState::InColumnList
507 } else {
508 ParseState::Start
509 }
510 }
511 }
512 }
513
514 fn get_suggestions_for_context(
515 &self,
516 context: &ParseState,
517 partial_word: &Option<String>,
518 query: &str,
519 ) -> Vec<String> {
520 let default_table = self
521 .schema
522 .get_first_table_name()
523 .unwrap_or("trade_deal".to_string());
524
525 let mut suggestions = match context {
526 ParseState::Start => vec!["SELECT".to_string()],
527 ParseState::AfterSelect => {
528 let mut cols = self
529 .schema
530 .get_columns(&default_table)
531 .into_iter()
532 .map(|col| quote_if_needed(&col))
533 .collect::<Vec<_>>();
534 cols.push("*".to_string());
535 cols
536 }
537 ParseState::InColumnList => {
538 let mut cols = self
539 .schema
540 .get_columns(&default_table)
541 .into_iter()
542 .map(|col| quote_if_needed(&col))
543 .collect::<Vec<_>>();
544 cols.push("FROM".to_string());
545 cols
546 }
547 ParseState::AfterFrom => self.schema.get_table_names(),
548 ParseState::AfterTable => {
549 vec!["WHERE".to_string(), "ORDER BY".to_string()]
550 }
551 ParseState::InWhere => {
552 let mut suggestions = self
554 .schema
555 .get_columns(&default_table)
556 .into_iter()
557 .map(|col| quote_if_needed(&col))
558 .collect::<Vec<_>>();
559
560 let add_keywords = if let Some(partial) = partial_word {
562 let partial_lower = partial.to_lowercase();
563 let matching_columns = suggestions
564 .iter()
565 .any(|col| col.to_lowercase().starts_with(&partial_lower));
566 !matching_columns } else {
568 true };
570
571 if add_keywords {
572 suggestions.extend(vec![
573 "AND".to_string(),
574 "OR".to_string(),
575 "IN".to_string(),
576 "ORDER BY".to_string(),
577 ]);
578 }
579
580 suggestions
581 }
582 ParseState::InOrderBy => {
583 let mut suggestions = Vec::new();
584
585 let selected_columns = self.extract_selected_columns(query, query.len());
587
588 if !selected_columns.is_empty() && !selected_columns.contains(&"*".to_string()) {
590 suggestions.extend(selected_columns);
591 } else {
592 suggestions.extend(
594 self.schema
595 .get_columns(&default_table)
596 .into_iter()
597 .map(|col| quote_if_needed(&col)),
598 );
599 }
600
601 suggestions.extend(vec!["ASC".to_string(), "DESC".to_string()]);
603 suggestions
604 }
605 _ => vec![],
606 };
607
608 if let Some(partial) = partial_word {
610 suggestions.retain(|suggestion| {
611 suggestion
612 .to_lowercase()
613 .starts_with(&partial.to_lowercase())
614 });
615 }
616
617 suggestions
618 }
619
620 fn extract_selected_columns(&self, query: &str, cursor_pos: usize) -> Vec<String> {
621 let mut selected_columns = Vec::new();
623
624 let query_upper = query.to_uppercase();
626 if let Some(select_pos) = query_upper.find("SELECT") {
627 let end_pos = query_upper
629 .find("FROM")
630 .unwrap_or(cursor_pos)
631 .min(cursor_pos);
632
633 if select_pos + 6 < end_pos {
635 let select_clause = &query[(select_pos + 6)..end_pos];
636
637 for part in select_clause.split(',') {
639 let trimmed = part.trim();
640 if !trimmed.is_empty() {
641 let col_name = if trimmed.starts_with('"') {
643 if let Some(close_quote_pos) = trimmed[1..].find('"') {
645 &trimmed[..close_quote_pos + 2]
647 } else {
648 trimmed
650 }
651 } else {
652 if let Some(space_pos) = trimmed.find(char::is_whitespace) {
654 &trimmed[..space_pos]
655 } else {
656 trimmed
657 }
658 };
659
660 selected_columns.push(col_name.to_string());
662 }
663 }
664 }
665 }
666
667 selected_columns
668 }
669
670 fn detect_method_call_context(
671 &self,
672 query_before_cursor: &str,
673 _cursor_pos: usize,
674 ) -> Option<(String, String)> {
675 if let Some(dot_pos) = query_before_cursor.rfind('.') {
681 let text_after_dot = &query_before_cursor[dot_pos + 1..];
684
685 if text_after_dot.contains(')')
688 && (text_after_dot.contains(" AND ")
689 || text_after_dot.contains(" OR ")
690 || text_after_dot.trim().ends_with(" AND")
691 || text_after_dot.trim().ends_with(" OR"))
692 {
693 return None; }
695
696 let before_dot = &query_before_cursor[..dot_pos];
698
699 let mut property_start = dot_pos;
701 let chars: Vec<char> = before_dot.chars().collect();
702
703 while property_start > 0 {
704 let char_pos = property_start - 1;
705 if char_pos < chars.len() {
706 let ch = chars[char_pos];
707 if ch.is_alphanumeric() || ch == '_' {
708 property_start -= 1;
709 } else {
710 break;
711 }
712 } else {
713 break;
714 }
715 }
716
717 if property_start < dot_pos {
718 let property_name = before_dot[property_start..].trim().to_string();
719
720 if let Some(property_type) = self.get_property_type(&property_name) {
722 return Some((property_name, property_type));
723 }
724 }
725 }
726
727 None
728 }
729
730 fn get_property_type(&self, property_name: &str) -> Option<String> {
731 let property_lower = property_name.to_lowercase();
735
736 let string_properties = [
738 "platformorderid",
739 "dealid",
740 "externalorderid",
741 "parentorderid",
742 "instrumentid",
743 "instrumentname",
744 "instrumenttype",
745 "isin",
746 "cusip",
747 "ticker",
748 "exchange",
749 "counterparty",
750 "counterpartyid",
751 "counterpartytype",
752 "counterpartycountry",
753 "trader",
754 "portfolio",
755 "strategy",
756 "desk",
757 "status",
758 "confirmationstatus",
759 "settlementstatus",
760 "allocationstatus",
761 "currency",
762 "side",
763 "producttype",
764 "venue",
765 "clearinghouse",
766 "prime",
767 "comments",
768 "book",
769 "source",
770 "sourcesystem",
771 ];
772
773 let numeric_properties = [
775 "price",
776 "quantity",
777 "notional",
778 "commission",
779 "accrual",
780 "netamount",
781 "accruedinterest",
782 "grossamount",
783 "settlementamount",
784 "fees",
785 "tax",
786 ];
787
788 let datetime_properties = [
790 "tradedate",
791 "settlementdate",
792 "createddate",
793 "modifieddate",
794 "valuedate",
795 "maturitydate",
796 "confirmationdate",
797 "executiondate",
798 "lastmodifieddate",
799 ];
800
801 if string_properties.contains(&property_lower.as_str()) {
802 Some("string".to_string())
803 } else if numeric_properties.contains(&property_lower.as_str()) {
804 Some("numeric".to_string())
805 } else if datetime_properties.contains(&property_lower.as_str()) {
806 Some("datetime".to_string())
807 } else {
808 Some("string".to_string())
810 }
811 }
812
813 fn find_safe_boundary(&self, s: &str, pos: usize) -> usize {
815 if pos >= s.len() {
816 return s.len();
817 }
818
819 if s.is_char_boundary(pos) {
821 return pos;
822 }
823
824 let mut safe_pos = pos;
826 while safe_pos > 0 && !s.is_char_boundary(safe_pos) {
827 safe_pos -= 1;
828 }
829 safe_pos
830 }
831
832 #[cfg(test)]
833 #[must_use]
834 pub fn test_extract_selected_columns(&self, query: &str, cursor_pos: usize) -> Vec<String> {
835 self.extract_selected_columns(query, cursor_pos)
836 }
837
838 fn get_string_method_suggestions(
839 &self,
840 property_type: &str,
841 partial_word: &Option<String>,
842 ) -> Vec<String> {
843 let mut suggestions = Vec::new();
844
845 match property_type {
846 "string" => {
847 let string_methods = vec![
851 "Contains('')",
852 "StartsWith('')",
853 "EndsWith('')",
854 "IndexOf('')",
855 "Substring(0, 5)",
856 "ToLower()",
857 "ToUpper()",
858 "Trim()",
859 "TrimStart()",
860 "TrimEnd()",
861 "IsNullOrEmpty()",
862 "Replace('', '')",
863 "Length()", ];
865
866 if let Some(partial) = partial_word {
867 let partial_lower = partial.to_lowercase();
868 for method in string_methods {
869 if method.to_lowercase().starts_with(&partial_lower) {
870 suggestions.push(method.to_string());
871 }
872 }
873 } else {
874 suggestions.extend(
875 string_methods
876 .into_iter()
877 .map(std::string::ToString::to_string),
878 );
879 }
880 }
881 "numeric" | "integer" | "float" | "decimal" => {
882 let numeric_string_methods = vec![
884 "Contains('')",
885 "StartsWith('')",
886 "EndsWith('')",
887 "ToString()",
888 "Length()", ];
891
892 if let Some(partial) = partial_word {
893 let partial_lower = partial.to_lowercase();
894 for method in numeric_string_methods {
895 if method.to_lowercase().starts_with(&partial_lower) {
896 suggestions.push(method.to_string());
897 }
898 }
899 } else {
900 suggestions.extend(
901 numeric_string_methods
902 .into_iter()
903 .map(std::string::ToString::to_string),
904 );
905 }
906 }
907 "datetime" => {
908 let datetime_methods = vec![
910 "Year()", "Month()", "Day()", "ToString(\"yyyy-MM-dd\")",
914 "AddDays(1)",
915 "Contains('')",
917 "StartsWith('')",
918 "EndsWith('')",
919 "Length()", ];
921
922 if let Some(partial) = partial_word {
923 let partial_lower = partial.to_lowercase();
924 for method in datetime_methods {
925 if method.to_lowercase().starts_with(&partial_lower) {
926 suggestions.push(method.to_string());
927 }
928 }
929 } else {
930 suggestions.extend(
931 datetime_methods
932 .into_iter()
933 .map(std::string::ToString::to_string),
934 );
935 }
936 }
937 _ => {
938 suggestions.push("ToString()".to_string());
940 }
941 }
942
943 suggestions
944 }
945}
946
947#[cfg(test)]
948mod tests {
949 use super::*;
950
951 fn create_test_parser() -> CursorAwareParser {
952 CursorAwareParser::new()
953 }
954
955 #[test]
956 fn test_basic_select_completion() {
957 let parser = create_test_parser();
958
959 let result = parser.get_completions("", 0);
961 println!("Context for empty query: {}", result.context);
962 assert_eq!(result.suggestions, vec!["SELECT"]);
963 assert!(result.context.contains("Start") || result.context.contains("Unknown"));
964
965 let result = parser.get_completions("SELECT ", 7);
967 println!("Context for 'SELECT ': {}", result.context);
968 assert!(result.suggestions.contains(&"*".to_string()));
969 assert!(result.suggestions.contains(&"dealId".to_string()));
970 assert!(result.context.contains("AfterSelect") || result.context.contains("SelectClause"));
971 }
972
973 #[test]
974 fn test_where_clause_completion() {
975 let parser = create_test_parser();
976
977 let query = "SELECT * FROM trade_deal WHERE ";
979 let result = parser.get_completions(query, query.len());
980 println!("Context for WHERE clause: {}", result.context);
981 assert!(result.suggestions.contains(&"dealId".to_string()));
982 assert!(result.suggestions.contains(&"platformOrderId".to_string()));
983 assert!(result.context.contains("InWhere") || result.context.contains("WhereClause"));
984 }
985
986 #[test]
987 fn test_method_call_detection() {
988 let parser = create_test_parser();
989
990 let query = "SELECT * FROM trade_deal WHERE platformOrderId.";
992 let result = parser.get_completions(query, query.len());
993 println!("Context for method call: {}", result.context);
994 println!("Suggestions: {:?}", result.suggestions);
995 assert!(result.suggestions.contains(&"Contains('')".to_string()));
996 assert!(result.suggestions.contains(&"StartsWith('')".to_string()));
997 assert!(result.context.contains("MethodCall") || result.context.contains("AfterColumn"));
998 }
999
1000 #[test]
1001 fn test_and_operator_context() {
1002 let parser = create_test_parser();
1003
1004 let query = "SELECT * FROM trade_deal WHERE allocationStatus.Contains(\"All\") AND ";
1006 let result = parser.get_completions(query, query.len());
1007 println!("Context after AND: {}", result.context);
1008 assert!(result.suggestions.contains(&"dealId".to_string()));
1009 assert!(result.suggestions.contains(&"platformOrderId".to_string()));
1010 assert!(
1011 result.context.contains("InWhere")
1012 || result.context.contains("AfterAND")
1013 || result.context.contains("WhereClause")
1014 );
1015 assert!(!result.context.contains("MethodCall"));
1016 }
1017
1018 #[test]
1019 fn test_and_operator_with_partial_word() {
1020 let parser = create_test_parser();
1021
1022 let query = "SELECT * FROM trade_deal WHERE allocationStatus.Contains(\"All\") AND p";
1024 let result = parser.get_completions(query, query.len());
1025
1026 assert!(result.suggestions.contains(&"platformOrderId".to_string()));
1028 assert!(result.suggestions.contains(&"price".to_string()));
1029 assert!(result.suggestions.contains(&"portfolio".to_string()));
1030
1031 assert!(!result.suggestions.contains(&"dealId".to_string()));
1033 assert!(!result.suggestions.contains(&"quantity".to_string()));
1034
1035 assert!(
1037 result.context.contains("InWhere")
1038 || result.context.contains("WhereClause")
1039 || result.context.contains("AfterAND")
1040 );
1041 assert!(!result.context.contains("MethodCall"));
1042
1043 assert!(result.context.contains("(partial: Some(\"p\"))"));
1045 }
1046
1047 #[test]
1048 fn test_or_operator_context() {
1049 let parser = create_test_parser();
1050
1051 let query = "SELECT * FROM trade_deal WHERE price > 100 OR ";
1053 let result = parser.get_completions(query, query.len());
1054 println!("Context after OR: {}", result.context);
1055 assert!(result.suggestions.contains(&"dealId".to_string()));
1056 assert!(
1057 result.context.contains("InWhere")
1058 || result.context.contains("AfterOR")
1059 || result.context.contains("WhereClause")
1060 );
1061 }
1062
1063 #[test]
1064 fn test_partial_word_extraction() {
1065 let parser = create_test_parser();
1066
1067 assert_eq!(
1069 parser.extract_word_at_cursor("SELECT deal", 11),
1070 Some("deal".to_string())
1071 );
1072 assert_eq!(
1073 parser.extract_word_at_cursor("WHERE p", 7),
1074 Some("p".to_string())
1075 );
1076 assert_eq!(
1077 parser.extract_word_at_cursor("AND platf", 9),
1078 Some("platf".to_string())
1079 );
1080
1081 assert_eq!(parser.extract_word_at_cursor("", 0), None);
1083 assert_eq!(parser.extract_word_at_cursor("SELECT ", 7), None);
1084 }
1085
1086 #[test]
1087 fn test_complex_query_with_multiple_conditions() {
1088 let parser = create_test_parser();
1089
1090 let query = "SELECT * FROM trade_deal WHERE platformOrderId.StartsWith(\"ABC\") AND price > 100 AND ";
1092 let result = parser.get_completions(query, query.len());
1093 println!("Context for complex query: {}", result.context);
1094 assert!(result.suggestions.contains(&"dealId".to_string()));
1095 assert!(
1096 result.context.contains("InWhere")
1097 || result.context.contains("AfterAND")
1098 || result.context.contains("WhereClause")
1099 );
1100 assert!(!result.context.contains("MethodCall"));
1101 }
1102
1103 #[test]
1104 fn test_in_clause_support() {
1105 let parser = create_test_parser();
1106
1107 let query = "SELECT * FROM trade_deal WHERE status IN ";
1109 let result = parser.get_completions(query, query.len());
1110 println!("Context after IN: {}", result.context);
1111 assert!(
1113 result.context.contains("InWhere")
1114 || result.context.contains("WhereClause")
1115 || result.context.contains("Unknown")
1116 );
1117 }
1118
1119 #[test]
1120 fn test_partial_method_name_completion() {
1121 let parser = create_test_parser();
1122
1123 let query = "SELECT * FROM trade_deal WHERE instrumentName.Con";
1125 let result = parser.get_completions(query, query.len());
1126 println!("Context for partial method: {}", result.context);
1127 println!("Suggestions: {:?}", result.suggestions);
1128
1129 assert!(result.context.contains("MethodCall") || result.context.contains("AfterColumn"));
1131 assert!(result.context.contains("(partial: Some(\"Con\"))"));
1132
1133 assert!(result.suggestions.contains(&"Contains('')".to_string()));
1135 assert!(!result.suggestions.contains(&"StartsWith('')".to_string())); }
1137
1138 #[test]
1139 fn test_partial_matching_quoted_identifier() {
1140 let parser = CursorAwareParser::new();
1141 let mut parser = parser;
1143 parser.update_single_table(
1144 "customers".to_string(),
1145 vec![
1146 "Index".to_string(),
1147 "Customer Id".to_string(), "First Name".to_string(), "Company".to_string(),
1150 ],
1151 );
1152
1153 let query = "SELECT customer";
1155 let result = parser.get_completions(query, query.len());
1156
1157 assert!(
1159 result.suggestions.iter().any(|s| s == "\"Customer Id\""),
1160 "Should suggest quoted Customer Id for partial \"customer\". Got: {:?}",
1161 result.suggestions
1162 );
1163 }
1164
1165 #[test]
1166 fn test_case_preservation_in_order_by() {
1167 let parser = CursorAwareParser::new();
1168 let mut parser = parser;
1169 parser.update_single_table(
1170 "customers".to_string(),
1171 vec!["Company".to_string(), "Country".to_string()],
1172 );
1173
1174 let query = "SELECT Company, Country FROM customers ORDER BY Com";
1176 let result = parser.get_completions(query, query.len());
1177
1178 assert!(
1180 result.suggestions.iter().any(|s| s == "Company"),
1181 "Should preserve case in ORDER BY suggestions. Got: {:?}",
1182 result.suggestions
1183 );
1184 }
1185
1186 #[test]
1187 fn test_extract_selected_columns_preserves_case() {
1188 let parser = CursorAwareParser::new();
1189
1190 let query = "SELECT Company, Country FROM customers";
1191 let columns = parser.test_extract_selected_columns(query, query.len());
1192
1193 assert_eq!(columns, vec!["Company", "Country"]);
1194 assert_ne!(
1195 columns,
1196 vec!["company", "country"],
1197 "Should preserve original case"
1198 );
1199 }
1200
1201 #[test]
1202 fn test_filtering_already_selected_columns() {
1203 let parser = CursorAwareParser::new();
1204 let mut parser = parser;
1205 parser.update_single_table(
1206 "customers".to_string(),
1207 vec![
1208 "Company".to_string(),
1209 "Country".to_string(),
1210 "Customer Id".to_string(),
1211 ],
1212 );
1213
1214 let query = "SELECT Company, ";
1218 let result = parser.get_completions(query, query.len());
1219
1220 assert!(
1221 result.suggestions.iter().any(|s| s == "Company"),
1222 "Should still suggest Company even though already selected"
1223 );
1224 assert!(
1225 result.suggestions.iter().any(|s| s == "Country"),
1226 "Should suggest Country"
1227 );
1228 assert!(
1229 result.suggestions.iter().any(|s| s == "\"Customer Id\""),
1230 "Should suggest Customer Id"
1231 );
1232 }
1233
1234 #[test]
1235 fn test_order_by_completion_with_quoted_columns() {
1236 let parser = CursorAwareParser::new();
1237 let mut parser = parser;
1238 parser.update_single_table(
1239 "customers".to_string(),
1240 vec![
1241 "City".to_string(),
1242 "Company".to_string(),
1243 "Country".to_string(),
1244 "Customer Id".to_string(),
1245 ],
1246 );
1247
1248 let query = r#"SELECT City,Company,Country,"Customer Id" FROM customers ORDER BY coun"#;
1250 let result = parser.get_completions(query, query.len());
1251
1252 assert_eq!(
1254 result.partial_word,
1255 Some("coun".to_string()),
1256 "Should extract 'coun' as partial, not something weird"
1257 );
1258
1259 assert!(
1261 result.suggestions.iter().any(|s| s == "Country"),
1262 "Should suggest Country for partial 'coun'. Got: {:?}",
1263 result.suggestions
1264 );
1265 }
1266
1267 #[test]
1268 fn test_order_by_quoted_partial_completion() {
1269 let parser = CursorAwareParser::new();
1270 let mut parser = parser;
1271 parser.update_single_table(
1272 "customers".to_string(),
1273 vec![
1274 "City".to_string(),
1275 "Company".to_string(),
1276 "Country".to_string(),
1277 "Customer Id".to_string(),
1278 ],
1279 );
1280
1281 let query =
1283 r#"select City,Company,Country,"Customer Id" from customers order by City, "Customer"#;
1284 let result = parser.get_completions(query, query.len());
1285
1286 assert_eq!(
1288 result.partial_word,
1289 Some("\"Customer".to_string()),
1290 "Should extract '\"Customer' as partial"
1291 );
1292
1293 assert!(
1295 result.suggestions.iter().any(|s| s == "\"Customer Id\""),
1296 "Should suggest properly quoted 'Customer Id' for partial '\"Customer'. Got: {:?}",
1297 result.suggestions
1298 );
1299
1300 assert!(
1302 !result.suggestions.iter().any(|s| s == "\"Customer"),
1303 "Should not have truncated suggestion '\"Customer'. Got: {:?}",
1304 result.suggestions
1305 );
1306 }
1307}