feat(lsp): refresh and settings for inlay hints and semantic tokens

README.md

@@ -20,7 +20,7 @@ This project provides a toolchain for Postgres development, built on Postgres' o
 
 The following features are available today:
 
-- [Autocompletion & Hover](https://pg-language-server.com/latest/features/editor_features/)
+- [Autocompletion, Hover, Inlay Hints & Semantic Highlighting](https://pg-language-server.com/latest/features/editor_features/)
 - [Syntax Diagnostics](https://pg-language-server.com/latest/features/syntax_diagnostics/)
 - [Type Checking](https://pg-language-server.com/latest/features/type_checking/)
 - [Formatting](https://pg-language-server.com/latest/features/formatting/)

crates/pgls_configuration/src/inlay_hints.rs

@@ -0,0 +1,20 @@
+use bpaf::Bpaf;
+use pgls_configuration_macros::{Merge, Partial};
+use serde::{Deserialize, Serialize};
+
+/// The configuration for inlay hints.
+#[derive(Clone, Debug, Deserialize, Eq, Partial, PartialEq, Serialize)]
+#[partial(derive(Bpaf, Clone, Eq, PartialEq, Merge))]
+#[partial(cfg_attr(feature = "schema", derive(schemars::JsonSchema)))]
+#[partial(serde(rename_all = "camelCase", default, deny_unknown_fields))]
+pub struct InlayHintsConfiguration {
+    /// if `false`, the language server sends no inlay hints. `true` by default
+    #[partial(bpaf(hide))]
+    pub enabled: bool,
+}
+
+impl Default for InlayHintsConfiguration {
+    fn default() -> Self {
+        Self { enabled: true }
+    }
+}

crates/pgls_configuration/src/lib.rs

@@ -10,11 +10,13 @@ pub mod database;
 pub mod diagnostics;
 pub mod files;
 pub mod format;
+pub mod inlay_hints;
 pub mod linter;
 pub mod migrations;
 pub mod pglinter;
 pub mod plpgsql_check;
 pub mod rules;
+pub mod semantic_tokens;
 pub mod splinter;
 pub mod typecheck;
 pub mod utils;
@@ -37,6 +39,9 @@ pub use format::{
     FormatConfiguration, FunctionArgumentGroups, IndentStyle as FormatIndentStyle, KeywordCase,
     PartialFormatConfiguration, partial_format_configuration,
 };
+use inlay_hints::{
+    InlayHintsConfiguration, PartialInlayHintsConfiguration, partial_inlay_hints_configuration,
+};
 pub use linter::{
     LinterConfiguration, LinterRuleSettings, PartialLinterConfiguration, Rules,
     partial_linter_configuration,
@@ -57,6 +62,10 @@ pub use rules::{
     AnalyzerGroup, RuleConfiguration, RuleFixConfiguration, RulePlainConfiguration, RuleSelector,
     RuleWithFixOptions, RuleWithOptions,
 };
+use semantic_tokens::{
+    PartialSemanticTokensConfiguration, SemanticTokensConfiguration,
+    partial_semantic_tokens_configuration,
+};
 use serde::{Deserialize, Serialize};
 use splinter::{
     PartialSplinterConfiguration, SplinterConfiguration, partial_splinter_configuration,
@@ -131,6 +140,20 @@ pub struct Configuration {
         bpaf(external(partial_database_configuration), optional, hide_usage)
     )]
     pub db: DatabaseConfiguration,
+
+    /// The configuration for inlay hints in editors
+    #[partial(
+        type,
+        bpaf(external(partial_inlay_hints_configuration), optional, hide_usage)
+    )]
+    pub inlay_hints: InlayHintsConfiguration,
+
+    /// The configuration for semantic highlighting in editors
+    #[partial(
+        type,
+        bpaf(external(partial_semantic_tokens_configuration), optional, hide_usage)
+    )]
+    pub semantic_tokens: SemanticTokensConfiguration,
 }
 
 impl PartialConfiguration {
@@ -197,6 +220,8 @@ impl PartialConfiguration {
                 conn_timeout_secs: Some(10),
                 disable_connection: Some(false),
             }),
+            inlay_hints: None,
+            semantic_tokens: None,
         }
     }
 

crates/pgls_configuration/src/semantic_tokens.rs

@@ -0,0 +1,21 @@
+use bpaf::Bpaf;
+use pgls_configuration_macros::{Merge, Partial};
+use serde::{Deserialize, Serialize};
+
+/// The configuration for semantic highlighting.
+#[derive(Clone, Debug, Deserialize, Eq, Partial, PartialEq, Serialize)]
+#[partial(derive(Bpaf, Clone, Eq, PartialEq, Merge))]
+#[partial(cfg_attr(feature = "schema", derive(schemars::JsonSchema)))]
+#[partial(serde(rename_all = "camelCase", default, deny_unknown_fields))]
+pub struct SemanticTokensConfiguration {
+    /// if `false`, the language server sends no semantic tokens, and editors keep the
+    /// highlighting of their grammar. `true` by default
+    #[partial(bpaf(hide))]
+    pub enabled: bool,
+}
+
+impl Default for SemanticTokensConfiguration {
+    fn default() -> Self {
+        Self { enabled: true }
+    }
+}

crates/pgls_lsp/src/configuration.rs

@@ -123,7 +123,9 @@ define_section_mask!(
     pglinter,
     typecheck,
     plpgsql_check,
-    db
+    db,
+    inlay_hints,
+    semantic_tokens
 );
 
 /// The sticky layer of configuration a client supplies at runtime.

