cqlsh_rs/
colorizer.rs

1//! CQL syntax colorization for the REPL prompt and output.
2//!
3//! Uses the unified CQL lexer (`cql_lexer`) for context-aware tokenization,
4//! which correctly distinguishes keywords from identifiers based on grammar
5//! position (e.g., USERS after FROM is an identifier, not a keyword).
6//! Also provides output coloring for query result values, headers, and errors
7//! matching Python cqlsh's color scheme.
8
9use crossterm::style::Stylize;
10
11use crate::cql_lexer::{self, TokenKind};
12use crate::driver::types::CqlValue;
13
14/// CQL syntax colorizer using ANSI escape codes.
15pub struct CqlColorizer {
16    enabled: bool,
17}
18
19impl CqlColorizer {
20    /// Create a new colorizer. If `enabled` is false, all methods return input unchanged.
21    pub fn new(enabled: bool) -> Self {
22        Self { enabled }
23    }
24
25    /// Returns whether colorization is enabled.
26    pub fn is_enabled(&self) -> bool {
27        self.enabled
28    }
29
30    /// Colorize a CQL result value matching Python cqlsh's color scheme.
31    ///
32    /// Color mapping:
33    /// - Text/Ascii → yellow bold
34    /// - Numeric/boolean/uuid/timestamp/date/time/duration/inet → green bold
35    /// - Blob → dark magenta (non-bold)
36    /// - Null → red bold
37    /// - Collection delimiters → blue bold, inner values colored by type
38    pub fn colorize_value(&self, value: &CqlValue) -> String {
39        if !self.enabled {
40            return value.to_string();
41        }
42        self.colorize_value_inner(value)
43    }
44
45    /// Colorize a column header (magenta bold, matching Python cqlsh default).
46    pub fn colorize_header(&self, name: &str) -> String {
47        if !self.enabled {
48            return name.to_string();
49        }
50        format!("{}", name.magenta().bold())
51    }
52
53    /// Colorize an error message (red bold, matching Python cqlsh).
54    pub fn colorize_error(&self, msg: &str) -> String {
55        if !self.enabled {
56            return msg.to_string();
57        }
58        format!("{}", msg.red().bold())
59    }
60
61    /// Colorize a warning message (red bold, matching Python cqlsh).
62    pub fn colorize_warning(&self, msg: &str) -> String {
63        self.colorize_error(msg)
64    }
65
66    /// Colorize the "Tracing session:" label (magenta bold).
67    pub fn colorize_trace_label(&self, label: &str) -> String {
68        if !self.enabled {
69            return label.to_string();
70        }
71        format!("{}", label.magenta().bold())
72    }
73
74    /// Colorize the cluster name in the welcome message (blue bold).
75    pub fn colorize_cluster_name(&self, name: &str) -> String {
76        if !self.enabled {
77            return name.to_string();
78        }
79        format!("{}", name.blue().bold())
80    }
81
82    /// Inner recursive value colorizer.
83    fn colorize_value_inner(&self, value: &CqlValue) -> String {
84        match value {
85            CqlValue::Ascii(s) | CqlValue::Text(s) => {
86                format!("{}", s.as_str().yellow().bold())
87            }
88            CqlValue::Int(_)
89            | CqlValue::BigInt(_)
90            | CqlValue::SmallInt(_)
91            | CqlValue::TinyInt(_)
92            | CqlValue::Float(_)
93            | CqlValue::Double(_)
94            | CqlValue::Decimal(_)
95            | CqlValue::Varint(_)
96            | CqlValue::Counter(_)
97            | CqlValue::Boolean(_)
98            | CqlValue::Uuid(_)
99            | CqlValue::TimeUuid(_)
100            | CqlValue::Timestamp(_)
101            | CqlValue::Date(_)
102            | CqlValue::Time(_)
103            | CqlValue::Duration { .. }
104            | CqlValue::Inet(_) => {
105                format!("{}", value.to_string().green().bold())
106            }
107            CqlValue::Blob(_) => {
108                format!("{}", value.to_string().dark_magenta())
109            }
110            CqlValue::Null => String::new(),
111            CqlValue::Unset => {
112                format!("{}", "<unset>".red().bold())
113            }
114            CqlValue::List(items) => {
115                let mut result = format!("{}", "[".blue().bold());
116                for (i, item) in items.iter().enumerate() {
117                    if i > 0 {
118                        result.push_str(&format!("{}", ", ".blue().bold()));
119                    }
120                    result.push_str(&self.colorize_collection_element(item));
121                }
122                result.push_str(&format!("{}", "]".blue().bold()));
123                result
124            }
125            CqlValue::Set(items) => {
126                let mut result = format!("{}", "{".blue().bold());
127                for (i, item) in items.iter().enumerate() {
128                    if i > 0 {
129                        result.push_str(&format!("{}", ", ".blue().bold()));
130                    }
131                    result.push_str(&self.colorize_collection_element(item));
132                }
133                result.push_str(&format!("{}", "}".blue().bold()));
134                result
135            }
136            CqlValue::Map(entries) => {
137                let mut result = format!("{}", "{".blue().bold());
138                for (i, (k, v)) in entries.iter().enumerate() {
139                    if i > 0 {
140                        result.push_str(&format!("{}", ", ".blue().bold()));
141                    }
142                    result.push_str(&self.colorize_collection_element(k));
143                    result.push_str(&format!("{}", ": ".blue().bold()));
144                    result.push_str(&self.colorize_collection_element(v));
145                }
146                result.push_str(&format!("{}", "}".blue().bold()));
147                result
148            }
149            CqlValue::Tuple(items) => {
150                let mut result = format!("{}", "(".blue().bold());
151                for (i, item) in items.iter().enumerate() {
152                    if i > 0 {
153                        result.push_str(&format!("{}", ", ".blue().bold()));
154                    }
155                    match item {
156                        Some(v) => result.push_str(&self.colorize_collection_element(v)),
157                        None => result.push_str(&format!("{}", "null".red().bold())),
158                    }
159                }
160                result.push_str(&format!("{}", ")".blue().bold()));
161                result
162            }
163            CqlValue::UserDefinedType { fields, .. } => {
164                let mut result = format!("{}", "{".blue().bold());
165                for (i, (name, val)) in fields.iter().enumerate() {
166                    if i > 0 {
167                        result.push_str(&format!("{}", ", ".blue().bold()));
168                    }
169                    // UDT field names are yellow (like text)
170                    result.push_str(&format!("{}", name.as_str().yellow().bold()));
171                    result.push_str(&format!("{}", ": ".blue().bold()));
172                    match val {
173                        Some(v) => result.push_str(&self.colorize_collection_element(v)),
174                        None => result.push_str(&format!("{}", "null".red().bold())),
175                    }
176                }
177                result.push_str(&format!("{}", "}".blue().bold()));
178                result
179            }
180        }
181    }
182
183    /// Colorize an element inside a collection, quoting strings like Display does.
184    fn colorize_collection_element(&self, value: &CqlValue) -> String {
185        match value {
186            CqlValue::Ascii(s) | CqlValue::Text(s) => {
187                // Inside collections, strings are quoted: 'value'
188                let quoted = format!("'{}'", s.replace('\'', "''"));
189                format!("{}", quoted.yellow().bold())
190            }
191            other => self.colorize_value_inner(other),
192        }
193    }
194
195    /// Colorize a line of CQL input for display.
196    ///
197    /// Uses the unified CQL lexer for context-aware tokenization. Colors:
198    /// - CQL keywords → bold blue
199    /// - String literals → green
200    /// - Numbers → cyan
201    /// - Comments → dark grey
202    /// - Identifiers, operators, punctuation → default (no color)
203    pub fn colorize_line(&self, line: &str) -> String {
204        if !self.enabled {
205            return line.to_string();
206        }
207
208        let tokens = cql_lexer::tokenize(line);
209        let mut result = String::with_capacity(line.len() * 2);
210
211        for token in &tokens {
212            match token.kind {
213                TokenKind::Keyword => {
214                    result.push_str(&format!("{}", token.text.as_str().blue().bold()));
215                }
216                TokenKind::StringLiteral | TokenKind::DollarStringLiteral => {
217                    result.push_str(&format!("{}", token.text.as_str().green()));
218                }
219                TokenKind::NumberLiteral => {
220                    result.push_str(&format!("{}", token.text.as_str().cyan()));
221                }
222                TokenKind::BlobLiteral => {
223                    result.push_str(&format!("{}", token.text.as_str().dark_magenta()));
224                }
225                TokenKind::UuidLiteral => {
226                    result.push_str(&format!("{}", token.text.as_str().green()));
227                }
228                TokenKind::BooleanLiteral => {
229                    result.push_str(&format!("{}", token.text.as_str().green().bold()));
230                }
231                TokenKind::LineComment | TokenKind::BlockComment => {
232                    result.push_str(&format!("{}", token.text.as_str().dark_grey()));
233                }
234                _ => {
235                    result.push_str(&token.text);
236                }
237            }
238        }
239
240        result
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247
248    #[test]
249    fn keywords_are_highlighted() {
250        let c = CqlColorizer::new(true);
251        let output = c.colorize_line("SELECT * FROM users");
252        assert!(output.contains("\x1b["), "should contain ANSI escape codes");
253        assert!(output.contains("SELECT"));
254        assert!(output.contains("FROM"));
255    }
256
257    #[test]
258    fn colorizer_disabled_returns_unchanged() {
259        let c = CqlColorizer::new(false);
260        let output = c.colorize_line("SELECT * FROM users");
261        assert_eq!(output, "SELECT * FROM users");
262    }
263
264    #[test]
265    fn string_literals_are_colored() {
266        let c = CqlColorizer::new(true);
267        let output = c.colorize_line("INSERT INTO t (a) VALUES ('hello')");
268        // 'hello' should be green (contains ANSI codes)
269        assert!(output.contains("\x1b["));
270        assert!(output.contains("hello"));
271    }
272
273    #[test]
274    fn numbers_are_colored() {
275        let c = CqlColorizer::new(true);
276        let output = c.colorize_line("SELECT * FROM t LIMIT 100");
277        assert!(output.contains("100"));
278    }
279
280    #[test]
281    fn comments_are_colored() {
282        let c = CqlColorizer::new(true);
283        let output = c.colorize_line("SELECT 1 -- test comment");
284        assert!(output.contains("test comment"));
285    }
286
287    #[test]
288    fn non_keywords_are_not_highlighted() {
289        let c = CqlColorizer::new(true);
290        let output = c.colorize_line("my_table");
291        // "my_table" is not a keyword, should not have ANSI codes
292        assert!(!output.contains("\x1b["));
293    }
294
295    #[test]
296    fn mixed_case_keywords() {
297        let c = CqlColorizer::new(true);
298        let output = c.colorize_line("select * from users");
299        assert!(
300            output.contains("\x1b["),
301            "lowercase keywords should also be highlighted"
302        );
303    }
304
305    #[test]
306    fn identifiers_after_from_not_highlighted() {
307        let c = CqlColorizer::new(true);
308        // "users" after FROM should NOT be highlighted as a keyword
309        let output = c.colorize_line("SELECT * FROM users");
310        // Extract the "users" portion — it should not contain ANSI codes
311        // The output is: <colored SELECT> <space> <*> <space> <colored FROM> <space> <users>
312        assert!(output.ends_with("users"));
313    }
314
315    #[test]
316    fn keyword_names_as_identifiers_not_highlighted() {
317        let c = CqlColorizer::new(true);
318        // KEY and SET after FROM should be identifiers, not keywords
319        let output_key = c.colorize_line("SELECT * FROM KEY");
320        assert!(
321            output_key.ends_with("KEY"),
322            "KEY after FROM should not be highlighted"
323        );
324        let output_set = c.colorize_line("SELECT * FROM SET");
325        assert!(
326            output_set.ends_with("SET"),
327            "SET after FROM should not be highlighted"
328        );
329    }
330
331    #[test]
332    fn qualified_name_after_dot_not_highlighted() {
333        let c = CqlColorizer::new(true);
334        let output = c.colorize_line("SELECT * FROM ks.users");
335        // "users" after dot should not be highlighted
336        assert!(output.ends_with("users"));
337    }
338
339    // --- Output coloring tests ---
340
341    #[test]
342    fn colorize_text_value_yellow() {
343        let c = CqlColorizer::new(true);
344        let output = c.colorize_value(&CqlValue::Text("hello".to_string()));
345        assert!(output.contains("\x1b["), "should contain ANSI codes");
346        assert!(output.contains("hello"));
347    }
348
349    #[test]
350    fn colorize_int_value_green() {
351        let c = CqlColorizer::new(true);
352        let output = c.colorize_value(&CqlValue::Int(42));
353        assert!(output.contains("\x1b["), "should contain ANSI codes");
354        assert!(output.contains("42"));
355    }
356
357    #[test]
358    fn colorize_null_value_empty() {
359        let c = CqlColorizer::new(true);
360        let output = c.colorize_value(&CqlValue::Null);
361        assert_eq!(output, "");
362    }
363
364    #[test]
365    fn colorize_blob_value_dark_magenta() {
366        let c = CqlColorizer::new(true);
367        let output = c.colorize_value(&CqlValue::Blob(vec![0xde, 0xad]));
368        assert!(output.contains("\x1b["), "should contain ANSI codes");
369        assert!(output.contains("dead"));
370    }
371
372    #[test]
373    fn colorize_list_with_blue_delimiters() {
374        let c = CqlColorizer::new(true);
375        let list = CqlValue::List(vec![CqlValue::Int(1), CqlValue::Int(2)]);
376        let output = c.colorize_value(&list);
377        assert!(output.contains("\x1b["), "should contain ANSI codes");
378    }
379
380    #[test]
381    fn colorize_value_disabled_returns_plain() {
382        let c = CqlColorizer::new(false);
383        let output = c.colorize_value(&CqlValue::Text("hello".to_string()));
384        assert_eq!(output, "hello");
385    }
386
387    #[test]
388    fn colorize_header_magenta() {
389        let c = CqlColorizer::new(true);
390        let output = c.colorize_header("name");
391        assert!(output.contains("\x1b["), "should contain ANSI codes");
392        assert!(output.contains("name"));
393    }
394
395    #[test]
396    fn colorize_error_red() {
397        let c = CqlColorizer::new(true);
398        let output = c.colorize_error("SyntaxException: bad input");
399        assert!(output.contains("\x1b["), "should contain ANSI codes");
400        assert!(output.contains("SyntaxException"));
401    }
402
403    #[test]
404    fn colorize_map_with_colored_elements() {
405        let c = CqlColorizer::new(true);
406        let map = CqlValue::Map(vec![(CqlValue::Text("key".to_string()), CqlValue::Int(42))]);
407        let output = c.colorize_value(&map);
408        assert!(output.contains("\x1b["), "should contain ANSI codes");
409    }
410
411    #[test]
412    fn colorize_udt_field_names_yellow() {
413        let c = CqlColorizer::new(true);
414        let udt = CqlValue::UserDefinedType {
415            keyspace: "ks".to_string(),
416            type_name: "my_type".to_string(),
417            fields: vec![
418                (
419                    "name".to_string(),
420                    Some(CqlValue::Text("Alice".to_string())),
421                ),
422                ("age".to_string(), Some(CqlValue::Int(30))),
423            ],
424        };
425        let output = c.colorize_value(&udt);
426        assert!(output.contains("\x1b["), "should contain ANSI codes");
427    }
428
429    #[test]
430    fn colorize_warning_same_as_error() {
431        let c = CqlColorizer::new(true);
432        let output = c.colorize_warning("Something bad");
433        assert!(output.contains("\x1b["));
434        assert!(output.contains("Something bad"));
435    }
436
437    #[test]
438    fn colorize_trace_label_colored() {
439        let c = CqlColorizer::new(true);
440        let output = c.colorize_trace_label("Tracing session:");
441        assert!(output.contains("\x1b["));
442        assert!(output.contains("Tracing session:"));
443    }
444
445    #[test]
446    fn colorize_cluster_name_colored() {
447        let c = CqlColorizer::new(true);
448        let output = c.colorize_cluster_name("Test Cluster");
449        assert!(output.contains("\x1b["));
450        assert!(output.contains("Test Cluster"));
451    }
452
453    #[test]
454    fn colorize_tuple_with_nulls() {
455        let c = CqlColorizer::new(true);
456        let tuple = CqlValue::Tuple(vec![
457            Some(CqlValue::Int(1)),
458            None,
459            Some(CqlValue::Text("x".to_string())),
460        ]);
461        let output = c.colorize_value(&tuple);
462        assert!(output.contains("\x1b["));
463        assert!(output.contains("null"));
464    }
465
466    #[test]
467    fn colorize_set_multiple_elements() {
468        let c = CqlColorizer::new(true);
469        let set = CqlValue::Set(vec![
470            CqlValue::Text("a".to_string()),
471            CqlValue::Text("b".to_string()),
472            CqlValue::Text("c".to_string()),
473        ]);
474        let output = c.colorize_value(&set);
475        assert!(output.contains("\x1b["));
476    }
477
478    #[test]
479    fn colorize_udt_with_none_fields() {
480        let c = CqlColorizer::new(true);
481        let udt = CqlValue::UserDefinedType {
482            keyspace: "ks".to_string(),
483            type_name: "t".to_string(),
484            fields: vec![
485                ("a".to_string(), None),
486                ("b".to_string(), Some(CqlValue::Int(1))),
487            ],
488        };
489        let output = c.colorize_value(&udt);
490        assert!(output.contains("null"));
491    }
492
493    #[test]
494    fn colorize_map_multi_elements() {
495        let c = CqlColorizer::new(true);
496        let map = CqlValue::Map(vec![
497            (CqlValue::Text("k1".to_string()), CqlValue::Int(1)),
498            (CqlValue::Text("k2".to_string()), CqlValue::Int(2)),
499        ]);
500        let output = c.colorize_value(&map);
501        assert!(output.contains("\x1b["));
502    }
503
504    #[test]
505    fn colorize_unset_value() {
506        let c = CqlColorizer::new(true);
507        let output = c.colorize_value(&CqlValue::Unset);
508        assert!(output.contains("\x1b["));
509        assert!(output.contains("unset"));
510    }
511
512    #[test]
513    fn colorize_disabled_returns_plain_for_collections() {
514        let c = CqlColorizer::new(false);
515        let list = CqlValue::List(vec![CqlValue::Int(1), CqlValue::Int(2)]);
516        let output = c.colorize_value(&list);
517        assert_eq!(output, "[1, 2]");
518
519        let tuple = CqlValue::Tuple(vec![Some(CqlValue::Int(1)), None]);
520        let output = c.colorize_value(&tuple);
521        assert_eq!(output, "(1, null)");
522
523        let udt = CqlValue::UserDefinedType {
524            keyspace: "ks".to_string(),
525            type_name: "t".to_string(),
526            fields: vec![("x".to_string(), Some(CqlValue::Int(5)))],
527        };
528        let output = c.colorize_value(&udt);
529        assert_eq!(output, "{x: 5}");
530    }
531
532    #[test]
533    fn colorize_disabled_cluster_name() {
534        let c = CqlColorizer::new(false);
535        assert_eq!(c.colorize_cluster_name("Test"), "Test");
536        assert_eq!(c.colorize_trace_label("Trace:"), "Trace:");
537        assert_eq!(c.colorize_warning("warn"), "warn");
538    }
539}
cqlsh_rs/colorizer.rs

cqlsh_rs/
colorizer.rs
