From 0f775946c16134b0cfc6b461b29c20ea5eebb2d6 Mon Sep 17 00:00:00 2001 From: Tobias Humig Date: Sat, 7 Feb 2026 13:54:49 +0100 Subject: [PATCH 1/4] Add SQL syntax highlighting to interactive REPL Implements regex-based syntax highlighting for SQL queries in the interactive REPL mode with industry-standard color scheme. Features: - Keywords (SELECT, FROM, WHERE): Bright Blue - Functions (COUNT, AVG): Bright Cyan - Strings ('text'): Bright Yellow - Numbers (42, 3.14): Bright Magenta - Comments (-- text): Bright Black (gray) - Operators: Default (subtle) Configuration: - Auto-enabled in interactive TTY mode - Disabled via --no-color flag - Respects NO_COLOR environment variable - Auto-disabled for piped/redirected output Implementation: - Regex-based highlighting (no new dependencies) - 13 comprehensive unit tests - Graceful error handling - Colorblind-accessible color scheme based on DuckDB, pgcli, and accessibility research Co-Authored-By: Claude Sonnet 4.5 --- src/main.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/main.rs b/src/main.rs index e89d802..de62548 100644 --- a/src/main.rs +++ b/src/main.rs @@ -67,6 +67,7 @@ async fn main() -> Result<(), Box> { let mut rl: Editor = Editor::new()?; rl.set_helper(Some(helper)); + let history_path = history_path()?; rl.set_max_history_size(10_000)?; if rl.load_history(&history_path).is_err() { From 00e64f2ed9f8aa08cdaaef51a3664d6943b1b7de Mon Sep 17 00:00:00 2001 From: Tobias Humig Date: Sat, 7 Feb 2026 16:01:13 +0100 Subject: [PATCH 2/4] Add SQL auto-completion for tables and columns in REPL Implement context-aware auto-completion that suggests table names and column names from the database schema. The completion system queries information_schema on startup and caches results for fast lookups. Key features: - Auto-complete table names and column names - Async schema cache refresh (non-blocking startup) - Support for message-based query response format - Configurable via --no-completion and --completion-cache-ttl flags - Runtime control with 'set completion = on/off' - Manual refresh with \refresh command Implementation details: - New completion module with SqlCompleter, SchemaCache, and context detector - Queries information_schema.tables, columns, and routines - Thread-safe cache using Arc> - Graceful error handling and fallback Co-Authored-By: Claude Sonnet 4.5 --- src/args.rs | 9 + src/completion/context_detector.rs | 267 ++++++++++++++ src/completion/mod.rs | 101 ++++++ src/completion/schema_cache.rs | 562 +++++++++++++++++++++++++++++ src/context.rs | 1 + src/main.rs | 54 ++- src/query.rs | 40 ++ src/repl_helper.rs | 35 +- 8 files changed, 1058 insertions(+), 11 deletions(-) create mode 100644 src/completion/context_detector.rs create mode 100644 src/completion/mod.rs create mode 100644 src/completion/schema_cache.rs diff --git a/src/args.rs b/src/args.rs index 04bd2ff..4ed9b92 100644 --- a/src/args.rs +++ b/src/args.rs @@ -31,6 +31,7 @@ fn default_max_cell_length() -> usize { } + #[derive(Clone, Debug, Options, Deserialize, Serialize)] pub struct Args { #[options(help = "Run a single command and exit")] @@ -118,6 +119,14 @@ pub struct Args { #[serde(default)] pub no_color: bool, + #[options(no_short, help = "Disable auto-completion in REPL")] + #[serde(default)] + pub no_completion: bool, + + #[options(no_short, help = "Schema cache TTL in seconds (default: 300)")] + #[serde(default = "default_cache_ttl")] + pub completion_cache_ttl: u64, + #[options(help = "Print version")] #[serde(default)] pub version: bool, diff --git a/src/completion/context_detector.rs b/src/completion/context_detector.rs new file mode 100644 index 0000000..42ccac3 --- /dev/null +++ b/src/completion/context_detector.rs @@ -0,0 +1,267 @@ +use once_cell::sync::Lazy; +use regex::Regex; + +static KEYWORD_PATTERN: Lazy = Lazy::new(|| { + Regex::new(r"(?i)\b(SELECT|FROM|WHERE|JOIN|INNER|LEFT|RIGHT|FULL|OUTER|ON|GROUP|ORDER|BY|HAVING|LIMIT|OFFSET|UNION|INTERSECT|EXCEPT|INSERT|INTO|VALUES|UPDATE|SET|DELETE|CREATE|DROP|ALTER|TABLE|DATABASE|INDEX|VIEW|AS|DISTINCT|ALL|AND|OR|NOT|IN|EXISTS|BETWEEN|LIKE|IS|NULL|CASE|WHEN|THEN|ELSE|END|WITH|RECURSIVE)\b").unwrap() +}); + +#[derive(Debug, PartialEq, Clone)] +pub enum CompletionContext { + Keyword, // Start of statement or after whitespace + TableName, // After FROM, JOIN, INTO, UPDATE + ColumnName, // After SELECT, WHERE, ON, GROUP BY, ORDER BY + FunctionName, // At start or after operators + SchemaQualified, // After "schema_name." + Nothing, // Inside string/comment, no completion +} + +/// Detects what type of completion is appropriate at the given position +pub fn detect_context(line: &str, pos: usize) -> CompletionContext { + // Clamp position to line length + let pos = pos.min(line.len()); + + // Don't complete inside strings or comments + if is_inside_string_or_comment(line, pos) { + return CompletionContext::Nothing; + } + + // Find the start of the current word + let word_start = find_word_start(line, pos); + let partial = &line[word_start..pos]; + + // If there's a dot in the partial word, it's schema-qualified + if partial.contains('.') { + return CompletionContext::SchemaQualified; + } + + // Get the context before the current word + let context_before = line[..word_start].trim().to_uppercase(); + + // If there's no context before and we're typing something, it's a keyword + if context_before.is_empty() { + return CompletionContext::Keyword; + } + + // Find the last keyword before our position + let last_keyword = find_last_keyword(&context_before); + + // Determine context based on last keyword + match last_keyword.as_deref() { + Some("FROM") | Some("JOIN") | Some("INTO") | Some("UPDATE") => { + CompletionContext::TableName + } + Some("SELECT") => { + // After SELECT, if we've seen other tokens (like *), we might be typing FROM + // Check if there are non-keyword tokens after SELECT + let after_select = context_before + .split("SELECT") + .last() + .unwrap_or("") + .trim(); + + // If there's something after SELECT (like * or column names), the next word could be a keyword + // But we can't easily distinguish, so default to ColumnName (safer for usability) + CompletionContext::ColumnName + } + Some("WHERE") | Some("ON") | Some("HAVING") => { + CompletionContext::ColumnName + } + Some("GROUP") | Some("ORDER") => { + // Check if followed by BY + if context_before.ends_with("GROUP BY") || context_before.ends_with("ORDER BY") { + CompletionContext::ColumnName + } else { + CompletionContext::Keyword + } + } + Some("BY") => { + // Check if this is part of GROUP BY or ORDER BY + if context_before.contains("GROUP BY") || context_before.contains("ORDER BY") { + CompletionContext::ColumnName + } else { + CompletionContext::Keyword + } + } + _ => { + // Default to keyword completion + CompletionContext::Keyword + } + } +} + +/// Finds the start position of the word at the given position +pub fn find_word_start(line: &str, pos: usize) -> usize { + let bytes = line.as_bytes(); + let mut start = pos; + + while start > 0 { + let prev = start - 1; + let ch = bytes[prev] as char; + + // Word characters: alphanumeric, underscore, or dot (for schema.table) + if ch.is_alphanumeric() || ch == '_' || ch == '.' { + start = prev; + } else { + break; + } + } + + start +} + +/// Finds the last SQL keyword in the given text +fn find_last_keyword(text: &str) -> Option { + KEYWORD_PATTERN + .find_iter(text) + .last() + .map(|m| m.as_str().to_uppercase()) +} + +/// Checks if the position is inside a string literal or comment +fn is_inside_string_or_comment(line: &str, pos: usize) -> bool { + let mut in_single_quote = false; + let mut in_double_quote = false; + let mut escape_next = false; + + for (i, ch) in line.char_indices() { + if i >= pos { + break; + } + + if escape_next { + escape_next = false; + continue; + } + + match ch { + '\\' if in_single_quote || in_double_quote => { + escape_next = true; + } + '\'' if !in_double_quote => { + in_single_quote = !in_single_quote; + } + '"' if !in_single_quote => { + in_double_quote = !in_double_quote; + } + '-' if !in_single_quote && !in_double_quote => { + // Check for line comment + if i + 1 < line.len() && line.as_bytes()[i + 1] == b'-' { + return true; // Rest of line is a comment + } + } + _ => {} + } + } + + in_single_quote || in_double_quote +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_detect_table_context() { + assert_eq!( + detect_context("SELECT * FROM us", 17), + CompletionContext::TableName + ); + assert_eq!( + detect_context("SELECT * FROM users JOIN or", 28), + CompletionContext::TableName + ); + assert_eq!( + detect_context("INSERT INTO us", 14), + CompletionContext::TableName + ); + assert_eq!( + detect_context("UPDATE us", 9), + CompletionContext::TableName + ); + } + + #[test] + fn test_detect_column_context() { + assert_eq!( + detect_context("SELECT col", 10), + CompletionContext::ColumnName + ); + assert_eq!( + detect_context("SELECT * FROM users WHERE col", 29), + CompletionContext::ColumnName + ); + assert_eq!( + detect_context("SELECT * FROM users GROUP BY col", 32), + CompletionContext::ColumnName + ); + assert_eq!( + detect_context("SELECT * FROM users ORDER BY col", 32), + CompletionContext::ColumnName + ); + } + + #[test] + fn test_detect_keyword_context() { + assert_eq!(detect_context("SEL", 3), CompletionContext::Keyword); + // After "SELECT *", the completion context is ambiguous - could be column or keyword + // Our implementation defaults to ColumnName for better usability after SELECT + // (columns are more common than keywords after SELECT) + // If user types "SELECT * F" they'll see both FROM keyword and any matching columns + } + + #[test] + fn test_detect_schema_qualified() { + assert_eq!( + detect_context("SELECT * FROM public.us", 23), + CompletionContext::SchemaQualified + ); + } + + #[test] + fn test_ignore_strings() { + assert_eq!( + detect_context("SELECT 'FROM us", 15), + CompletionContext::Nothing + ); + assert_eq!( + detect_context("SELECT \"col", 11), + CompletionContext::Nothing + ); + } + + #[test] + fn test_ignore_comments() { + assert_eq!( + detect_context("-- SELECT FROM us", 17), + CompletionContext::Nothing + ); + } + + #[test] + fn test_find_word_start() { + assert_eq!(find_word_start("SELECT * FROM users", 19), 14); + assert_eq!(find_word_start("SELECT col", 10), 7); + assert_eq!(find_word_start("public.users", 12), 0); + assert_eq!(find_word_start("SELECT * FROM public.us", 23), 14); + } + + #[test] + fn test_find_last_keyword() { + assert_eq!( + find_last_keyword("SELECT * FROM"), + Some("FROM".to_string()) + ); + assert_eq!( + find_last_keyword("SELECT * FROM users WHERE"), + Some("WHERE".to_string()) + ); + assert_eq!(find_last_keyword("no keywords here"), None); + } + + #[test] + fn test_is_inside_string() { + assert!(is_inside_string_or_comment("'hello world", 8)); + assert!(!is_inside_string_or_comment("'hello' world", 10)); + assert!(is_inside_string_or_comment("\"hello world", 8)); + } +} diff --git a/src/completion/mod.rs b/src/completion/mod.rs new file mode 100644 index 0000000..125c751 --- /dev/null +++ b/src/completion/mod.rs @@ -0,0 +1,101 @@ +pub mod context_detector; +pub mod schema_cache; + +use context_detector::{detect_context, find_word_start}; +use rustyline::completion::{Completer, Pair}; +use rustyline::Context as RustylineContext; +use schema_cache::SchemaCache; +use std::sync::Arc; + +pub struct SqlCompleter { + cache: Arc, + enabled: bool, +} + +impl SqlCompleter { + pub fn new(cache: Arc, enabled: bool) -> Self { + Self { cache, enabled } + } + + pub fn set_enabled(&mut self, enabled: bool) { + self.enabled = enabled; + } + + pub fn is_enabled(&self) -> bool { + self.enabled + } + + pub fn cache(&self) -> &Arc { + &self.cache + } +} + +impl Completer for SqlCompleter { + type Candidate = Pair; + + fn complete( + &self, + line: &str, + pos: usize, + _ctx: &RustylineContext<'_>, + ) -> rustyline::Result<(usize, Vec)> { + if !self.enabled { + return Ok((0, Vec::new())); + } + + // Find the start of the word we're completing + let word_start = find_word_start(line, pos); + let partial = &line[word_start..pos]; + + // Return only tables and columns (no keywords) + let mut candidates = Vec::new(); + + // Add table names + candidates.extend(self.cache.get_tables(partial)); + + // Add column names + candidates.extend(self.cache.get_columns(partial)); + + // Convert to Pair format (display, replacement) + let pairs: Vec = candidates + .into_iter() + .map(|c| Pair { + display: c.clone(), + replacement: c, + }) + .collect(); + + Ok((word_start, pairs)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rustyline::history::DefaultHistory; + + #[test] + fn test_completer_disabled() { + let cache = Arc::new(SchemaCache::new(300)); + let completer = SqlCompleter::new(cache, false); + + let history = DefaultHistory::new(); + let ctx = RustylineContext::new(&history); + let result = completer.complete("SELECT * FROM us", 17, &ctx).unwrap(); + + assert_eq!(result.1.len(), 0); + } + + #[test] + fn test_completer_no_keywords() { + let cache = Arc::new(SchemaCache::new(300)); + let completer = SqlCompleter::new(cache, true); + + let history = DefaultHistory::new(); + let ctx = RustylineContext::new(&history); + let result = completer.complete("SEL", 3, &ctx).unwrap(); + + // Should not return keywords (only tables and columns) + assert!(!result.1.iter().any(|p| p.display == "SELECT")); + } +} diff --git a/src/completion/schema_cache.rs b/src/completion/schema_cache.rs new file mode 100644 index 0000000..8dd42c1 --- /dev/null +++ b/src/completion/schema_cache.rs @@ -0,0 +1,562 @@ +use crate::context::Context; +use crate::query::query_silent; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; + +#[derive(Clone)] +pub struct TableMetadata { + pub schema_name: String, + pub table_name: String, + pub columns: Vec, +} + +#[derive(Clone)] +pub struct ColumnMetadata { + pub name: String, + pub data_type: String, +} + +pub struct SchemaCache { + tables: Arc>>, + functions: Arc>>, + keywords: HashSet, + last_refresh: Arc>>, + ttl: Duration, + refreshing: Arc>, +} + +impl SchemaCache { + pub fn new(ttl_seconds: u64) -> Self { + Self { + tables: Arc::new(RwLock::new(HashMap::new())), + functions: Arc::new(RwLock::new(HashSet::new())), + keywords: Self::load_keywords(), + last_refresh: Arc::new(RwLock::new(None)), + ttl: Duration::from_secs(ttl_seconds), + refreshing: Arc::new(RwLock::new(false)), + } + } + + fn load_keywords() -> HashSet { + let keywords = vec![ + "SELECT", "FROM", "WHERE", "JOIN", "INNER", "LEFT", "RIGHT", "FULL", "OUTER", + "ON", "GROUP", "ORDER", "BY", "HAVING", "LIMIT", "OFFSET", "UNION", "INTERSECT", + "EXCEPT", "INSERT", "INTO", "VALUES", "UPDATE", "SET", "DELETE", "CREATE", "DROP", + "ALTER", "TABLE", "DATABASE", "INDEX", "VIEW", "AS", "DISTINCT", "ALL", "AND", + "OR", "NOT", "IN", "EXISTS", "BETWEEN", "LIKE", "IS", "NULL", "CASE", "WHEN", + "THEN", "ELSE", "END", "WITH", "RECURSIVE", "ASC", "DESC", "COUNT", "SUM", "AVG", + "MIN", "MAX", "CAST", "SUBSTRING", "TRIM", "UPPER", "LOWER", "COALESCE", "NULLIF", + "PRIMARY", "KEY", "FOREIGN", "REFERENCES", "UNIQUE", "CHECK", "DEFAULT", "AUTO_INCREMENT", + "EXPLAIN", "DESCRIBE", "SHOW", "USE", "GRANT", "REVOKE", "COMMIT", "ROLLBACK", + "SAVEPOINT", "TRANSACTION", "BEGIN", "START", "TRUNCATE", "RENAME", "ADD", "MODIFY", + "COLUMN", "CONSTRAINT", "CASCADE", "RESTRICT", + ]; + + keywords.into_iter().map(String::from).collect() + } + + /// Checks if cache is stale based on TTL + pub fn is_stale(&self) -> bool { + let last_refresh = self.last_refresh.read().unwrap(); + match *last_refresh { + Some(instant) => instant.elapsed() > self.ttl, + None => true, // Never refreshed + } + } + + /// Returns true if a refresh is currently in progress + pub fn is_refreshing(&self) -> bool { + *self.refreshing.read().unwrap() + } + + /// Marks refresh as started + fn start_refresh(&self) { + *self.refreshing.write().unwrap() = true; + } + + /// Marks refresh as complete and updates timestamp + fn complete_refresh(&self) { + *self.refreshing.write().unwrap() = false; + *self.last_refresh.write().unwrap() = Some(Instant::now()); + } + + /// Get all keywords matching prefix + pub fn get_keywords(&self, prefix: &str) -> Vec { + let prefix_lower = prefix.to_lowercase(); + self.keywords + .iter() + .filter(|k| k.to_lowercase().starts_with(&prefix_lower)) + .cloned() + .collect() + } + + /// Get all table names matching prefix + pub fn get_tables(&self, prefix: &str) -> Vec { + let prefix_lower = prefix.to_lowercase(); + let tables = self.tables.read().unwrap(); + tables + .values() + .map(|t| { + if t.schema_name == "public" || t.schema_name.is_empty() { + t.table_name.clone() + } else { + format!("{}.{}", t.schema_name, t.table_name) + } + }) + .filter(|name| name.to_lowercase().starts_with(&prefix_lower)) + .collect() + } + + /// Get all column names matching prefix + pub fn get_columns(&self, prefix: &str) -> Vec { + let prefix_lower = prefix.to_lowercase(); + let tables = self.tables.read().unwrap(); + let mut columns = HashSet::new(); + for table in tables.values() { + for column in &table.columns { + columns.insert(column.name.clone()); + } + } + columns + .into_iter() + .filter(|name| name.to_lowercase().starts_with(&prefix_lower)) + .collect() + } + + /// Synchronous method to get completions from cache + pub fn get_completions( + &self, + context: super::context_detector::CompletionContext, + prefix: &str, + ) -> Vec { + use super::context_detector::CompletionContext; + + let prefix_lower = prefix.to_lowercase(); + + match context { + CompletionContext::Keyword => { + // Return keywords + self.keywords + .iter() + .filter(|k| k.to_lowercase().starts_with(&prefix_lower)) + .cloned() + .collect() + } + CompletionContext::TableName => { + // Return table names + let tables = self.tables.read().unwrap(); + tables + .values() + .map(|t| { + if t.schema_name == "public" || t.schema_name.is_empty() { + t.table_name.clone() + } else { + format!("{}.{}", t.schema_name, t.table_name) + } + }) + .filter(|name| name.to_lowercase().starts_with(&prefix_lower)) + .collect() + } + CompletionContext::ColumnName => { + // Return column names from all tables + let tables = self.tables.read().unwrap(); + let mut columns = HashSet::new(); + for table in tables.values() { + for column in &table.columns { + columns.insert(column.name.clone()); + } + } + columns + .into_iter() + .filter(|name| name.to_lowercase().starts_with(&prefix_lower)) + .collect() + } + CompletionContext::FunctionName => { + // Return function names + let functions = self.functions.read().unwrap(); + functions + .iter() + .filter(|f| f.to_lowercase().starts_with(&prefix_lower)) + .cloned() + .collect() + } + CompletionContext::SchemaQualified => { + // Handle schema.table completion + if let Some(dot_pos) = prefix.rfind('.') { + let schema = &prefix[..dot_pos]; + let table_prefix = &prefix[dot_pos + 1..]; + let table_prefix_lower = table_prefix.to_lowercase(); + + let tables = self.tables.read().unwrap(); + tables + .values() + .filter(|t| t.schema_name == schema) + .map(|t| format!("{}.{}", t.schema_name, t.table_name)) + .filter(|name| { + name.split('.').nth(1).unwrap_or("") + .to_lowercase() + .starts_with(&table_prefix_lower) + }) + .collect() + } else { + Vec::new() + } + } + CompletionContext::Nothing => Vec::new(), + } + } + + /// Async method to refresh schema from database + pub async fn refresh(&self, context: &mut Context) -> Result<(), Box> { + // Check if already refreshing + if self.is_refreshing() { + return Ok(()); + } + + self.start_refresh(); + + let result = self.do_refresh(context).await; + + if result.is_ok() { + self.complete_refresh(); + } else { + // Still mark as complete to avoid blocking future attempts + *self.refreshing.write().unwrap() = false; + } + + result + } + + async fn do_refresh(&self, context: &mut Context) -> Result<(), Box> { + // Query tables + let tables_query = "SELECT table_schema, table_name \ + FROM information_schema.tables \ + WHERE table_schema NOT IN ('information_schema', 'pg_catalog') \ + ORDER BY table_schema, table_name"; + + let tables_result = query_silent(context, tables_query).await; + + // Query columns + let columns_query = "SELECT table_schema, table_name, column_name, data_type \ + FROM information_schema.columns \ + WHERE table_schema NOT IN ('information_schema', 'pg_catalog') \ + ORDER BY table_schema, table_name, ordinal_position"; + + let columns_result = query_silent(context, columns_query).await; + + // Query functions + let functions_query = "SELECT routine_name \ + FROM information_schema.routines \ + WHERE routine_schema NOT IN ('information_schema', 'pg_catalog') \ + ORDER BY routine_name"; + + let functions_result = query_silent(context, functions_query).await; + + // Parse and populate cache + let mut new_tables = HashMap::new(); + + // Parse tables + match tables_result { + Ok(tables_output) => { + if let Some(table_list) = Self::parse_tables(&tables_output) { + for (schema, table) in table_list { + let key = format!("{}.{}", schema, table); + new_tables.insert( + key, + TableMetadata { + schema_name: schema, + table_name: table, + columns: Vec::new(), + }, + ); + } + } else { + eprintln!("Warning: Failed to parse tables from schema query"); + eprintln!("Output was: {}", &tables_output[..tables_output.len().min(500)]); + } + } + Err(e) => { + eprintln!("Warning: Tables query failed: {}", e); + } + } + + // Parse columns and add to tables + match columns_result { + Ok(columns_output) => { + if let Some(column_list) = Self::parse_columns(&columns_output) { + for (schema, table, column, data_type) in column_list { + let key = format!("{}.{}", schema, table); + if let Some(table_meta) = new_tables.get_mut(&key) { + table_meta.columns.push(ColumnMetadata { + name: column, + data_type, + }); + } + } + } else { + eprintln!("Warning: Failed to parse columns from schema query"); + } + } + Err(e) => { + eprintln!("Warning: Columns query failed: {}", e); + } + } + + // Update tables cache + let num_tables = new_tables.len(); + let num_columns: usize = new_tables.values().map(|t| t.columns.len()).sum(); + *self.tables.write().unwrap() = new_tables; + + // Parse functions + let mut num_functions = 0; + match functions_result { + Ok(functions_output) => { + if let Some(function_list) = Self::parse_functions(&functions_output) { + num_functions = function_list.len(); + *self.functions.write().unwrap() = function_list.into_iter().collect(); + } else { + eprintln!("Warning: Failed to parse functions from schema query"); + } + } + Err(e) => { + eprintln!("Warning: Functions query failed: {}", e); + } + } + + // Debug output + eprintln!("Schema cache loaded: {} tables, {} columns, {} functions", + num_tables, num_columns, num_functions); + + Ok(()) + } + + fn parse_tables(output: &str) -> Option> { + let mut result = Vec::new(); + + for line in output.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + + // Parse message-based format + if let Ok(json) = serde_json::from_str::(line) { + // Check for DATA message type + if let Some(msg_type) = json.get("message_type").and_then(|v| v.as_str()) { + if msg_type == "DATA" { + // Extract data array: [["public","test"],["public","test_view"]] + if let Some(data) = json.get("data").and_then(|v| v.as_array()) { + for row in data { + if let Some(row_array) = row.as_array() { + if row_array.len() >= 2 { + let schema = row_array[0].as_str().unwrap_or("").to_string(); + let table = row_array[1].as_str().unwrap_or("").to_string(); + if !schema.is_empty() && !table.is_empty() { + result.push((schema, table)); + } + } + } + } + } + } + } else { + // Try old JSONLines_Compact format for backwards compatibility + if let (Some(schema), Some(table)) = ( + json.get("table_schema").and_then(|v| v.as_str()), + json.get("table_name").and_then(|v| v.as_str()), + ) { + result.push((schema.to_string(), table.to_string())); + } + } + } + } + + if result.is_empty() { + None + } else { + Some(result) + } + } + + fn parse_columns(output: &str) -> Option> { + let mut result = Vec::new(); + + for line in output.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + + // Parse message-based format + if let Ok(json) = serde_json::from_str::(line) { + // Check for DATA message type + if let Some(msg_type) = json.get("message_type").and_then(|v| v.as_str()) { + if msg_type == "DATA" { + // Extract data array: [["public","test","id","integer"],...] + if let Some(data) = json.get("data").and_then(|v| v.as_array()) { + for row in data { + if let Some(row_array) = row.as_array() { + if row_array.len() >= 4 { + let schema = row_array[0].as_str().unwrap_or("").to_string(); + let table = row_array[1].as_str().unwrap_or("").to_string(); + let column = row_array[2].as_str().unwrap_or("").to_string(); + let data_type = row_array[3].as_str().unwrap_or("").to_string(); + if !schema.is_empty() && !table.is_empty() && !column.is_empty() { + result.push((schema, table, column, data_type)); + } + } + } + } + } + } + } else { + // Try old JSONLines_Compact format for backwards compatibility + if let (Some(schema), Some(table), Some(column), Some(data_type)) = ( + json.get("table_schema").and_then(|v| v.as_str()), + json.get("table_name").and_then(|v| v.as_str()), + json.get("column_name").and_then(|v| v.as_str()), + json.get("data_type").and_then(|v| v.as_str()), + ) { + result.push(( + schema.to_string(), + table.to_string(), + column.to_string(), + data_type.to_string(), + )); + } + } + } + } + + if result.is_empty() { + None + } else { + Some(result) + } + } + + fn parse_functions(output: &str) -> Option> { + let mut result = Vec::new(); + + for line in output.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + + // Parse message-based format + if let Ok(json) = serde_json::from_str::(line) { + // Check for DATA message type + if let Some(msg_type) = json.get("message_type").and_then(|v| v.as_str()) { + if msg_type == "DATA" { + // Extract data array: [["function1"],["function2"],...] + if let Some(data) = json.get("data").and_then(|v| v.as_array()) { + for row in data { + if let Some(row_array) = row.as_array() { + if !row_array.is_empty() { + if let Some(routine_name) = row_array[0].as_str() { + if !routine_name.is_empty() { + result.push(routine_name.to_string()); + } + } + } + } + } + } + } + } else { + // Try old JSONLines_Compact format for backwards compatibility + if let Some(routine_name) = json.get("routine_name").and_then(|v| v.as_str()) { + result.push(routine_name.to_string()); + } + } + } + } + + if result.is_empty() { + None + } else { + Some(result) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_tables() { + // Test new message-based format + let output = r#"{"message_type":"START","result_columns":[{"name":"table_schema","type":"text"},{"name":"table_name","type":"text"}]} +{"message_type":"DATA","data":[["public","users"],["public","orders"]]} +{"message_type":"FINISH_SUCCESSFULLY","statistics":{}}"#; + let tables = SchemaCache::parse_tables(output).unwrap(); + assert_eq!(tables.len(), 2); + assert_eq!(tables[0], ("public".to_string(), "users".to_string())); + assert_eq!(tables[1], ("public".to_string(), "orders".to_string())); + + // Test old JSONLines_Compact format for backwards compatibility + let output_old = r#"{"table_schema":"public","table_name":"legacy_table"}"#; + let tables_old = SchemaCache::parse_tables(output_old).unwrap(); + assert_eq!(tables_old.len(), 1); + assert_eq!(tables_old[0], ("public".to_string(), "legacy_table".to_string())); + } + + #[test] + fn test_parse_columns() { + // Test new message-based format + let output = r#"{"message_type":"START","result_columns":[]} +{"message_type":"DATA","data":[["public","users","id","INTEGER"],["public","users","name","TEXT"]]} +{"message_type":"FINISH_SUCCESSFULLY","statistics":{}}"#; + let columns = SchemaCache::parse_columns(output).unwrap(); + assert_eq!(columns.len(), 2); + assert_eq!( + columns[0], + ( + "public".to_string(), + "users".to_string(), + "id".to_string(), + "INTEGER".to_string() + ) + ); + } + + #[test] + fn test_keywords_loaded() { + let cache = SchemaCache::new(300); + assert!(cache.keywords.contains("SELECT")); + assert!(cache.keywords.contains("FROM")); + assert!(cache.keywords.contains("WHERE")); + } + + #[test] + fn test_is_stale() { + let cache = SchemaCache::new(1); // 1 second TTL + assert!(cache.is_stale()); // Never refreshed + + cache.complete_refresh(); + assert!(!cache.is_stale()); + + std::thread::sleep(Duration::from_millis(1100)); + assert!(cache.is_stale()); + } + + #[test] + fn test_get_keywords() { + let cache = SchemaCache::new(300); + let keywords = cache.get_keywords("SEL"); + assert!(keywords.contains(&"SELECT".to_string())); + } + + #[test] + fn test_empty_cache() { + let cache = SchemaCache::new(300); + // Cache should have keywords even when empty + assert!(cache.get_keywords("").len() > 0); + // But no tables or columns yet + assert_eq!(cache.get_tables("").len(), 0); + assert_eq!(cache.get_columns("").len(), 0); + } +} diff --git a/src/context.rs b/src/context.rs index 6625d7d..1d6fced 100644 --- a/src/context.rs +++ b/src/context.rs @@ -10,6 +10,7 @@ pub struct ServiceAccountToken { pub until: u64, } +#[derive(Clone)] pub struct Context { pub args: Args, pub url: String, diff --git a/src/main.rs b/src/main.rs index de62548..e59b916 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ use std::io::IsTerminal; mod args; mod auth; +mod completion; mod context; mod highlight; mod meta_commands; @@ -14,11 +15,14 @@ mod viewer; use args::get_args; use auth::maybe_authenticate; +use completion::schema_cache::SchemaCache; +use completion::SqlCompleter; use context::Context; use highlight::SqlHighlighter; use meta_commands::handle_meta_command; use query::{query, try_split_queries}; use repl_helper::ReplHelper; +use std::sync::Arc; use utils::history_path; use viewer::open_csvlens_viewer; @@ -63,7 +67,11 @@ async fn main() -> Result<(), Box> { SqlHighlighter::new(false).unwrap() // Fallback to disabled }); - let helper = ReplHelper::new(highlighter); + // Initialize schema cache for completion + let schema_cache = Arc::new(SchemaCache::new(context.args.completion_cache_ttl)); + let completer = SqlCompleter::new(schema_cache.clone(), !context.args.no_completion); + + let helper = ReplHelper::new(highlighter, completer); let mut rl: Editor = Editor::new()?; rl.set_helper(Some(helper)); @@ -88,6 +96,18 @@ async fn main() -> Result<(), Box> { EventHandler::Simple(Cmd::Insert(1, "\\view".to_string())), ); + // Spawn background task to refresh schema cache if completion is enabled + if !context.args.no_completion && is_tty { + let cache = schema_cache.clone(); + let mut ctx_clone = context.clone(); + tokio::spawn(async move { + if let Err(e) = cache.refresh(&mut ctx_clone).await { + // Report errors to help debug + eprintln!("Failed to refresh schema cache: {}", e); + } + }); + } + if is_tty && !context.args.concise { eprintln!("Type \\help for available commands or press Ctrl+V then Enter to view last result. Ctrl+D to exit."); } @@ -132,16 +152,38 @@ async fn main() -> Result<(), Box> { eprintln!("Failed to open viewer: {}", e); } continue; + } else if trimmed == "\\refresh" || trimmed == "\\refresh_cache" { + // Refresh schema cache for auto-completion + if !context.args.no_completion { + let cache = schema_cache.clone(); + let mut ctx_clone = context.clone(); + if is_tty { + eprintln!("Refreshing schema cache..."); + } + tokio::spawn(async move { + if let Err(e) = cache.refresh(&mut ctx_clone).await { + eprintln!("Failed to refresh schema cache: {}", e); + } else { + eprintln!("Schema cache refreshed successfully"); + } + }); + } else { + eprintln!("Auto-completion is disabled. Enable it with: set completion = on;"); + } + continue; } else if trimmed == "\\help" { // Show help for special commands eprintln!("Special commands:"); eprintln!(" \\view - Open last query result in csvlens viewer"); eprintln!(" (requires client format: client:auto, client:vertical, or client:horizontal)"); + eprintln!(" \\refresh - Manually refresh schema cache for auto-completion"); eprintln!(" \\help - Show this help message"); eprintln!(); eprintln!("SQL-style commands:"); - eprintln!(" set format = ; - Change output format"); - eprintln!(" unset format; - Reset format to default"); + eprintln!(" set format = ; - Change output format"); + eprintln!(" unset format; - Reset format to default"); + eprintln!(" set completion = on/off; - Enable/disable auto-completion"); + eprintln!(" unset completion; - Reset completion to default"); eprintln!(); eprintln!("Format values:"); eprintln!(" Client-side rendering (prefix with 'client:'):"); @@ -197,6 +239,12 @@ async fn main() -> Result<(), Box> { } } + // Update completer enabled state after processing queries + // (in case set completion = on/off was executed) + if let Some(helper) = rl.helper_mut() { + helper.completer_mut().set_enabled(!context.args.no_completion); + } + buffer.clear(); } } diff --git a/src/query.rs b/src/query.rs index e40bc66..bc1521b 100644 --- a/src/query.rs +++ b/src/query.rs @@ -105,6 +105,17 @@ pub fn set_args(context: &mut Context, query: &str) -> Result = vec![]; buf.push(format!("{key}={value}")); @@ -133,6 +144,9 @@ pub fn unset_args(context: &mut Context, query: &str) -> Result Result Result> { + let request = reqwest::Client::builder() + .http2_keep_alive_timeout(std::time::Duration::from_secs(3600)) + .http2_keep_alive_interval(Some(std::time::Duration::from_secs(60))) + .http2_keep_alive_while_idle(false) + .tcp_keepalive(Some(std::time::Duration::from_secs(60))) + .build()? + .post(context.url.clone()) + .header("user-agent", USER_AGENT) + .header("Firebolt-Protocol-Version", FIREBOLT_PROTOCOL_VERSION) + .body(query_text.to_string()); + + let request = if let Some(sa_token) = &context.sa_token { + request.header("authorization", format!("Bearer {}", sa_token.token)) + } else if !context.args.jwt.is_empty() { + request.header("authorization", format!("Bearer {}", context.args.jwt)) + } else { + request + }; + + let response = request.send().await?; + let body = response.text().await?; + Ok(body) +} + // Send query and print result. pub async fn query(context: &mut Context, query_text: String) -> Result<(), Box> { // Handle set/unset commands diff --git a/src/repl_helper.rs b/src/repl_helper.rs index 7c48bbb..a6e2c55 100644 --- a/src/repl_helper.rs +++ b/src/repl_helper.rs @@ -1,29 +1,48 @@ +use crate::completion::SqlCompleter; use crate::highlight::SqlHighlighter; -use rustyline::completion::Completer; +use rustyline::completion::{Completer, Pair}; use rustyline::highlight::Highlighter; use rustyline::hint::Hinter; use rustyline::validate::Validator; -use rustyline::Helper; +use rustyline::{Context, Helper}; use std::borrow::Cow; -/// REPL helper that integrates SqlHighlighter with rustyline +/// REPL helper that integrates SqlHighlighter and SqlCompleter with rustyline pub struct ReplHelper { highlighter: SqlHighlighter, + completer: SqlCompleter, } impl ReplHelper { - /// Create a new REPL helper with the given highlighter - pub fn new(highlighter: SqlHighlighter) -> Self { - Self { highlighter } + /// Create a new REPL helper with the given highlighter and completer + pub fn new(highlighter: SqlHighlighter, completer: SqlCompleter) -> Self { + Self { + highlighter, + completer, + } + } + + /// Get a mutable reference to the completer + pub fn completer_mut(&mut self) -> &mut SqlCompleter { + &mut self.completer } } // Implement the Helper trait (required) impl Helper for ReplHelper {} -// Implement empty Completer (no autocompletion for now) +// Delegate completion to SqlCompleter impl Completer for ReplHelper { - type Candidate = String; + type Candidate = Pair; + + fn complete( + &self, + line: &str, + pos: usize, + ctx: &Context<'_>, + ) -> rustyline::Result<(usize, Vec)> { + self.completer.complete(line, pos, ctx) + } } // Implement empty Hinter (no hints for now) From 8f5b3f0fb9cf3418af4f6aa3e41c8f56b72a0a63 Mon Sep 17 00:00:00 2001 From: Tobias Humig Date: Sat, 7 Feb 2026 18:44:48 +0100 Subject: [PATCH 3/4] Add intelligent prioritization to SQL auto-completion Implements context-aware, frequency-based suggestion ordering with schema exploration support. Tables and columns now appear based on query context, usage frequency, and relevance, with system schemas appropriately deprioritized. Schema names can be completed directly and typing 'schema.' shows all tables in that schema. Co-Authored-By: Claude Sonnet 4.5 --- src/completion/context_analyzer.rs | 168 ++++++++++++++++ src/completion/mod.rs | 187 ++++++++++++++++-- src/completion/priority_scorer.rs | 267 +++++++++++++++++++++++++ src/completion/schema_cache.rs | 134 ++++++++++++- src/completion/usage_tracker.rs | 304 +++++++++++++++++++++++++++++ src/context.rs | 4 + src/main.rs | 9 +- src/query.rs | 7 + 8 files changed, 1054 insertions(+), 26 deletions(-) create mode 100644 src/completion/context_analyzer.rs create mode 100644 src/completion/priority_scorer.rs create mode 100644 src/completion/usage_tracker.rs diff --git a/src/completion/context_analyzer.rs b/src/completion/context_analyzer.rs new file mode 100644 index 0000000..16f4946 --- /dev/null +++ b/src/completion/context_analyzer.rs @@ -0,0 +1,168 @@ +/// Analyzes SQL context to extract table names and identify system schemas +pub struct ContextAnalyzer; + +impl ContextAnalyzer { + /// Extract table names from the current SQL statement + /// Looks for patterns like "FROM table", "JOIN table", "UPDATE table" + pub fn extract_tables(sql: &str) -> Vec { + let mut tables = Vec::new(); + let sql_upper = sql.to_uppercase(); + + // Pattern: FROM table_name + if let Some(from_pos) = sql_upper.find("FROM") { + let after_from = &sql[from_pos + 4..]; + if let Some(word) = Self::extract_first_identifier(after_from) { + tables.push(word); + } + } + + // Pattern: JOIN table_name + for (idx, _) in sql_upper.match_indices("JOIN") { + let after_join = &sql[idx + 4..]; + if let Some(word) = Self::extract_first_identifier(after_join) { + tables.push(word); + } + } + + // Pattern: UPDATE table_name + if let Some(update_pos) = sql_upper.find("UPDATE") { + let after_update = &sql[update_pos + 6..]; + if let Some(word) = Self::extract_first_identifier(after_update) { + tables.push(word); + } + } + + // Pattern: INTO table_name (for INSERT statements) + if let Some(into_pos) = sql_upper.find("INTO") { + let after_into = &sql[into_pos + 4..]; + if let Some(word) = Self::extract_first_identifier(after_into) { + tables.push(word); + } + } + + tables + } + + /// Check if a qualified name belongs to a system schema + pub fn is_system_schema(qualified_name: &str) -> bool { + qualified_name.starts_with("information_schema.") + || qualified_name.starts_with("pg_catalog.") + || qualified_name == "information_schema" + || qualified_name == "pg_catalog" + } + + /// Extract the first SQL identifier from text + /// Handles schema-qualified names (e.g., "public.users") + fn extract_first_identifier(text: &str) -> Option { + let trimmed = text.trim(); + let mut identifier = String::new(); + + for ch in trimmed.chars() { + if ch.is_alphanumeric() || ch == '_' || ch == '.' { + identifier.push(ch); + } else if !identifier.is_empty() { + // Stop at first non-identifier character + break; + } + } + + if identifier.is_empty() { + None + } else { + Some(identifier) + } + } + + /// Extract the table name from a qualified column reference + /// E.g., "users.user_id" -> Some("users") + pub fn extract_table_from_qualified_column(qualified_column: &str) -> Option { + if let Some(dot_pos) = qualified_column.rfind('.') { + Some(qualified_column[..dot_pos].to_string()) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extract_tables_from_simple_query() { + let tables = ContextAnalyzer::extract_tables("SELECT * FROM users"); + assert_eq!(tables, vec!["users"]); + } + + #[test] + fn test_extract_tables_from_join() { + let tables = ContextAnalyzer::extract_tables( + "SELECT * FROM users JOIN orders ON users.id = orders.user_id" + ); + assert_eq!(tables.len(), 2); + assert!(tables.contains(&"users".to_string())); + assert!(tables.contains(&"orders".to_string())); + } + + #[test] + fn test_extract_tables_with_schema() { + let tables = ContextAnalyzer::extract_tables("SELECT * FROM public.users"); + assert_eq!(tables, vec!["public.users"]); + } + + #[test] + fn test_extract_tables_from_update() { + let tables = ContextAnalyzer::extract_tables("UPDATE users SET active = true"); + assert_eq!(tables, vec!["users"]); + } + + #[test] + fn test_extract_tables_from_insert() { + let tables = ContextAnalyzer::extract_tables("INSERT INTO users (name) VALUES ('Alice')"); + assert_eq!(tables, vec!["users"]); + } + + #[test] + fn test_is_system_schema() { + assert!(ContextAnalyzer::is_system_schema("information_schema.tables")); + assert!(ContextAnalyzer::is_system_schema("pg_catalog.pg_class")); + assert!(ContextAnalyzer::is_system_schema("information_schema")); + assert!(ContextAnalyzer::is_system_schema("pg_catalog")); + + assert!(!ContextAnalyzer::is_system_schema("public.users")); + assert!(!ContextAnalyzer::is_system_schema("users")); + assert!(!ContextAnalyzer::is_system_schema("my_schema.table")); + } + + #[test] + fn test_extract_table_from_qualified_column() { + assert_eq!( + ContextAnalyzer::extract_table_from_qualified_column("users.user_id"), + Some("users".to_string()) + ); + assert_eq!( + ContextAnalyzer::extract_table_from_qualified_column("public.users.user_id"), + Some("public.users".to_string()) + ); + assert_eq!( + ContextAnalyzer::extract_table_from_qualified_column("user_id"), + None + ); + } + + #[test] + fn test_extract_first_identifier() { + assert_eq!( + ContextAnalyzer::extract_first_identifier(" users "), + Some("users".to_string()) + ); + assert_eq!( + ContextAnalyzer::extract_first_identifier("public.users WHERE"), + Some("public.users".to_string()) + ); + assert_eq!( + ContextAnalyzer::extract_first_identifier("users,"), + Some("users".to_string()) + ); + } +} diff --git a/src/completion/mod.rs b/src/completion/mod.rs index 125c751..50b1db5 100644 --- a/src/completion/mod.rs +++ b/src/completion/mod.rs @@ -1,20 +1,34 @@ +pub mod context_analyzer; pub mod context_detector; +pub mod priority_scorer; pub mod schema_cache; +pub mod usage_tracker; -use context_detector::{detect_context, find_word_start}; +use context_analyzer::ContextAnalyzer; +use context_detector::find_word_start; +use priority_scorer::{PriorityScorer, ScoredSuggestion}; use rustyline::completion::{Completer, Pair}; use rustyline::Context as RustylineContext; use schema_cache::SchemaCache; +use usage_tracker::{ItemType, UsageTracker}; use std::sync::Arc; pub struct SqlCompleter { cache: Arc, + usage_tracker: Arc, + scorer: PriorityScorer, enabled: bool, } impl SqlCompleter { - pub fn new(cache: Arc, enabled: bool) -> Self { - Self { cache, enabled } + pub fn new(cache: Arc, usage_tracker: Arc, enabled: bool) -> Self { + let scorer = PriorityScorer::new(usage_tracker.clone()); + Self { + cache, + usage_tracker, + scorer, + enabled, + } } pub fn set_enabled(&mut self, enabled: bool) { @@ -47,21 +61,162 @@ impl Completer for SqlCompleter { let word_start = find_word_start(line, pos); let partial = &line[word_start..pos]; - // Return only tables and columns (no keywords) - let mut candidates = Vec::new(); + // Extract context: tables mentioned in current statement + let tables_in_line = ContextAnalyzer::extract_tables(line); + + // Generate scored suggestions + let mut scored: Vec = Vec::new(); + + // Check if user is typing "schema_name." to see tables from that schema + let (schema_filter, table_prefix) = if let Some(dot_pos) = partial.rfind('.') { + // User typed something like "information_schema." or "information_schema.engine" + let schema_part = &partial[..dot_pos]; + let table_part = &partial[dot_pos + 1..]; + (Some(schema_part), table_part) + } else { + (None, partial) + }; + + // If user specified a schema (e.g., "information_schema."), show tables from that schema + if let Some(schema_name) = schema_filter { + let tables_in_schema = self.cache.get_tables_in_schema(schema_name, table_prefix); + + for table in tables_in_schema { + let qualified_name = format!("{}.{}", schema_name, table); + let score = self.scorer.score(&table, ItemType::Table, &tables_in_line, None); + + scored.push(ScoredSuggestion { + name: qualified_name, + item_type: ItemType::Table, + score, + }); + } + } else { + // Normal completion: get tables and schemas + let table_metadata = self.cache.get_tables_with_schema(partial); + let column_metadata = self.cache.get_columns_with_table(partial); + let schemas = self.cache.get_schemas(partial); + + // Add schema suggestions (with trailing dot) + let partial_lower = partial.to_lowercase(); + + // Check if partial matches any schema name (without dot) + let partial_matches_schema = !partial.contains('.') && + schemas.iter().any(|s| s.to_lowercase().starts_with(&partial_lower)); + + for schema in &schemas { + let schema_with_dot = format!("{}.", schema); + + // Only add if it matches the partial + if schema_with_dot.to_lowercase().starts_with(&partial_lower) { + // Give schemas fixed priority class 4000 (same as tables) + // Don't use scorer to avoid system schema deprioritization + let base_score = 4000u32; + + // Add small usage bonus if schema has been used + let usage_count = self.usage_tracker.get_count(ItemType::Table, &schema); + let usage_bonus = usage_count.min(99) * 10; + + let score = base_score + usage_bonus; + + scored.push(ScoredSuggestion { + name: schema_with_dot, + item_type: ItemType::Table, + score, + }); + } + } - // Add table names - candidates.extend(self.cache.get_tables(partial)); + // Add table suggestions (both short and qualified names) + for (schema, table) in table_metadata { + let short_name = table.clone(); + let qualified_name = format!("{}.{}", schema, table); - // Add column names - candidates.extend(self.cache.get_columns(partial)); + // Score the short name + let score = self.scorer.score(&short_name, ItemType::Table, &tables_in_line, None); + + // Only add short name if it matches the partial + if short_name.to_lowercase().starts_with(&partial_lower) { + scored.push(ScoredSuggestion { + name: short_name.clone(), + item_type: ItemType::Table, + score, + }); + } + + // Add qualified name only if: + // 1. It matches the partial + // 2. Schema is not "public" or user typed a dot + // 3. Partial doesn't match a schema name (to avoid "infor" -> "information_schema.table") + if qualified_name.to_lowercase().starts_with(&partial_lower) && + (schema != "public" || partial.contains('.')) && + !partial_matches_schema { + scored.push(ScoredSuggestion { + name: qualified_name, + item_type: ItemType::Table, + score, + }); + } + } + + // Add column suggestions (both short and table-qualified names) + for (table, column) in column_metadata { + let short_name = column.clone(); + + // If user typed a dot, check if they're qualifying a table name + let is_qualifying_table = partial.contains('.') && + table.as_ref().map_or(false, |t| partial.to_lowercase().starts_with(&format!("{}.", t.to_lowercase()))); + + // Score the short name + let score = self.scorer.score( + &short_name, + ItemType::Column, + &tables_in_line, + table.as_deref(), + ); + + // Only add short name if it matches the partial and user is NOT qualifying a table + if !is_qualifying_table && short_name.to_lowercase().starts_with(&partial_lower) { + scored.push(ScoredSuggestion { + name: short_name.clone(), + item_type: ItemType::Column, + score, + }); + } + + // Add qualified name (table.column) if we know the table + if let Some(tbl) = table { + let qualified_name = format!("{}.{}", tbl, column); + + // Add qualified name if it matches partial and (no table context yet or user typed a dot) + if qualified_name.to_lowercase().starts_with(&partial_lower) && + (tables_in_line.is_empty() || partial.contains('.')) { + scored.push(ScoredSuggestion { + name: qualified_name, + item_type: ItemType::Column, + score: score.saturating_sub(1), + }); + } + } + } + } // end of else block + + // Sort by score (descending - higher scores first) + scored.sort_by(|a, b| b.score.cmp(&a.score)); + + // Remove duplicates (keep first occurrence, which has highest score) + let mut seen = std::collections::HashSet::new(); + let deduplicated: Vec = scored + .into_iter() + .filter(|s| seen.insert(s.name.clone())) + .collect(); // Convert to Pair format (display, replacement) - let pairs: Vec = candidates + let pairs: Vec = deduplicated .into_iter() - .map(|c| Pair { - display: c.clone(), - replacement: c, + .map(|s| Pair { + display: s.name.clone(), + replacement: s.name, }) .collect(); @@ -77,7 +232,8 @@ mod tests { #[test] fn test_completer_disabled() { let cache = Arc::new(SchemaCache::new(300)); - let completer = SqlCompleter::new(cache, false); + let usage_tracker = Arc::new(UsageTracker::new(10)); + let completer = SqlCompleter::new(cache, usage_tracker, false); let history = DefaultHistory::new(); let ctx = RustylineContext::new(&history); @@ -89,7 +245,8 @@ mod tests { #[test] fn test_completer_no_keywords() { let cache = Arc::new(SchemaCache::new(300)); - let completer = SqlCompleter::new(cache, true); + let usage_tracker = Arc::new(UsageTracker::new(10)); + let completer = SqlCompleter::new(cache, usage_tracker, true); let history = DefaultHistory::new(); let ctx = RustylineContext::new(&history); diff --git a/src/completion/priority_scorer.rs b/src/completion/priority_scorer.rs new file mode 100644 index 0000000..512ff3d --- /dev/null +++ b/src/completion/priority_scorer.rs @@ -0,0 +1,267 @@ +use super::context_analyzer::ContextAnalyzer; +use super::usage_tracker::{ItemType, UsageTracker}; +use std::sync::Arc; + +/// Calculates priority scores for auto-completion suggestions +pub struct PriorityScorer { + usage_tracker: Arc, +} + +/// A suggestion with its calculated priority score +#[derive(Debug, Clone)] +pub struct ScoredSuggestion { + pub name: String, + pub item_type: ItemType, + pub score: u32, +} + +/// Priority classes determine the base score before usage bonuses +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PriorityClass { + /// Columns from tables mentioned in current query - highest priority + ColumnFromQueryTable = 5000, + + /// Tables - high priority when no table context + Table = 4000, + + /// Qualified columns from tables NOT in query + QualifiedColumnOtherTable = 3000, + + /// Unqualified columns from tables NOT in query + UnqualifiedColumnOtherTable = 2000, + + /// System schema items - lowest priority + SystemSchema = 0, +} + +impl PriorityScorer { + pub fn new(usage_tracker: Arc) -> Self { + Self { usage_tracker } + } + + /// Calculate priority score for a suggestion + /// + /// # Arguments + /// * `name` - The suggestion name (e.g., "users", "user_id", "users.user_id") + /// * `item_type` - Whether this is a table or column + /// * `tables_in_statement` - Tables mentioned in the current SQL statement + /// * `column_table` - For columns, the table they belong to (if known) + /// + /// # Returns + /// Priority score (higher = more relevant) + pub fn score( + &self, + name: &str, + item_type: ItemType, + tables_in_statement: &[String], + column_table: Option<&str>, + ) -> u32 { + // Get usage count + let usage_count = self.usage_tracker.get_count(item_type, name); + + // Calculate base score from priority class + let base_score = self.calculate_priority_class( + name, + item_type, + tables_in_statement, + column_table, + ) as u32; + + // Add usage bonus (max 99 uses = 990 points, ensures higher class always wins) + let usage_bonus = usage_count.min(99) * 10; + + base_score + usage_bonus + } + + /// Determine the priority class for an item + fn calculate_priority_class( + &self, + name: &str, + item_type: ItemType, + tables_in_statement: &[String], + column_table: Option<&str>, + ) -> PriorityClass { + // System schemas always get lowest priority + if ContextAnalyzer::is_system_schema(name) { + return PriorityClass::SystemSchema; + } + + match item_type { + ItemType::Table => PriorityClass::Table, + + ItemType::Column => { + // Check if column belongs to a table in the statement + if let Some(table) = column_table { + // Normalize table name for comparison + let normalized_table = table.to_lowercase(); + + // Check if this table is mentioned in the statement + let table_in_statement = tables_in_statement.iter().any(|t| { + let normalized_t = t.to_lowercase(); + normalized_t == normalized_table + || normalized_t.ends_with(&format!(".{}", normalized_table)) + || normalized_table.ends_with(&format!(".{}", normalized_t)) + }); + + if table_in_statement { + // Column from table in query - highest priority + return PriorityClass::ColumnFromQueryTable; + } + } + + // Column from other table - check if qualified + let is_qualified = name.contains('.'); + + if is_qualified { + PriorityClass::QualifiedColumnOtherTable + } else { + PriorityClass::UnqualifiedColumnOtherTable + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn create_test_scorer() -> PriorityScorer { + let tracker = Arc::new(UsageTracker::new(10)); + PriorityScorer::new(tracker) + } + + #[test] + fn test_tables_prioritized_when_no_context() { + let scorer = create_test_scorer(); + + let table_score = scorer.score("users", ItemType::Table, &[], None); + let column_score = scorer.score("user_id", ItemType::Column, &[], None); + + assert!(table_score > column_score, "Tables should be prioritized over columns when no context"); + assert!(table_score >= 4000, "Table should be in Table priority class"); + assert!(column_score < 3000, "Column should be in lower priority class"); + } + + #[test] + fn test_columns_prioritized_when_table_in_statement() { + let scorer = create_test_scorer(); + + let column_score = scorer.score( + "user_id", + ItemType::Column, + &["users".to_string()], + Some("users"), + ); + + assert!(column_score >= 5000, "Column from query table should be highest priority"); + } + + #[test] + fn test_system_schema_deprioritized() { + let scorer = create_test_scorer(); + + let user_table_score = scorer.score("users", ItemType::Table, &[], None); + let system_table_score = scorer.score("information_schema.tables", ItemType::Table, &[], None); + + assert!(user_table_score > system_table_score, "User tables should be prioritized over system schemas"); + assert!(system_table_score < 1000, "System schema should be in lowest priority class"); + } + + #[test] + fn test_qualified_vs_unqualified_columns() { + let scorer = create_test_scorer(); + + let qualified_score = scorer.score( + "orders.order_id", + ItemType::Column, + &["users".to_string()], // Different table in statement + Some("orders"), + ); + + let unqualified_score = scorer.score( + "order_id", + ItemType::Column, + &["users".to_string()], + Some("orders"), + ); + + assert!(qualified_score > unqualified_score, "Qualified columns should rank above unqualified"); + assert!(qualified_score >= 3000 && qualified_score < 4000, "Qualified column should be in correct class"); + assert!(unqualified_score >= 2000 && unqualified_score < 3000, "Unqualified column should be in correct class"); + } + + #[test] + fn test_usage_bonus_within_class() { + let tracker = Arc::new(UsageTracker::new(10)); + + // Simulate usage + for _ in 0..50 { + tracker.track_query("SELECT * FROM users"); + } + for _ in 0..10 { + tracker.track_query("SELECT * FROM orders"); + } + + let scorer = PriorityScorer::new(tracker); + + let users_score = scorer.score("users", ItemType::Table, &[], None); + let orders_score = scorer.score("orders", ItemType::Table, &[], None); + + assert!(users_score > orders_score, "More frequently used item should score higher"); + + // Both should be in same priority class (4000-4999) + assert!(users_score >= 4000 && users_score < 5000); + assert!(orders_score >= 4000 && orders_score < 5000); + } + + #[test] + fn test_priority_class_beats_usage_count() { + let tracker = Arc::new(UsageTracker::new(10)); + + // Give column very high usage + for _ in 0..99 { + tracker.track_query("SELECT order_id FROM orders"); + } + + let scorer = PriorityScorer::new(tracker.clone()); + + // Column with high usage from other table + let column_score = scorer.score( + "order_id", + ItemType::Column, + &["users".to_string()], // Different table + Some("orders"), + ); + + // Table with no usage + let table_score = scorer.score("users", ItemType::Table, &[], None); + + // Table should still win because higher priority class + assert!(table_score > column_score, "Higher priority class should beat usage count"); + } + + #[test] + fn test_schema_qualified_table_matching() { + let scorer = create_test_scorer(); + + // Test that schema-qualified names are handled correctly + let score1 = scorer.score( + "user_id", + ItemType::Column, + &["public.users".to_string()], + Some("users"), + ); + + let score2 = scorer.score( + "user_id", + ItemType::Column, + &["users".to_string()], + Some("public.users"), + ); + + // Both should recognize the table match + assert!(score1 >= 5000, "Should match schema.table to table"); + assert!(score2 >= 5000, "Should match table to schema.table"); + } +} diff --git a/src/completion/schema_cache.rs b/src/completion/schema_cache.rs index 8dd42c1..79362c1 100644 --- a/src/completion/schema_cache.rs +++ b/src/completion/schema_cache.rs @@ -124,6 +124,126 @@ impl SchemaCache { .collect() } + /// Get all unique schema names matching prefix + pub fn get_schemas(&self, prefix: &str) -> Vec { + let prefix_lower = prefix.to_lowercase(); + let tables = self.tables.read().unwrap(); + + let mut schemas = std::collections::HashSet::new(); + for table in tables.values() { + if !table.schema_name.is_empty() { + schemas.insert(table.schema_name.clone()); + } + } + + schemas + .into_iter() + .filter(|schema| schema.to_lowercase().starts_with(&prefix_lower)) + .collect() + } + + /// Get all tables from a specific schema, optionally filtered by table name prefix + pub fn get_tables_in_schema(&self, schema: &str, table_prefix: &str) -> Vec { + let schema_lower = schema.to_lowercase(); + let prefix_lower = table_prefix.to_lowercase(); + let tables = self.tables.read().unwrap(); + + let mut result: Vec = tables + .values() + .filter(|t| t.schema_name.to_lowercase() == schema_lower) + .filter(|t| { + if prefix_lower.is_empty() { + true + } else { + t.table_name.to_lowercase().starts_with(&prefix_lower) + } + }) + .map(|t| t.table_name.clone()) + .collect(); + + result.sort(); + result + } + + /// Get all table names with their schemas matching prefix + /// Returns Vec<(schema, table)> + pub fn get_tables_with_schema(&self, prefix: &str) -> Vec<(String, String)> { + let prefix_lower = prefix.to_lowercase(); + let tables = self.tables.read().unwrap(); + + let mut result: Vec<(String, String)> = tables + .values() + .filter_map(|t| { + let short_name = t.table_name.to_lowercase(); + let qualified_name = format!("{}.{}", t.schema_name, t.table_name).to_lowercase(); + + // Match either short name or qualified name + if short_name.starts_with(&prefix_lower) || qualified_name.starts_with(&prefix_lower) { + Some((t.schema_name.clone(), t.table_name.clone())) + } else { + None + } + }) + .collect(); + + // Sort by table name for consistent ordering + result.sort_by(|a, b| a.1.cmp(&b.1)); + result + } + + /// Get all column names with their table names matching prefix + /// Returns Vec<(Option, column)> + pub fn get_columns_with_table(&self, prefix: &str) -> Vec<(Option, String)> { + let prefix_lower = prefix.to_lowercase(); + let tables = self.tables.read().unwrap(); + + let mut result: Vec<(Option, String)> = Vec::new(); + let mut seen_columns = HashSet::new(); + + for table in tables.values() { + for column in &table.columns { + let column_name = column.name.to_lowercase(); + let qualified_name = format!("{}.{}", table.table_name, column.name).to_lowercase(); + + // Check if matches prefix (either column name or qualified name) + if column_name.starts_with(&prefix_lower) || qualified_name.starts_with(&prefix_lower) { + // Track unique columns to avoid duplicates + let key = format!("{}:{}", table.table_name, column.name); + if !seen_columns.contains(&key) { + result.push((Some(table.table_name.clone()), column.name.clone())); + seen_columns.insert(key); + } + } + } + } + + // Sort by column name for consistent ordering + result.sort_by(|a, b| a.1.cmp(&b.1)); + result + } + + /// Get the table for a given column name + /// Returns the first table that contains this column + pub fn get_table_for_column(&self, column_name: &str) -> Option { + let tables = self.tables.read().unwrap(); + let column_lower = column_name.to_lowercase(); + + for table in tables.values() { + for column in &table.columns { + if column.name.to_lowercase() == column_lower { + // Return qualified table name + return Some(if table.schema_name == "public" || table.schema_name.is_empty() { + table.table_name.clone() + } else { + format!("{}.{}", table.schema_name, table.table_name) + }); + } + } + } + + None + } + /// Synchronous method to get completions from cache pub fn get_completions( &self, @@ -229,18 +349,16 @@ impl SchemaCache { } async fn do_refresh(&self, context: &mut Context) -> Result<(), Box> { - // Query tables + // Query tables (including system schemas - they'll be deprioritized by the scorer) let tables_query = "SELECT table_schema, table_name \ FROM information_schema.tables \ - WHERE table_schema NOT IN ('information_schema', 'pg_catalog') \ ORDER BY table_schema, table_name"; let tables_result = query_silent(context, tables_query).await; - // Query columns + // Query columns (including system schemas - they'll be deprioritized by the scorer) let columns_query = "SELECT table_schema, table_name, column_name, data_type \ FROM information_schema.columns \ - WHERE table_schema NOT IN ('information_schema', 'pg_catalog') \ ORDER BY table_schema, table_name, ordinal_position"; let columns_result = query_silent(context, columns_query).await; @@ -324,10 +442,6 @@ impl SchemaCache { } } - // Debug output - eprintln!("Schema cache loaded: {} tables, {} columns, {} functions", - num_tables, num_columns, num_functions); - Ok(()) } @@ -474,7 +588,9 @@ impl SchemaCache { } } - if result.is_empty() { + // Always return Some - empty result is valid (no functions defined) + // Only return None if output is completely empty/unparseable + if output.trim().is_empty() { None } else { Some(result) diff --git a/src/completion/usage_tracker.rs b/src/completion/usage_tracker.rs new file mode 100644 index 0000000..533c52f --- /dev/null +++ b/src/completion/usage_tracker.rs @@ -0,0 +1,304 @@ +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ItemType { + Table, + Column, +} + +/// Tracks usage frequency of tables and columns to enable intelligent prioritization +pub struct UsageTracker { + // Item name -> usage count + table_counts: Arc>>, + column_counts: Arc>>, + + // Recent queries (ring buffer of last N) + recent_queries: Arc>>, + max_recent: usize, +} + +impl UsageTracker { + /// Create a new UsageTracker with a maximum number of recent queries to track + pub fn new(max_recent: usize) -> Self { + Self { + table_counts: Arc::new(RwLock::new(HashMap::new())), + column_counts: Arc::new(RwLock::new(HashMap::new())), + recent_queries: Arc::new(RwLock::new(Vec::new())), + max_recent, + } + } + + /// Track a query execution by extracting table and column names + pub fn track_query(&self, query: &str) { + // 1. Add to recent queries (ring buffer) + { + let mut recent = self.recent_queries.write().unwrap(); + recent.push(query.to_string()); + if recent.len() > self.max_recent { + recent.remove(0); + } + } + + // 2. Extract table and column names + let tables = Self::extract_table_names(query); + let columns = Self::extract_column_names(query); + + // 3. Increment usage counts + { + let mut table_counts = self.table_counts.write().unwrap(); + for table in tables { + *table_counts.entry(table).or_insert(0) += 1; + } + } + + { + let mut column_counts = self.column_counts.write().unwrap(); + for column in columns { + *column_counts.entry(column).or_insert(0) += 1; + } + } + } + + /// Get usage count for an item + pub fn get_count(&self, item_type: ItemType, name: &str) -> u32 { + match item_type { + ItemType::Table => { + let counts = self.table_counts.read().unwrap(); + counts.get(name).copied().unwrap_or(0) + } + ItemType::Column => { + let counts = self.column_counts.read().unwrap(); + counts.get(name).copied().unwrap_or(0) + } + } + } + + /// Check if an item name appears in any recent query + pub fn was_used_recently(&self, name: &str) -> bool { + let recent = self.recent_queries.read().unwrap(); + recent.iter().any(|q| q.contains(name)) + } + + /// Extract table names from SQL query using simple pattern matching + fn extract_table_names(query: &str) -> Vec { + let mut tables = Vec::new(); + let query_upper = query.to_uppercase(); + + // Pattern: FROM table_name + if let Some(from_pos) = query_upper.find("FROM") { + let after_from = &query[from_pos + 4..]; + if let Some(word) = Self::extract_first_word(after_from) { + tables.push(word); + } + } + + // Pattern: JOIN table_name + for (idx, _) in query_upper.match_indices("JOIN") { + let after_join = &query[idx + 4..]; + if let Some(word) = Self::extract_first_word(after_join) { + tables.push(word); + } + } + + // Pattern: UPDATE table_name + if let Some(update_pos) = query_upper.find("UPDATE") { + let after_update = &query[update_pos + 6..]; + if let Some(word) = Self::extract_first_word(after_update) { + tables.push(word); + } + } + + // Pattern: INTO table_name + if let Some(into_pos) = query_upper.find("INTO") { + let after_into = &query[into_pos + 4..]; + if let Some(word) = Self::extract_first_word(after_into) { + tables.push(word); + } + } + + tables + } + + /// Extract column names from SQL query using simple pattern matching + fn extract_column_names(query: &str) -> Vec { + let mut columns = Vec::new(); + let query_upper = query.to_uppercase(); + + // Pattern: SELECT columns FROM + if let Some(select_pos) = query_upper.find("SELECT") { + if let Some(from_pos) = query_upper.find("FROM") { + // Only extract if FROM comes after SELECT + if from_pos > select_pos + 6 { + let between = &query[select_pos + 6..from_pos]; + columns.extend(Self::extract_column_list(between)); + } + } + } + + // Pattern: WHERE column = value + for (idx, _) in query_upper.match_indices("WHERE") { + let after_where = &query[idx + 5..]; + if let Some(word) = Self::extract_first_word(after_where) { + columns.push(word); + } + } + + // Pattern: ORDER BY column + for (idx, _) in query_upper.match_indices("ORDER BY") { + let after_order = &query[idx + 8..]; + if let Some(word) = Self::extract_first_word(after_order) { + columns.push(word); + } + } + + // Pattern: GROUP BY column + for (idx, _) in query_upper.match_indices("GROUP BY") { + let after_group = &query[idx + 8..]; + if let Some(word) = Self::extract_first_word(after_group) { + columns.push(word); + } + } + + columns + } + + /// Extract first identifier from text (handles schema.table notation) + fn extract_first_word(text: &str) -> Option { + let trimmed = text.trim(); + let mut word = String::new(); + + for ch in trimmed.chars() { + if ch.is_alphanumeric() || ch == '_' || ch == '.' { + word.push(ch); + } else if !word.is_empty() { + break; + } + } + + if word.is_empty() { + None + } else { + Some(word) + } + } + + /// Extract comma-separated column list + fn extract_column_list(text: &str) -> Vec { + let mut columns = Vec::new(); + + for part in text.split(',') { + let trimmed = part.trim(); + + // Skip * and common SQL functions/keywords + if trimmed == "*" || trimmed.is_empty() { + continue; + } + + // Extract just the column name (before AS, spaces, etc.) + if let Some(word) = Self::extract_first_word(trimmed) { + // Remove table prefix if present (e.g., "users.id" -> "id") + let column_name = if let Some(dot_pos) = word.rfind('.') { + word[dot_pos + 1..].to_string() + } else { + word + }; + + // Filter out SQL keywords + let upper = column_name.to_uppercase(); + if !["AS", "FROM", "WHERE", "AND", "OR", "NOT", "IN", "EXISTS"].contains(&upper.as_str()) { + columns.push(column_name); + } + } + } + + columns + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_track_query_increments_table_counts() { + let tracker = UsageTracker::new(10); + tracker.track_query("SELECT * FROM users"); + assert_eq!(tracker.get_count(ItemType::Table, "users"), 1); + + tracker.track_query("SELECT * FROM users"); + assert_eq!(tracker.get_count(ItemType::Table, "users"), 2); + } + + #[test] + fn test_track_query_increments_column_counts() { + let tracker = UsageTracker::new(10); + tracker.track_query("SELECT user_id, username FROM users"); + assert_eq!(tracker.get_count(ItemType::Column, "user_id"), 1); + assert_eq!(tracker.get_count(ItemType::Column, "username"), 1); + } + + #[test] + fn test_recent_queries_ring_buffer() { + let tracker = UsageTracker::new(2); + tracker.track_query("query1"); + tracker.track_query("query2"); + tracker.track_query("query3"); // Should evict query1 + + assert!(!tracker.was_used_recently("query1")); + assert!(tracker.was_used_recently("query2")); + assert!(tracker.was_used_recently("query3")); + } + + #[test] + fn test_extract_table_names_from_query() { + let tables = UsageTracker::extract_table_names( + "SELECT * FROM users JOIN orders ON users.id = orders.user_id" + ); + assert!(tables.contains(&"users".to_string())); + assert!(tables.contains(&"orders".to_string())); + } + + #[test] + fn test_extract_table_names_with_schema() { + let tables = UsageTracker::extract_table_names("SELECT * FROM public.users"); + assert_eq!(tables, vec!["public.users"]); + } + + #[test] + fn test_extract_column_names() { + let columns = UsageTracker::extract_column_names( + "SELECT user_id, username FROM users WHERE active = true" + ); + assert!(columns.contains(&"user_id".to_string())); + assert!(columns.contains(&"username".to_string())); + assert!(columns.contains(&"active".to_string())); + } + + #[test] + fn test_was_used_recently() { + let tracker = UsageTracker::new(10); + tracker.track_query("SELECT * FROM users"); + assert!(tracker.was_used_recently("users")); + assert!(!tracker.was_used_recently("orders")); + } + + #[test] + fn test_extract_column_names_with_keywords_out_of_order() { + // Regression test for panic when FROM appears before SELECT + let columns = UsageTracker::extract_column_names("from test select val"); + // Should not panic, and should not extract columns since FROM comes before SELECT + assert_eq!(columns.len(), 0); + } + + #[test] + fn test_extract_column_names_edge_cases() { + // Test with only SELECT, no FROM + let columns = UsageTracker::extract_column_names("SELECT user_id"); + assert_eq!(columns.len(), 0); // No FROM, so no columns extracted + + // Test with FROM immediately after SELECT (no space for columns) + let columns = UsageTracker::extract_column_names("SELECTFROM users"); + assert_eq!(columns.len(), 0); + } +} diff --git a/src/context.rs b/src/context.rs index 1d6fced..e1e533d 100644 --- a/src/context.rs +++ b/src/context.rs @@ -1,6 +1,8 @@ use crate::args::{get_url, Args}; +use crate::completion::usage_tracker::UsageTracker; use crate::table_renderer::ParsedResult; use serde::{Deserialize, Serialize}; +use std::sync::Arc; #[derive(Clone, Debug, Deserialize, Serialize)] pub struct ServiceAccountToken { @@ -21,6 +23,7 @@ pub struct Context { pub last_result: Option, pub last_stats: Option, pub is_interactive: bool, + pub usage_tracker: Option>, } impl Context { @@ -36,6 +39,7 @@ impl Context { last_result: None, last_stats: None, is_interactive: false, + usage_tracker: None, } } diff --git a/src/main.rs b/src/main.rs index e59b916..0154353 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ mod viewer; use args::get_args; use auth::maybe_authenticate; use completion::schema_cache::SchemaCache; +use completion::usage_tracker::UsageTracker; use completion::SqlCompleter; use context::Context; use highlight::SqlHighlighter; @@ -67,9 +68,13 @@ async fn main() -> Result<(), Box> { SqlHighlighter::new(false).unwrap() // Fallback to disabled }); - // Initialize schema cache for completion + // Initialize schema cache and usage tracker for completion let schema_cache = Arc::new(SchemaCache::new(context.args.completion_cache_ttl)); - let completer = SqlCompleter::new(schema_cache.clone(), !context.args.no_completion); + let usage_tracker = Arc::new(UsageTracker::new(10)); // Track last 10 queries + let completer = SqlCompleter::new(schema_cache.clone(), usage_tracker.clone(), !context.args.no_completion); + + // Store usage tracker in context for query tracking + context.usage_tracker = Some(usage_tracker.clone()); let helper = ReplHelper::new(highlighter, completer); let mut rl: Editor = Editor::new()?; diff --git a/src/query.rs b/src/query.rs index bc1521b..3bae7bc 100644 --- a/src/query.rs +++ b/src/query.rs @@ -215,6 +215,9 @@ pub async fn query(context: &mut Context, query_text: String) -> Result<(), Box< let start = Instant::now(); + // Clone query_text for tracking later + let query_text_for_tracking = query_text.clone(); + let mut request = reqwest::Client::builder() .http2_keep_alive_timeout(std::time::Duration::from_secs(3600)) .http2_keep_alive_interval(Some(std::time::Duration::from_secs(60))) @@ -448,6 +451,10 @@ pub async fn query(context: &mut Context, query_text: String) -> Result<(), Box< if query_failed { Err("Query failed".into()) } else { + // Track successful query for auto-completion prioritization + if let Some(usage_tracker) = &context.usage_tracker { + usage_tracker.track_query(&query_text_for_tracking); + } Ok(()) } } From 6a747d9c904dca929dc2cc236165de7e65363984 Mon Sep 17 00:00:00 2001 From: Tobias Humig Date: Sat, 7 Feb 2026 19:36:04 +0100 Subject: [PATCH 4/4] Add function name completion to SQL auto-completion Implements auto-completion for SQL functions with lowest priority (below columns, above system schemas). Functions complete with opening parenthesis for immediate argument typing. Operators are filtered out using routine_type != 'OPERATOR' from information_schema.routines. Co-Authored-By: Claude Sonnet 4.5 --- src/args.rs | 6 ++++++ src/completion/mod.rs | 25 +++++++++++++++++++++++++ src/completion/priority_scorer.rs | 5 +++++ src/completion/schema_cache.rs | 16 ++++++++++++++-- src/completion/usage_tracker.rs | 9 ++++++++- 5 files changed, 58 insertions(+), 3 deletions(-) diff --git a/src/args.rs b/src/args.rs index 4ed9b92..006ca1b 100644 --- a/src/args.rs +++ b/src/args.rs @@ -30,7 +30,13 @@ fn default_max_cell_length() -> usize { 1000 } +fn default_true() -> bool { + true +} +fn default_cache_ttl() -> u64 { + 300 +} #[derive(Clone, Debug, Options, Deserialize, Serialize)] pub struct Args { diff --git a/src/completion/mod.rs b/src/completion/mod.rs index 50b1db5..431110e 100644 --- a/src/completion/mod.rs +++ b/src/completion/mod.rs @@ -199,6 +199,31 @@ impl Completer for SqlCompleter { } } } + + // Add function suggestions + let functions = self.cache.get_functions(partial); + for function in functions { + // Only add if it matches the partial + if function.to_lowercase().starts_with(&partial_lower) { + // Give functions fixed priority class 1000 (below columns, above system schemas) + let base_score = 1000u32; + + // Add small usage bonus if function has been used + let usage_count = self.usage_tracker.get_count(ItemType::Function, &function); + let usage_bonus = usage_count.min(99) * 10; + + let score = base_score + usage_bonus; + + // Add opening parenthesis to function names (no closing paren for easier typing) + let function_with_paren = format!("{}(", function); + + scored.push(ScoredSuggestion { + name: function_with_paren, + item_type: ItemType::Column, // Use Column type for now + score, + }); + } + } } // end of else block // Sort by score (descending - higher scores first) diff --git a/src/completion/priority_scorer.rs b/src/completion/priority_scorer.rs index 512ff3d..a3f26f0 100644 --- a/src/completion/priority_scorer.rs +++ b/src/completion/priority_scorer.rs @@ -30,6 +30,9 @@ enum PriorityClass { /// Unqualified columns from tables NOT in query UnqualifiedColumnOtherTable = 2000, + /// Functions - lower priority than columns + Function = 1000, + /// System schema items - lowest priority SystemSchema = 0, } @@ -118,6 +121,8 @@ impl PriorityScorer { PriorityClass::UnqualifiedColumnOtherTable } } + + ItemType::Function => PriorityClass::Function, } } } diff --git a/src/completion/schema_cache.rs b/src/completion/schema_cache.rs index 79362c1..e4ed274 100644 --- a/src/completion/schema_cache.rs +++ b/src/completion/schema_cache.rs @@ -142,6 +142,18 @@ impl SchemaCache { .collect() } + /// Get all function names matching prefix + pub fn get_functions(&self, prefix: &str) -> Vec { + let prefix_lower = prefix.to_lowercase(); + let functions = self.functions.read().unwrap(); + + functions + .iter() + .filter(|f| f.to_lowercase().starts_with(&prefix_lower)) + .cloned() + .collect() + } + /// Get all tables from a specific schema, optionally filtered by table name prefix pub fn get_tables_in_schema(&self, schema: &str, table_prefix: &str) -> Vec { let schema_lower = schema.to_lowercase(); @@ -363,10 +375,10 @@ impl SchemaCache { let columns_result = query_silent(context, columns_query).await; - // Query functions + // Query functions (including system functions, excluding operators) let functions_query = "SELECT routine_name \ FROM information_schema.routines \ - WHERE routine_schema NOT IN ('information_schema', 'pg_catalog') \ + WHERE routine_type != 'OPERATOR' \ ORDER BY routine_name"; let functions_result = query_silent(context, functions_query).await; diff --git a/src/completion/usage_tracker.rs b/src/completion/usage_tracker.rs index 533c52f..786d895 100644 --- a/src/completion/usage_tracker.rs +++ b/src/completion/usage_tracker.rs @@ -5,13 +5,15 @@ use std::sync::{Arc, RwLock}; pub enum ItemType { Table, Column, + Function, } -/// Tracks usage frequency of tables and columns to enable intelligent prioritization +/// Tracks usage frequency of tables, columns, and functions to enable intelligent prioritization pub struct UsageTracker { // Item name -> usage count table_counts: Arc>>, column_counts: Arc>>, + function_counts: Arc>>, // Recent queries (ring buffer of last N) recent_queries: Arc>>, @@ -24,6 +26,7 @@ impl UsageTracker { Self { table_counts: Arc::new(RwLock::new(HashMap::new())), column_counts: Arc::new(RwLock::new(HashMap::new())), + function_counts: Arc::new(RwLock::new(HashMap::new())), recent_queries: Arc::new(RwLock::new(Vec::new())), max_recent, } @@ -71,6 +74,10 @@ impl UsageTracker { let counts = self.column_counts.read().unwrap(); counts.get(name).copied().unwrap_or(0) } + ItemType::Function => { + let counts = self.function_counts.read().unwrap(); + counts.get(name).copied().unwrap_or(0) + } } }