Skip to content

Commit b0336dc

Browse files
juleswritescodejuleswritescode
committed
intermediary
1 parent a331f41 commit b0336dc

File tree

5 files changed

+143
-30
lines changed

5 files changed

+143
-30
lines changed

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,7 @@ use super::QueryTryFrom;
99
static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
1010
static QUERY_STR: &str = r#"
1111
[
12-
(field
13-
(field_qualifier)?
14-
(column_identifier)
15-
) @reference
12+
(column_reference) @reference
1613
1714
(parameter) @parameter
1815
]

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

Lines changed: 25 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,10 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
99
static QUERY_STR: &str = r#"
1010
(select_expression
1111
(term
12-
(field
13-
(field_qualifier
14-
(object_reference) @alias
15-
"."
16-
)?
17-
(column_identifier) @column
12+
(object_reference
13+
first: (any_identifier) @first
14+
second: (any_identifier)? @second
15+
third: (any_identifier)? @third
1816
)
1917
)
2018
","?
@@ -26,6 +24,8 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
2624

2725
#[derive(Debug)]
2826
pub struct SelectColumnMatch<'a> {
27+
#[allow(unused)]
28+
pub(crate) schema: Option<tree_sitter::Node<'a>>,
2929
pub(crate) alias: Option<tree_sitter::Node<'a>>,
3030
pub(crate) column: tree_sitter::Node<'a>,
3131
}
@@ -78,6 +78,7 @@ impl<'a> Query<'a> for SelectColumnMatch<'a> {
7878
if m.captures.len() == 1 {
7979
let capture = m.captures[0].node;
8080
to_return.push(QueryResult::SelectClauseColumns(SelectColumnMatch {
81+
schema: None,
8182
alias: None,
8283
column: capture,
8384
}));
@@ -88,6 +89,19 @@ impl<'a> Query<'a> for SelectColumnMatch<'a> {
8889
let column = m.captures[1].node;
8990

9091
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),
91105
alias: Some(alias),
92106
column,
93107
}));
@@ -141,7 +155,8 @@ select
141155
u.id,
142156
u.email,
143157
cs.user_settings,
144-
cs.client_id
158+
cs.client_id,
159+
public.client_settings.id
145160
from
146161
auth.users u
147162
join public.client_settings cs
@@ -176,5 +191,8 @@ from
176191

