Skip to content

Commit c8c6b1b

Browse files
juleswritescodejuleswritescode
committed
very very cool
1 parent e6452af commit c8c6b1b

File tree

6 files changed

+95
-97
lines changed

6 files changed

+95
-97
lines changed

‎crates/pgt_treesitter/src/queries/helper.rs‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
use tree_sitter::Node;
2+
3+
pub(crate) static OBJECT_REFERENCE_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
4+
static QUERY_STR: &str = r#"
5+
(object_reference
6+
object_reference_1of1: (any_identifier) @tail
7+
)
8+
(object_reference
9+
object_reference_1of2: (any_identifier) @head
10+
object_reference_2of2: (any_identifier) @tail
11+
)
12+
(object_reference
13+
object_reference_1of1: (any_identifier) @head
14+
object_reference_1of2: (any_identifier) @middle
15+
object_reference_1of3: (any_identifier) @tail
16+
)
17+
)
18+
"#;
19+
tree_sitter::Query::new(&pgt_treesitter_grammar::LANGUAGE.into(), QUERY_STR)
20+
.expect("Invalid TS Query")
21+
});
22+
23+
pub(crate) fn object_reference_query<'a>(
24+
node: Node<'a>,
25+
stmt: &'a str,
26+
) -> Option<(Option<Node<'a>>, Option<Node<'a>>, Node<'a>)> {
27+
let mut cursor = tree_sitter::QueryCursor::new();
28+
let matches = cursor.matches(&TS_QUERY, root_node, stmt.as_bytes());
29+
30+
assert!(
31+
matches.len() <= 1,
32+
"Please pass a single `object_reference` node into the `object_reference_query`!"
33+
);
34+
35+
if matches[0].len() == 0 {
36+
None
37+
} else if matches[0].captures.len() == 1 {
38+
Some((None, None, m.captures[0].node))
39+
} else if matches[0].captures.len() == 2 {
40+
Some((None, m.captures[0].node, m.captures[1].node))
41+
} else if matches[0].captures.len() == 3 {
42+
Some((
43+
Some(m.captures[0].node),
44+
Some(m.captures[1].node),
45+
m.captures[2].node,
46+
))
47+
} else {
48+
None
49+
}
50+
}

‎crates/pgt_treesitter/src/queries/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
mod helper;
12
mod insert_columns;
23
mod parameters;
34
mod relations;

‎crates/pgt_treesitter/src/queries/relations.rs‎

Lines changed: 9 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::sync::LazyLock;
22

3-
use crate::queries::{Query, QueryResult};
3+
use crate::queries::{Query, QueryResult, helper::object_reference_query};
44

55
use tree_sitter::StreamingIterator;
66

@@ -9,27 +9,15 @@ use super::QueryTryFrom;
99
static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
1010
static QUERY_STR: &str = r#"
1111
(relation
12-
(object_reference
13-
object_reference_first: (any_identifier) @first
14-
object_reference_second: (any_identifier)? @second
15-
object_reference_third: (any_identifier)? @third
16-
)
12+
(object_reference) @ref
1713
)
1814
(insert
19-
(object_reference
20-
object_reference_first: (any_identifier) @first
21-
object_reference_second: (any_identifier)? @second
22-
object_reference_third: (any_identifier)? @third
23-
)
15+
(object_reference) @ref
2416
)
2517
(alter_table
2618
(keyword_alter)
2719
(keyword_table)
28-
(object_reference
29-
object_reference_first: (any_identifier) @first
30-
object_reference_second: (any_identifier)? @second
31-
object_reference_third: (any_identifier)? @third
32-
)
20+
(object_reference) @ref
3321
)
3422
"#;
3523
tree_sitter::Query::new(&pgt_treesitter_grammar::LANGUAGE.into(), QUERY_STR)
@@ -87,23 +75,11 @@ impl<'a> Query<'a> for RelationMatch<'a> {
8775
let mut to_return = vec![];
8876

8977
matches.for_each(|m| {
90-
if m.captures.len() == 1 {
91-
let capture = m.captures[0].node;
92-
to_return.push(QueryResult::Relation(RelationMatch {
93-
schema: None,
94-
table: capture,
95-
}));
96-
}
97-
98-
if m.captures.len() == 2 {
99-
let schema = m.captures[0].node;
100-
let table = m.captures[1].node;
101-
102-
to_return.push(QueryResult::Relation(RelationMatch {
103-
schema: Some(schema),
104-
table,
105-
}));
106-
}
78+
m.captures.iter().for_each(|capture| {
79+
if let Some((_, schema, table)) = object_reference_query(capture.node, stmt) {
80+
to_return.push(RelationMatch { schema, table })
81+
}
82+
});
10783
});
10884

10985
to_return

‎crates/pgt_treesitter/src/queries/select_columns.rs‎

Lines changed: 11 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::sync::LazyLock;
22

3-
use crate::queries::{Query, QueryResult};
3+
use crate::queries::{Query, QueryResult, helper::object_reference_query};
44
use tree_sitter::StreamingIterator;
55

