fix: sql fn params

Author: psteinroe

crates/pgt_typecheck/src/lib.rs

@@ -1,18 +1,23 @@
 mod diagnostics;
+mod typed_identifier;
 
-pub use diagnostics::TypecheckDiagnostic;
 use diagnostics::create_type_error;
+pub use diagnostics::TypecheckDiagnostic;
 use pgt_text_size::TextRange;
 use sqlx::postgres::PgDatabaseError;
 pub use sqlx::postgres::PgSeverity;
 use sqlx::{Executor, PgPool};
+use typed_identifier::{apply_identifiers, 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 cst: &'a tree_sitter::Node<'a>,
+    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.cst,
+        params.sql,
+    );
+
+    let res = conn.prepare(prepared).await;
 
     match res {
         Ok(_) => Ok(None),

crates/pgt_typecheck/src/typed_identifier.rs

@@ -0,0 +1,18 @@
+#[derive(Debug)]
+pub struct TypedIdentifier {
+    pub schema: Option<String>,
+    pub relation: String,
+    pub name: String,
+    pub type_: String,
+}
+
+/// Applies the identifiers to the SQL string by replacing them with their default values.
+pub fn apply_identifiers<'a>(
+    identifiers: Vec<TypedIdentifier>,
+    schema_cache: &'a pgt_schema_cache::SchemaCache,
+    cst: &'a tree_sitter::Node<'a>,
+    sql: &'a str,
+) -> &'a str {
+    // TODO
+    sql
+}

crates/pgt_workspace/src/workspace/server/document.rs

@@ -34,6 +34,13 @@ impl Document {
         }
     }
 
+    pub fn statement_content(&self, id: &StatementId) -> Option<&str> {
+        self.positions
+            .iter()
+            .find(|(statement_id, _)| statement_id == id)
+            .map(|(_, range)| &self.content[*range])
+    }
+
     /// Returns true if there is at least one fatal error in the diagnostics
     ///
     /// A fatal error is a scan error that prevents the document from being used

crates/pgt_workspace/src/workspace/server/parsed_document.rs

@@ -12,7 +12,7 @@ use super::{
     change::StatementChange,
     document::{Document, StatementIterator},
     pg_query::PgQueryStore,
-    sql_function::SQLFunctionBodyStore,
+    sql_function::{SQLFunctionBodyStore, SQLFunctionSignature},
     statement_identifier::StatementId,
     tree_sitter::TreeSitterStore,
 };
@@ -274,6 +274,7 @@ impl<'a> StatementMapper<'a> for AsyncDiagnosticsMapper {
         String,
         Option<pgt_query_ext::NodeEnum>,
         Arc<tree_sitter::Tree>,
+        Option<Arc<SQLFunctionSignature>>,
     );
 
     fn map(
@@ -293,7 +294,26 @@ impl<'a> StatementMapper<'a> for AsyncDiagnosticsMapper {
 
         let cst_result = parser.cst_db.get_or_cache_tree(&id, &content_owned);
 
-        (id, range, content_owned, ast_option, cst_result)
+        let sql_fn_sig = id
+            .parent()
+            .and_then(|root| {
+                let c = parser.doc.statement_content(&root)?;
+                Some((root, c))
+            })
+            .and_then(|(root, c)| {
+                let ast_option = parser
+                    .ast_db
+                    .get_or_cache_ast(&root, c)
+                    .as_ref()
+                    .clone()
+                    .ok();
+
+                let ast_option = ast_option.as_ref()?;
+
+                parser.sql_fn_db.get_function_signature(&root, ast_option)
+            });
+
+        (id, range, content_owned, ast_option, cst_result, sql_fn_sig)
     }
 }
 
@@ -413,7 +433,7 @@ mod tests {
 
     #[test]
     fn sql_function_body() {
-        let input = "CREATE FUNCTION add(integer, integer) RETURNS integer
+        let input = "CREATE FUNCTION add(test0 integer, test1 integer) RETURNS integer
     AS 'select $1 + $2;'
     LANGUAGE SQL
     IMMUTABLE

crates/pgt_workspace/src/workspace/server/sql_function.rs

@@ -5,6 +5,18 @@ use pgt_text_size::TextRange;
 
 use super::statement_identifier::StatementId;
 
+#[derive(Debug, Clone)]
+pub struct SQLFunctionArgs {
+    pub name: Option<String>,
+    pub type_: (Option<String>, String),
+}
+
+#[derive(Debug, Clone)]
+pub struct SQLFunctionSignature {
+    pub name: (Option<String>, String),
+    pub args: Vec<SQLFunctionArgs>,
+}
+
 #[derive(Debug, Clone)]
 pub struct SQLFunctionBody {
     pub range: TextRange,
@@ -13,11 +25,33 @@ pub struct SQLFunctionBody {
 
 pub struct SQLFunctionBodyStore {
     db: DashMap<StatementId, Option<Arc<SQLFunctionBody>>>,
+    sig_db: DashMap<StatementId, Option<Arc<SQLFunctionSignature>>>,
 }
 
 impl SQLFunctionBodyStore {
     pub fn new() -> SQLFunctionBodyStore {
-        SQLFunctionBodyStore { db: DashMap::new() }
+        SQLFunctionBodyStore {
+            db: DashMap::new(),
+            sig_db: DashMap::new(),
+        }
+    }
+
+    pub fn get_function_signature(
+        &self,
+        statement: &StatementId,
+        ast: &pgt_query_ext::NodeEnum,
+    ) -> Option<Arc<SQLFunctionSignature>> {
+        // First check if we already have this statement cached
+        if let Some(existing) = self.sig_db.get(statement).map(|x| x.clone()) {
+            return existing;
+        }
+
+        // If not cached, try to extract it from the AST
+        let fn_sig = get_sql_fn_signature(ast).map(Arc::new);
+
+        // Cache the result and return it
+        self.sig_db.insert(statement.clone(), fn_sig.clone());
+        fn_sig
     }
 
     pub fn get_function_body(
@@ -48,6 +82,48 @@ impl SQLFunctionBodyStore {
     }
 }
 
+/// Extracts SQL function signature from a CreateFunctionStmt node.
+fn get_sql_fn_signature(ast: &pgt_query_ext::NodeEnum) -> Option<SQLFunctionSignature> {
+    let create_fn = match ast {
+        pgt_query_ext::NodeEnum::CreateFunctionStmt(cf) => cf,
+        _ => return None,
+    };
+
+    println!("create_fn: {:?}", create_fn);
+
+    // Extract language from function options
+    let language = find_option_value(create_fn, "language")?;
+
+    // Only process SQL functions
+    if language != "sql" {
+        return None;
+    }
+
+    let fn_name = parse_name(&create_fn.funcname)?;
+
+    // we return None if anything is not expected
+    let mut fn_args = Vec::new();
+    for arg in &create_fn.parameters {
+        if let Some(pgt_query_ext::NodeEnum::FunctionParameter(node)) = &arg.node {
+            let arg_name = (!node.name.is_empty()).then_some(node.name.clone());
+
+            let type_name = parse_name(&node.arg_type.as_ref().unwrap().names)?;
+
+            fn_args.push(SQLFunctionArgs {
+                name: arg_name,
+                type_: type_name,
+            });
+        } else {
+            return None;
+        }
+    }
+
+    Some(SQLFunctionSignature {
+        name: fn_name,
+        args: fn_args,
+    })
+}
+
 /// Extracts SQL function body and its text range from a CreateFunctionStmt node.
 /// Returns None if the function is not an SQL function or if the body can't be found.
 fn get_sql_fn(ast: &pgt_query_ext::NodeEnum, content: &str) -> Option<SQLFunctionBody> {
@@ -56,6 +132,8 @@ fn get_sql_fn(ast: &pgt_query_ext::NodeEnum, content: &str) -> Option<SQLFunctio
         _ => return None,
     };
 
+    println!("create_fn: {:?}", create_fn);
+
     // Extract language from function options
     let language = find_option_value(create_fn, "language")?;
 
@@ -120,3 +198,19 @@ fn find_option_value(
             }
         })
 }
+
+fn parse_name(nodes: &Vec<pgt_query_ext::protobuf::Node>) -> Option<(Option<String>, String)> {
+    let names = nodes
+        .iter()
+        .map(|n| match &n.node {
+            Some(pgt_query_ext::NodeEnum::String(s)) => Some(s.sval.clone()),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+
+    match names.as_slice() {
+        [Some(schema), Some(name)] => Some((Some(schema.clone()), name.clone())),
+        [Some(name)] => Some((None, name.clone())),
+        _ => None,
+    }
+}

crates/pgt_workspace/src/workspace/server/statement_identifier.rs

@@ -57,6 +57,21 @@ impl StatementId {
             StatementId::Child(s) => s.inner,
         }
     }
+
+    pub fn is_root(&self) -> bool {
+        matches!(self, StatementId::Root(_))
+    }
+
+    pub fn is_child(&self) -> bool {
+        matches!(self, StatementId::Child(_))
+    }
+
+    pub fn parent(&self) -> Option<StatementId> {
+        match self {
+            StatementId::Root(_) => None,
+            StatementId::Child(id) => Some(StatementId::Root(id.clone())),
+        }
+    }
 }
 
 /// Helper struct to generate unique statement ids