From 62aac65220f41ec6b92631d9a1d30aa92f7d9199 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 12 Apr 2025 21:13:24 +0200 Subject: [PATCH 1/5] feat: annotations --- Cargo.lock | 1 + crates/pgt_workspace/Cargo.toml | 1 + crates/pgt_workspace/src/workspace/server.rs | 1 + .../src/workspace/server/annotation.rs | 88 +++++++++++++++++++ .../src/workspace/server/change.rs | 6 +- .../src/workspace/server/parsed_document.rs | 4 + .../workspace/server/statement_identifier.rs | 13 ++- 7 files changed, 108 insertions(+), 6 deletions(-) create mode 100644 crates/pgt_workspace/src/workspace/server/annotation.rs diff --git a/Cargo.lock b/Cargo.lock index 79ec52f0..4a1f6ea5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2765,6 +2765,7 @@ dependencies = [ "pgt_console", "pgt_diagnostics", "pgt_fs", + "pgt_lexer", "pgt_query_ext", "pgt_schema_cache", "pgt_statement_splitter", diff --git a/crates/pgt_workspace/Cargo.toml b/crates/pgt_workspace/Cargo.toml index 7df42b19..5f598b2d 100644 --- a/crates/pgt_workspace/Cargo.toml +++ b/crates/pgt_workspace/Cargo.toml @@ -25,6 +25,7 @@ pgt_configuration = { workspace = true } pgt_console = { workspace = true } pgt_diagnostics = { workspace = true } pgt_fs = { workspace = true, features = ["serde"] } +pgt_lexer = { workspace = true } pgt_query_ext = { workspace = true } pgt_schema_cache = { workspace = true } pgt_statement_splitter = { workspace = true } diff --git a/crates/pgt_workspace/src/workspace/server.rs b/crates/pgt_workspace/src/workspace/server.rs index 27f5e8be..5e33bc27 100644 --- a/crates/pgt_workspace/src/workspace/server.rs +++ b/crates/pgt_workspace/src/workspace/server.rs @@ -43,6 +43,7 @@ use super::{ pub use statement_identifier::StatementId; mod analyser; +mod annotation; mod async_helper; mod change; mod db_connection; diff --git a/crates/pgt_workspace/src/workspace/server/annotation.rs b/crates/pgt_workspace/src/workspace/server/annotation.rs new file mode 100644 index 00000000..0d5a8686 --- /dev/null +++ b/crates/pgt_workspace/src/workspace/server/annotation.rs @@ -0,0 +1,88 @@ +use std::sync::Arc; + +use dashmap::DashMap; +use pgt_lexer::{SyntaxKind, WHITESPACE_TOKENS}; + +use super::statement_identifier::StatementId; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StatementAnnotations { + ends_with_semicolon: bool, +} + +pub struct AnnotationStore { + db: DashMap>>, +} + +impl AnnotationStore { + pub fn new() -> AnnotationStore { + AnnotationStore { db: DashMap::new() } + } + + #[allow(unused)] + pub fn get_annotations( + &self, + statement: &StatementId, + content: &str, + ) -> Option> { + if let Some(existing) = self.db.get(statement).map(|x| x.clone()) { + return existing; + } + + // we swallow the error here because the lexing within the document would have already + // thrown and we wont even get here if that happened. + let annotations = pgt_lexer::lex(content).ok().map(|tokens| { + let ends_with_semicolon = tokens + .iter() + .rev() + .find(|token| !WHITESPACE_TOKENS.contains(&token.kind)) + .is_some_and(|token| token.kind == SyntaxKind::Ascii59); + + Arc::new(StatementAnnotations { + ends_with_semicolon, + }) + }); + + self.db.insert(statement.clone(), None); + annotations + } + + #[allow(unused)] + pub fn clear_statement(&self, id: &StatementId) { + self.db.remove(id); + + if let Some(child_id) = id.get_child_id() { + self.db.remove(&child_id); + } + } +} + +#[cfg(test)] +mod tests { + use crate::workspace::StatementId; + + use super::AnnotationStore; + + #[test] + fn annotates_correctly() { + let store = AnnotationStore::new(); + + let test_cases = [ + ("SELECT * FROM foo", false), + ("SELECT * FROM foo;", true), + ("SELECT * FROM foo ;", true), + ("SELECT * FROM foo ; ", true), + ("SELECT * FROM foo ;\n", true), + ("SELECT * FROM foo\n", false), + ]; + + for (idx, (content, expected)) in test_cases.iter().enumerate() { + let statement_id = StatementId::Root(idx.into()); + + let annotations = store.get_annotations(&statement_id, content); + + assert!(annotations.is_some()); + assert_eq!(annotations.unwrap().ends_with_semicolon, *expected); + } + } +} diff --git a/crates/pgt_workspace/src/workspace/server/change.rs b/crates/pgt_workspace/src/workspace/server/change.rs index afe0eb64..69c68189 100644 --- a/crates/pgt_workspace/src/workspace/server/change.rs +++ b/crates/pgt_workspace/src/workspace/server/change.rs @@ -409,7 +409,7 @@ mod tests { use pgt_diagnostics::Diagnostic; use pgt_text_size::TextRange; - use crate::workspace::{ChangeFileParams, ChangeParams, server::statement_identifier::root_id}; + use crate::workspace::{ChangeFileParams, ChangeParams}; use pgt_fs::PgTPath; @@ -886,14 +886,14 @@ mod tests { assert_eq!( changed[2], StatementChange::Added(AddedStatement { - stmt: StatementId::Root(root_id(2)), + stmt: StatementId::Root(2.into()), text: "select id,test from users".to_string() }) ); assert_eq!( changed[3], StatementChange::Added(AddedStatement { - stmt: StatementId::Root(root_id(3)), + stmt: StatementId::Root(3.into()), text: "select 1;".to_string() }) ); diff --git a/crates/pgt_workspace/src/workspace/server/parsed_document.rs b/crates/pgt_workspace/src/workspace/server/parsed_document.rs index a110fb1f..affc6fe3 100644 --- a/crates/pgt_workspace/src/workspace/server/parsed_document.rs +++ b/crates/pgt_workspace/src/workspace/server/parsed_document.rs @@ -8,6 +8,7 @@ use pgt_text_size::{TextRange, TextSize}; use crate::workspace::ChangeFileParams; use super::{ + annotation::AnnotationStore, change::StatementChange, document::{Document, StatementIterator}, pg_query::PgQueryStore, @@ -24,6 +25,7 @@ pub struct ParsedDocument { ast_db: PgQueryStore, cst_db: TreeSitterStore, sql_fn_db: SQLFunctionBodyStore, + annoation_db: AnnotationStore, } impl ParsedDocument { @@ -33,6 +35,7 @@ impl ParsedDocument { let cst_db = TreeSitterStore::new(); let ast_db = PgQueryStore::new(); let sql_fn_db = SQLFunctionBodyStore::new(); + let annoation_db = AnnotationStore::new(); doc.iter().for_each(|(stmt, _, content)| { cst_db.add_statement(&stmt, content); @@ -44,6 +47,7 @@ impl ParsedDocument { ast_db, cst_db, sql_fn_db, + annoation_db, } } diff --git a/crates/pgt_workspace/src/workspace/server/statement_identifier.rs b/crates/pgt_workspace/src/workspace/server/statement_identifier.rs index 0739fb2f..66f177aa 100644 --- a/crates/pgt_workspace/src/workspace/server/statement_identifier.rs +++ b/crates/pgt_workspace/src/workspace/server/statement_identifier.rs @@ -6,9 +6,16 @@ pub struct RootId { inner: usize, } -#[cfg(test)] -pub fn root_id(inner: usize) -> RootId { - RootId { inner } +impl From for usize { + fn from(val: RootId) -> Self { + val.inner + } +} + +impl From for RootId { + fn from(inner: usize) -> Self { + RootId { inner } + } } #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] From 457940000aeb538a33b807fb703e6a7178ed9782 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Steinr=C3=B6tter?= Date: Sat, 12 Apr 2025 21:56:25 +0200 Subject: [PATCH 2/5] Update crates/pgt_workspace/src/workspace/server/parsed_document.rs Co-authored-by: Julian Domke <68325451+juleswritescode@users.noreply.github.com> --- crates/pgt_workspace/src/workspace/server/parsed_document.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/pgt_workspace/src/workspace/server/parsed_document.rs b/crates/pgt_workspace/src/workspace/server/parsed_document.rs index affc6fe3..b0e87080 100644 --- a/crates/pgt_workspace/src/workspace/server/parsed_document.rs +++ b/crates/pgt_workspace/src/workspace/server/parsed_document.rs @@ -25,7 +25,7 @@ pub struct ParsedDocument { ast_db: PgQueryStore, cst_db: TreeSitterStore, sql_fn_db: SQLFunctionBodyStore, - annoation_db: AnnotationStore, + annotation_db: AnnotationStore, } impl ParsedDocument { From 783bb102572cdd197c2fb481a28dcc5e93b7d5f5 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 12 Apr 2025 21:57:46 +0200 Subject: [PATCH 3/5] fix: only impl From in test --- .../pgt_workspace/src/workspace/server/statement_identifier.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/pgt_workspace/src/workspace/server/statement_identifier.rs b/crates/pgt_workspace/src/workspace/server/statement_identifier.rs index 66f177aa..8c02814d 100644 --- a/crates/pgt_workspace/src/workspace/server/statement_identifier.rs +++ b/crates/pgt_workspace/src/workspace/server/statement_identifier.rs @@ -6,12 +6,14 @@ pub struct RootId { inner: usize, } +#[cfg(test)] impl From for usize { fn from(val: RootId) -> Self { val.inner } } +#[cfg(test)] impl From for RootId { fn from(inner: usize) -> Self { RootId { inner } From b9a4dee5cc73e96e79971db8d16b4dc478206651 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 12 Apr 2025 22:17:48 +0200 Subject: [PATCH 4/5] fix: typo --- crates/pgt_workspace/src/workspace/server/parsed_document.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/pgt_workspace/src/workspace/server/parsed_document.rs b/crates/pgt_workspace/src/workspace/server/parsed_document.rs index b0e87080..ceb8388e 100644 --- a/crates/pgt_workspace/src/workspace/server/parsed_document.rs +++ b/crates/pgt_workspace/src/workspace/server/parsed_document.rs @@ -35,7 +35,7 @@ impl ParsedDocument { let cst_db = TreeSitterStore::new(); let ast_db = PgQueryStore::new(); let sql_fn_db = SQLFunctionBodyStore::new(); - let annoation_db = AnnotationStore::new(); + let annotation_db = AnnotationStore::new(); doc.iter().for_each(|(stmt, _, content)| { cst_db.add_statement(&stmt, content); @@ -47,7 +47,7 @@ impl ParsedDocument { ast_db, cst_db, sql_fn_db, - annoation_db, + annotation_db, } } From 027e4ab91ef8e5005a4eaa438a7274e85974b21b Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 12 Apr 2025 22:19:12 +0200 Subject: [PATCH 5/5] fix: add missing clear_statement calls --- crates/pgt_workspace/src/workspace/server/annotation.rs | 1 - crates/pgt_workspace/src/workspace/server/parsed_document.rs | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/pgt_workspace/src/workspace/server/annotation.rs b/crates/pgt_workspace/src/workspace/server/annotation.rs index 0d5a8686..321dd3ac 100644 --- a/crates/pgt_workspace/src/workspace/server/annotation.rs +++ b/crates/pgt_workspace/src/workspace/server/annotation.rs @@ -47,7 +47,6 @@ impl AnnotationStore { annotations } - #[allow(unused)] pub fn clear_statement(&self, id: &StatementId) { self.db.remove(id); diff --git a/crates/pgt_workspace/src/workspace/server/parsed_document.rs b/crates/pgt_workspace/src/workspace/server/parsed_document.rs index ceb8388e..2b64d24a 100644 --- a/crates/pgt_workspace/src/workspace/server/parsed_document.rs +++ b/crates/pgt_workspace/src/workspace/server/parsed_document.rs @@ -73,6 +73,7 @@ impl ParsedDocument { self.cst_db.remove_statement(s); self.ast_db.clear_statement(s); self.sql_fn_db.clear_statement(s); + self.annotation_db.clear_statement(s); } StatementChange::Modified(s) => { tracing::debug!( @@ -88,6 +89,7 @@ impl ParsedDocument { self.cst_db.modify_statement(s); self.ast_db.clear_statement(&s.old_stmt); self.sql_fn_db.clear_statement(&s.old_stmt); + self.annotation_db.clear_statement(&s.old_stmt); } } }