diff --git a/src/args.rs b/src/args.rs index 04bd2ff..006ca1b 100644 --- a/src/args.rs +++ b/src/args.rs @@ -30,6 +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 { @@ -118,6 +125,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_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/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..431110e --- /dev/null +++ b/src/completion/mod.rs @@ -0,0 +1,283 @@ +pub mod context_analyzer; +pub mod context_detector; +pub mod priority_scorer; +pub mod schema_cache; +pub mod usage_tracker; + +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, 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) { + 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]; + + // 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 suggestions (both short and qualified names) + for (schema, table) in table_metadata { + let short_name = table.clone(); + let qualified_name = format!("{}.{}", schema, table); + + // 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), + }); + } + } + } + + // 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) + 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 = deduplicated + .into_iter() + .map(|s| Pair { + display: s.name.clone(), + replacement: s.name, + }) + .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 usage_tracker = Arc::new(UsageTracker::new(10)); + let completer = SqlCompleter::new(cache, usage_tracker, 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 usage_tracker = Arc::new(UsageTracker::new(10)); + let completer = SqlCompleter::new(cache, usage_tracker, 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/priority_scorer.rs b/src/completion/priority_scorer.rs new file mode 100644 index 0000000..a3f26f0 --- /dev/null +++ b/src/completion/priority_scorer.rs @@ -0,0 +1,272 @@ +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, + + /// Functions - lower priority than columns + Function = 1000, + + /// 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 + } + } + + ItemType::Function => PriorityClass::Function, + } + } +} + +#[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 new file mode 100644 index 0000000..e4ed274 --- /dev/null +++ b/src/completion/schema_cache.rs @@ -0,0 +1,690 @@ +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() + } + + /// 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 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(); + 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, + 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 (including system schemas - they'll be deprioritized by the scorer) + let tables_query = "SELECT table_schema, table_name \ + FROM information_schema.tables \ + ORDER BY table_schema, table_name"; + + let tables_result = query_silent(context, tables_query).await; + + // 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 \ + ORDER BY table_schema, table_name, ordinal_position"; + + let columns_result = query_silent(context, columns_query).await; + + // Query functions (including system functions, excluding operators) + let functions_query = "SELECT routine_name \ + FROM information_schema.routines \ + WHERE routine_type != 'OPERATOR' \ + 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); + } + } + + 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()); + } + } + } + } + + // 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) + } + } +} + +#[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/completion/usage_tracker.rs b/src/completion/usage_tracker.rs new file mode 100644 index 0000000..786d895 --- /dev/null +++ b/src/completion/usage_tracker.rs @@ -0,0 +1,311 @@ +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ItemType { + Table, + Column, + Function, +} + +/// 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>>, + 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())), + function_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) + } + ItemType::Function => { + let counts = self.function_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 6625d7d..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 { @@ -10,6 +12,7 @@ pub struct ServiceAccountToken { pub until: u64, } +#[derive(Clone)] pub struct Context { pub args: Args, pub url: String, @@ -20,6 +23,7 @@ pub struct Context { pub last_result: Option, pub last_stats: Option, pub is_interactive: bool, + pub usage_tracker: Option>, } impl Context { @@ -35,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 e89d802..0154353 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,15 @@ 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; 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,10 +68,19 @@ async fn main() -> Result<(), Box> { SqlHighlighter::new(false).unwrap() // Fallback to disabled }); - let helper = ReplHelper::new(highlighter); + // Initialize schema cache and usage tracker for completion + let schema_cache = Arc::new(SchemaCache::new(context.args.completion_cache_ttl)); + 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()?; 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() { @@ -87,6 +101,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."); } @@ -131,16 +157,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:'):"); @@ -196,6 +244,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..3bae7bc 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 @@ -175,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))) @@ -408,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(()) } } 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)