From 059fefac7b540f01e8716a2acdd81e3ac9b53899 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Fri, 23 Oct 2026 13:55:59 +0000 Subject: [PATCH 1/4] fix(workspace): stop the statement iterator from overflowing the stack ParseIterator::next recursed once for every statement the filter skipped, so hover, code actions and completions overflowed the 2 MiB stack of the tokio workers on files with a few thousand statements and aborted the server. It is a loop now. Hover and code actions only look at the top-level statement at the cursor and its SQL function body instead of filtering every statement of the file. --- crates/pgls_workspace/src/workspace/server.rs | 17 +--- .../src/workspace/server.tests.rs | 76 +++++++++++++++++ .../src/workspace/server/document.rs | 84 +++++++++++++++---- 3 files changed, 145 insertions(+), 32 deletions(-) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index b186db9f2..831235944 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -14,9 +14,7 @@ use async_helper::run_async; use connection_manager::ConnectionManager; #[cfg(feature = "db")] use document::TypecheckDiagnosticsMapper; -use document::{ - CursorPositionFilter, DefaultMapper, Document, ExecuteStatementMapper, FormatStatementMapper, -}; +use document::{DefaultMapper, Document, ExecuteStatementMapper, FormatStatementMapper}; #[cfg(feature = "db")] use futures::{StreamExt, TryStreamExt, stream}; #[cfg(feature = "db")] @@ -1006,10 +1004,7 @@ impl Workspace for WorkspaceServer { }; let mut actions: Vec = parser - .iter_with_filter( - DefaultMapper, - CursorPositionFilter::new(params.cursor_position), - ) + .iter_at(DefaultMapper, params.cursor_position) .map(|(stmt, _, txt)| { let title = format!( "Execute Statement: {}...", @@ -1758,13 +1753,7 @@ impl Workspace for WorkspaceServer { }; // The innermost statement: the body of a SQL function rather than its `CREATE FUNCTION`. - match doc - .iter_with_filter( - WithCSTandASTMapper, - CursorPositionFilter::new(params.position), - ) - .last() - { + match doc.iter_at(WithCSTandASTMapper, params.position).last() { Some((stmt_id, range, ts_tree, maybe_ast)) => { let position_in_stmt = params.position - range.start(); let references = self.references_at(doc, &snapshot, params.position); diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 8f0e857a0..533f52252 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2122,6 +2122,82 @@ async fn test_hover_shows_resolved_references(test_db: PgPool) { assert_eq!(blocks.len(), 2, "{blocks:?}"); } +/// Hover and code actions on a file with thousands of statements, on a thread with the 2 MiB +/// stack of the language server's tokio workers. They used to recurse once per statement before +/// or after the cursor and overflow the stack, which aborts the server. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn hover_and_code_actions_on_long_files(test_db: PgPool) { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let workspace = get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + + let path = PgLSPath::new("test.sql"); + let content = format!( + "create table long_file_items (id int);\n{}", + "select id from long_file_items;\n".repeat(5_000) + ); + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.clone(), + version: 1, + }) + .expect("Unable to open test file"); + + let first = u32::try_from(content.find("id from").unwrap()).unwrap(); + let last = u32::try_from(content.rfind("id from").unwrap()).unwrap(); + std::thread::scope(|scope| { + std::thread::Builder::new() + .stack_size(2 * 1024 * 1024) + .spawn_scoped(scope, || { + for position in [first, last] { + let blocks = workspace + .on_hover(OnHoverParams { + path: path.clone(), + position: position.into(), + }) + .expect("Unable to hover") + .markdown_blocks; + assert_eq!(blocks.len(), 1, "{blocks:?}"); + assert!( + blocks[0].contains("public.long_file_items.id"), + "{blocks:?}" + ); + + let actions = workspace + .pull_code_actions(crate::features::code_actions::CodeActionsParams { + path: path.clone(), + cursor_position: position.into(), + only: vec![], + skip: vec![], + }) + .expect("Unable to pull code actions") + .actions; + let executes: Vec<_> = actions + .iter() + .filter(|action| action.title.starts_with("Execute Statement")) + .collect(); + assert_eq!(executes.len(), 1, "{actions:#?}"); + } + }) + .unwrap() + .join() + .unwrap(); + }); +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_inlay_hints_follow_migration_sections(test_db: PgPool) { let mut conf = PartialConfiguration::init(); diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 4dd5c91c5..87eee4774 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -92,6 +92,27 @@ impl Document { ParseIterator::new(self, mapper, filter) } + /// The statements that contain `position`, outermost first: the top-level statement, and the + /// body of a SQL function if `position` is in it. Unlike filtering all statements with a + /// [`CursorPositionFilter`], this doesn't look at the other statements of the document. + pub fn iter_at<'a, M>( + &'a self, + mapper: M, + position: TextSize, + ) -> ParseIterator<'a, M, CursorPositionFilter> + where + M: StatementMapper<'a>, + { + // The ranges are ordered and don't overlap, so only the first one that ends after + // `position` can contain it. + let first = self.ranges.partition_point(|range| range.end() <= position); + let ranges = &self.ranges[first..self.ranges.len().min(first + 1)]; + ParseIterator { + ranges: ranges.iter(), + ..ParseIterator::new(self, mapper, CursorPositionFilter::new(position)) + } + } + #[allow(dead_code)] pub fn count(&self) -> usize { self.iter(DefaultMapper).count() @@ -136,21 +157,18 @@ where type Item = M::Output; fn next(&mut self) -> Option { - // First check if we have any pending sub-statements to process - if let Some((id, range, content)) = self.pending_sub_statements.pop() { - if self.filter.predicate(&id, &range, content.as_str()) { - return Some(self.mapper.map(self.parser, id, range)); + // A loop rather than recursion: the filter can skip thousands of statements in a row. + loop { + // First check if we have any pending sub-statements to process + if let Some((id, range, content)) = self.pending_sub_statements.pop() { + if self.filter.predicate(&id, &range, content.as_str()) { + return Some(self.mapper.map(self.parser, id, range)); + } + continue; } - // If the sub-statement doesn't pass the filter, continue to the next item - return self.next(); - } - - // Process the next top-level statement - let next_range = self.ranges.next(); - - if let Some(range) = next_range { - // If we should include sub-statements and this statement has an AST + // Process the next top-level statement + let range = self.ranges.next()?; let content = &self.parser.content[*range]; let root_id = StatementId::new(content); @@ -171,12 +189,7 @@ where if self.filter.predicate(&root_id, range, content) { return Some(self.mapper.map(self.parser, root_id, *range)); } - - // If the current statement doesn't pass the filter, try the next one - return self.next(); } - - None } } @@ -702,6 +715,41 @@ $$ LANGUAGE plpgsql;"; assert_eq!(results2[0].2, "INSERT INTO users VALUES (1);"); } + #[test] + fn iter_at_yields_the_statements_at_the_position() { + let input = + "select 1;\ncreate function f() returns int language sql as 'select 2;';\nselect 3;"; + let d = Document::new(input.to_string(), 1); + let at = |needle: &str| { + let position = TextSize::try_from(input.find(needle).unwrap()).unwrap(); + d.iter_at(DefaultMapper, position) + .map(|(_, _, sql)| sql) + .collect::>() + }; + + assert_eq!(at("1;"), ["select 1;"]); + // The `CREATE FUNCTION`, then its body. + assert_eq!( + at("2;"), + [ + "create function f() returns int language sql as 'select 2;';", + "select 2;" + ] + ); + assert_eq!( + at("returns"), + ["create function f() returns int language sql as 'select 2;';"] + ); + assert_eq!(at("select 3"), ["select 3;"]); + // Between two statements, and past the end. + assert!(at("\ncreate").is_empty()); + assert!( + d.iter_at(DefaultMapper, TextSize::of(input)) + .next() + .is_none() + ); + } + #[test] fn test_id_filter() { let input = "SELECT 1; SELECT 2;"; From 2e0067085ce77c512af20a2098193a106469dc1f Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 24 Oct 2026 06:15:31 +0000 Subject: [PATCH 2/4] perf(workspace): keep every statement of a document cached and walk once for hover The parse and tree-sitter caches of a document held 1,000 statements, so on bigger files every walk over the document parsed every statement again. They now keep the statements of the document and drop the ones that are gone when the content changes. Hover walked the statements before the cursor twice: once to resolve the references, and once more in snapshot_before, which also deep-cloned every AST. It now builds the snapshot from the catalog of its own walk, skips it when the hovered references are all to objects the file didn't change, and caches it by the text before the statement. --- Cargo.lock | 23 +-- crates/pgls_workspace/Cargo.toml | 1 - crates/pgls_workspace/src/workspace/server.rs | 186 +++++++++++------- .../src/workspace/server.tests.rs | 77 ++++++++ .../src/workspace/server/annotation.rs | 18 +- .../src/workspace/server/document.rs | 36 ++++ .../src/workspace/server/pg_query.rs | 29 +-- .../src/workspace/server/tree_sitter.rs | 22 ++- 8 files changed, 263 insertions(+), 129 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9860413af..647215b80 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1160,12 +1160,6 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" -[[package]] -name = "foldhash" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" - [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1473,7 +1467,7 @@ checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ "allocator-api2", "equivalent", - "foldhash 0.1.5", + "foldhash", ] [[package]] @@ -1481,11 +1475,6 @@ name = "hashbrown" version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" -dependencies = [ - "allocator-api2", - "equivalent", - "foldhash 0.2.0", -] [[package]] name = "hashlink" @@ -1989,15 +1978,6 @@ dependencies = [ "logos-codegen", ] -[[package]] -name = "lru" -version = "0.16.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" -dependencies = [ - "hashbrown 0.16.1", -] - [[package]] name = "ls-types" version = "0.0.6" @@ -2990,7 +2970,6 @@ dependencies = [ "globset", "ignore", "json-strip-comments 3.1.0", - "lru", "pgls_analyse", "pgls_analyser", "pgls_catalog", diff --git a/crates/pgls_workspace/Cargo.toml b/crates/pgls_workspace/Cargo.toml index 280b3733c..bfdb1e9e5 100644 --- a/crates/pgls_workspace/Cargo.toml +++ b/crates/pgls_workspace/Cargo.toml @@ -13,7 +13,6 @@ version = "0.0.0" [dependencies] globset = "0.4.16" -lru = "0.16.3" ignore = { workspace = true } json-strip-comments = { workspace = true } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 831235944..8612f6cff 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -1,4 +1,5 @@ use std::{ + borrow::Borrow, collections::HashMap, fs, hash::{DefaultHasher, Hash, Hasher}, @@ -29,8 +30,8 @@ use pgls_analyser::{ Analyser, AnalyserConfig, AnalyserParams, AnalysisResult, Direction, FileKind, LinterDiagnostic, LinterOptions, MigrationSection, RuleFix, }; -use pgls_catalog::resolve::{Reference, ResolveParams, resolve}; -use pgls_catalog::{Catalog, CatalogBase, Session, TransactionMode}; +use pgls_catalog::resolve::{ColumnSource, Reference, ResolveParams, resolve}; +use pgls_catalog::{Catalog, CatalogBase, Origin, Session, TransactionMode}; use std::borrow::Cow; use pgls_catalog::Snapshot; @@ -109,13 +110,22 @@ pub struct WorkspaceServer { /// The indexed database snapshot the catalog starts from catalog_base: Mutex>>, - /// The last snapshot built for completions and hover, by the statements it includes - snapshot_before: Mutex)>>, + /// The last snapshot built for completions and hover + snapshot_before: Mutex>, #[cfg(feature = "db")] connection: ConnectionManager, } +/// A database snapshot with the changes of the statements before a statement. +struct SnapshotBefore { + /// The database snapshot it is built from. + base: Arc, + /// The text that the statement sees. + key: u64, + snapshot: Arc, +} + /// The `Workspace` object is long-lived, so we want it to be able to cross /// unwind boundaries. /// In return, we have to make sure operations on the workspace either do not @@ -463,14 +473,15 @@ impl WorkspaceServer { /// and the session follow the statements, so that each statement sees the objects created /// earlier in the file. Each migration section gets a session of its own, and the changes of /// a down section don't carry into the next section. `visit` gets every parsed statement with - /// the catalog and the session as they are before it. + /// the catalog and the session as they are before it. Returns the catalog after the walked + /// statements. fn walk_statements( &self, doc: &Document, snapshot: &Arc, end: TextSize, mut visit: impl FnMut(&pgls_analyser::AnalysableStatement, &Catalog, &Session), - ) { + ) -> Catalog { let (search_path, sections) = self .workspaces() .settings() @@ -541,33 +552,33 @@ impl WorkspaceServer { session.apply(&stmt.root); } } + catalog } - /// The references whose span contains `position`, in the innermost statement at `position` - /// (the body of a SQL function rather than its `CREATE FUNCTION`). - fn references_at( + /// For hover on the statement `id` at `range`: the references whose span contains + /// `position`, and the database snapshot as the statement sees it. + fn hover_references( &self, doc: &Document, - snapshot: &Arc, + id: &StatementId, + range: TextRange, position: TextSize, - ) -> Vec { - let mut innermost: Option<(TextRange, Vec)> = None; - self.walk_statements( + snapshot: Arc, + ) -> (Vec, Arc) { + let mut found = None; + let catalog = self.walk_statements( doc, - snapshot, + &snapshot, position + TextSize::from(1), |stmt, catalog, session| { let Some(sql) = stmt.sql.as_deref() else { return; }; - if !stmt.range.contains(position) - || innermost - .as_ref() - .is_some_and(|(range, _)| !range.contains_range(stmt.range)) - { + if stmt.range != range { return; } - let references = resolve(ResolveParams { + let position = position - range.start(); + let (references, others): (Vec<_>, Vec<_>) = resolve(ResolveParams { stmt: &stmt.root, catalog, search_path: session.search_path(), @@ -575,22 +586,30 @@ impl WorkspaceServer { sql: Some(sql), collect_references: true, }) - .references; - innermost = Some((stmt.range, references)); + .references + .into_iter() + .partition(|reference| { + reference + .span() + .is_some_and(|span| span.contains_inclusive(position)) + }); + // Objects that the file didn't change are the same in the database snapshot. + let snapshot = if in_database(&references, &others) { + Arc::clone(&snapshot) + } else { + self.snapshot_before(doc, id, range, Arc::clone(&snapshot), || catalog) + }; + found = Some((references, snapshot)); }, ); - let Some((range, references)) = innermost else { - return Vec::new(); - }; - let position = position - range.start(); - references - .into_iter() - .filter(|reference| { - reference - .span() - .is_some_and(|span| span.contains_inclusive(position)) - }) - .collect() + // A statement that doesn't parse is not visited. It changes nothing, so it sees the + // catalog after the walk. + found.unwrap_or_else(|| { + ( + Vec::new(), + self.snapshot_before(doc, id, range, snapshot, || catalog), + ) + }) } /// The indexed database snapshot, built once per snapshot. @@ -604,53 +623,41 @@ impl WorkspaceServer { base } - /// The database snapshot as the statement at `statement_start` sees it: with the objects - /// that the statements before it create, change, or drop. - fn snapshot_before( + /// The database snapshot as the statement `id` at `range` sees it: with the objects that the + /// statements before it create, change, or drop. `catalog` returns the catalog that the + /// statement sees, and is only called when the snapshot is not cached. + fn snapshot_before>( &self, doc: &Document, - statement_start: TextSize, + id: &StatementId, + range: TextRange, snapshot: Arc, + catalog: impl FnOnce() -> C, ) -> Arc { - let statements: Vec<_> = doc - .iter(ExecuteStatementMapper) - .take_while(|(_, range, _, _)| range.end() <= statement_start) - .filter(|(id, _, _, _)| !id.is_child()) - .filter_map(|(_, _, sql, ast)| Some((sql, ast?))) - .collect(); - if statements.is_empty() { - return snapshot; - } - - let search_path = self - .workspaces() - .settings() - .map(|settings| settings.typecheck.search_path.clone()) - .unwrap_or_default(); - + // The statement sees the text before it, and the body of a SQL function also sees its + // `CREATE FUNCTION`. let mut hasher = DefaultHasher::new(); - Arc::as_ptr(&snapshot).hash(&mut hasher); - search_path.hash(&mut hasher); - for (sql, _) in &statements { - sql.hash(&mut hasher); - } + doc.get_document_content()[..usize::from(range.start())].hash(&mut hasher); + id.parent_content().hash(&mut hasher); let key = hasher.finish(); let mut cached = self.snapshot_before.lock().unwrap(); - if let Some((cached_key, cached_snapshot)) = cached.as_ref() - && *cached_key == key + if let Some(cached) = cached.as_ref() + && cached.key == key + && Arc::ptr_eq(&cached.base, &snapshot) { - return Arc::clone(cached_snapshot); + return Arc::clone(&cached.snapshot); } - let mut catalog = Catalog::new(Some(self.catalog_base(&snapshot))); - let mut session = Session::new(snapshot.expand_search_path(&search_path)); - for (_, ast) in &statements { - catalog.apply(ast, session.search_path()); - session.apply(ast); - } - let result = catalog.snapshot().unwrap_or(snapshot); - *cached = Some((key, Arc::clone(&result))); + let result = catalog() + .borrow() + .snapshot() + .unwrap_or_else(|| Arc::clone(&snapshot)); + *cached = Some(SnapshotBefore { + base: snapshot, + key, + snapshot: Arc::clone(&result), + }); result } @@ -913,6 +920,8 @@ impl Workspace for WorkspaceServer { params.vcs_base_path, params.gitignore_matches.as_slice(), )?; + // The search path and the migration sections change what a statement sees. + *self.snapshot_before.lock().unwrap() = None; Ok(()) } @@ -1707,7 +1716,10 @@ impl Workspace for WorkspaceServer { } Some((id, range, cst)) => { let position = params.position - range.start(); - let snapshot = self.snapshot_before(parsed_doc, range.start(), snapshot); + let snapshot = + self.snapshot_before(parsed_doc, &id, range, Arc::clone(&snapshot), || { + self.walk_statements(parsed_doc, &snapshot, range.start(), |_, _, _| {}) + }); let items = pgls_completions::complete(pgls_completions::CompletionParams { position, @@ -1756,8 +1768,8 @@ impl Workspace for WorkspaceServer { match doc.iter_at(WithCSTandASTMapper, params.position).last() { Some((stmt_id, range, ts_tree, maybe_ast)) => { let position_in_stmt = params.position - range.start(); - let references = self.references_at(doc, &snapshot, params.position); - let snapshot = self.snapshot_before(doc, range.start(), snapshot); + let (references, snapshot) = + self.hover_references(doc, &stmt_id, range, params.position, snapshot); let markdown_blocks = pgls_hover::on_hover(pgls_hover::OnHoverParams { ts_tree: &ts_tree, @@ -1844,6 +1856,34 @@ impl Workspace for WorkspaceServer { } } +/// Whether hover finds what `references` resolve to in the database snapshot: objects that the +/// file didn't change, and columns of CTEs and subqueries. `others` are the other references of +/// the statement. +fn in_database(references: &[Reference], others: &[Reference]) -> bool { + // The statement references each relation that it takes columns from. + let unchanged_relation = |schema: &str, name: &str| { + references.iter().chain(others).any(|reference| { + matches!( + reference, + Reference::Relation { relation, .. } + if relation.schema == schema + && relation.name == name + && relation.origin == Origin::Database + ) + }) + }; + !references.is_empty() + && references.iter().all(|reference| match reference { + Reference::Relation { relation, .. } => relation.origin == Origin::Database, + Reference::Function { function, .. } => function.origin == Origin::Database, + Reference::Column { + source: ColumnSource::Relation { schema, name }, + .. + } => unchanged_relation(schema, name), + Reference::Column { .. } => true, + }) +} + /// The migration sections of `doc`, like the up and down migrations of dbmate and goose. fn document_sections(settings: &Settings, doc: &Document) -> Vec { settings diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 533f52252..e1255917a 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2122,6 +2122,83 @@ async fn test_hover_shows_resolved_references(test_db: PgPool) { assert_eq!(blocks.len(), 2, "{blocks:?}"); } +/// Hover shows the objects as the statement sees them: created earlier in the file, changed by +/// the file, or unchanged in the database, also after the file changes. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn hover_follows_the_ddl_before_the_statement(test_db: PgPool) { + test_db + .execute( + "create table hover_db_items (id int); + create table hover_db_tags (id int);", + ) + .await + .expect("Failed to setup test database"); + + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + ..Default::default() + }), + ..Default::default() + }); + let workspace = get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + + let path = PgLSPath::new("test.sql"); + let statements = "alter table hover_db_items add column extra int;\n\ + select id from hover_file_items;\n\ + select extra from hover_db_items;\n\ + select id from hover_db_tags;\n"; + let hover = |content: &str, version: i32, needle: &str| { + if version == 1 { + workspace.open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version, + }) + } else { + workspace.change_file(crate::workspace::ChangeFileParams { + path: path.clone(), + content: content.into(), + version, + }) + } + .expect("Unable to update test file"); + let offset = content.find(needle).expect("needle in content") + 2; + let blocks = workspace + .on_hover(OnHoverParams { + path: path.clone(), + position: u32::try_from(offset).unwrap().into(), + }) + .expect("Unable to hover") + .markdown_blocks; + assert_eq!(blocks.len(), 1, "{blocks:?}"); + blocks.into_iter().next().unwrap() + }; + + let content = format!("create table hover_file_items (id int);\n{statements}"); + let block = hover(&content, 1, "hover_file_items;"); + assert!(block.contains("- id: int4"), "{block}"); + assert!(!block.contains("- name"), "{block}"); + // A table of the database that the file changed. + let block = hover(&content, 1, "hover_db_items;"); + assert!(block.contains("- extra: int4"), "{block}"); + // A table of the database that the file didn't change. + let block = hover(&content, 1, "hover_db_tags;"); + assert!(block.contains("public.hover_db_tags"), "{block}"); + + // The statements before the hovered one changed. + let content = format!("create table hover_file_items (id int, name text);\n{statements}"); + let block = hover(&content, 2, "hover_file_items;"); + assert!(block.contains("- name: text"), "{block}"); +} + /// Hover and code actions on a file with thousands of statements, on a thread with the 2 MiB /// stack of the language server's tokio workers. They used to recurse once per statement before /// or after the cursor and overflow the stack, which aborts the server. diff --git a/crates/pgls_workspace/src/workspace/server/annotation.rs b/crates/pgls_workspace/src/workspace/server/annotation.rs index 3f2bb43ef..756deac2d 100644 --- a/crates/pgls_workspace/src/workspace/server/annotation.rs +++ b/crates/pgls_workspace/src/workspace/server/annotation.rs @@ -1,20 +1,17 @@ -use std::num::NonZeroUsize; use std::sync::{Arc, Mutex}; -use lru::LruCache; use pgls_lexer::SyntaxKind; +use rustc_hash::FxHashMap; use super::statement_identifier::StatementId; -const DEFAULT_CACHE_SIZE: usize = 1000; - #[derive(Debug, Clone, PartialEq, Eq)] pub struct StatementAnnotations { ends_with_semicolon: bool, } pub struct AnnotationStore { - db: Mutex>>, + db: Mutex>>, } const WHITESPACE_TOKENS: [SyntaxKind; 6] = [ @@ -29,12 +26,15 @@ const WHITESPACE_TOKENS: [SyntaxKind; 6] = [ impl AnnotationStore { pub fn new() -> AnnotationStore { AnnotationStore { - db: Mutex::new(LruCache::new( - NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(), - )), + db: Mutex::default(), } } + /// Drops the annotations of the statements for which `keep` returns `false`. + pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { + self.db.lock().unwrap().retain(|id, _| keep(id)); + } + #[allow(unused)] pub fn get_annotations( &self, @@ -58,7 +58,7 @@ impl AnnotationStore { ends_with_semicolon, }); - cache.put(statement_id.clone(), annotations.clone()); + cache.insert(statement_id.clone(), annotations.clone()); annotations } diff --git a/crates/pgls_workspace/src/workspace/server/document.rs b/crates/pgls_workspace/src/workspace/server/document.rs index 87eee4774..36716aa17 100644 --- a/crates/pgls_workspace/src/workspace/server/document.rs +++ b/crates/pgls_workspace/src/workspace/server/document.rs @@ -5,6 +5,7 @@ use pgls_diagnostics::{Diagnostic, DiagnosticExt, serde::Diagnostic as SDiagnost use pgls_query_ext::diagnostics::SyntaxDiagnostic; use pgls_suppressions::Suppressions; use pgls_text_size::{TextRange, TextSize}; +use rustc_hash::FxHashSet; use super::{ annotation::AnnotationStore, @@ -56,6 +57,19 @@ impl Document { self.ranges = ranges; self.diagnostics = diagnostics; self.suppressions = Suppressions::from(self.content.as_str()); + + // Keep what is cached for the statements that are still in the document. The bodies + // of SQL functions are kept with their function. + let statements: FxHashSet<&str> = self + .ranges + .iter() + .map(|range| &self.content[*range]) + .collect(); + let is_present = + |id: &StatementId| statements.contains(id.parent_content().unwrap_or(id.content())); + self.ast_db.retain(is_present); + self.cst_db.retain(is_present); + self.annotation_db.retain(is_present); } pub fn suppressions(&self) -> &Suppressions { @@ -750,6 +764,28 @@ $$ LANGUAGE plpgsql;"; ); } + #[test] + fn keeps_the_parsed_statements_that_are_still_in_the_document() { + let mut d = Document::new("select 1;\nselect 2;".to_string(), 1); + let first = StatementId::new("select 1;"); + let second = StatementId::new("select 2;"); + let first_ast = d.ast_db.get_or_cache_ast(&first); + let second_ast = d.ast_db.get_or_cache_ast(&second); + let first_tree = d.cst_db.get_or_cache_tree(&first); + + d.update_content("select 1;\nselect 3;".to_string(), 2); + + assert!(Arc::ptr_eq(&first_ast, &d.ast_db.get_or_cache_ast(&first))); + assert!(Arc::ptr_eq( + &first_tree, + &d.cst_db.get_or_cache_tree(&first) + )); + assert!(!Arc::ptr_eq( + &second_ast, + &d.ast_db.get_or_cache_ast(&second) + )); + } + #[test] fn test_id_filter() { let input = "SELECT 1; SELECT 2;"; diff --git a/crates/pgls_workspace/src/workspace/server/pg_query.rs b/crates/pgls_workspace/src/workspace/server/pg_query.rs index a17f3eaba..80f8c64e2 100644 --- a/crates/pgls_workspace/src/workspace/server/pg_query.rs +++ b/crates/pgls_workspace/src/workspace/server/pg_query.rs @@ -1,33 +1,34 @@ -use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock, Mutex}; -use lru::LruCache; pub use pgls_lexer::{convert_to_positional_params, convert_to_positional_params_with_metadata}; use pgls_query_ext::diagnostics::*; use pgls_text_size::TextRange; use regex::Regex; +use rustc_hash::FxHashMap; use super::statement_identifier::StatementId; -const DEFAULT_CACHE_SIZE: usize = 1000; - +/// The parse results of the statements of a document. They are kept until the statement is +/// no longer in the document, see [`PgQueryStore::retain`]. pub struct PgQueryStore { - ast_db: Mutex>>>, - plpgsql_db: Mutex>>, + ast_db: Mutex>>>, + plpgsql_db: Mutex>>, } impl PgQueryStore { pub fn new() -> PgQueryStore { PgQueryStore { - ast_db: Mutex::new(LruCache::new( - NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(), - )), - plpgsql_db: Mutex::new(LruCache::new( - NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(), - )), + ast_db: Mutex::default(), + plpgsql_db: Mutex::default(), } } + /// Drops the results of the statements for which `keep` returns `false`. + pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { + self.ast_db.lock().unwrap().retain(|id, _| keep(id)); + self.plpgsql_db.lock().unwrap().retain(|id, _| keep(id)); + } + pub fn get_or_cache_ast( &self, statement: &StatementId, @@ -47,7 +48,7 @@ impl PgQueryStore { }) }), ); - cache.put(statement.clone(), r.clone()); + cache.insert(statement.clone(), r.clone()); r } @@ -90,7 +91,7 @@ impl PgQueryStore { }) .map_err(|e| SyntaxDiagnostic::new(e.to_string(), Some(range))); - cache.put(statement.clone(), r.clone()); + cache.insert(statement.clone(), r.clone()); Some(r) } diff --git a/crates/pgls_workspace/src/workspace/server/tree_sitter.rs b/crates/pgls_workspace/src/workspace/server/tree_sitter.rs index 2a63971a7..ccbb89881 100644 --- a/crates/pgls_workspace/src/workspace/server/tree_sitter.rs +++ b/crates/pgls_workspace/src/workspace/server/tree_sitter.rs @@ -1,14 +1,13 @@ -use std::num::NonZeroUsize; use std::sync::{Arc, Mutex}; -use lru::LruCache; +use rustc_hash::FxHashMap; use super::statement_identifier::StatementId; -const DEFAULT_CACHE_SIZE: usize = 1000; - +/// The tree-sitter trees of the statements of a document. They are kept until the statement is +/// no longer in the document, see [`TreeSitterStore::retain`]. pub struct TreeSitterStore { - db: Mutex>>, + db: Mutex>>, parser: Mutex, } @@ -20,17 +19,20 @@ impl TreeSitterStore { .expect("Error loading sql language"); TreeSitterStore { - db: Mutex::new(LruCache::new( - NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap(), - )), + db: Mutex::default(), parser: Mutex::new(parser), } } + /// Drops the trees of the statements for which `keep` returns `false`. + pub fn retain(&self, keep: impl Fn(&StatementId) -> bool) { + self.db.lock().unwrap().retain(|id, _| keep(id)); + } + pub fn get_or_cache_tree(&self, statement: &StatementId) -> Arc { // First check cache { - let mut cache = self.db.lock().unwrap(); + let cache = self.db.lock().unwrap(); if let Some(existing) = cache.get(statement) { return existing.clone(); } @@ -48,7 +50,7 @@ impl TreeSitterStore { if let Some(existing) = cache.get(statement) { return existing.clone(); } - cache.put(statement.clone(), tree.clone()); + cache.insert(statement.clone(), tree.clone()); } tree From 1adefada37f6273ef0f1ac794f9fb91d439efde0 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 26 Oct 2026 08:40:27 +0000 Subject: [PATCH 3/4] Format --- crates/pgls_workspace/src/workspace/server.tests.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 70f1c8a08..997cab9e9 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2309,7 +2309,8 @@ async fn hover_follows_the_ddl_before_the_statement(test_db: PgPool) { }), ..Default::default() }); - let (workspace, project_key) = get_test_workspace(Some(conf)).expect("Unable to create test workspace"); + let (workspace, project_key) = + get_test_workspace(Some(conf)).expect("Unable to create test workspace"); let path = PgLSPath::new("test.sql"); let statements = "alter table hover_db_items add column extra int;\n\ From 1681e8a714adef7550659484533337fec49bb02b Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 26 Oct 2026 10:36:24 +0000 Subject: [PATCH 4/4] Hover: types, enum literals and parameters in the database check --- crates/pgls_workspace/src/workspace/server.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index c88d6a923..0be21bbc6 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -2099,6 +2099,11 @@ fn in_database(references: &[Reference], others: &[Reference]) -> bool { .. } => unchanged_relation(schema, name), Reference::Column { .. } => true, + Reference::Type { type_info, .. } | Reference::EnumLiteral { type_info, .. } => { + type_info.origin == Origin::Database + } + // Parameters belong to a function of the file. + Reference::Parameter { .. } => false, // The columns of a `*` can come from objects that the file changed. Reference::Star { .. } => false, })