fix: sql fn params

crates/pgt_completions/src/context.rs

@@ -171,7 +171,9 @@ impl<'a> CompletionContext<'a> {
                         table_alias_match.get_table(sql),
                     );
                 }
-            };
+
+                _ => {}
+            }
         }
     }
 

crates/pgt_schema_cache/src/lib.rs

@@ -16,3 +16,4 @@ pub use functions::{Behavior, Function, FunctionArg, FunctionArgs};
 pub use schema_cache::SchemaCache;
 pub use schemas::Schema;
 pub use tables::{ReplicaIdentity, Table};
+pub use types::{PostgresType, PostgresTypeAttribute};

crates/pgt_schema_cache/src/types.rs

@@ -6,13 +6,13 @@ use crate::schema_cache::SchemaCacheItem;
 
 #[derive(Debug, Clone, Default)]
 pub struct TypeAttributes {
-    attrs: Vec<PostgresTypeAttribute>,
+    pub attrs: Vec<PostgresTypeAttribute>,
 }
 
 #[derive(Debug, Clone, Default, Deserialize)]
 pub struct PostgresTypeAttribute {
-    name: String,
-    type_id: i64,
+    pub name: String,
+    pub type_id: i64,
 }
 
 impl From<Option<JsonValue>> for TypeAttributes {

crates/pgt_treesitter_queries/src/lib.rs

@@ -70,7 +70,7 @@ mod tests {
 
     use crate::{
         TreeSitterQueriesExecutor,
-        queries::{RelationMatch, TableAliasMatch},
+        queries::{ParameterMatch, RelationMatch, TableAliasMatch},
     };
 
     #[test]
@@ -207,11 +207,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
@@ -255,4 +255,33 @@ 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_path(sql), "v_test");
+
+        assert_eq!(results[1].get_path(sql), "fn_name.custom_type.v_test2");
+
+        assert_eq!(results[2].get_path(sql), "$3");
+
+        assert_eq!(results[3].get_path(sql), "custom_type.v_test3");
+    }
 }

crates/pgt_treesitter_queries/src/queries/mod.rs

@@ -1,12 +1,15 @@
+mod parameters;
 mod relations;
 mod table_aliases;
 
+pub use parameters::*;
 pub use relations::*;
 pub use table_aliases::*;
 
 #[derive(Debug)]
 pub enum QueryResult<'a> {
     Relation(RelationMatch<'a>),
+    Parameter(ParameterMatch<'a>),
     TableAliases(TableAliasMatch<'a>),
 }
 
@@ -23,6 +26,12 @@ impl QueryResult<'_> {
 
                 start >= range.start_point && end <= range.end_point
             }
+            Self::Parameter(pm) => {
+                let node_range = pm.node.range();
+
+                node_range.start_point >= range.start_point
+                    && node_range.end_point <= range.end_point
+            }
             QueryResult::TableAliases(m) => {
                 let start = m.table.start_position();
                 let end = m.alias.end_position();

crates/pgt_treesitter_queries/src/queries/parameters.rs

@@ -0,0 +1,82 @@
+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) node: tree_sitter::Node<'a>,
+}
+
+impl ParameterMatch<'_> {
+    pub fn get_path(&self, sql: &str) -> String {
+        self.node
+            .utf8_text(sql.as_bytes())
+            .expect("Failed to get path from ParameterMatch")
+            .to_string()
+    }
+
+    pub fn get_range(&self) -> tree_sitter::Range {
+        self.node.range()
+    }
+
+    pub fn get_byte_range(&self) -> std::ops::Range<usize> {
+        let range = self.node.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;
+                }
+
+                Some(QueryResult::Parameter(ParameterMatch {
+                    node: captures[0].node,
+                }))
+            })
+            .collect()
+    }
+}

crates/pgt_typecheck/Cargo.toml

@@ -12,15 +12,16 @@ version              = "0.0.0"
 
 
 [dependencies]
-pgt_console.workspace      = true
-pgt_diagnostics.workspace  = true
-pgt_query_ext.workspace    = true
-pgt_schema_cache.workspace = true
-pgt_text_size.workspace    = true
-sqlx.workspace             = true
-tokio.workspace            = true
-tree-sitter.workspace      = true
-tree_sitter_sql.workspace  = true
+pgt_console.workspace            = true
+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
+tree_sitter_sql.workspace        = true
 
 [dev-dependencies]
 insta.workspace          = true

crates/pgt_typecheck/src/diagnostics.rs

@@ -97,21 +97,26 @@ impl Advices for TypecheckAdvices {
 pub(crate) fn create_type_error(
     pg_err: &PgDatabaseError,
     ts: &tree_sitter::Tree,
+    positions_valid: bool,
 ) -> TypecheckDiagnostic {
     let position = pg_err.position().and_then(|pos| match pos {
         sqlx::postgres::PgErrorPosition::Original(pos) => Some(pos - 1),
         _ => None,
     });
 
     let range = position.and_then(|pos| {
-        ts.root_node()
-            .named_descendant_for_byte_range(pos, pos)
-            .map(|node| {
-                TextRange::new(
-                    node.start_byte().try_into().unwrap(),
-                    node.end_byte().try_into().unwrap(),
-                )
-            })
+        if positions_valid {
+            ts.root_node()
+                .named_descendant_for_byte_range(pos, pos)
+                .map(|node| {
+                    TextRange::new(
+                        node.start_byte().try_into().unwrap(),
+                        node.end_byte().try_into().unwrap(),
+                    )
+                })
+        } else {
+            None
+        }
     });
 
     let severity = match pg_err.severity() {

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};
+use typed_identifier::apply_identifiers;
+pub use typed_identifier::{IdentifierType, TypedIdentifier};
 
 #[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,13 +56,24 @@ pub async fn check_sql(
     // each typecheck operation.
     conn.close_on_drop();
 
-    let res = conn.prepare(params.sql).await;
+    let (prepared, positions_valid) = apply_identifiers(
+        params.identifiers,
+        params.schema_cache,
+        params.tree,
+        params.sql,
+    );
+
+    let res = conn.prepare(&prepared).await;
 
     match res {
         Ok(_) => Ok(None),
         Err(sqlx::Error::Database(err)) => {
             let pg_err = err.downcast_ref::<PgDatabaseError>();
-            Ok(Some(create_type_error(pg_err, params.tree)))
+            Ok(Some(create_type_error(
+                pg_err,
+                params.tree,
+                positions_valid,
+            )))
         }
         Err(err) => Err(err),
     }