Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 24aee5a

Browse files
feat(hover): hover on columns
1 parent 339ceb3 commit 24aee5a

File tree

6 files changed

+357
-53
lines changed

6 files changed

+357
-53
lines changed
Lines changed: 245 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
use pgt_schema_cache::{Column, Table};
2+
use pgt_treesitter::context::TreesitterContext;
3+
4+
pub(crate) trait ContextualPriority {
5+
fn relevance_score(&self, ctx: &TreesitterContext) -> f32;
6+
}
7+
8+
impl ContextualPriority for Column {
9+
fn relevance_score(&self, ctx: &TreesitterContext) -> f32 {
10+
let mut score = 0.0;
11+
12+
// high score if we match the specific alias or table being referenced in the cursor context
13+
if let Some(table_or_alias) = ctx.schema_or_alias_name.as_ref() {
14+
if table_or_alias == self.table_name.as_str() {
15+
score += 250.0;
16+
} else if let Some(table_name) = ctx.mentioned_table_aliases.get(table_or_alias) {
17+
if table_name == self.table_name.as_str() {
18+
score += 250.0;
19+
}
20+
}
21+
}
22+
23+
// medium score if the current column maps to any of the query's mentioned
24+
// "(schema.)table" combinations
25+
for (schema_opt, tables) in &ctx.mentioned_relations {
26+
if tables.contains(&self.table_name) {
27+
if schema_opt.as_deref() == Some(&self.schema_name) {
28+
score += 150.0;
29+
} else {
30+
score += 100.0;
31+
}
32+
}
33+
}
34+
35+
if self.schema_name == "public" && score == 0.0 {
36+
score += 10.0;
37+
}
38+
39+
if self.is_primary_key && score == 0.0 {
40+
score += 5.0;
41+
}
42+
43+
score
44+
}
45+
}
46+
47+
impl ContextualPriority for Table {
48+
fn relevance_score(&self, ctx: &TreesitterContext) -> f32 {
49+
let mut score = 0.0;
50+
51+
// Highest priority: table explicitly mentioned in the query
52+
for (schema_opt, tables) in &ctx.mentioned_relations {
53+
if tables.contains(&self.name) {
54+
// Extra points if schema also matches
55+
if schema_opt.as_deref() == Some(&self.schema) {
56+
score += 200.0;
57+
} else {
58+
score += 150.0;
59+
}
60+
}
61+
}
62+
63+
// Check if table is mentioned via alias
64+
if ctx
65+
.mentioned_table_aliases
66+
.values()
67+
.any(|table_name| *table_name == self.name)
68+
{
69+
score += 140.0;
70+
}
71+
72+
// Medium priority: same schema as mentioned tables
73+
if ctx
74+
.mentioned_relations
75+
.keys()
76+
.any(|schema| schema.as_deref() == Some(&self.schema))
77+
{
78+
score += 50.0;
79+
}
80+
81+
// Lower priority: public schema
82+
if self.schema == "public" {
83+
score += 10.0;
84+
}
85+
86+
score
87+
}
88+
}
89+
90+
/// Will first sort the items by a score and then filter out items with a score gap algorithm.
91+
///
92+
/// `[200, 180, 150, 140]` => all items are returned
93+
///
94+
/// `[200, 180, 15, 10]` => first two items are returned
95+
///
96+
/// `[200, 30, 20, 10]` => only first item is returned
97+
pub(crate) fn prioritize_by_context<T: ContextualPriority + std::fmt::Debug>(
98+
items: Vec<T>,
99+
ctx: &TreesitterContext,
100+
) -> Vec<T> {
101+
let mut scored: Vec<_> = items
102+
.into_iter()
103+
.map(|item| {
104+
let score = item.relevance_score(ctx);
105+
(item, score)
106+
})
107+
.collect();
108+
109+
scored.sort_by(|(_, score_a), (_, score_b)| score_b.partial_cmp(score_a).unwrap());
110+
111+
if scored.is_empty() {
112+
return vec![];
113+
}
114+
115+
// always include the top result
116+
let top_result = scored.remove(0);
117+
let mut results = vec![top_result.0];
118+
let mut prev_score = top_result.1;
119+
120+
// include additional results until we hit a significant score gap of 30%
121+
for (item, score) in scored.into_iter() {
122+
let gap = prev_score - score;
123+
if gap > prev_score * 0.3 {
124+
break;
125+
}
126+
results.push(item);
127+
prev_score = score;
128+
}
129+
130+
results
131+
}
132+
133+
#[cfg(test)]
134+
mod tests {
135+
use pgt_test_utils::QueryWithCursorPosition;
136+
use pgt_text_size::TextSize;
137+
138+
use super::*;
139+
140+
fn create_test_column(schema: &str, table: &str, name: &str, is_pk: bool) -> Column {
141+
Column {
142+
name: name.to_string(),
143+
table_name: table.to_string(),
144+
table_oid: 1,
145+
class_kind: pgt_schema_cache::ColumnClassKind::OrdinaryTable,
146+
schema_name: schema.to_string(),
147+
type_id: 23,
148+
type_name: Some("integer".to_string()),
149+
is_nullable: false,
150+
is_primary_key: is_pk,
151+
is_unique: is_pk,
152+
default_expr: None,
153+
varchar_length: None,
154+
comment: None,
155+
}
156+
}
157+
158+
fn create_test_context(query: QueryWithCursorPosition) -> TreesitterContext<'static> {
159+
use pgt_treesitter::TreeSitterContextParams;
160+
161+
let (pos, sql) = query.get_text_and_position();
162+
163+
let mut parser = tree_sitter::Parser::new();
164+
parser.set_language(tree_sitter_sql::language()).unwrap();
165+
let tree = parser.parse(sql.clone(), None).unwrap();
166+
167+
// Leak some stuff so test setup is easier
168+
let leaked_tree: &'static tree_sitter::Tree = Box::leak(Box::new(tree));
169+
let leaked_sql: &'static String = Box::leak(Box::new(sql));
170+
171+
let position = TextSize::new(pos.try_into().unwrap());
172+
173+
let ctx = pgt_treesitter::context::TreesitterContext::new(TreeSitterContextParams {
174+
position,
175+
text: leaked_sql,
176+
tree: leaked_tree,
177+
});
178+
179+
ctx
180+
}
181+
182+
#[test]
183+
fn column_scoring_prioritizes_mentioned_tables() {
184+
let query = format!(
185+
"select id{} from auth.users",
186+
QueryWithCursorPosition::cursor_marker()
187+
);
188+
189+
let ctx = create_test_context(query.into());
190+
191+
let auth_users_id = create_test_column("auth", "users", "id", true);
192+
let public_posts_id = create_test_column("public", "posts", "id", false);
193+
let public_users_id = create_test_column("public", "users", "id", false);
194+
195+
let columns = vec![auth_users_id, public_posts_id, public_users_id];
196+
let result = prioritize_by_context(columns, &ctx);
197+
198+
assert_eq!(result.len(), 1);
199+
assert_eq!(result[0].table_name, "users");
200+
assert_eq!(result[0].schema_name, "auth");
201+
assert_eq!(result[0].is_primary_key, true);
202+
}
203+
204+
#[test]
205+
fn column_scoring_prioritizes_mentioned_table_names() {
206+
let query = format!(
207+
"select users.id{} from ",
208+
QueryWithCursorPosition::cursor_marker()
209+
);
210+
211+
let ctx = create_test_context(query.into());
212+
213+
let videos_id = create_test_column("public", "videos", "id", false);
214+
let posts_id = create_test_column("public", "posts", "id", false);
215+
let users_id = create_test_column("public", "users", "id", false);
216+
217+
let columns = vec![videos_id, posts_id, users_id];
218+
let result = prioritize_by_context(columns, &ctx);
219+
220+
assert_eq!(result.len(), 1);
221+
assert_eq!(result[0].table_name, "users");
222+
}
223+
224+
#[test]
225+
fn column_scoring_prioritizes_mentioned_tables_with_aliases() {
226+
let query = format!(
227+
"select p.id{} as post_id from auth.users u join public.posts p on u.id = p.user_id;",
228+
QueryWithCursorPosition::cursor_marker()
229+
);
230+
231+
let ctx = create_test_context(query.into());
232+
233+
let auth_users_id = create_test_column("auth", "users", "id", true);
234+
let public_posts_id = create_test_column("public", "posts", "id", false);
235+
let public_users_id = create_test_column("public", "users", "id", false);
236+
237+
let columns = vec![auth_users_id, public_posts_id, public_users_id];
238+
let result = prioritize_by_context(columns, &ctx);
239+
240+
assert_eq!(result.len(), 1);
241+
assert_eq!(result[0].table_name, "posts");
242+
assert_eq!(result[0].schema_name, "public");
243+
assert_eq!(result[0].is_primary_key, false);
244+
}
245+
}

