Skip to content

Commit 896bfb1

Browse files
psteinroepsteinroe
committed
progress
1 parent e032c3f commit 896bfb1

File tree

6 files changed

+223
-211
lines changed

6 files changed

+223
-211
lines changed

‎crates/pgt_schema_cache/src/types.rs

Lines changed: 58 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,13 @@ use crate::schema_cache::SchemaCacheItem;
66

77
#[derive(Debug, Clone, Default)]
88
pub struct TypeAttributes {
9-
attrs: Vec<PostgresTypeAttribute>,
9+
pub attrs: Vec<PostgresTypeAttribute>,
1010
}
1111

1212
#[derive(Debug, Clone, Default, Deserialize)]
1313
pub struct PostgresTypeAttribute {
14-
name: String,
15-
type_id: i64,
14+
pub name: String,
15+
pub type_id: i64,
1616
}
1717

1818
impl From<Option<JsonValue>> for TypeAttributes {
@@ -56,3 +56,58 @@ impl SchemaCacheItem for PostgresType {
5656
.await
5757
}
5858
}
59+
60+
#[cfg(test)]
61+
mod tests {
62+
use pgt_test_utils::test_database::get_new_test_db;
63+
use sqlx::Executor;
64+
65+
use crate::{schema_cache::SchemaCacheItem, types::PostgresType};
66+
67+
#[tokio::test]
68+
async fn test_types() {
69+
let setup = r#"
70+
CREATE TYPE "public"."priority" AS ENUM (
71+
'critical',
72+
'high',
73+
'default',
74+
'low',
75+
'very_low'
76+
);
77+
78+
CREATE TYPE complex AS (
79+
r double precision,
80+
i double precision
81+
);
82+
"#;
83+
84+
let test_db = get_new_test_db().await;
85+
86+
test_db
87+
.execute(setup)
88+
.await
89+
.expect("Failed to setup test database");
90+
91+
let types = PostgresType::load(&test_db).await.unwrap();
92+
93+
let enum_type = types.iter().find(|t| t.name == "priority");
94+
let comp_type = types.iter().find(|t| t.name == "complex");
95+
96+
println!("{:?}", enum_type);
97+
// search for type id
98+
println!("{:?}", comp_type);
99+
100+
comp_type.and_then(|t| {
101+
t.attributes.attrs.iter().for_each(|a| {
102+
let typ = types.iter().find(|t| t.id == a.type_id);
103+
println!(
104+
"{}: {} - {:?}",
105+
a.name,
106+
a.type_id,
107+
typ.as_ref().map(|t| t.name.clone())
108+
);
109+
});
110+
Some(())
111+
});
112+
}
113+
}

‎crates/pgt_treesitter_queries/src/lib.rs

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ impl<'a> Iterator for QueryResultIter<'a> {
6969
mod tests {
7070

7171
use crate::{
72-
queries::{Field, ParameterMatch, RelationMatch}, TreeSitterQueriesExecutor
72+
TreeSitterQueriesExecutor,
73+
queries::{Field, ParameterMatch, RelationMatch},
7374
};
7475

7576
#[test]
@@ -208,20 +209,12 @@ on sq1.id = pt.id;
208209

209210
assert_eq!(results.len(), 4);
210211

211-
assert_eq!(results[0].get_root(sql), None);
212-
assert_eq!(results[0].get_path(sql), None);
213-
assert_eq!(results[0].get_field(sql), Field::Text("v_test".to_string()));
212+
assert_eq!(results[0].get_path(sql), "v_test");
214213

215-
assert_eq!(results[1].get_root(sql), Some("fn_name".into()));
216-
assert_eq!(results[1].get_path(sql), Some("custom_type".into()));
217-
assert_eq!(results[1].get_field(sql), Field::Text("v_test2".to_string()));
214+
assert_eq!(results[1].get_path(sql), "fn_name.custom_type._test2");
218215

219-
assert_eq!(results[2].get_root(sql), None);
220-
assert_eq!(results[2].get_path(sql), None);
221-
assert_eq!(results[2].get_field(sql), Field::Parameter(3));
216+
assert_eq!(results[2].get_path(sql), "$3");
222217

223-
assert_eq!(results[3].get_root(sql), None);
224-
assert_eq!(results[3].get_path(sql), Some("custom_type".into()));
225-
assert_eq!(results[3].get_field(sql), Field::Text("v_test3".to_string()));
218+
assert_eq!(results[3].get_path(sql), "custom_type.v_test3");
226219
}
227220
}

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

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,10 @@ impl QueryResult<'_> {
2424
start >= range.start_point && end <= range.end_point
2525
}
2626
Self::Parameter(pm) => {
27-
let start = match pm.root {
28-
Some(s) => s.start_position(),
29-
None => pm.path.as_ref().unwrap().start_position(),
30-
};
27+
let node_range = pm.node.range();
3128

32-
let end = pm.field.end_position();
33-
34-
start >= range.start_point && end <= range.end_point
29+
node_range.start_point >= range.start_point
30+
&& node_range.end_point <= range.end_point
3531
}
3632
}
3733
}

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