177192
assert_eq!(results[3].get_alias(sql), Some("cs".into()));
178193
assert_eq!(results[3].get_column(sql), "client_id");
194+
195+
assert_eq!(results[4].get_alias(sql), Some("client_settings".into()));
196+
assert_eq!(results[4].get_column(sql), "id");
179197
}
180198
}

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

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,9 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
1111
(where
1212
(binary_expression
1313
(binary_expression
14-
(field
15-
(field_qualifier
16-
(object_reference) @alias
17-
"."
18-
)?
14+
(column_reference
15+
(schema_identifier)? @schema
16+
(table_identifier)? @table
1917
(column_identifier) @column
2018
)
2119
)
@@ -28,6 +26,8 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
2826

2927
#[derive(Debug)]
3028
pub struct WhereColumnMatch<'a> {
29+
#[allow(unused)]
30+
pub(crate) schema: Option<tree_sitter::Node<'a>>,
3131
pub(crate) alias: Option<tree_sitter::Node<'a>>,
3232
pub(crate) column: tree_sitter::Node<'a>,
3333
}
@@ -80,6 +80,7 @@ impl<'a> Query<'a> for WhereColumnMatch<'a> {
8080
if m.captures.len() == 1 {
8181
let capture = m.captures[0].node;
8282
to_return.push(QueryResult::WhereClauseColumns(WhereColumnMatch {
83+
schema: None,
8384
alias: None,
8485
column: capture,
8586
}));
@@ -90,6 +91,19 @@ impl<'a> Query<'a> for WhereColumnMatch<'a> {
9091
let column = m.captures[1].node;
9192

9293
to_return.push(QueryResult::WhereClauseColumns(WhereColumnMatch {
94+
schema: None,
95+
alias: Some(alias),
96+
column,
97+
}));
98+
}
99+
100+
if m.captures.len() == 2 {
101+
let schema = m.captures[0].node;
102+
let alias = m.captures[1].node;
103+
let column = m.captures[2].node;
104+
105+
to_return.push(QueryResult::WhereClauseColumns(WhereColumnMatch {
106+
schema: Some(schema),
93107
alias: Some(alias),
94108
column,
95109
}));

‎crates/pgt_treesitter_grammar/grammar.js‎

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2383,16 +2383,6 @@ module.exports = grammar({
23832383
optional(seq($.keyword_for, $.keyword_columns)),
23842384
optional(seq($.keyword_cache, $.keyword_metadata)),
23852385
optional($.keyword_noscan)
2386-
),
2387-
// Impala
2388-
seq(
2389-
$.keyword_compute,
2390-
optional($.keyword_incremental),
2391-
$.keyword_stats,
2392-
$.object_reference,
2393-
optional(
2394-
choice(paren_list(repeat1($.field), false), $._partition_spec)
2395-
)
23962386
)
23972387
),
23982388

@@ -3221,7 +3211,7 @@ module.exports = grammar({
32213211
1,
32223212
choice(
32233213
$.literal,
3224-
$.column_reference,
3214+
$.object_reference,
32253215
$.parameter,
32263216
$.list,
32273217
$.case,
@@ -3553,8 +3543,8 @@ module.exports = grammar({
35533543
$._sql_parameter,
35543544
seq("`", $._identifier, "`")
35553545
),
3556-
_sql_parameter: (_) => /[:$@?][a-zA-Z_][0-9a-zA-Z_]+/,
3557-
_identifier: (_) => /[a-zA-Z_][0-9a-zA-Z_]+/,
3546+
_sql_parameter: (_) => /[:$@?][a-zA-Z_][0-9a-zA-Z_]*/,
3547+
_identifier: (_) => /[a-zA-Z_][0-9a-zA-Z_]*/,
35583548

35593549
_anything: (_) => /\S+/,
35603550
},

‎crates/pgt_treesitter_grammar/tests/grammar_tests.rs‎

Lines changed: 95 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
use insta::assert_snapshot;
22

33
fn printed_tree(sql: &str) -> String {
4+
let sql = sql.trim();
45
let mut result = String::new();
56
let mut parser = tree_sitter::Parser::new();
67
let _ = parser.set_language(&pgt_treesitter_grammar::LANGUAGE.into());
7-
let tree = parser.parse(sql.trim(), None).expect("Unable to parse!");
8+
let tree = parser.parse(sql, None).expect("Unable to parse!");
89
pgt_test_utils::print_ts_tree(&tree.root_node(), sql, 0, &mut result);
910
result
1011
}
@@ -102,3 +103,96 @@ fn test_2() {
102103
; [54..55] ';'
103104
");
104105
}
106+
107+
#[test]
108+
fn test_3() {
109+
let sql = r#"
110+
select
111+
u.id,
112+
u.email,
113+
cs.user_settings,
114+
cs.client_id
115+
from
116+
auth.users u
117+
join public.client_settings cs
118+
on u.id = cs.user_id;
119+
120+
"#;
121+
122+
assert_snapshot!(printed_tree(sql), @r"
123+
program [0..25] 'select u.id, u.email, cs.user_settings, cs.client_id from auth.users u join public.client_settings cs on u.id = cs.user_id;'
124+
statement [0..24] 'select u.id, u.email, cs.user_settings, cs.client_id from auth.users u join public.client_settings cs on u.id = cs.user_id'
125+
select [0..16] 'select u.id, u.email, cs.user_settings, cs.client_id'
126+
keyword_select [0..6] 'select'
127+
select_expression [4..16] 'u.id, u.email, cs.user_settings, cs.client_id'
128+
term [4..8] 'u.id'
129+
column_reference [4..8] 'u.id'
130+
table_identifier [4..5] 'u'
131+
. [5..6] '.'
132+
column_identifier [6..8] 'id'
133+
, [8..9] ','
134+
term [4..11] 'u.email'
135+
column_reference [4..11] 'u.email'
136+
table_identifier [4..5] 'u'
137+
. [5..6] '.'
138+
column_identifier [6..11] 'email'
139+
, [11..12] ','
140+
term [4..20] 'cs.user_settings'
141+
column_reference [4..20] 'cs.user_settings'
142+
table_identifier [4..6] 'cs'
143+
. [6..7] '.'
144+
column_identifier [7..20] 'user_settings'
145+
, [20..21] ','
146+
term [4..16] 'cs.client_id'
147+
column_reference [4..16] 'cs.client_id'
148+
table_identifier [4..6] 'cs'
149+
. [6..7] '.'
150+
column_identifier [7..16] 'client_id'
151+
from [0..24] 'from auth.users u join public.client_settings cs on u.id = cs.user_id'
152+
keyword_from [0..4] 'from'
153+
relation [4..16] 'auth.users u'
154+
object_reference [4..14] 'auth.users'
155+
any_identifier [4..8] 'auth'
156+
. [8..9] '.'
157+
any_identifier [9..14] 'users'
158+
any_identifier [15..16] 'u'
159+
join [4..24] 'join public.client_settings cs on u.id = cs.user_id'
160+
keyword_join [4..8] 'join'
161+
relation [9..34] 'public.client_settings cs'
162+
object_reference [9..31] 'public.client_settings'
163+
any_identifier [9..15] 'public'
164+
. [15..16] '.'
165+
any_identifier [16..31] 'client_settings'
166+
any_identifier [32..34] 'cs'
167+
keyword_on [4..6] 'on'
168+
binary_expression [7..24] 'u.id = cs.user_id'
169+
column_reference [7..11] 'u.id'
170+
table_identifier [7..8] 'u'
171+
. [8..9] '.'
172+
column_identifier [9..11] 'id'
173+
= [12..13] '='
174+
column_reference [14..24] 'cs.user_id'
175+
table_identifier [14..16] 'cs'
176+
. [16..17] '.'
177+
column_identifier [17..24] 'user_id'
178+
; [24..25] ';'
179+
");
180+
}
181+
182+
#[test]
183+
fn test_4() {
184+
let sql = r#"select "auth".REPLACED_TOKEN"#;
185+
186+
assert_snapshot!(printed_tree(sql), @r#"
187+
program [0..28] 'select "auth".REPLACED_TOKEN'
188+
statement [0..28] 'select "auth".REPLACED_TOKEN'
189+
select [0..28] 'select "auth".REPLACED_TOKEN'
190+
keyword_select [0..6] 'select'
191+
select_expression [7..28] '"auth".REPLACED_TOKEN'
192+
term [7..28] '"auth".REPLACED_TOKEN'
193+
object_reference [7..28] '"auth".REPLACED_TOKEN'
194+
any_identifier [7..13] '"auth"'
195+
. [13..14] '.'
196+
any_identifier [14..28] 'REPLACED_TOKEN'
197+
"#);
198+
}

0 commit comments

Comments
 (0)