add ts query

Author: psteinroe

crates/pgt_treesitter_queries/src/lib.rs

@@ -68,7 +68,10 @@ impl<'a> Iterator for QueryResultIter<'a> {
 #[cfg(test)]
 mod tests {
 
-    use crate::{TreeSitterQueriesExecutor, queries::RelationMatch};
+    use crate::{
+        TreeSitterQueriesExecutor,
+        queries::{ParameterMatch, RelationMatch},
+    };
 
     #[test]
     fn finds_all_relations_and_ignores_functions() {
@@ -137,11 +140,11 @@ where
 select
   *
 from (
-    select * 
+    select *
     from (
       select *
       from private.something
-    ) as sq2 
+    ) as sq2
     join private.tableau pt1
     on sq2.id = pt1.id
   ) as sq1
@@ -185,4 +188,41 @@ on sq1.id = pt.id;
         assert_eq!(results[0].get_schema(sql), Some("private".into()));
         assert_eq!(results[0].get_table(sql), "something");
     }
+
+    #[test]
+    fn extracts_parameters() {
+        let sql = r#"select v_test + fn_name.custom_type.v_test2 + $3 + custom_type.v_test3;"#;
+
+        let mut parser = tree_sitter::Parser::new();
+        parser.set_language(tree_sitter_sql::language()).unwrap();
+
+        let tree = parser.parse(sql, None).unwrap();
+
+        let mut executor = TreeSitterQueriesExecutor::new(tree.root_node(), sql);
+
+        executor.add_query_results::<ParameterMatch>();
+
+        let results: Vec<&ParameterMatch> = executor
+            .get_iter(None)
+            .filter_map(|q| q.try_into().ok())
+            .collect();
+
+        assert_eq!(results.len(), 4);
+
+        assert_eq!(results[0].get_root(sql), None);
+        assert_eq!(results[0].get_path(sql), None);
+        assert_eq!(results[0].get_field(sql), "v_test");
+
+        assert_eq!(results[1].get_root(sql), Some("fn_name".into()));
+        assert_eq!(results[1].get_path(sql), Some("custom_type".into()));
+        assert_eq!(results[1].get_field(sql), "v_test2");
+
+        assert_eq!(results[2].get_root(sql), None);
+        assert_eq!(results[2].get_path(sql), None);
+        assert_eq!(results[2].get_field(sql), "$3");
+
+        assert_eq!(results[3].get_root(sql), None);
+        assert_eq!(results[3].get_path(sql), Some("custom_type".into()));
+        assert_eq!(results[3].get_field(sql), "v_test3");
+    }
 }

crates/pgt_treesitter_queries/src/queries/mod.rs

@@ -1,10 +1,13 @@
+mod parameters;
 mod relations;
 
+pub use parameters::*;
 pub use relations::*;
 
 #[derive(Debug)]
 pub enum QueryResult<'a> {
     Relation(RelationMatch<'a>),
+    Parameter(ParameterMatch<'a>),
 }
 
 impl QueryResult<'_> {
@@ -18,6 +21,16 @@ impl QueryResult<'_> {
 
                 let end = rm.table.end_position();
 
+                start >= range.start_point && end <= range.end_point
+            }
+            Self::Parameter(pm) => {
+                let start = match pm.root {
+                    Some(s) => s.start_position(),
+                    None => pm.path.as_ref().unwrap().start_position(),
+                };
+
+                let end = pm.field.end_position();
+
                 start >= range.start_point && end <= range.end_point
             }
         }

crates/pgt_treesitter_queries/src/queries/parameters.rs