‎crates/pgt_hover/src/hovered_item.rs

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
use crate::{contextual_priority::ContextualPriority, to_markdown::ToHoverMarkdown};
2+
3+
/// Mapper type that will be used for filtering and turning data to markdown.
4+
#[derive(Debug)]
5+
pub(crate) enum HoverItem<'a> {
6+
Table(&'a pgt_schema_cache::Table),
7+
Column(&'a pgt_schema_cache::Column),
8+
}
9+
10+
impl<'a> From<&'a pgt_schema_cache::Table> for HoverItem<'a> {
11+
fn from(value: &'a pgt_schema_cache::Table) -> Self {
12+
HoverItem::Table(value)
13+
}
14+
}
15+
16+
impl<'a> From<&'a pgt_schema_cache::Column> for HoverItem<'a> {
17+
fn from(value: &'a pgt_schema_cache::Column) -> Self {
18+
HoverItem::Column(value)
19+
}
20+
}
21+
22+
impl<'a> ContextualPriority for HoverItem<'a> {
23+
fn relevance_score(&self, ctx: &pgt_treesitter::TreesitterContext) -> f32 {
24+
match self {
25+
HoverItem::Table(table) => table.relevance_score(ctx),
26+
HoverItem::Column(column) => column.relevance_score(ctx),
27+
}
28+
}
29+
}
30+
31+
impl<'a> ToHoverMarkdown for HoverItem<'a> {
32+
fn hover_headline<W: std::fmt::Write>(&self, writer: &mut W) -> Result<(), std::fmt::Error> {
33+
match self {
34+
HoverItem::Table(table) => ToHoverMarkdown::hover_headline(*table, writer),
35+
HoverItem::Column(column) => ToHoverMarkdown::hover_headline(*column, writer),
36+
}
37+
}
38+
39+
fn hover_body<W: std::fmt::Write>(&self, writer: &mut W) -> Result<bool, std::fmt::Error> {
40+
match self {
41+
HoverItem::Table(table) => ToHoverMarkdown::hover_body(*table, writer),
42+
HoverItem::Column(column) => ToHoverMarkdown::hover_body(*column, writer),
43+
}
44+
}
45+
46+
fn hover_footer<W: std::fmt::Write>(&self, writer: &mut W) -> Result<bool, std::fmt::Error> {
47+
match self {
48+
HoverItem::Table(table) => ToHoverMarkdown::hover_footer(*table, writer),
49+
HoverItem::Column(column) => ToHoverMarkdown::hover_footer(*column, writer),
50+
}
51+
}
52+
}

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /