May 22, 2025 · Apr 22, 2025 · Apr 25, 2025 · Apr 25, 2025 · Apr 25, 2025 · Apr 29, 2025
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/pgt_completions/src/context/mod.rs b/crates/pgt_completions/src/context/mod.rs
                            .insert(Some(WrappingClause::Select), new);
                    }
                }
                _ => {}
            };
        }
    }
diff --git a/crates/pgt_schema_cache/src/lib.rs b/crates/pgt_schema_cache/src/lib.rs
 pub use schemas::Schema;
 pub use tables::{ReplicaIdentity, Table};
 pub use triggers::{Trigger, TriggerAffected, TriggerEvent};
 pub use types::{PostgresType, PostgresTypeAttribute};
diff --git a/crates/pgt_schema_cache/src/types.rs b/crates/pgt_schema_cache/src/types.rs

 #[derive(Debug, Clone, Default)]
 pub struct TypeAttributes {
    attrs: Vec<PostgresTypeAttribute>,
 pubattrs: Vec<PostgresTypeAttribute>,
 }

 #[derive(Debug, Clone, Default, Deserialize)]
 pub struct PostgresTypeAttribute {
    name: String,
    type_id: i64,
 pubname: String,
 pubtype_id: i64,
 }

 impl From<Option<JsonValue>> for TypeAttributes {
diff --git a/crates/pgt_treesitter_queries/src/lib.rs b/crates/pgt_treesitter_queries/src/lib.rs

    use crate::{
        TreeSitterQueriesExecutor,
        queries::{RelationMatch, TableAliasMatch},
        queries::{ParameterMatch,RelationMatch, TableAliasMatch},
    };

    #[test]
 select
  *
 from (
    select *
    select *
    from (
      select *
      from private.something
    ) as sq2
    ) as sq2
    join private.tableau pt1
    on sq2.id = pt1.id
  ) as sq1
        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");
    }
 }
diff --git a/crates/pgt_treesitter_queries/src/queries/mod.rs b/crates/pgt_treesitter_queries/src/queries/mod.rs
 mod parameters;
 mod relations;
 mod select_columns;
 mod table_aliases;

 pub use parameters::*;
 pub use relations::*;
 pub use select_columns::*;
 pub use table_aliases::*;

 #[derive(Debug)]
 pub enum QueryResult<'a> {
    Relation(RelationMatch<'a>),
    Parameter(ParameterMatch<'a>),
    TableAliases(TableAliasMatch<'a>),
    SelectClauseColumns(SelectColumnMatch<'a>),
 }

                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();
diff --git a/crates/pgt_treesitter_queries/src/queries/parameters.rs b/crates/pgt_treesitter_queries/src/queries/parameters.rs
 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()
    }
 }
diff --git a/crates/pgt_typecheck/Cargo.toml b/crates/pgt_typecheck/Cargo.toml


 [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
diff --git a/crates/pgt_typecheck/src/diagnostics.rs b/crates/pgt_typecheck/src/diagnostics.rs
 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() {
diff --git a/crates/pgt_typecheck/src/lib.rs b/crates/pgt_typecheck/src/lib.rs
 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)]
    // 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),
    }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -270,6 +270,7 @@ impl<'a> CompletionContext<'a> {
		.insert(Some(WrappingClause::Select), new);
		}
		}
		_ => {}
		};
		}
		}
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -19,3 +19,4 @@ pub use schema_cache::SchemaCache;
		pub use schemas::Schema;
		pub use tables::{ReplicaIdentity, Table};
		pub use triggers::{Trigger, TriggerAffected, TriggerEvent};
		pub use types::{PostgresType, PostgresTypeAttribute};
Original file line number	Diff line number	Diff line change
Expand Up		@@ -6,13 +6,13 @@ use crate::schema_cache::SchemaCacheItem;

		#[derive(Debug, Clone, Default)]
		pub struct TypeAttributes {
		attrs: Vec<PostgresTypeAttribute>,
		pubattrs: Vec<PostgresTypeAttribute>,
		}

		#[derive(Debug, Clone, Default, Deserialize)]
		pub struct PostgresTypeAttribute {
		name: String,
		type_id: i64,
		pubname: String,
		pubtype_id: i64,
		}

		impl From<Option<JsonValue>> for TypeAttributes {
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -70,7 +70,7 @@ mod tests {

		use crate::{
		TreeSitterQueriesExecutor,
		queries::{RelationMatch, TableAliasMatch},
		queries::{ParameterMatch,RelationMatch, TableAliasMatch},
		};

		#[test]
Expand DownExpand Up		@@ -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
Expand DownExpand Up		@@ -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");
		}
		}
Original file line number	Diff line number	Diff line change
		@@ -1,14 +1,17 @@
		mod parameters;
		mod relations;
		mod select_columns;
		mod table_aliases;

		pub use parameters::*;
		pub use relations::*;
		pub use select_columns::*;
		pub use table_aliases::*;

		#[derive(Debug)]
		pub enum QueryResult<'a> {
		Relation(RelationMatch<'a>),
		Parameter(ParameterMatch<'a>),
		TableAliases(TableAliasMatch<'a>),
		SelectClauseColumns(SelectColumnMatch<'a>),
		}
Expand All		@@ -26,6 +29,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();
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -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()
		}
		}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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() {
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -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)]
Expand DownExpand Up		@@ -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),
		}
Expand Down