crates/pgls_lsp/src/handlers/code_actions.rs

@@ -332,6 +332,8 @@ pub async fn execute_command(
         }
         "pgls.invalidateSchemaCache" => {
             session.workspace.invalidate_schema_cache(true)?;
+            // Everything that depends on the schema is out of date.
+            session.update_all_diagnostics().await;
 
             session
                 .client

crates/pgls_lsp/src/handlers/inlay_hints.rs

@@ -29,7 +29,10 @@ pub(crate) fn inlay_hints(
     }) {
         Ok(result) => result,
         Err(e) => match e {
-            WorkspaceError::DatabaseConnectionError(_) => return Ok(None),
+            WorkspaceError::DatabaseConnectionError(_) => {
+                session.editor_features_without_database();
+                return Ok(None);
+            }
             _ => return Err(e.into()),
         },
     };

crates/pgls_lsp/src/handlers/semantic_tokens.rs

@@ -89,7 +89,10 @@ fn semantic_tokens(
         }) {
         Ok(result) => result,
         Err(e) => match e {
-            WorkspaceError::DatabaseConnectionError(_) => return Ok(None),
+            WorkspaceError::DatabaseConnectionError(_) => {
+                session.editor_features_without_database();
+                return Ok(None);
+            }
             _ => return Err(e.into()),
         },
     };

crates/pgls_lsp/src/session.rs

@@ -154,6 +154,11 @@ pub(crate) struct Session {
     /// before a new task is spawned. This ensures that only one analysis run
     /// fires at the end of a rapid-fire burst rather than one per keystroke.
     diagnostic_debounce: Arc<Mutex<FxHashMap<ls_types::Uri, AbortHandle>>>,
+
+    /// Set when inlay hints or semantic tokens were requested while the database was
+    /// unreachable. The next diagnostics run that finds it reachable asks the client to request
+    /// them again.
+    editor_features_without_database: AtomicBool,
 }
 
 /// The parameters provided by the client in the "initialize" request
@@ -263,6 +268,7 @@ impl Session {
             #[cfg(test)]
             max_running_transitions: std::sync::atomic::AtomicUsize::new(0),
             diagnostic_debounce: Arc::new(Mutex::new(FxHashMap::default())),
+            editor_features_without_database: AtomicBool::new(false),
         }
     }
 
@@ -511,9 +517,64 @@ impl Session {
             .publish_diagnostics(url, diagnostics, Some(captured_version))
             .await;
 
+        if self
+            .editor_features_without_database
+            .load(Ordering::Relaxed)
+        {
+            let connection = self.workspace.check_database_connection(
+                features::diagnostics::CheckDatabaseConnectionParams {
+                    project_key: self.project_key(&pgls_path),
+                },
+            );
+            if connection.is_ok_and(|connection| connection.error.is_none())
+                && self
+                    .editor_features_without_database
+                    .swap(false, Ordering::Relaxed)
+            {
+                self.refresh_editor_features();
+            }
+        }
+
         Ok(())
     }
 
+    /// Records that inlay hints or semantic tokens were requested while the database was
+    /// unreachable, so that the client is asked to request them again once it is reachable.
+    pub(crate) fn editor_features_without_database(&self) {
+        self.editor_features_without_database
+            .store(true, Ordering::Relaxed);
+    }
+
+    /// Asks the client to request inlay hints and semantic tokens again, if it supports that.
+    /// Called whenever their results may have changed without an edit: the configuration was
+    /// reloaded, the schema changed, or the database became reachable.
+    pub(crate) fn refresh_editor_features(&self) {
+        let workspace = self
+            .client_capabilities()
+            .and_then(|capabilities| capabilities.workspace.as_ref());
+        let inlay_hints = workspace
+            .and_then(|workspace| workspace.inlay_hint.as_ref())
+            .and_then(|inlay_hint| inlay_hint.refresh_support)
+            .unwrap_or(false);
+        let semantic_tokens = workspace
+            .and_then(|workspace| workspace.semantic_tokens.as_ref())
+            .and_then(|semantic_tokens| semantic_tokens.refresh_support)
+            .unwrap_or(false);
+        if !inlay_hints && !semantic_tokens {
+            return;
+        }
+        // The client answers these requests asynchronously, so don't wait for it.
+        let client = self.client.clone();
+        tokio::spawn(async move {
+            if inlay_hints && let Err(err) = client.inlay_hint_refresh().await {
+                error!("Failed to refresh inlay hints: {err}");
+            }
+            if semantic_tokens && let Err(err) = client.semantic_tokens_refresh().await {
+                error!("Failed to refresh semantic tokens: {err}");
+            }
+        });
+    }
+
     /// Schedules a debounced diagnostic update for the given URL.
     ///
     /// If a diagnostic task for `url` is already pending it is aborted first,
@@ -558,7 +619,8 @@ impl Session {
         }
     }
 
-    /// Updates diagnostics for every [`Document`] in this [`Session`]
+    /// Updates diagnostics for every [`Document`] in this [`Session`], and asks the client to
+    /// request inlay hints and semantic tokens again.
     pub(crate) async fn update_all_diagnostics(&self) {
         let mut futures: FuturesUnordered<_> = self
             .documents
@@ -573,6 +635,10 @@ impl Session {
                 error!("Error while updating diagnostics: {}", e);
             }
         }
+
+        // Whatever invalidated the diagnostics of every document may have changed the inlay
+        // hints and semantic tokens too.
+        self.refresh_editor_features();
     }
 
     /// Get a [`Document`] matching the provided [`ls_types::Uri`]