@@ -0,0 +1,124 @@
+use std::sync::LazyLock;
+
+use crate::{Query, QueryResult};
+
+use super::QueryTryFrom;
+
+static TS_QUERY: LazyLock<tree_sitter::Query> = LazyLock::new(|| {
+    static QUERY_STR: &str = r#"
+[
+  (field
+    (identifier)) @reference
+  (field
+    (object_reference)
+    "." (identifier)) @reference
+  (parameter) @parameter
+]
+"#;
+    tree_sitter::Query::new(tree_sitter_sql::language(), QUERY_STR).expect("Invalid TS Query")
+});
+
+#[derive(Debug)]
+pub struct ParameterMatch<'a> {
+    pub(crate) root: Option<tree_sitter::Node<'a>>,
+    pub(crate) path: Option<tree_sitter::Node<'a>>,
+
+    pub(crate) field: tree_sitter::Node<'a>,
+}
+
+impl ParameterMatch<'_> {
+    pub fn get_root(&self, sql: &str) -> Option<String> {
+        let str = self
+            .root
+            .as_ref()?
+            .utf8_text(sql.as_bytes())
+            .expect("Failed to get schema from RelationMatch");
+
+        Some(str.to_string())
+    }
+
+    pub fn get_path(&self, sql: &str) -> Option<String> {
+        let str = self
+            .path
+            .as_ref()?
+            .utf8_text(sql.as_bytes())
+            .expect("Failed to get table from RelationMatch");
+
+        Some(str.to_string())
+    }
+
+    pub fn get_field(&self, sql: &str) -> String {
+        self.field
+            .utf8_text(sql.as_bytes())
+            .expect("Failed to get table from RelationMatch")
+            .to_string()
+    }
+}
+
+impl<'a> TryFrom<&'a QueryResult<'a>> for &'a ParameterMatch<'a> {
+    type Error = String;
+
+    fn try_from(q: &'a QueryResult<'a>) -> Result<Self, Self::Error> {
+        match q {
+            QueryResult::Parameter(r) => Ok(r),
+
+            #[allow(unreachable_patterns)]
+            _ => Err("Invalid QueryResult type".into()),
+        }
+    }
+}
+
+impl<'a> QueryTryFrom<'a> for ParameterMatch<'a> {
+    type Ref = &'a ParameterMatch<'a>;
+}
+
+impl<'a> Query<'a> for ParameterMatch<'a> {
+    fn execute(root_node: tree_sitter::Node<'a>, stmt: &'a str) -> Vec<crate::QueryResult<'a>> {
+        let mut cursor = tree_sitter::QueryCursor::new();
+
+        let matches = cursor.matches(&TS_QUERY, root_node, stmt.as_bytes());
+
+        matches
+            .filter_map(|m| {
+                let captures = m.captures;
+
+                // We expect exactly one capture for a parameter
+                if captures.len() != 1 {
+                    return None;
+                }
+
+                let field = captures[0].node;
+                let text = match field.utf8_text(stmt.as_bytes()) {
+                    Ok(t) => t,
+                    Err(_) => return None,
+                };
+                let parts: Vec<&str> = text.split('.').collect();
+
+                let param_match = match parts.len() {
+                    // Simple field: field_name
+                    1 => ParameterMatch {
+                        root: None,
+                        path: None,
+                        field,
+                    },
+                    // Table qualified: table.field_name
+                    2 => ParameterMatch {
+                        root: None,
+                        path: field.named_child(0),
+                        field: field.named_child(1)?,
+                    },
+                    // Fully qualified: schema.table.field_name
+                    3 => ParameterMatch {
+                        root: field.named_child(0).and_then(|n| n.named_child(0)),
+                        path: field.named_child(0).and_then(|n| n.named_child(1)),
+                        field: field.named_child(1)?,
+                    },
+                    // Unexpected number of parts
+                    _ => return None,
+                };
+
+                Some(QueryResult::Parameter(param_match))
+            })
+            .collect()
+    }
+}

crates/pgt_typecheck/src/lib.rs

@@ -1,13 +1,13 @@
 mod diagnostics;
 mod typed_identifier;
 
-use diagnostics::create_type_error;
 pub use diagnostics::TypecheckDiagnostic;
-pub use typed_identifier::TypedIdentifier;
+use diagnostics::create_type_error;
 use pgt_text_size::TextRange;
 use sqlx::postgres::PgDatabaseError;
 pub use sqlx::postgres::PgSeverity;
 use sqlx::{Executor, PgPool};
