diff --git a/crates/pgls_completions/src/providers/helper.rs b/crates/pgls_completions/src/providers/helper.rs index 834c41a36..0da2d9298 100644 --- a/crates/pgls_completions/src/providers/helper.rs +++ b/crates/pgls_completions/src/providers/helper.rs @@ -45,7 +45,7 @@ pub(crate) fn with_schema_or_alias( item_name: &str, schema_or_alias_name: Option<&str>, ) -> String { - let is_already_prefixed_with_schema_name = ctx.schema_or_alias_name.is_some(); + let is_already_prefixed_with_schema_name = ctx.has_any_qualifier(); let with_quotes = node_text_surrounded_by_quotes(ctx); let single_leading_quote = only_leading_quote(ctx); diff --git a/crates/pgls_completions/src/relevance/filtering.rs b/crates/pgls_completions/src/relevance/filtering.rs index 899b7fd38..072214080 100644 --- a/crates/pgls_completions/src/relevance/filtering.rs +++ b/crates/pgls_completions/src/relevance/filtering.rs @@ -225,7 +225,7 @@ impl CompletionFilter<'_> { (ctx.matches_ancestor_history(&[ "grantable_on_table", "object_reference", - ]) && ctx.schema_or_alias_name.is_none()) + ]) && !ctx.has_any_qualifier()) || ctx.matches_ancestor_history(&["grantable_on_all"]) } @@ -308,18 +308,25 @@ impl CompletionFilter<'_> { } fn check_mentioned_schema_or_alias(&self, ctx: &TreesitterContext) -> Option<()> { - if ctx.schema_or_alias_name.is_none() { - return Some(()); - } - - let schema_or_alias = ctx.schema_or_alias_name.as_ref().unwrap().replace('"', ""); + let tail_qualifier = match ctx.tail_qualifier_sanitized() { + Some(q) => q, + None => return Some(()), // no qualifier = this check passes + }; let matches = match self.data { - CompletionRelevanceData::Table(table) => table.schema == schema_or_alias, - CompletionRelevanceData::Function(f) => f.schema == schema_or_alias, - CompletionRelevanceData::Column(col) => ctx - .get_mentioned_table_for_alias(&schema_or_alias) - .is_some_and(|t| t == &col.table_name), + CompletionRelevanceData::Table(table) => table.schema == tail_qualifier, + CompletionRelevanceData::Function(f) => f.schema == tail_qualifier, + CompletionRelevanceData::Column(col) => { + let table = ctx + .get_mentioned_table_for_alias(&tail_qualifier) + .unwrap_or(&tail_qualifier); + + if let Some(schema) = ctx.head_qualifier_sanitized() { + col.schema_name == schema.as_str() && col.table_name == table.as_str() + } else { + col.table_name == table.as_str() + } + } // we should never allow schema suggestions if there already was one. CompletionRelevanceData::Schema(_) => false, diff --git a/crates/pgls_completions/src/relevance/scoring.rs b/crates/pgls_completions/src/relevance/scoring.rs index d41ade208..46f761fbb 100644 --- a/crates/pgls_completions/src/relevance/scoring.rs +++ b/crates/pgls_completions/src/relevance/scoring.rs @@ -79,7 +79,7 @@ impl CompletionScore<'_> { }; let has_mentioned_tables = ctx.has_any_mentioned_relations(); - let has_mentioned_schema = ctx.schema_or_alias_name.is_some(); + let has_qualifier = ctx.has_any_qualifier(); self.score += match self.data { CompletionRelevanceData::Table(_) => match clause_type { @@ -122,13 +122,13 @@ impl CompletionScore<'_> { _ => -15, }, CompletionRelevanceData::Schema(_) => match clause_type { - WrappingClause::From if !has_mentioned_schema => 15, - WrappingClause::Join { .. } if !has_mentioned_schema => 15, - WrappingClause::Update if !has_mentioned_schema => 15, - WrappingClause::Delete if !has_mentioned_schema => 15, - WrappingClause::AlterPolicy if !has_mentioned_schema => 15, - WrappingClause::DropPolicy if !has_mentioned_schema => 15, - WrappingClause::CreatePolicy if !has_mentioned_schema => 15, + WrappingClause::From if !has_qualifier => 15, + WrappingClause::Join { .. } if !has_qualifier => 15, + WrappingClause::Update if !has_qualifier => 15, + WrappingClause::Delete if !has_qualifier => 15, + WrappingClause::AlterPolicy if !has_qualifier => 15, + WrappingClause::DropPolicy if !has_qualifier => 15, + WrappingClause::CreatePolicy if !has_qualifier => 15, _ => -50, }, CompletionRelevanceData::Policy(_) => match clause_type { @@ -149,15 +149,15 @@ impl CompletionScore<'_> { Some(wn) => wn, }; - let has_mentioned_schema = ctx.schema_or_alias_name.is_some(); + let has_qualifier = ctx.has_any_qualifier(); let has_node_text = ctx .get_node_under_cursor_content() .is_some_and(|txt| !sanitization::is_sanitized_token(txt.as_str())); self.score += match self.data { CompletionRelevanceData::Table(_) => match wrapping_node { - WrappingNode::Relation if has_mentioned_schema => 15, - WrappingNode::Relation if !has_mentioned_schema => 10, + WrappingNode::Relation if has_qualifier => 15, + WrappingNode::Relation if !has_qualifier => 10, WrappingNode::BinaryExpression => 5, _ => -50, }, @@ -172,8 +172,8 @@ impl CompletionScore<'_> { _ => -15, }, CompletionRelevanceData::Schema(_) => match wrapping_node { - WrappingNode::Relation if !has_mentioned_schema && !has_node_text => 15, - WrappingNode::Relation if !has_mentioned_schema && has_node_text => 0, + WrappingNode::Relation if !has_qualifier && !has_node_text => 15, + WrappingNode::Relation if !has_qualifier && has_node_text => 0, _ => -50, }, CompletionRelevanceData::Policy(_) => 0, @@ -191,17 +191,30 @@ impl CompletionScore<'_> { } fn check_matches_schema(&mut self, ctx: &TreesitterContext) { - let schema_name = match ctx.schema_or_alias_name.as_ref() { - None => return, - Some(n) => n.replace('"', ""), + let schema_from_qualifier = match self.data { + CompletionRelevanceData::Table(_) | CompletionRelevanceData::Function(_) => { + ctx.tail_qualifier_sanitized() + } + + CompletionRelevanceData::Column(_) | CompletionRelevanceData::Policy(_) => { + ctx.head_qualifier_sanitized() + } + + CompletionRelevanceData::Schema(_) | CompletionRelevanceData::Role(_) => None, }; + if schema_from_qualifier.is_none() { + return; + } + + let schema_from_qualifier = schema_from_qualifier.unwrap(); + let data_schema = match self.get_schema_name() { Some(s) => s, None => return, }; - if schema_name == data_schema { + if schema_from_qualifier == data_schema { self.score += 25; } else { self.score -= 10; diff --git a/crates/pgls_hover/src/hoverables/column.rs b/crates/pgls_hover/src/hoverables/column.rs index cecc75a51..f0a3ee0e6 100644 --- a/crates/pgls_hover/src/hoverables/column.rs +++ b/crates/pgls_hover/src/hoverables/column.rs @@ -74,12 +74,19 @@ impl ContextualPriority for Column { let mut score = 0.0; // high score if we match the specific alias or table being referenced in the cursor context - if let Some(table_or_alias) = ctx.schema_or_alias_name.as_ref() { - if table_or_alias.replace('"', "") == self.table_name.as_str() { + + if let Some(table_or_alias) = ctx.tail_qualifier_sanitized() { + let table = ctx + .get_mentioned_table_for_alias(&table_or_alias) + .unwrap_or(&table_or_alias); + + if table == self.table_name.as_str() { score += 250.0; - } else if let Some(table_name) = ctx.get_mentioned_table_for_alias(table_or_alias) { - if table_name == self.table_name.as_str() { - score += 250.0; + } + + if let Some(schema) = ctx.head_qualifier_sanitized() { + if schema == self.schema_name.as_str() { + score += 50.0; } } } diff --git a/crates/pgls_hover/src/hovered_node.rs b/crates/pgls_hover/src/hovered_node.rs index 878257919..cd48686b2 100644 --- a/crates/pgls_hover/src/hovered_node.rs +++ b/crates/pgls_hover/src/hovered_node.rs @@ -35,7 +35,7 @@ impl HoveredNode { } Some(HoveredNode::Table(( - ctx.schema_or_alias_name.clone(), + ctx.identifier_qualifiers.1.clone(), node_content, ))) } @@ -52,14 +52,14 @@ impl HoveredNode { }) => { Some(HoveredNode::Table(( - ctx.schema_or_alias_name.clone(), + ctx.identifier_qualifiers.1.clone(), node_content, ))) } "column_identifier" => Some(HoveredNode::Column(( None, - ctx.schema_or_alias_name.clone(), + ctx.identifier_qualifiers.1.clone(), node_content, ))), @@ -67,7 +67,7 @@ impl HoveredNode { if ctx.matches_ancestor_history(&["invocation", "object_reference"]) => { Some(HoveredNode::Function(( - ctx.schema_or_alias_name.clone(), + ctx.identifier_qualifiers.1.clone(), node_content, ))) } @@ -103,7 +103,7 @@ impl HoveredNode { { let sanitized = node_content.replace(['(', ')'], ""); Some(HoveredNode::PostgresType(( - ctx.schema_or_alias_name.clone(), + ctx.identifier_qualifiers.1.clone(), sanitized, ))) } @@ -114,7 +114,7 @@ impl HoveredNode { } "grant_table" => Some(HoveredNode::Table(( - ctx.schema_or_alias_name.clone(), + ctx.identifier_qualifiers.1.clone(), node_content, ))), diff --git a/crates/pgls_treesitter/src/context/mod.rs b/crates/pgls_treesitter/src/context/mod.rs index 5f47b4aaf..9788a6eb2 100644 --- a/crates/pgls_treesitter/src/context/mod.rs +++ b/crates/pgls_treesitter/src/context/mod.rs @@ -101,25 +101,16 @@ pub struct TreesitterContext<'a> { pub text: &'a str, pub position: usize, - /// If the cursor is on a node that uses dot notation - /// to specify an alias or schema, this will hold the schema's or - /// alias's name. + /// Tuple containing up to two qualifiers for identifier-node under the cursor: (head, tail) /// - /// Here, `auth` is a schema name: - /// ```sql - /// select * from auth.users; - /// ``` - /// - /// Here, `u` is an alias name: - /// ```sql - /// select - /// * - /// from - /// auth.users u - /// left join identities i - /// on u.id = i.user_id; - /// ``` - pub schema_or_alias_name: Option, + /// The qualifiers represent different "parents" based on the context, for example: + /// - `column` -> (None, None) + /// - `table.column` -> (None, Some("table")) + /// - `alias.column` -> (None, Some("alias")) + /// - `schema.table` -> (None, Some("schema")) + /// - `schema.table.column` -> (Some("schema"), Some("table")) + /// - `table` -> (None, None) + pub identifier_qualifiers: (Option, Option), pub wrapping_clause_type: Option>, @@ -140,7 +131,7 @@ impl<'a> TreesitterContext<'a> { text: params.text, position: usize::from(params.position), node_under_cursor: None, - schema_or_alias_name: None, + identifier_qualifiers: (None, None), wrapping_clause_type: None, wrapping_node_kind: None, wrapping_statement_range: None, @@ -333,10 +324,20 @@ impl<'a> TreesitterContext<'a> { let content = self.get_ts_node_content(¤t_node); if let Some(txt) = content { let parts: Vec<&str> = txt.split('.').collect(); - // we do not want to set it if we're on the schema or alias node itself - let is_on_schema_node = start + parts[0].len() >= self.position; - if parts.len() == 2 && !is_on_schema_node { - self.schema_or_alias_name = Some(parts[0].to_string()); + // we do not want to set it if we're on the first qualifier itself + let is_on_first_part = start + parts[0].len() >= self.position; + + if parts.len() == 2 && !is_on_first_part { + self.identifier_qualifiers = (None, Some(parts[0].to_string())); + } else if parts.len() == 3 && !is_on_first_part { + let is_on_second_part = + start + parts[0].len() + 1 + parts[1].len() >= self.position; + if !is_on_second_part { + self.identifier_qualifiers = + (Some(parts[0].to_string()), Some(parts[1].to_string())); + } else { + self.identifier_qualifiers = (None, Some(parts[0].to_string())); + } } } } @@ -778,6 +779,37 @@ impl<'a> TreesitterContext<'a> { pub fn has_mentioned_columns(&self) -> bool { !self.mentioned_columns.is_empty() } + + /// Returns the head qualifier (leftmost), sanitized (quotes removed) + /// For `schema.table.`: returns `Some("schema")` + /// For `table.`: returns `None` + pub fn head_qualifier_sanitized(&self) -> Option { + self.identifier_qualifiers + .0 + .as_ref() + .map(|s| s.replace('"', "")) + } + + /// Returns the tail qualifier (rightmost), sanitized (quotes removed) + /// For `schema.table.`: returns `Some("table")` + /// For `table.`: returns `Some("table")` + pub fn tail_qualifier_sanitized(&self) -> Option { + self.identifier_qualifiers + .1 + .as_ref() + .map(|s| s.replace('"', "")) + } + + /// Returns true if there is at least one qualifier present + pub fn has_any_qualifier(&self) -> bool { + match self.identifier_qualifiers { + (Some(_), Some(_)) => true, + (None, Some(_)) => true, + (None, None) => false, + + (Some(_), None) => unreachable!(), + } + } } #[cfg(test)] @@ -926,7 +958,7 @@ mod tests { let ctx = TreesitterContext::new(params); assert_eq!( - ctx.schema_or_alias_name, + ctx.identifier_qualifiers.1, expected_schema.map(|f| f.to_string()) ); }