Skip to content

Commit d34f31e

Browse files
juleswritescodejuleswritescode
authored
refactor: replace schema_or_alias
1 parent 65cc053 commit d34f31e

File tree

6 files changed

+123
-64
lines changed

6 files changed

+123
-64
lines changed

‎crates/pgls_completions/src/providers/helper.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ pub(crate) fn with_schema_or_alias(
4545
item_name: &str,
4646
schema_or_alias_name: Option<&str>,
4747
) -> String {
48-
let is_already_prefixed_with_schema_name = ctx.schema_or_alias_name.is_some();
48+
let is_already_prefixed_with_schema_name = ctx.has_any_qualifier();
4949

5050
let with_quotes = node_text_surrounded_by_quotes(ctx);
5151
let single_leading_quote = only_leading_quote(ctx);

‎crates/pgls_completions/src/relevance/filtering.rs‎

Lines changed: 18 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ impl CompletionFilter<'_> {
225225
(ctx.matches_ancestor_history(&[
226226
"grantable_on_table",
227227
"object_reference",
228-
]) && ctx.schema_or_alias_name.is_none())
228+
]) && !ctx.has_any_qualifier())
229229
|| ctx.matches_ancestor_history(&["grantable_on_all"])
230230
}
231231

@@ -308,18 +308,25 @@ impl CompletionFilter<'_> {
308308
}
309309

