From c98a67cd33bc9a600ad15b92c85dfff178aecb0b Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 30 Oct 2025 11:58:30 +0100 Subject: [PATCH 1/5] feat(lsp): add code action to invalidate the schema cache --- crates/pgls_lsp/src/handlers/code_actions.rs | 29 +++- crates/pgls_lsp/src/server.rs | 1 + crates/pgls_lsp/tests/server.rs | 156 +++++++++++++++++- .../src/features/code_actions.rs | 1 + crates/pgls_workspace/src/workspace.rs | 8 + crates/pgls_workspace/src/workspace/client.rs | 4 + crates/pgls_workspace/src/workspace/server.rs | 13 ++ .../workspace/server/schema_cache_manager.rs | 13 ++ 8 files changed, 222 insertions(+), 3 deletions(-) diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 5fd1546d9..04d8b243d 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -53,6 +53,21 @@ pub fn get_actions( .map(|reason| CodeActionDisabled { reason }), ..Default::default() }), + CommandActionCategory::InvalidateSchemaCache => Some(CodeAction { + title: title.clone(), + kind: Some(lsp_types::CodeActionKind::EMPTY), + command: Some({ + Command { + title: title.clone(), + command: command_id, + arguments: None, + } + }), + disabled: action + .disabled_reason + .map(|reason| CodeActionDisabled { reason }), + ..Default::default() + }), } } @@ -68,7 +83,8 @@ pub fn get_actions( pub fn command_id(command: &CommandActionCategory) -> String { match command { - CommandActionCategory::ExecuteStatement(_) => "pgt.executeStatement".into(), + CommandActionCategory::ExecuteStatement(_) => "pgls.executeStatement".into(), + CommandActionCategory::InvalidateSchemaCache => "pgls.invalidateSchemaCache".into(), } } @@ -80,7 +96,7 @@ pub async fn execute_command( let command = params.command; match command.as_str() { - "pgt.executeStatement" => { + "pgls.executeStatement" => { let statement_id = serde_json::from_value::( params.arguments[0].clone(), )?; @@ -105,7 +121,16 @@ pub async fn execute_command( Ok(None) } + "pgls.invalidateSchemaCache" => { + session.workspace.invalidate_schema_cache(true)?; + session + .client + .show_message(MessageType::INFO, "Schema cache invalidated") + .await; + + Ok(None) + } any => Err(anyhow!(format!("Unknown command: {}", any))), } } diff --git a/crates/pgls_lsp/src/server.rs b/crates/pgls_lsp/src/server.rs index 18f38f007..b5bc92ac2 100644 --- a/crates/pgls_lsp/src/server.rs +++ b/crates/pgls_lsp/src/server.rs @@ -461,6 +461,7 @@ impl ServerFactory { workspace_method!(builder, get_completions); workspace_method!(builder, register_project_folder); workspace_method!(builder, unregister_project_folder); + workspace_method!(builder, invalidate_schema_cache); let (service, socket) = builder.finish(); ServerConnection { socket, service } diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index e4798c536..b02a4b3ea 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -916,7 +916,7 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { .find_map(|action_or_cmd| match action_or_cmd { lsp::CodeActionOrCommand::CodeAction(code_action) => { let command = code_action.command.as_ref(); - if command.is_some_and(|cmd| &cmd.command == "pgt.executeStatement") { + if command.is_some_and(|cmd| &cmd.command == "pgls.executeStatement") { let command = command.unwrap(); let arguments = command.arguments.as_ref().unwrap().clone(); Some((command.command.clone(), arguments)) @@ -952,6 +952,160 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { Ok(()) } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_invalidate_schema_cache(test_db: PgPool) -> Result<()> { + let factory = ServerFactory::default(); + let mut fs = MemoryFileSystem::default(); + + let database = test_db + .connect_options() + .get_database() + .unwrap() + .to_string(); + let host = test_db.connect_options().get_host().to_string(); + + // Setup: Create a table with only id column (no name column yet) + let setup = r#" + create table public.users ( + id serial primary key + ); + "#; + + test_db + .execute(setup) + .await + .expect("Failed to setup test database"); + + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + database: Some(database), + host: Some(host), + ..Default::default() + }), + ..Default::default() + }); + + fs.insert( + url!("postgres-language-server.jsonc") + .to_file_path() + .unwrap(), + serde_json::to_string_pretty(&conf).unwrap(), + ); + + let (service, client) = factory + .create_with_fs(None, DynRef::Owned(Box::new(fs))) + .into_inner(); + + let (stream, sink) = client.split(); + let mut server = Server::new(service); + + let (sender, mut receiver) = channel(CHANNEL_BUFFER_SIZE); + let reader = tokio::spawn(client_handler(stream, sink, sender)); + + server.initialize().await?; + server.initialized().await?; + + server.load_configuration().await?; + + // Open a document that references a non-existent 'name' column + let doc_content = "select name from public.users;"; + server.open_document(doc_content).await?; + + // Wait for typecheck diagnostics showing column doesn't exist + let got_error = tokio::time::timeout(Duration::from_secs(5), async { + loop { + match receiver.next().await { + Some(ServerNotification::PublishDiagnostics(msg)) => { + if msg + .diagnostics + .iter() + .any(|d| d.message.contains("column \"name\" does not exist")) + { + return true; + } + } + _ => continue, + } + } + }) + .await + .is_ok(); + + assert!( + got_error, + "Expected typecheck error for non-existent column 'name'" + ); + + // Add the missing column to the database + let alter_table = r#" + alter table public.users + add column name text; + "#; + + test_db + .execute(alter_table) + .await + .expect("Failed to add column to table"); + + // Invalidate the schema cache (all = false for current connection only) + server + .request::("pgt/invalidate_schema_cache", "_invalidate_cache", false) + .await?; + + // Change the document slightly to trigger re-analysis + server + .change_document( + 1, + vec![TextDocumentContentChangeEvent { + range: Some(Range { + start: Position { + line: 0, + character: 30, + }, + end: Position { + line: 0, + character: 30, + }, + }), + range_length: Some(0), + text: " ".to_string(), + }], + ) + .await?; + + // Wait for diagnostics to clear (no typecheck error anymore) + let error_cleared = tokio::time::timeout(Duration::from_secs(5), async { + loop { + match receiver.next().await { + Some(ServerNotification::PublishDiagnostics(msg)) => { + // Check that there's no typecheck error for the column + let has_column_error = msg + .diagnostics + .iter() + .any(|d| d.message.contains("column \"name\" does not exist")); + if !has_column_error { + return true; + } + } + _ => continue, + } + } + }) + .await + .is_ok(); + + assert!( + error_cleared, + "Expected typecheck error to be cleared after schema cache invalidation" + ); + + server.shutdown().await?; + reader.abort(); + + Ok(()) +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_issue_281(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); diff --git a/crates/pgls_workspace/src/features/code_actions.rs b/crates/pgls_workspace/src/features/code_actions.rs index 55a3d0ca0..fb16e49d2 100644 --- a/crates/pgls_workspace/src/features/code_actions.rs +++ b/crates/pgls_workspace/src/features/code_actions.rs @@ -48,6 +48,7 @@ pub struct CommandAction { #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub enum CommandActionCategory { ExecuteStatement(StatementId), + InvalidateSchemaCache, } #[derive(Debug, serde::Serialize, serde::Deserialize)] diff --git a/crates/pgls_workspace/src/workspace.rs b/crates/pgls_workspace/src/workspace.rs index dda370a48..d82ca5058 100644 --- a/crates/pgls_workspace/src/workspace.rs +++ b/crates/pgls_workspace/src/workspace.rs @@ -158,6 +158,14 @@ pub trait Workspace: Send + Sync + RefUnwindSafe { &self, params: ExecuteStatementParams, ) -> Result; + + /// Invalidate the schema cache. + /// + /// # Arguments + /// * `all` - If true, clears all cached schemas. If false, clears only the current connection's cache. + /// + /// The schema will be reloaded lazily on the next operation that requires it. + fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError>; } /// Convenience function for constructing a server instance of [Workspace] diff --git a/crates/pgls_workspace/src/workspace/client.rs b/crates/pgls_workspace/src/workspace/client.rs index 70f7c20a3..431e029fd 100644 --- a/crates/pgls_workspace/src/workspace/client.rs +++ b/crates/pgls_workspace/src/workspace/client.rs @@ -168,4 +168,8 @@ where ) -> Result { self.request("pgt/on_hover", params) } + + fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError> { + self.request("pgt/invalidate_schema_cache", all) + } } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index f7ae3225d..43c790d53 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -424,6 +424,19 @@ impl Workspace for WorkspaceServer { }) } + fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError> { + if all { + self.schema_cache.clear_all(); + } else { + // Only clear current connection if one exists + if let Some(pool) = self.get_current_connection() { + self.schema_cache.clear(&pool); + } + // If no connection, nothing to clear - just return Ok + } + Ok(()) + } + #[ignored_path(path=¶ms.path)] fn pull_diagnostics( &self, diff --git a/crates/pgls_workspace/src/workspace/server/schema_cache_manager.rs b/crates/pgls_workspace/src/workspace/server/schema_cache_manager.rs index 9786946bf..fe7ee41c4 100644 --- a/crates/pgls_workspace/src/workspace/server/schema_cache_manager.rs +++ b/crates/pgls_workspace/src/workspace/server/schema_cache_manager.rs @@ -46,4 +46,17 @@ impl SchemaCacheManager { schemas.insert(key, schema_cache.clone()); Ok(schema_cache) } + + /// Clear the schema cache for a specific connection + pub fn clear(&self, pool: &PgPool) { + let key: ConnectionKey = pool.into(); + let mut schemas = self.schemas.write().unwrap(); + schemas.remove(&key); + } + + /// Clear all schema caches + pub fn clear_all(&self) { + let mut schemas = self.schemas.write().unwrap(); + schemas.clear(); + } } From 9842753ff1e75e98bf4783b4aab3aea951acd572 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 30 Oct 2025 12:03:04 +0100 Subject: [PATCH 2/5] progress --- crates/pgls_lsp/tests/server.rs | 136 ++++++++++++++++---------------- 1 file changed, 70 insertions(+), 66 deletions(-) diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index b02a4b3ea..fd88774b6 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -1000,7 +1000,7 @@ async fn test_invalidate_schema_cache(test_db: PgPool) -> Result<()> { let (stream, sink) = client.split(); let mut server = Server::new(service); - let (sender, mut receiver) = channel(CHANNEL_BUFFER_SIZE); + let (sender, _receiver) = channel(CHANNEL_BUFFER_SIZE); let reader = tokio::spawn(client_handler(stream, sink, sender)); server.initialize().await?; @@ -1008,33 +1008,46 @@ async fn test_invalidate_schema_cache(test_db: PgPool) -> Result<()> { server.load_configuration().await?; - // Open a document that references a non-existent 'name' column - let doc_content = "select name from public.users;"; + // Open a document to get completions from + let doc_content = "select from public.users;"; server.open_document(doc_content).await?; - // Wait for typecheck diagnostics showing column doesn't exist - let got_error = tokio::time::timeout(Duration::from_secs(5), async { - loop { - match receiver.next().await { - Some(ServerNotification::PublishDiagnostics(msg)) => { - if msg - .diagnostics - .iter() - .any(|d| d.message.contains("column \"name\" does not exist")) - { - return true; - } - } - _ => continue, - } - } - }) - .await - .is_ok(); + // Get completions before adding the column - 'name' should NOT be present + let completions_before = server + .get_completion(CompletionParams { + work_done_progress_params: WorkDoneProgressParams::default(), + partial_result_params: PartialResultParams::default(), + context: None, + text_document_position: TextDocumentPositionParams { + text_document: TextDocumentIdentifier { + uri: url!("document.sql"), + }, + position: Position { + line: 0, + character: 7, + }, + }, + }) + .await? + .unwrap(); + + let items_before = match completions_before { + CompletionResponse::Array(ref a) => a, + CompletionResponse::List(ref l) => &l.items, + }; + + let has_name_before = items_before.iter().any(|item| { + item.label == "name" + && item.label_details.as_ref().is_some_and(|d| { + d.description + .as_ref() + .is_some_and(|desc| desc.contains("public.users")) + }) + }); assert!( - got_error, - "Expected typecheck error for non-existent column 'name'" + !has_name_before, + "Column 'name' should not be in completions before it's added to the table" ); // Add the missing column to the database @@ -1053,51 +1066,42 @@ async fn test_invalidate_schema_cache(test_db: PgPool) -> Result<()> { .request::("pgt/invalidate_schema_cache", "_invalidate_cache", false) .await?; - // Change the document slightly to trigger re-analysis - server - .change_document( - 1, - vec![TextDocumentContentChangeEvent { - range: Some(Range { - start: Position { - line: 0, - character: 30, - }, - end: Position { - line: 0, - character: 30, - }, - }), - range_length: Some(0), - text: " ".to_string(), - }], - ) - .await?; + // Get completions after invalidating cache - 'name' should NOW be present + let completions_after = server + .get_completion(CompletionParams { + work_done_progress_params: WorkDoneProgressParams::default(), + partial_result_params: PartialResultParams::default(), + context: None, + text_document_position: TextDocumentPositionParams { + text_document: TextDocumentIdentifier { + uri: url!("document.sql"), + }, + position: Position { + line: 0, + character: 7, + }, + }, + }) + .await? + .unwrap(); - // Wait for diagnostics to clear (no typecheck error anymore) - let error_cleared = tokio::time::timeout(Duration::from_secs(5), async { - loop { - match receiver.next().await { - Some(ServerNotification::PublishDiagnostics(msg)) => { - // Check that there's no typecheck error for the column - let has_column_error = msg - .diagnostics - .iter() - .any(|d| d.message.contains("column \"name\" does not exist")); - if !has_column_error { - return true; - } - } - _ => continue, - } - } - }) - .await - .is_ok(); + let items_after = match completions_after { + CompletionResponse::Array(ref a) => a, + CompletionResponse::List(ref l) => &l.items, + }; + + let has_name_after = items_after.iter().any(|item| { + item.label == "name" + && item.label_details.as_ref().is_some_and(|d| { + d.description + .as_ref() + .is_some_and(|desc| desc.contains("public.users")) + }) + }); assert!( - error_cleared, - "Expected typecheck error to be cleared after schema cache invalidation" + has_name_after, + "Column 'name' should be in completions after schema cache invalidation" ); server.shutdown().await?; From 2133bdaa95e39f5ef3c9836703cd41d6bde2e09a Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 30 Oct 2025 12:06:11 +0100 Subject: [PATCH 3/5] progress --- crates/pgls_workspace/src/workspace/server.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 43c790d53..7d6f2f2d0 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -358,7 +358,7 @@ impl Workspace for WorkspaceServer { None => Some("Statement execution not allowed against database.".into()), }; - let actions = parser + let mut actions: Vec = parser .iter_with_filter( DefaultMapper, CursorPositionFilter::new(params.cursor_position), @@ -379,6 +379,21 @@ impl Workspace for WorkspaceServer { }) .collect(); + // Add reload schema cache action + let reload_disabled_reason = if self.get_current_connection().is_some() { + None + } else { + Some("No database connection available.".into()) + }; + + actions.push(CodeAction { + title: "Reload Schema Cache".into(), + kind: CodeActionKind::Command(CommandAction { + category: CommandActionCategory::InvalidateSchemaCache, + }), + disabled_reason: reload_disabled_reason, + }); + Ok(CodeActionsResult { actions }) } From ba3defd4594e458ed35beccb691621f2e9c1f9de Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 30 Oct 2025 12:10:27 +0100 Subject: [PATCH 4/5] progress --- crates/pgls_lsp/src/handlers/code_actions.rs | 2 +- crates/pgls_workspace/src/workspace/server.rs | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 04d8b243d..388191a96 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -122,7 +122,7 @@ pub async fn execute_command( Ok(None) } "pgls.invalidateSchemaCache" => { - session.workspace.invalidate_schema_cache(true)?; + session.workspace.invalidate_schema_cache(false)?; session .client diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 7d6f2f2d0..a3a388681 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -379,19 +379,18 @@ impl Workspace for WorkspaceServer { }) .collect(); - // Add reload schema cache action - let reload_disabled_reason = if self.get_current_connection().is_some() { + let invalidate_disabled_reason = if self.get_current_connection().is_some() { None } else { Some("No database connection available.".into()) }; actions.push(CodeAction { - title: "Reload Schema Cache".into(), + title: "Invalidate Schema Cache".into(), kind: CodeActionKind::Command(CommandAction { category: CommandActionCategory::InvalidateSchemaCache, }), - disabled_reason: reload_disabled_reason, + disabled_reason: invalidate_disabled_reason, }); Ok(CodeActionsResult { actions }) From 92c681560152dfa214a7ae366632eb5ea7ce89cb Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 30 Oct 2025 12:16:14 +0100 Subject: [PATCH 5/5] progress --- crates/pgls_lsp/src/handlers/code_actions.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 388191a96..04d8b243d 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -122,7 +122,7 @@ pub async fn execute_command( Ok(None) } "pgls.invalidateSchemaCache" => { - session.workspace.invalidate_schema_cache(false)?; + session.workspace.invalidate_schema_cache(true)?; session .client