fix: sql fn params

crates/pgt_treesitter_queries/src/lib.rs

@@ -68,7 +68,9 @@ impl<'a> Iterator for QueryResultIter<'a> {
 #[cfg(test)]
 mod tests {
 
-    use crate::{TreeSitterQueriesExecutor, queries::RelationMatch};
+    use crate::{
+        queries::{Field, ParameterMatch, RelationMatch}, TreeSitterQueriesExecutor
+    };
 
     #[test]
     fn finds_all_relations_and_ignores_functions() {
@@ -137,11 +139,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 +187,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), Field::Text("v_test".to_string()));
+
+        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), Field::Text("v_test2".to_string()));
+
+        assert_eq!(results[2].get_root(sql), None);
+        assert_eq!(results[2].get_path(sql), None);
+        assert_eq!(results[2].get_field(sql), Field::Parameter(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), Field::Text("v_test3".to_string()));
+    }
 }

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,149 @@
+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>,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum Field {
+    Text(String),
+    Parameter(usize),
+}
+
+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) -> Field {
+        let text = self
+            .field
+            .utf8_text(sql.as_bytes())
+            .expect("Failed to get field from RelationMatch");
+
+        if let Some(stripped) = text.strip_prefix('$') {
+            return Field::Parameter(
+                stripped
+                    .parse::<usize>()
+                    .expect("Failed to parse parameter"),
+            );
+        }
+
+        Field::Text(text.to_string())
+    }
+
+    pub fn get_range(&self) -> tree_sitter::Range {
+        self.field.range()
+    }
+
+    pub fn get_byte_range(&self) -> std::ops::Range<usize> {
+        let range = self.field.range();
+        range.start_byte..range.end_byte
+    }
+}
+
+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/Cargo.toml

@@ -17,6 +17,7 @@ pgt_diagnostics.workspace  = true
 pgt_query_ext.workspace    = true
 pgt_schema_cache.workspace = true
 pgt_text_size.workspace    = true
+pgt_treesitter_queries.workspace = true
 sqlx.workspace             = true
 tokio.workspace            = true
 tree-sitter.workspace      = true

crates/pgt_typecheck/src/lib.rs

@@ -1,18 +1,23 @@
 mod diagnostics;
+mod typed_identifier;
 
 pub use diagnostics::TypecheckDiagnostic;
 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)]
 pub struct TypecheckParams<'a> {
     pub conn: &'a PgPool,
     pub sql: &'a str,
     pub ast: &'a pgt_query_ext::NodeEnum,
     pub tree: &'a tree_sitter::Tree,
+    pub schema_cache: &'a pgt_schema_cache::SchemaCache,
+    pub identifiers: Vec<TypedIdentifier>,
 }
 
 #[derive(Debug, Clone)]
@@ -51,7 +56,14 @@ pub async fn check_sql(
     // each typecheck operation.
     conn.close_on_drop();
 
-    let res = conn.prepare(params.sql).await;
+    let prepared = apply_identifiers(
+        params.identifiers,
+        params.schema_cache,
+        params.tree,
+        params.sql,
+    );
+
+    let res = conn.prepare(prepared).await;
 
     match res {
         Ok(_) => Ok(None),