310310
fn check_mentioned_schema_or_alias(&self, ctx: &TreesitterContext) -> Option<()> {
311-
if ctx.schema_or_alias_name.is_none() {
312-
return Some(());
313-
}
314-
315-
let schema_or_alias = ctx.schema_or_alias_name.as_ref().unwrap().replace('"', "");
311+
let tail_qualifier = match ctx.tail_qualifier_sanitized() {
312+
Some(q) => q,
313+
None => return Some(()), // no qualifier = this check passes
314+
};
316315

317316
let matches = match self.data {
318-
CompletionRelevanceData::Table(table) => table.schema == schema_or_alias,
319-
CompletionRelevanceData::Function(f) => f.schema == schema_or_alias,
320-
CompletionRelevanceData::Column(col) => ctx
321-
.get_mentioned_table_for_alias(&schema_or_alias)
322-
.is_some_and(|t| t == &col.table_name),
317+
CompletionRelevanceData::Table(table) => table.schema == tail_qualifier,
318+
CompletionRelevanceData::Function(f) => f.schema == tail_qualifier,
319+
CompletionRelevanceData::Column(col) => {
320+
let table = ctx
321+
.get_mentioned_table_for_alias(&tail_qualifier)
322+
.unwrap_or(&tail_qualifier);
323+
324+
if let Some(schema) = ctx.head_qualifier_sanitized() {
325+
col.schema_name == schema.as_str() && col.table_name == table.as_str()
326+
} else {
327+
col.table_name == table.as_str()
328+
}
329+
}
323330

324331
// we should never allow schema suggestions if there already was one.
325332
CompletionRelevanceData::Schema(_) => false,

‎crates/pgls_completions/src/relevance/scoring.rs‎

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ impl CompletionScore<'_> {
7979
};
8080

8181
let has_mentioned_tables = ctx.has_any_mentioned_relations();
82-
let has_mentioned_schema = ctx.schema_or_alias_name.is_some();
82+
let has_qualifier = ctx.has_any_qualifier();
8383

8484
self.score += match self.data {
8585
CompletionRelevanceData::Table(_) => match clause_type {
@@ -122,13 +122,13 @@ impl CompletionScore<'_> {
122122
_ => -15,
123123
},
124124
CompletionRelevanceData::Schema(_) => match clause_type {
125-
WrappingClause::From if !has_mentioned_schema => 15,
126-
WrappingClause::Join { .. } if !has_mentioned_schema => 15,
127-
WrappingClause::Update if !has_mentioned_schema => 15,
128-
WrappingClause::Delete if !has_mentioned_schema => 15,
129-
WrappingClause::AlterPolicy if !has_mentioned_schema => 15,
130-
WrappingClause::DropPolicy if !has_mentioned_schema => 15,
131-
WrappingClause::CreatePolicy if !has_mentioned_schema => 15,
125+
WrappingClause::From if !has_qualifier => 15,
126+
WrappingClause::Join { .. } if !has_qualifier => 15,
127+
WrappingClause::Update if !has_qualifier => 15,
128+
WrappingClause::Delete if !has_qualifier => 15,
129+
WrappingClause::AlterPolicy if !has_qualifier => 15,
130+
WrappingClause::DropPolicy if !has_qualifier => 15,
131+
WrappingClause::CreatePolicy if !has_qualifier => 15,
132132
_ => -50,
133133
},
134134
CompletionRelevanceData::Policy(_) => match clause_type {
@@ -149,15 +149,15 @@ impl CompletionScore<'_> {
149149
Some(wn) => wn,
150150
};
151151

152-
let has_mentioned_schema = ctx.schema_or_alias_name.is_some();
152+
let has_qualifier = ctx.has_any_qualifier();
153153
let has_node_text = ctx
154154
.get_node_under_cursor_content()
155155
.is_some_and(|txt| !sanitization::is_sanitized_token(txt.as_str()));
156156

157157
self.score += match self.data {
158158
CompletionRelevanceData::Table(_) => match wrapping_node {
159-
WrappingNode::Relation if has_mentioned_schema => 15,
160-
WrappingNode::Relation if !has_mentioned_schema => 10,
159+
WrappingNode::Relation if has_qualifier => 15,
160+
WrappingNode::Relation if !has_qualifier => 10,
161161
WrappingNode::BinaryExpression => 5,
162162
_ => -50,
163163
},
@@ -172,8 +172,8 @@ impl CompletionScore<'_> {
172172
_ => -15,
173173
},
174174
CompletionRelevanceData::Schema(_) => match wrapping_node {
175-
WrappingNode::Relation if !has_mentioned_schema && !has_node_text => 15,
176-
WrappingNode::Relation if !has_mentioned_schema && has_node_text => 0,
175+
WrappingNode::Relation if !has_qualifier && !has_node_text => 15,
176+
WrappingNode::Relation if !has_qualifier && has_node_text => 0,
177177
_ => -50,
178178
},
179179
CompletionRelevanceData::Policy(_) => 0,
@@ -191,17 +191,30 @@ impl CompletionScore<'_> {
191191
}
192192

193193
fn check_matches_schema(&mut self, ctx: &TreesitterContext) {
194-
let schema_name = match ctx.schema_or_alias_name.as_ref() {
195-
None => return,
196-
Some(n) => n.replace('"', ""),
194+
let schema_from_qualifier = match self.data {
195+
CompletionRelevanceData::Table(_) | CompletionRelevanceData::Function(_) => {
196+
ctx.tail_qualifier_sanitized()
197+
}
198+
199+
CompletionRelevanceData::Column(_) | CompletionRelevanceData::Policy(_) => {
200+
ctx.head_qualifier_sanitized()
201+
}
202+
203+
CompletionRelevanceData::Schema(_) | CompletionRelevanceData::Role(_) => None,
197204
};
198205

206+
if schema_from_qualifier.is_none() {
207+
return;
208+
}
209+
210+
let schema_from_qualifier = schema_from_qualifier.unwrap();
211+
199212
let data_schema = match self.get_schema_name() {
200213
Some(s) => s,
201214
None => return,
202215
};
203216

204-
if schema_name == data_schema {
217+
if schema_from_qualifier == data_schema {
205218
self.score += 25;
206219
} else {
207220
self.score -= 10;

‎crates/pgls_hover/src/hoverables/column.rs‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -74,12 +74,19 @@ impl ContextualPriority for Column {
7474
let mut score = 0.0;
7575

7676
// high score if we match the specific alias or table being referenced in the cursor context
77-
if let Some(table_or_alias) = ctx.schema_or_alias_name.as_ref() {
78-
if table_or_alias.replace('"', "") == self.table_name.as_str() {
77+
78+
if let Some(table_or_alias) = ctx.tail_qualifier_sanitized() {
79+
let table = ctx
80+
.get_mentioned_table_for_alias(&table_or_alias)
81+
.unwrap_or(&table_or_alias);
82+
83+
if table == self.table_name.as_str() {
7984
score += 250.0;
80-
} else if let Some(table_name) = ctx.get_mentioned_table_for_alias(table_or_alias) {
81-
if table_name == self.table_name.as_str() {
82-
score += 250.0;
85+
}
86+
87+
if let Some(schema) = ctx.head_qualifier_sanitized() {
88+
if schema == self.schema_name.as_str() {
89+
score += 50.0;
8390
}
8491
}
8592
}

‎crates/pgls_hover/src/hovered_node.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ impl HoveredNode {
3535
}
3636

3737
Some(HoveredNode::Table((
38-
ctx.schema_or_alias_name.clone(),
38+
ctx.identifier_qualifiers.1.clone(),
3939
node_content,
4040
)))
4141
}
@@ -52,22 +52,22 @@ impl HoveredNode {
5252
}) =>
5353
{
5454
Some(HoveredNode::Table((
55-
ctx.schema_or_alias_name.clone(),
55+
ctx.identifier_qualifiers.1.clone(),
5656
node_content,
5757
)))
5858
}
5959

6060
"column_identifier" => Some(HoveredNode::Column((
6161
None,
62-
ctx.schema_or_alias_name.clone(),
62+
ctx.identifier_qualifiers.1.clone(),
6363
node_content,
6464
))),
6565

6666
"any_identifier"
6767
if ctx.matches_ancestor_history(&["invocation", "object_reference"]) =>
6868
{
6969
Some(HoveredNode::Function((
70-
ctx.schema_or_alias_name.clone(),
70+
ctx.identifier_qualifiers.1.clone(),
7171
node_content,
7272
)))
7373
}
@@ -103,7 +103,7 @@ impl HoveredNode {
103103
{
104104
let sanitized = node_content.replace(['(', ')'], "");
105105
Some(HoveredNode::PostgresType((
106-
ctx.schema_or_alias_name.clone(),
106+
ctx.identifier_qualifiers.1.clone(),
107107
sanitized,
108108
)))
109109
}
@@ -114,7 +114,7 @@ impl HoveredNode {
114114
}
115115

116116
"grant_table" => Some(HoveredNode::Table((
117-
ctx.schema_or_alias_name.clone(),
117+
ctx.identifier_qualifiers.1.clone(),
118118
node_content,
119119
))),
120120

‎crates/pgls_treesitter/src/context/mod.rs‎

Lines changed: 56 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -101,25 +101,16 @@ pub struct TreesitterContext<'a> {
101101
pub text: &'a str,
102102
pub position: usize,
103103

104-
/// If the cursor is on a node that uses dot notation
105-
/// to specify an alias or schema, this will hold the schema's or
106-
/// alias's name.
104+
/// Tuple containing up to two qualifiers for identifier-node under the cursor: (head, tail)
107105
///
108-
/// Here, `auth` is a schema name:
109-
/// ```sql
110-
/// select * from auth.users;
111-
/// ```
112-
///
113-
/// Here, `u` is an alias name:
114-
/// ```sql
115-
/// select
116-
/// *
117-
/// from
118-
/// auth.users u
119-
/// left join identities i
120-
/// on u.id = i.user_id;
121-
/// ```
122-
pub schema_or_alias_name: Option<String>,
106+
/// The qualifiers represent different "parents" based on the context, for example:
107+
/// - `column` -> (None, None)
108+
/// - `table.column` -> (None, Some("table"))
109+
/// - `alias.column` -> (None, Some("alias"))
110+
/// - `schema.table` -> (None, Some("schema"))
111+
/// - `schema.table.column` -> (Some("schema"), Some("table"))
112+
/// - `table` -> (None, None)
113+
pub identifier_qualifiers: (Option<String>, Option<String>),
123114

124115
pub wrapping_clause_type: Option<WrappingClause<'a>>,
125116

@@ -140,7 +131,7 @@ impl<'a> TreesitterContext<'a> {
140131
text: params.text,
141132
position: usize::from(params.position),
142133
node_under_cursor: None,
143-
schema_or_alias_name: None,
134+
identifier_qualifiers: (None, None),
144135
wrapping_clause_type: None,
145136
wrapping_node_kind: None,
146137
wrapping_statement_range: None,
@@ -333,10 +324,20 @@ impl<'a> TreesitterContext<'a> {
333324
let content = self.get_ts_node_content(&current_node);
334325
if let Some(txt) = content {
335326
let parts: Vec<&str> = txt.split('.').collect();
336-
// we do not want to set it if we're on the schema or alias node itself
337-
let is_on_schema_node = start + parts[0].len() >= self.position;
338-
if parts.len() == 2 && !is_on_schema_node {
339-
self.schema_or_alias_name = Some(parts[0].to_string());
327+
// we do not want to set it if we're on the first qualifier itself
328+
let is_on_first_part = start + parts[0].len() >= self.position;
329+
330+
if parts.len() == 2 && !is_on_first_part {
331+
self.identifier_qualifiers = (None, Some(parts[0].to_string()));
332+
} else if parts.len() == 3 && !is_on_first_part {
333+
let is_on_second_part =
334+
start + parts[0].len() + 1 + parts[1].len() >= self.position;
335+
if !is_on_second_part {
336+
self.identifier_qualifiers =
337+
(Some(parts[0].to_string()), Some(parts[1].to_string()));
338+
} else {
339+
self.identifier_qualifiers = (None, Some(parts[0].to_string()));
340+
}
340341
}
341342
}
342343
}
@@ -778,6 +779,37 @@ impl<'a> TreesitterContext<'a> {
778779
pub fn has_mentioned_columns(&self) -> bool {
779780
!self.mentioned_columns.is_empty()
780781
}
782+
783+
/// Returns the head qualifier (leftmost), sanitized (quotes removed)
784+
/// For `schema.table.<column>`: returns `Some("schema")`
785+
/// For `table.<column>`: returns `None`
786+
pub fn head_qualifier_sanitized(&self) -> Option<String> {
787+
self.identifier_qualifiers
788+
.0
789+
.as_ref()
790+
.map(|s| s.replace('"', ""))
791+
}
792+
793+
/// Returns the tail qualifier (rightmost), sanitized (quotes removed)
794+
/// For `schema.table.<column>`: returns `Some("table")`
795+
/// For `table.<column>`: returns `Some("table")`
796+
pub fn tail_qualifier_sanitized(&self) -> Option<String> {
797+
self.identifier_qualifiers
798+
.1
799+
.as_ref()
800+
.map(|s| s.replace('"', ""))
801+
}
802+
803+
/// Returns true if there is at least one qualifier present
804+
pub fn has_any_qualifier(&self) -> bool {
805+
match self.identifier_qualifiers {
806+
(Some(_), Some(_)) => true,
807+
(None, Some(_)) => true,
808+
(None, None) => false,
809+
810+
(Some(_), None) => unreachable!(),
811+
}
812+
}
781813
}
782814

783815
#[cfg(test)]
@@ -926,7 +958,7 @@ mod tests {
926958
let ctx = TreesitterContext::new(params);
927959

928960
assert_eq!(
929-
ctx.schema_or_alias_name,
961+
ctx.identifier_qualifiers.1,
930962
expected_schema.map(|f| f.to_string())
931963
);
932964
}

0 commit comments

Comments
 (0)