+pub use typed_identifier::TypedIdentifier;
 use typed_identifier::apply_identifiers;
 
 #[derive(Debug)]

crates/pgt_typecheck/src/typed_identifier.rs

@@ -1,24 +1,17 @@
-#[derive(Debug)]
-pub struct Type {
-    pub schema: Option<String>,
-    pub name: String,
-    pub oid: i32,
-}
-
 #[derive(Debug)]
 pub struct TypedIdentifier {
     pub schema: Option<String>,
     pub relation: Option<String>,
     pub name: String,
-    pub type_: Type,
+    pub type_: (Option<String>, String),
 }
 
 impl TypedIdentifier {
     pub fn new(
         schema: Option<String>,
         relation: Option<String>,
         name: String,
-        type_: Type,
+        type_: (Option<String>, String),
     ) -> Self {
         TypedIdentifier {
             schema,
@@ -44,5 +37,107 @@ pub fn apply_identifiers<'a>(
     println!("Applying identifiers to SQL: {}", sql);
     println!("Identifiers: {:?}", identifiers);
     println!("CST: {:#?}", cst);
+
     sql
 }
+
+#[cfg(test)]
+mod tests {
+    use pgt_test_utils::test_database::get_new_test_db;
+    use sqlx::Executor;
+
+    #[tokio::test]
+    async fn test_apply_identifiers() {
+        let input = "select v_test + fn_name.custom_type.v_test2 + $3 + test.field;";
+
+        let test_db = get_new_test_db().await;
+
+        let mut parser = tree_sitter::Parser::new();
+        parser
+            .set_language(tree_sitter_sql::language())
+            .expect("Error loading sql language");
+
+        let schema_cache = pgt_schema_cache::SchemaCache::load(&test_db)
+            .await
+            .expect("Failed to load Schema Cache");
+
+        let root = pgt_query_ext::parse(input).unwrap();
+        let tree = parser.parse(input, None).unwrap();
+
+        println!("Parsed SQL: {:?}", root);
+        println!("Parsed CST: {:?}", tree);
+
+        // let mut parameters = Vec::new();
+
+        enum Parameter {
+            Identifier {
+                range: (usize, usize),
+                name: (Option<String>, String),
+            },
+            Parameter {
+                range: (usize, usize),
+                idx: usize,
+            },
+        }
+
+        let mut c = tree.walk();
+
+        'outer: loop {
+            // 0. Add the current node to the map.
+            println!("Current node: {:?}", c.node());
+            match c.node().kind() {
+                "identifier" => {
+                    println!(
+                        "Found identifier: {:?}",
+                        c.node().utf8_text(input.as_bytes()).unwrap()
+                    );
+                }
+                "parameter" => {
+                    println!(
+                        "Found parameter: {:?}",
+                        c.node().utf8_text(input.as_bytes()).unwrap()
+                    );
+                }
+                "object_reference" => {
+                    println!(
+                        "Found object reference: {:?}",
+                        c.node().utf8_text(input.as_bytes()).unwrap()
+                    );
+
+                    // let source = self.text;
+                    // ts_node.utf8_text(source.as_bytes()).ok().map(|txt| {
+                    //     if SanitizedCompletionParams::is_sanitized_token(txt) {
+                    //         NodeText::Replaced
+                    //     } else {
+                    //         NodeText::Original(txt)
+                    //     }
+                    // })
+                }
+                _ => {}
+            }
+
+            // 1. Go to its child and continue.
+            if c.goto_first_child() {
+                continue 'outer;
+            }
+
+            // 2. We've reached a leaf (node without a child). We will go to a sibling.
+            if c.goto_next_sibling() {
+                continue 'outer;
+            }
+
+            // 3. If there are no more siblings, we need to go back up.
+            'inner: loop {
+                // 4. Check if we've reached the root node. If so, we're done.
+                if !c.goto_parent() {
+                    break 'outer;
+                }
+                // 5. Go to the previous node's sibling.
+                if c.goto_next_sibling() {
+                    // And break out of the inner loop.
+                    break 'inner;
+                }
+            }
+        }
+    }
+}