From f4bff3f2bb5996b86bf8a15cdd812cfe7964b246 Mon Sep 17 00:00:00 2001
From: Matthew Turner <matthew.m.turner@outlook.com>
Date: Thu, 19 Sep 2024 23:55:41 -0400
Subject: [PATCH] Syntax highlight for keywords

---
 Cargo.lock                      |  1 +
 Cargo.toml                      |  2 +-
 src/app/state/tabs/flightsql.rs |  6 +++++-
 src/app/state/tabs/sql.rs       | 28 +++++++++++++++++++++++++++-
 4 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 46cf05e..7971f2c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3987,6 +3987,7 @@ checksum = "29c07084342a575cea919eea996b9658a358c800b03d435df581c1d7c60e065a"
 dependencies = [
  "crossterm",
  "ratatui",
+ "regex",
  "unicode-width",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index b222dd5..542523f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -39,7 +39,7 @@ tokio-util = "0.7.10"
 toml = "0.8.12"
 tonic = { version = "0.11.0", optional = true }
 tui-logger = {version = "0.12", features = ["tracing-support"]}
-tui-textarea = "0.6.1"
+tui-textarea = { version = "0.6.1", features = ["search"] }
 url = { version = "2.5.2", optional = true }
 
 [dev-dependencies]
diff --git a/src/app/state/tabs/flightsql.rs b/src/app/state/tabs/flightsql.rs
index a49d92c..74aee20 100644
--- a/src/app/state/tabs/flightsql.rs
+++ b/src/app/state/tabs/flightsql.rs
@@ -25,6 +25,7 @@ use ratatui::style::Style;
 use ratatui::widgets::TableState;
 use tui_textarea::TextArea;
 
+use crate::app::state::tabs::sql;
 use crate::execution::ExecutionStats;
 
 #[derive(Clone, Debug)]
@@ -115,7 +116,8 @@ impl<'app> FlightSQLTabState<'app> {
         // TODO: Enable vim mode from config?
         let mut textarea = TextArea::new(empty_text);
         textarea.set_style(Style::default().fg(tailwind::WHITE));
-
+        textarea.set_search_pattern(sql::keyword_regex()).unwrap();
+        textarea.set_search_style(sql::keyword_style());
         Self {
             editor: textarea,
             editor_editable: false,
@@ -152,6 +154,8 @@ impl<'app> FlightSQLTabState<'app> {
     pub fn clear_editor(&mut self) {
         let mut textarea = TextArea::new(vec!["".to_string()]);
         textarea.set_style(Style::default().fg(tailwind::WHITE));
+        textarea.set_search_pattern(sql::keyword_regex()).unwrap();
+        textarea.set_search_style(sql::keyword_style());
         self.editor = textarea;
     }
 
diff --git a/src/app/state/tabs/sql.rs b/src/app/state/tabs/sql.rs
index 7bc633d..d4cafc1 100644
--- a/src/app/state/tabs/sql.rs
+++ b/src/app/state/tabs/sql.rs
@@ -19,9 +19,10 @@ use core::cell::RefCell;
 use std::time::Duration;
 
 use datafusion::arrow::array::RecordBatch;
+use datafusion::sql::sqlparser::keywords;
 use ratatui::crossterm::event::KeyEvent;
 use ratatui::style::palette::tailwind;
-use ratatui::style::Style;
+use ratatui::style::{Modifier, Style};
 use ratatui::widgets::TableState;
 use tui_textarea::TextArea;
 
@@ -101,6 +102,27 @@ impl Query {
     }
 }
 
+pub fn get_keywords() -> Vec<String> {
+    keywords::ALL_KEYWORDS
+        .iter()
+        .map(|k| k.to_string())
+        .collect()
+}
+
+pub fn keyword_regex() -> String {
+    format!(
+        "(?i)(^|[^a-zA-Z0-9\'\"`._]*?)({})($|[^a-zA-Z0-9\'\"`._]*)",
+        get_keywords().join("|")
+    )
+}
+
+pub fn keyword_style() -> Style {
+    Style::default()
+        .bg(tailwind::BLACK)
+        .fg(tailwind::YELLOW.c100)
+        .add_modifier(Modifier::BOLD)
+}
+
 #[derive(Debug, Default)]
 pub struct SQLTabState<'app> {
     editor: TextArea<'app>,
@@ -115,6 +137,8 @@ impl<'app> SQLTabState<'app> {
         // TODO: Enable vim mode from config?
         let mut textarea = TextArea::new(empty_text);
         textarea.set_style(Style::default().fg(tailwind::WHITE));
+        textarea.set_search_pattern(keyword_regex()).unwrap();
+        textarea.set_search_style(keyword_style());
         Self {
             editor: textarea,
             editor_editable: false,
@@ -151,6 +175,8 @@ impl<'app> SQLTabState<'app> {
     pub fn clear_editor(&mut self) {
         let mut textarea = TextArea::new(vec!["".to_string()]);
         textarea.set_style(Style::default().fg(tailwind::WHITE));
+        textarea.set_search_pattern(keyword_regex()).unwrap();
+        textarea.set_search_style(keyword_style());
         self.editor = textarea;
     }