66
use super::QueryTryFrom;
@@ -9,11 +9,7 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
99
static QUERY_STR: &str = r#"
1010
(select_expression
1111
(term
12-
(object_reference
13-
object_reference_first: (any_identifier) @first
14-
object_reference_second: (any_identifier)? @second
15-
object_reference_third: (any_identifier)? @third
16-
)
12+
(object_reference) @ref
1713
)
1814
","?
1915
)
@@ -75,37 +71,15 @@ impl<'a> Query<'a> for SelectColumnMatch<'a> {
7571
let mut to_return = vec![];
7672

7773
matches.for_each(|m| {
78-
if m.captures.len() == 1 {
79-
let capture = m.captures[0].node;
80-
to_return.push(QueryResult::SelectClauseColumns(SelectColumnMatch {
81-
schema: None,
82-
alias: None,
83-
column: capture,
84-
}));
85-
}
86-
87-
if m.captures.len() == 2 {
88-
let alias = m.captures[0].node;
89-
let column = m.captures[1].node;
90-
91-
to_return.push(QueryResult::SelectClauseColumns(SelectColumnMatch {
92-
schema: None,
93-
alias: Some(alias),
94-
column,
95-
}));
96-
}
97-
98-
if m.captures.len() == 3 {
99-
let schema = m.captures[0].node;
100-
let alias = m.captures[1].node;
101-
let column = m.captures[2].node;
102-
103-
to_return.push(QueryResult::SelectClauseColumns(SelectColumnMatch {
104-
schema: Some(schema),
105-
alias: Some(alias),
106-
column,
107-
}));
108-
}
74+
m.captures.iter().for_each(|capture| {
75+
if let Some((schema, alias, column)) = object_reference_query(capture.node, stmt) {
76+
to_return.push(SelectColumnMatch {
77+
schema,
78+
alias,
79+
column,
80+
});
81+
}
82+
});
10983
});
11084

11185
to_return

‎crates/pgt_treesitter/src/queries/table_aliases.rs‎

Lines changed: 18 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,7 @@ use super::QueryTryFrom;
88
static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
99
static QUERY_STR: &str = r#"
1010
(relation
11-
(object_reference
12-
object_reference_first: (any_identifier) @first
13-
object_reference_second: (any_identifier)? @second
14-
)
11+
(object_reference) @ref
1512
(keyword_as)?
1613
(any_identifier) @alias
1714
)
@@ -77,27 +74,27 @@ impl<'a> Query<'a> for TableAliasMatch<'a> {
7774
let mut to_return = vec![];
7875

7976
matches.for_each(|m| {
80-
if m.captures.len() == 3 {
81-
let schema = m.captures[0].node;
82-
let table = m.captures[1].node;
83-
let alias = m.captures[2].node;
84-
85-
to_return.push(QueryResult::TableAliases(TableAliasMatch {
86-
table,
87-
alias,
88-
schema: Some(schema),
89-
}));
77+
if m.captures.len() == 1 {
78+
let obj_ref = m.captures[0].node;
79+
if let Some((_, schema, table)) = object_reference_query(obj_ref, stmt) {
80+
to_return.push(QueryResult::TableAliases(TableAliasMatch {
81+
schema,
82+
table,
83+
alias: None,
84+
}));
85+
}
9086
}
9187

9288
if m.captures.len() == 2 {
93-
let table = m.captures[0].node;
89+
let obj_ref = m.captures[0].node;
9490
let alias = m.captures[1].node;
95-
96-
to_return.push(QueryResult::TableAliases(TableAliasMatch {
97-
table,
98-
alias,
99-
schema: None,
100-
}));
91+
if let Some((_, schema, table)) = object_reference_query(obj_ref, stmt) {
92+
to_return.push(QueryResult::TableAliases(TableAliasMatch {
93+
schema,
94+
table,
95+
alias,
96+
}));
97+
}
10198
}
10299
});
103100

‎crates/pgt_treesitter_grammar/grammar.js‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3490,18 +3490,18 @@ module.exports = grammar({
34903490
object_reference: ($) =>
34913491
choice(
34923492
seq(
3493-
field("object_reference_first", $.any_identifier),
3493+
field("object_reference_1of3", $.any_identifier),
34943494
".",
3495-
field("object_reference_second", $.any_identifier),
3495+
field("object_reference_2of3", $.any_identifier),
34963496
".",
3497-
field("object_reference_third", $.any_identifier)
3497+
field("object_reference_3of3", $.any_identifier)
34983498
),
34993499
seq(
3500-
field("object_reference_first", $.any_identifier),
3500+
field("object_reference_1of2", $.any_identifier),
35013501
".",
3502-
field("object_reference_second", $.any_identifier)
3502+
field("object_reference_2of2", $.any_identifier)
35033503
),
3504-
field("object_reference_first", $.any_identifier)
3504+
field("object_reference_1of1", $.any_identifier)
35053505
),
35063506

35073507
type_reference: ($) =>

0 commit comments

Comments
 (0)