feat(lsp): add code action to invalidate the schema cache

Author: psteinroe

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::<pgls_workspace::workspace::StatementId>(
                 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))),
     }
 }

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 }

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::<bool, ()>("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();

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)]

crates/pgls_workspace/src/workspace.rs

@@ -158,6 +158,14 @@ pub trait Workspace: Send + Sync + RefUnwindSafe {
         &self,
         params: ExecuteStatementParams,
     ) -> Result<ExecuteStatementResult, WorkspaceError>;
+
+    /// 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]

crates/pgls_workspace/src/workspace/client.rs

@@ -168,4 +168,8 @@ where
     ) -> Result<crate::features::on_hover::OnHoverResult, WorkspaceError> {
         self.request("pgt/on_hover", params)
     }
+
+    fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError> {
+        self.request("pgt/invalidate_schema_cache", all)
+    }
 }

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=&params.path)]
     fn pull_diagnostics(
         &self,

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();
+    }
 }