Lines changed: 10 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -20,62 +20,23 @@ static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
2020

2121
#[derive(Debug)]
2222
pub struct ParameterMatch<'a> {
23-
pub(crate) root: Option<tree_sitter::Node<'a>>,
24-
pub(crate) path: Option<tree_sitter::Node<'a>>,
25-
26-
pub(crate) field: tree_sitter::Node<'a>,
27-
}
28-
29-
#[derive(Debug, PartialEq)]
30-
pub enum Field {
31-
Text(String),
32-
Parameter(usize),
23+
pub(crate) node: tree_sitter::Node<'a>,
3324
}
3425

3526
impl ParameterMatch<'_> {
36-
pub fn get_root(&self, sql: &str) -> Option<String> {
37-
let str = self
38-
.root
39-
.as_ref()?
40-
.utf8_text(sql.as_bytes())
41-
.expect("Failed to get schema from RelationMatch");
42-
43-
Some(str.to_string())
44-
}
45-
46-
pub fn get_path(&self, sql: &str) -> Option<String> {
47-
let str = self
48-
.path
49-
.as_ref()?
27+
pub fn get_path(&self, sql: &str) -> String {
28+
self.node
5029
.utf8_text(sql.as_bytes())
51-
.expect("Failed to get table from RelationMatch");
52-
53-
Some(str.to_string())
54-
}
55-
56-
pub fn get_field(&self, sql: &str) -> Field {
57-
let text = self
58-
.field
59-
.utf8_text(sql.as_bytes())
60-
.expect("Failed to get field from RelationMatch");
61-
62-
if let Some(stripped) = text.strip_prefix('$') {
63-
return Field::Parameter(
64-
stripped
65-
.parse::<usize>()
66-
.expect("Failed to parse parameter"),
67-
);
68-
}
69-
70-
Field::Text(text.to_string())
30+
.expect("Failed to get path from ParameterMatch")
31+
.to_string()
7132
}
7233

7334
pub fn get_range(&self) -> tree_sitter::Range {
74-
self.field.range()
35+
self.node.range()
7536
}
7637

7738
pub fn get_byte_range(&self) -> std::ops::Range<usize> {
78-
let range = self.field.range();
39+
let range = self.node.range();
7940
range.start_byte..range.end_byte
8041
}
8142
}
@@ -112,37 +73,9 @@ impl<'a> Query<'a> for ParameterMatch<'a> {
11273
return None;
11374
}
11475

115-
let field = captures[0].node;
116-
let text = match field.utf8_text(stmt.as_bytes()) {
117-
Ok(t) => t,
118-
Err(_) => return None,
119-
};
120-
let parts: Vec<&str> = text.split('.').collect();
121-
122-
let param_match = match parts.len() {
123-
// Simple field: field_name
124-
1 => ParameterMatch {
125-
root: None,
126-
path: None,
127-
field,
128-
},
129-
// Table qualified: table.field_name
130-
2 => ParameterMatch {
131-
root: None,
132-
path: field.named_child(0),
133-
field: field.named_child(1)?,
134-
},
135-
// Fully qualified: schema.table.field_name
136-
3 => ParameterMatch {
137-
root: field.named_child(0).and_then(|n| n.named_child(0)),
138-
path: field.named_child(0).and_then(|n| n.named_child(1)),
139-
field: field.named_child(1)?,
140-
},
141-
// Unexpected number of parts
142-
_ => return None,
143-
};
144-
145-
Some(QueryResult::Parameter(param_match))
76+
Some(QueryResult::Parameter(ParameterMatch {
77+
node: captures[0].node,
78+
}))
14679
})
14780
.collect()
14881
}

‎crates/pgt_typecheck/Cargo.toml

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,16 @@ version = "0.0.0"
1212

1313

1414
[dependencies]
15-
pgt_console.workspace = true
16-
pgt_diagnostics.workspace = true
17-
pgt_query_ext.workspace = true
18-
pgt_schema_cache.workspace = true
19-
pgt_text_size.workspace = true
15+
pgt_console.workspace = true
16+
pgt_diagnostics.workspace = true
17+
pgt_query_ext.workspace = true
18+
pgt_schema_cache.workspace = true
19+
pgt_text_size.workspace = true
2020
pgt_treesitter_queries.workspace = true
21-
sqlx.workspace = true
22-
tokio.workspace = true
23-
tree-sitter.workspace = true
24-
tree_sitter_sql.workspace = true
21+
sqlx.workspace = true
22+
tokio.workspace = true
23+
tree-sitter.workspace = true
24+
tree_sitter_sql.workspace = true
2525

2626
[dev-dependencies]
2727
insta.workspace = true

0 commit comments

Comments
 (0)