From 255f14687c957370fe2630c470c39131d9da1e1c Mon Sep 17 00:00:00 2001 From: psteinroe Date: Thu, 22 Oct 2026 18:14:58 +0000 Subject: [PATCH 1/2] feat(lsp): add semantic highlighting from resolved references Provide textDocument/semanticTokens (full and range) for the names that the catalog resolver resolves to exactly one object: schemas, relations, columns, functions, types, SQL function parameters, and string literals coerced to an enum type. The resolver now also records type, parameter, and enum literal references when references are collected. Closes #874 --- AGENTS.md | 1 + ARCHITECTURE.md | 6 +- Cargo.lock | 12 + Cargo.toml | 1 + .../pgls_catalog/src/resolve/expr/a_expr.rs | 59 ++++- .../src/resolve/expr/assignment.rs | 9 +- crates/pgls_catalog/src/resolve/expr/mod.rs | 103 +++++--- .../src/resolve/expr/type_cast.rs | 3 + crates/pgls_catalog/src/resolve/mod.rs | 40 ++- .../src/resolve/nodes/column_ref.rs | 13 +- .../src/resolve/nodes/insert_stmt.rs | 25 +- crates/pgls_catalog/src/resolve/nodes/mod.rs | 12 +- .../src/resolve/nodes/on_conflict_clause.rs | 1 + .../src/resolve/nodes/type_name.rs | 9 +- .../src/resolve/nodes/update_stmt.rs | 1 + .../src/resolve/reference_tests.rs | 161 +++++++++++- crates/pgls_catalog/src/resolve/resolver.rs | 73 +++++- crates/pgls_catalog/src/resolve/span.rs | 72 +++++- crates/pgls_lsp/src/capabilities.rs | 13 +- crates/pgls_lsp/src/handlers.rs | 1 + .../pgls_lsp/src/handlers/semantic_tokens.rs | 125 +++++++++ crates/pgls_lsp/src/server.rs | 23 ++ crates/pgls_lsp/tests/server.rs | 95 +++++++ crates/pgls_semantic_tokens/Cargo.toml | 25 ++ crates/pgls_semantic_tokens/src/lib.rs | 239 ++++++++++++++++++ .../tests/semantic_tokens_tests.rs | 153 +++++++++++ crates/pgls_workspace/Cargo.toml | 1 + crates/pgls_workspace/src/features/mod.rs | 1 + .../src/features/semantic_tokens.rs | 63 +++++ crates/pgls_workspace/src/workspace.rs | 8 + crates/pgls_workspace/src/workspace/client.rs | 7 + crates/pgls_workspace/src/workspace/server.rs | 76 ++++++ .../src/workspace/server.tests.rs | 83 ++++++ docs/features/editor_features.md | 16 ++ 34 files changed, 1461 insertions(+), 69 deletions(-) create mode 100644 crates/pgls_lsp/src/handlers/semantic_tokens.rs create mode 100644 crates/pgls_semantic_tokens/Cargo.toml create mode 100644 crates/pgls_semantic_tokens/src/lib.rs create mode 100644 crates/pgls_semantic_tokens/tests/semantic_tokens_tests.rs create mode 100644 crates/pgls_workspace/src/features/semantic_tokens.rs diff --git a/AGENTS.md b/AGENTS.md index eb7e6b560..0dcb92fcc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -103,6 +103,7 @@ The project uses a modular Rust workspace with crates prefixed with `pgls_`: **Features:** - `pgls_completions` - Autocompletion engine - `pgls_hover` - Hover information provider +- `pgls_semantic_tokens` - Semantic highlighting of resolved names - `pgls_analyser` & `pgls_analyse` - Linting and analysis framework - `pgls_typecheck` - EXPLAIN fallback for typecheck - `pgls_catalog` - Database snapshot, the file's DDL on top of it, name resolution, and type inference diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 77b17f821..84d0a1f37 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -63,7 +63,7 @@ Tree-sitter integration used by features (e.g. completions and hover) that need ##### `crates/pgls_catalog` -The database catalog. `Snapshot` is the schema of the connected database (tables, columns, functions, types, ...), loaded by introspection queries or from JSON. `Catalog` applies the DDL of the current file on top of it without touching the database, and the name resolver checks statements against the result. On request, the resolver also returns the functions, columns, and relations that names resolve to with certainty, for features like inlay hints. Lookups answer found, missing, or unknown, so the typecheck rules only report what is certainly wrong. +The database catalog. `Snapshot` is the schema of the connected database (tables, columns, functions, types, ...), loaded by introspection queries or from JSON. `Catalog` applies the DDL of the current file on top of it without touching the database, and the name resolver checks statements against the result. On request, the resolver also returns what names resolve to with certainty (functions, columns, relations, types, SQL function parameters, and string literals coerced to enums), for features like inlay hints and semantic tokens. Lookups answer found, missing, or unknown, so the typecheck rules only report what is certainly wrong. `typing/` ports Postgres' type rules (coercion, common types, polymorphic types, and function and operator overload resolution) from the latest supported Postgres, and links the Postgres function each port follows. The resolver infers the type of every expression with them. @@ -107,6 +107,10 @@ Handles `-- pgls-ignore` style rule suppression comments. Autocompletion and hover providers. They operate on the catalog as of the statement at the cursor plus that statement and its parse results, and are intentionally free of any language-server flavour so they can be reused from the CLI. Hover shows the function overload, column, or relation that the name resolver resolves the hovered name to, and falls back to looking objects up by name when there is no unique reference. +##### `crates/pgls_inlay_hints`, `crates/pgls_semantic_tokens` + +Inlay hints and semantic tokens of a single statement, built from the references of the name resolver. The workspace walks the file like the analyser does, so that each statement sees the objects created earlier in the file. + #### Diagnostics and utilities - `crates/pgls_diagnostics` (+ `_categories`, `_macros`) — error and warning reporting model. diff --git a/Cargo.lock b/Cargo.lock index 9860413af..c75c255fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2821,6 +2821,17 @@ dependencies = [ "ureq", ] +[[package]] +name = "pgls_semantic_tokens" +version = "0.0.0" +dependencies = [ + "pgls_catalog", + "pgls_query", + "pgls_test_utils", + "pgls_text_size", + "sqlx", +] + [[package]] name = "pgls_splinter" version = "0.0.0" @@ -3009,6 +3020,7 @@ dependencies = [ "pgls_pretty_print", "pgls_query", "pgls_query_ext", + "pgls_semantic_tokens", "pgls_splinter", "pgls_statement_splitter", "pgls_suppressions", diff --git a/Cargo.toml b/Cargo.toml index a8613ff9e..b4bf4547e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -87,6 +87,7 @@ pgls_pretty_print_codegen = { path = "./crates/pgls_pretty_print_codegen", ver pgls_query = { path = "./crates/pgls_query", version = "0.0.0" } pgls_query_ext = { path = "./crates/pgls_query_ext", version = "0.0.0" } pgls_query_macros = { path = "./crates/pgls_query_macros", version = "0.0.0" } +pgls_semantic_tokens = { path = "./crates/pgls_semantic_tokens", version = "0.0.0" } pgls_splinter = { path = "./crates/pgls_splinter", version = "0.0.0" } pgls_statement_splitter = { path = "./crates/pgls_statement_splitter", version = "0.0.0" } pgls_suppressions = { path = "./crates/pgls_suppressions", version = "0.0.0" } diff --git a/crates/pgls_catalog/src/resolve/expr/a_expr.rs b/crates/pgls_catalog/src/resolve/expr/a_expr.rs index 9b0d4a6a2..4da97bc6a 100644 --- a/crates/pgls_catalog/src/resolve/expr/a_expr.rs +++ b/crates/pgls_catalog/src/resolve/expr/a_expr.rs @@ -7,6 +7,8 @@ use super::*; pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr) -> Option { { use protobuf::AExprKind as Kind; + let lexpr = n.lexpr.as_deref().and_then(|x| x.node.as_ref()); + let rexpr = n.rexpr.as_deref().and_then(|x| x.node.as_ref()); match Kind::try_from(n.kind).ok()? { // Port of `transformAExprDistinct`; scalar cases look up `=` like `make_distinct_op`: // https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_expr.c#L1031 @@ -22,8 +24,10 @@ pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr .as_deref() .and_then(|x| x.node.as_ref()) .and_then(|x| infer_expr(r, x)); - if let (Some(left), Some(right)) = (left.as_ref(), right.as_ref()) { - check_operator(r, "=", Some(left), right, n.location); + if let (Some(left), Some(right)) = (left.as_ref(), right.as_ref()) + && let Some(operator) = check_operator(r, "=", Some(left), right, n.location) + { + record_enum_operands(r, &operator, (lexpr, left), (rexpr, right)); } return named(r, "bool"); } @@ -106,7 +110,18 @@ pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr .collect::>(); match select_common_type(r.catalog, &all) { Selection::Match(Type::Named(common)) => { - check_operator(r, name, Some(left), &Type::Named(common), n.location); + check_operator( + r, + name, + Some(left), + &Type::Named(common.clone()), + n.location, + ); + for item in + std::iter::once(lexpr).chain(items.iter().copied().map(Some)) + { + r.record_enum_literal(item, &common); + } for (item, ty) in items.iter().zip(&types) { if matches!(item, NodeEnum::ColumnRef(_)) { check_operator(r, name, Some(left), ty, n.location); @@ -119,8 +134,10 @@ pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr _ => return named(r, "bool"), } } - for ty in &types { - check_operator(r, name, Some(left), ty, n.location); + for (item, ty) in items.iter().zip(&types) { + if let Some(operator) = check_operator(r, name, Some(left), ty, n.location) { + record_enum_operands(r, &operator, (lexpr, left), (Some(item), ty)); + } } return named(r, "bool"); } @@ -149,11 +166,24 @@ pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr Kind::AexprNotBetween | Kind::AexprNotBetweenSym => ("<", ">"), _ => (">=", "<="), }; - check_operator(r, lower, Some(left), low, n.location); - check_operator(r, upper, Some(left), high, n.location); + let low_expr = list.items[0].node.as_ref(); + let high_expr = list.items[1].node.as_ref(); + let mut comparisons = + vec![(lower, low_expr, low), (upper, high_expr, high)]; if matches!(kind, Kind::AexprBetweenSym | Kind::AexprNotBetweenSym) { - check_operator(r, lower, Some(left), high, n.location); - check_operator(r, upper, Some(left), low, n.location); + comparisons.extend([(lower, high_expr, high), (upper, low_expr, low)]); + } + for (operator, bound_expr, bound) in comparisons { + if let Some(operator) = + check_operator(r, operator, Some(left), bound, n.location) + { + record_enum_operands( + r, + &operator, + (lexpr, left), + (bound_expr, bound), + ); + } } } } @@ -214,7 +244,9 @@ pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr .node .as_ref() .and_then(|x| infer_expr(r, x))?; - check_operator(r, "=", Some(&left), &right, n.location); + if let Some(operator) = check_operator(r, "=", Some(&left), &right, n.location) { + record_enum_operands(r, &operator, (lexpr, &left), (rexpr, &right)); + } return common_types(r, &[left, right]); } _ => {} @@ -240,7 +272,12 @@ pub(super) fn infer_a_expr(r: &mut Resolver<'_>, n: &pgls_query::protobuf::AExpr r.search_path, ); match selection { - Selection::Match(o) => o.result, + Selection::Match(o) => { + if let Some(left) = left.as_ref() { + record_enum_operands(r, &o.operator, (lexpr, left), (rexpr, &right)); + } + o.result + } Selection::NoMatch | Selection::Ambiguous => { let failure = if matches!(selection, Selection::NoMatch) { crate::resolve::MatchFailure::NoMatch diff --git a/crates/pgls_catalog/src/resolve/expr/assignment.rs b/crates/pgls_catalog/src/resolve/expr/assignment.rs index a0ca041c1..3d29ae0f1 100644 --- a/crates/pgls_catalog/src/resolve/expr/assignment.rs +++ b/crates/pgls_catalog/src/resolve/expr/assignment.rs @@ -1,18 +1,25 @@ +use pgls_query::NodeEnum; + use super::super::resolver::Resolver; use crate::typing::{CoercionContext, Decision, Type, can_coerce}; /// Reports a value that can't be assigned to a column of INSERT or UPDATE, like /// [`transformAssignedExpr`] does with an assignment coercion. Unknown literals are coerced at -/// runtime and never reported. +/// runtime and never reported. `expr` is the value, if it is in the statement itself rather than +/// the output of a query. /// /// [`transformAssignedExpr`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_target.c#L455 pub(crate) fn check_assignment( r: &mut Resolver<'_>, column: &str, column_type: Type, + expr: Option<&NodeEnum>, expr_type: Option, location: i32, ) { + if let Type::Named(target) = &column_type { + r.record_enum_literal(expr, target); + } let Some(expr_type) = expr_type else { return }; if matches!(expr_type, Type::UnknownLiteral) { return; diff --git a/crates/pgls_catalog/src/resolve/expr/mod.rs b/crates/pgls_catalog/src/resolve/expr/mod.rs index 90f8e0914..43c89af4a 100644 --- a/crates/pgls_catalog/src/resolve/expr/mod.rs +++ b/crates/pgls_catalog/src/resolve/expr/mod.rs @@ -29,8 +29,8 @@ use crate::{ OperatorKind, lookup::Lookup, typing::{ - CallArg, CoercionContext, Decision, Selection, Type, can_coerce, normalize_type_name, - select_common_type, select_function, select_operator, + CallArg, CoercionContext, Decision, Selection, Type, TypeId, can_coerce, + normalize_type_name, select_common_type, select_function, select_operator, }, }; pub(super) use assignment::check_assignment; @@ -137,7 +137,8 @@ fn common_types(r: &Resolver<'_>, ts: &[Type]) -> Option { } /// Checks the `=` that `IS DISTINCT FROM`, `NULLIF` and simple `CASE` compare with, like -/// [`make_op`] (through [`oper`]) does. Reports only proven selection failures. +/// [`make_op`] (through [`oper`]) does. Reports only proven selection failures. Returns the +/// selected operator. /// /// [`make_op`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_oper.c#L660 /// [`oper`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_oper.c#L370 @@ -147,7 +148,7 @@ fn check_operator( left: Option<&Type>, right: &Type, location: i32, -) { +) -> Option { let selection = select_operator( r.catalog, None, @@ -158,34 +159,72 @@ fn check_operator( r.search_path, ); let failure = match selection { - Selection::NoMatch => Some(crate::resolve::MatchFailure::NoMatch), - Selection::Ambiguous => Some(crate::resolve::MatchFailure::Ambiguous), - Selection::Match(operator_result) => { - let Lookup::Found(boolean_info) = - r.catalog.type_(Some("pg_catalog"), "bool", r.search_path) - else { - return; - }; - let Some(boolean) = boolean_info.id else { - return; - }; - if operator_result.result.as_ref() != Some(&Type::Named(boolean)) { - return; - } - None - } - Selection::Unknown => None, + Selection::NoMatch => crate::resolve::MatchFailure::NoMatch, + Selection::Ambiguous => crate::resolve::MatchFailure::Ambiguous, + Selection::Match(operator_result) => return Some(operator_result.operator), + Selection::Unknown => return None, }; - if let Some(failure) = failure { - let span = r.operator_span(location); - r.report_with_span( - crate::resolve::FindingKind::OperatorMismatch { - operator: operator.to_owned(), - left: left.cloned(), - right: right.clone(), - failure, - }, - span, - ); + let span = r.operator_span(location); + r.report_with_span( + crate::resolve::FindingKind::OperatorMismatch { + operator: operator.to_owned(), + left: left.cloned(), + right: right.clone(), + failure, + }, + span, + ); + None +} + +/// The polymorphic types that an operand of an unknown literal resolves to the type of the other +/// operand, like `anyenum` in the comparison operators of enums. +const POLYMORPHIC_SCALAR_TYPES: &[&str] = &[ + "anyelement", + "anynonarray", + "anyenum", + "anycompatible", + "anycompatiblenonarray", +]; + +/// Records a string literal operand of `operator` that is coerced to the enum type of the other +/// operand. Postgres first assumes that an unknown literal has the type of the other operand +/// ([`binary_oper_exact`]), and the polymorphic operators of enums resolve it to that type +/// ([`enforce_generic_type_consistency`]). +/// +/// [`binary_oper_exact`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_oper.c#L262 +/// [`enforce_generic_type_consistency`]: https://github.com/postgres/postgres/blob/REL_18_6/src/backend/parser/parse_coerce.c#L2109 +fn record_enum_operands( + r: &mut Resolver<'_>, + operator: &crate::OperatorInfo, + left: (Option<&NodeEnum>, &Type), + right: (Option<&NodeEnum>, &Type), +) { + if !r.collects_references() { + return; + } + match (left.1, right.1) { + (Type::UnknownLiteral, Type::Named(other)) + if takes_type_of_other(r, operator.left.as_ref(), other) => + { + r.record_enum_literal(left.0, other); + } + (Type::Named(other), Type::UnknownLiteral) + if takes_type_of_other(r, operator.right.as_ref(), other) => + { + r.record_enum_literal(right.0, other); + } + _ => {} } } + +/// Whether an operand declared as `declared` takes the type `other` of the other operand. +fn takes_type_of_other(r: &Resolver<'_>, declared: Option<&TypeId>, other: &TypeId) -> bool { + let Some(declared) = declared else { + return false; + }; + declared == other + || r.catalog.type_by_id(declared).found().is_some_and(|info| { + info.schema == "pg_catalog" && POLYMORPHIC_SCALAR_TYPES.contains(&info.name.as_str()) + }) +} diff --git a/crates/pgls_catalog/src/resolve/expr/type_cast.rs b/crates/pgls_catalog/src/resolve/expr/type_cast.rs index 7def4bd4a..97efc0b63 100644 --- a/crates/pgls_catalog/src/resolve/expr/type_cast.rs +++ b/crates/pgls_catalog/src/resolve/expr/type_cast.rs @@ -17,6 +17,9 @@ pub(super) fn infer_type_cast( .as_ref() .and_then(|x| infer_expr(r, x)); let (id, _) = normalize_type_name(r.catalog, n.type_name.as_ref()?, r.search_path); + if let Some(id) = id.as_ref() { + r.record_enum_literal(n.arg.as_deref().and_then(|arg| arg.node.as_ref()), id); + } if let (Some(from), Some(to)) = (from, id.as_ref()) && !matches!(from, Type::UnknownLiteral) && matches!( diff --git a/crates/pgls_catalog/src/resolve/mod.rs b/crates/pgls_catalog/src/resolve/mod.rs index 44abe6d6e..4e5fc415d 100644 --- a/crates/pgls_catalog/src/resolve/mod.rs +++ b/crates/pgls_catalog/src/resolve/mod.rs @@ -204,6 +204,36 @@ pub enum Reference { span: Option, relation: crate::RelationInfo, }, + /// The type of a cast, like `int` in `x::int` or `date` in `date 'today'`. + Type { + /// The location of the type name in the statement. + location: i32, + /// The span of the (possibly qualified) type name, without modifiers like `(10)` or `[]`. + /// For multi-word types like `double precision`, only the first word. + span: Option, + type_info: crate::TypeInfo, + }, + /// A parameter of the SQL function whose body is resolved, by name (`id`, `fn.id`) or by + /// number (`$1`). + Parameter { + /// The location of the parameter reference in the statement. + location: i32, + /// The span of the whole parameter reference, including the function name qualifier. + span: Option, + /// The index of the parameter in [`FunctionContext::params`]. + index: usize, + }, + /// A string literal that is coerced to an enum type, like `'happy'` in `mood = 'happy'`. + /// Only recorded where the enum type is certain: casts, operators whose other operand is the + /// enum, function parameters of the enum type, and values assigned to enum columns in INSERT + /// and UPDATE. + EnumLiteral { + /// The location of the literal in the statement. + location: i32, + /// The span of the literal, including its quotes. + span: Option, + type_info: crate::TypeInfo, + }, } impl Reference { @@ -211,7 +241,10 @@ impl Reference { match self { Self::Function { location, .. } | Self::Column { location, .. } - | Self::Relation { location, .. } => *location, + | Self::Relation { location, .. } + | Self::Type { location, .. } + | Self::Parameter { location, .. } + | Self::EnumLiteral { location, .. } => *location, } } @@ -219,7 +252,10 @@ impl Reference { match self { Self::Function { span, .. } | Self::Column { span, .. } - | Self::Relation { span, .. } => *span, + | Self::Relation { span, .. } + | Self::Type { span, .. } + | Self::Parameter { span, .. } + | Self::EnumLiteral { span, .. } => *span, } } } diff --git a/crates/pgls_catalog/src/resolve/nodes/column_ref.rs b/crates/pgls_catalog/src/resolve/nodes/column_ref.rs index a03c50bdf..a0da0840a 100644 --- a/crates/pgls_catalog/src/resolve/nodes/column_ref.rs +++ b/crates/pgls_catalog/src/resolve/nodes/column_ref.rs @@ -75,7 +75,15 @@ fn resolve_unqualified_column(r: &mut Resolver, name: &str, location: i32) { ColumnLookup::NotFound => { // A whole-row reference to a FROM item, or a parameter of the SQL function. let is_item = find_item(&r.levels, name).is_some(); - if is_item || r.function_param(name).is_some() || has_opaque_items(&r.levels) { + let opaque = has_opaque_items(&r.levels); + if let Some(index) = r.function_param_index(name) { + // A column of an item with unknown columns would take precedence. + if !is_item && !opaque { + r.record_parameter(index, location); + } + return; + } + if is_item || opaque { return; } r.report( @@ -115,6 +123,9 @@ fn resolve_qualified_column(r: &mut Resolver, qualifier: &str, name: &str, locat if let Some(function) = r.function { // `function_name.parameter` if function.function_name == qualifier { + if let Some(index) = r.function_param_index(name) { + r.record_parameter(index, location); + } return; } // A field of a composite parameter. diff --git a/crates/pgls_catalog/src/resolve/nodes/insert_stmt.rs b/crates/pgls_catalog/src/resolve/nodes/insert_stmt.rs index 6c93bf053..df71b8a9d 100644 --- a/crates/pgls_catalog/src/resolve/nodes/insert_stmt.rs +++ b/crates/pgls_catalog/src/resolve/nodes/insert_stmt.rs @@ -74,14 +74,17 @@ fn resolve_insert(r: &mut Resolver, n: &InsertStmt) -> Columns { }) .collect() }; - let check = - |r: &mut Resolver, value: Option, target: Option<&TypedColumn>, location: i32| { - if let Some(target) = target - && let Some(column_type) = target.ty.clone() - { - check_assignment(r, &target.name, column_type, value, location); - } - }; + let check = |r: &mut Resolver, + expr: Option<&NodeEnum>, + value: Option, + target: Option<&TypedColumn>, + location: i32| { + if let Some(target) = target + && let Some(column_type) = target.ty.clone() + { + check_assignment(r, &target.name, column_type, expr, value, location); + } + }; match values_rows(source) { // Each row of `INSERT ... VALUES` is coerced to the target columns on its own. Some(rows) => { @@ -91,8 +94,8 @@ fn resolve_insert(r: &mut Resolver, n: &InsertStmt) -> Columns { let Some(value) = value.node.as_ref() else { continue; }; - let value = infer_expr(r, value); - check(r, value, *target, location); + let value_type = infer_expr(r, value); + check(r, Some(value), value_type, *target, location); } } } @@ -104,7 +107,7 @@ fn resolve_insert(r: &mut Resolver, n: &InsertStmt) -> Columns { .enumerate() { let location = select_target_location(source, index).unwrap_or_default(); - check(r, value.ty.clone(), *target, location); + check(r, None, value.ty.clone(), *target, location); } } } diff --git a/crates/pgls_catalog/src/resolve/nodes/mod.rs b/crates/pgls_catalog/src/resolve/nodes/mod.rs index 2888f0213..ceec888d1 100644 --- a/crates/pgls_catalog/src/resolve/nodes/mod.rs +++ b/crates/pgls_catalog/src/resolve/nodes/mod.rs @@ -116,11 +116,21 @@ pub(super) fn resolve_node_enum(r: &mut Resolver, node: &NodeEnum) -> Columns { resolve_opt(r, &n.end_offset); } NodeEnum::List(n) => resolve_list(r, &n.items), + // `$1` in the body of a SQL function. + NodeEnum::ParamRef(n) => { + let params = r.function.map_or(0, |function| function.params.len()); + if let Some(index) = usize::try_from(n.number) + .ok() + .and_then(|number| number.checked_sub(1)) + .filter(|&index| index < params) + { + r.record_parameter(index, n.location); + } + } // Field names in indirection (`(row).field`) and leaves. NodeEnum::String(_) | NodeEnum::AStar(_) | NodeEnum::AConst(_) - | NodeEnum::ParamRef(_) | NodeEnum::SqlvalueFunction(_) | NodeEnum::SetToDefault(_) | NodeEnum::CurrentOfExpr(_) => {} diff --git a/crates/pgls_catalog/src/resolve/nodes/on_conflict_clause.rs b/crates/pgls_catalog/src/resolve/nodes/on_conflict_clause.rs index 157444747..a21ce747b 100644 --- a/crates/pgls_catalog/src/resolve/nodes/on_conflict_clause.rs +++ b/crates/pgls_catalog/src/resolve/nodes/on_conflict_clause.rs @@ -42,6 +42,7 @@ pub(super) fn resolve_on_conflict_clause(r: &mut Resolver, n: &OnConflictClause, r, &assignment.name, column_type, + Some(expr), expr_type, assignment .val diff --git a/crates/pgls_catalog/src/resolve/nodes/type_name.rs b/crates/pgls_catalog/src/resolve/nodes/type_name.rs index 0b0c2372b..87a65cb47 100644 --- a/crates/pgls_catalog/src/resolve/nodes/type_name.rs +++ b/crates/pgls_catalog/src/resolve/nodes/type_name.rs @@ -2,7 +2,7 @@ use pgls_query::protobuf::TypeName; use super::{Resolver, string::string_values}; use crate::lookup::Lookup; -use crate::resolve::FindingKind; +use crate::resolve::{FindingKind, Reference}; /// Types whose input is the name of a catalog object, which we don't resolve. const OBJECT_NAME_TYPES: &[&str] = &[ @@ -49,6 +49,13 @@ pub(super) fn resolve_type_name(r: &mut Resolver, n: &TypeName) { if OBJECT_NAME_TYPES.contains(&name) { r.depends_on_file(); } + if r.collects_references() { + r.record(Reference::Type { + location: n.location, + span: r.name_span(n.location), + type_info, + }); + } } Lookup::Missing => r.report( FindingKind::UnknownType { diff --git a/crates/pgls_catalog/src/resolve/nodes/update_stmt.rs b/crates/pgls_catalog/src/resolve/nodes/update_stmt.rs index 766ced257..baa0508fb 100644 --- a/crates/pgls_catalog/src/resolve/nodes/update_stmt.rs +++ b/crates/pgls_catalog/src/resolve/nodes/update_stmt.rs @@ -50,6 +50,7 @@ fn resolve_update(r: &mut Resolver, n: &UpdateStmt) -> Columns { r, &assignment.name, column_type, + Some(expr), expr_type, assignment .val diff --git a/crates/pgls_catalog/src/resolve/reference_tests.rs b/crates/pgls_catalog/src/resolve/reference_tests.rs index 4434b2b61..2f043af05 100644 --- a/crates/pgls_catalog/src/resolve/reference_tests.rs +++ b/crates/pgls_catalog/src/resolve/reference_tests.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use sqlx::{Executor, PgPool}; -use super::{ColumnSource, Reference, ResolveParams, resolve}; +use super::{ColumnSource, FunctionContext, FunctionParam, Reference, ResolveParams, resolve}; use crate::{ Catalog, CatalogBase, Snapshot, typing::{Type, format_type}, @@ -19,6 +19,10 @@ const SETUP: &str = r#" language sql as 'select greeting || name'; create function amb(a int, b text) returns int language sql as 'select 1'; create function amb(c text, d int) returns int language sql as 'select 1'; + create type mood as enum ('happy', 'sad'); + create table people (id int8, current mood); + create function cheer(m mood) returns mood language sql as 'select m'; + create function cheer_all(variadic moods mood[]) returns int language sql as 'select 1'; "#; async fn catalog(test_db: &PgPool) -> (Catalog, Vec) { @@ -31,6 +35,16 @@ async fn catalog(test_db: &PgPool) -> (Catalog, Vec) { /// The references of `sql`, each as the referencing text and what it resolves to. fn references(catalog: &Catalog, search_path: &[String], sql: &str) -> Vec { + function_references_in(catalog, search_path, sql, None) +} + +/// The references of `sql` as the body of `function`. +fn function_references_in( + catalog: &Catalog, + search_path: &[String], + sql: &str, + function: Option<&FunctionContext>, +) -> Vec { let stmt = pgls_query::parse(sql) .expect("valid sql") .into_root() @@ -39,7 +53,7 @@ fn references(catalog: &Catalog, search_path: &[String], sql: &str) -> Vec Vec { format!("relation {}.{}", relation.schema, relation.name) } + Reference::Type { type_info, .. } => { + format!("type {}.{}", type_info.schema, type_info.name) + } + Reference::Parameter { index, .. } => format!("parameter {index}"), + Reference::EnumLiteral { type_info, .. } => { + format!("enum literal {}.{}", type_info.schema, type_info.name) + } }; format!("{text} => {target}") }) @@ -127,6 +148,7 @@ async fn function_references(test_db: PgPool) { "select pick(1), public.pick('a'::text)" ), [ + "text => type pg_catalog.text", "pick => function public.pick(integer) <- i", "public.pick => function public.pick(text) <- t", ] @@ -289,6 +311,141 @@ async fn relation_references(test_db: PgPool) { ); } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn type_references(test_db: PgPool) { + let (catalog, search_path) = catalog(&test_db).await; + assert_eq!( + references( + &catalog, + &search_path, + "select 1::int, cast(2 as pg_catalog.int8), '{}'::text[], 'sad'::mood, x::nope" + ), + [ + "int => type pg_catalog.int4", + "pg_catalog.int8 => type pg_catalog.int8", + "text => type pg_catalog.text", + "mood => type public.mood", + "'sad' => enum literal public.mood", + ] + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn enum_literal_references(test_db: PgPool) { + let (catalog, search_path) = catalog(&test_db).await; + assert_eq!( + references( + &catalog, + &search_path, + "select id from people where current = 'happy' or 'sad' <> current \ + or current in ('happy', 'sad') or current between 'happy' and 'sad' \ + or current is distinct from 'sad' or nullif(current, 'sad') is null" + ), + [ + "people => relation public.people", + "id => column public.people.id: bigint (int8)", + "current => column public.people.current: public.mood (mood)", + "current => column public.people.current: public.mood (mood)", + "current => column public.people.current: public.mood (mood)", + "current => column public.people.current: public.mood (mood)", + "current => column public.people.current: public.mood (mood)", + "current => column public.people.current: public.mood (mood)", + "'happy' => enum literal public.mood", + "'sad' => enum literal public.mood", + "'happy' => enum literal public.mood", + "'sad' => enum literal public.mood", + "'happy' => enum literal public.mood", + "'sad' => enum literal public.mood", + "'sad' => enum literal public.mood", + "'sad' => enum literal public.mood", + ] + ); + assert_eq!( + references( + &catalog, + &search_path, + "insert into people (current) values ('happy'), ('sad')" + ), + [ + "people => relation public.people", + "'happy' => enum literal public.mood", + "'sad' => enum literal public.mood", + ] + ); + assert_eq!( + references( + &catalog, + &search_path, + "update people set current = 'sad' where id = 1" + ), + [ + "people => relation public.people", + "'sad' => enum literal public.mood", + "id => column public.people.id: bigint (int8)", + ] + ); + assert_eq!( + references( + &catalog, + &search_path, + "select cheer('happy'), cheer_all('happy', 'sad')" + ), + [ + "'happy' => enum literal public.mood", + "cheer => function public.cheer(public.mood) <- m", + "'happy' => enum literal public.mood", + "'sad' => enum literal public.mood", + "cheer_all => function public.cheer_all(public.mood[]) <- moods, moods", + ] + ); + // Literals that are not coerced to an enum are no enum literals. + assert_eq!( + references( + &catalog, + &search_path, + "select 'happy', current::text = 'happy' from people" + ), + [ + "people => relation public.people", + "current => column public.people.current: public.mood (mood)", + "text => type pg_catalog.text", + ] + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn parameter_references(test_db: PgPool) { + let (catalog, search_path) = catalog(&test_db).await; + let param = |name: &str| FunctionParam { + name: Some(name.to_owned()), + type_schema: None, + type_name: "int8".to_owned(), + is_array: false, + }; + let function = FunctionContext { + function_name: "lookup".to_owned(), + params: vec![param("user_id"), param("name")], + }; + assert_eq!( + function_references_in( + &catalog, + &search_path, + "select title from posts where user_id = $1 or id = lookup.user_id or id = name", + Some(&function), + ), + [ + "posts => relation public.posts", + "title => column public.posts.title: text (text)", + "user_id => column public.posts.user_id: bigint (int8)", + "$1 => parameter 0", + "id => column public.posts.id: bigint (int8)", + "lookup.user_id => parameter 0", + "id => column public.posts.id: bigint (int8)", + "name => parameter 1", + ] + ); +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn references_are_only_collected_on_request(test_db: PgPool) { let (catalog, search_path) = catalog(&test_db).await; diff --git a/crates/pgls_catalog/src/resolve/resolver.rs b/crates/pgls_catalog/src/resolve/resolver.rs index eb4fd48d0..e230f8c39 100644 --- a/crates/pgls_catalog/src/resolve/resolver.rs +++ b/crates/pgls_catalog/src/resolve/resolver.rs @@ -4,7 +4,7 @@ use pgls_text_size::{TextRange, TextSize}; use pgls_query::{ NodeEnum, - protobuf::{CoercionForm, FuncCall}, + protobuf::{self, CoercionForm, FuncCall}, }; use super::{ @@ -12,8 +12,9 @@ use super::{ scope::{Cte, Level}, span, }; -use crate::FunctionInfo; use crate::lookup::{CatalogView, Lookup, Origin}; +use crate::typing::TypeId; +use crate::{FunctionInfo, TypeKind}; pub(super) struct Resolver<'a> { pub catalog: &'a dyn CatalogView, @@ -73,10 +74,16 @@ impl<'a> Resolver<'a> { /// The parameter of the SQL function with this name. pub fn function_param(&self, name: &str) -> Option<&'a FunctionParam> { + let index = self.function_param_index(name)?; + self.function?.params.get(index) + } + + /// The index of the parameter of the SQL function with this name. + pub fn function_param_index(&self, name: &str) -> Option { self.function? .params .iter() - .find(|param| param.name.as_deref() == Some(name)) + .position(|param| param.name.as_deref() == Some(name)) } /// SQL function bodies are only checked against the database if their parameters have @@ -157,6 +164,23 @@ impl<'a> Resolver<'a> { }; arguments.push(ResolvedCallArgument { parameter, named }); } + for (arg, argument) in call.args.iter().zip(&arguments) { + let value = match arg.node.as_ref() { + Some(NodeEnum::NamedArgExpr(named)) => { + named.arg.as_deref().and_then(|arg| arg.node.as_ref()) + } + value => value, + }; + let is_variadic = expanded && argument.parameter + 1 == parameters.len(); + let target = if is_variadic { + signature.variadic_element.as_ref() + } else { + parameters[argument.parameter].ty.as_ref() + }; + if let Some(target) = target { + self.record_enum_literal(value, target); + } + } self.record(Reference::Function { location: call.location, span: self.name_span(call.location), @@ -165,6 +189,44 @@ impl<'a> Resolver<'a> { }); } + /// Records that the parameter `index` of the SQL function is referenced at `location`, when + /// references are being collected. + pub fn record_parameter(&mut self, index: usize, location: i32) { + if !self.collects_references() { + return; + } + self.record(Reference::Parameter { + location, + span: self.name_span(location), + index, + }); + } + + /// Records `node` if it is a string literal that is coerced to `target`, and `target` is an + /// enum type. + pub fn record_enum_literal(&mut self, node: Option<&NodeEnum>, target: &TypeId) { + if !self.collects_references() { + return; + } + let Some(NodeEnum::AConst(constant)) = node else { + return; + }; + if !matches!(constant.val, Some(protobuf::a_const::Val::Sval(_))) { + return; + } + let Lookup::Found(type_info) = self.catalog.type_by_id(target) else { + return; + }; + if type_info.kind != Some(TypeKind::Enum) { + return; + } + self.record(Reference::EnumLiteral { + location: constant.location, + span: self.literal_span(constant.location), + type_info, + }); + } + /// Marks the statement as depending on something that is not an unchanged database object. pub fn depends_on_file(&mut self) { self.database_only = false; @@ -205,6 +267,11 @@ impl<'a> Resolver<'a> { self.span_to(location, span::cast_end) } + /// The span of the string constant starting at `location`. + pub fn literal_span(&self, location: i32) -> Option { + self.span_to(location, span::literal_end) + } + /// The span of the first token of an expression starting at `location`. pub fn expression_span(&self, location: i32) -> Option { self.span_to(location, |sql, start| { diff --git a/crates/pgls_catalog/src/resolve/span.rs b/crates/pgls_catalog/src/resolve/span.rs index 93f93d8e2..c60e501a0 100644 --- a/crates/pgls_catalog/src/resolve/span.rs +++ b/crates/pgls_catalog/src/resolve/span.rs @@ -142,6 +142,57 @@ pub(super) fn cast_end(sql: &str, start: usize) -> Option { Some(end) } +/// The end of the string constant starting at `start`: `'...'`, with an `E`, `B`, `X`, `N`, or +/// `U&` prefix, or `$tag$...$tag$`. Constants continued on a new line (`'a'\n'b'`) end after +/// their last part. +pub(super) fn literal_end(sql: &str, start: usize) -> Option { + let bytes = sql.as_bytes(); + let tail = sql.get(start..)?; + if let Some(after_dollar) = tail.strip_prefix('$') { + let tag_len = after_dollar.find('$')? + 2; + let tag = &tail[..tag_len]; + let tag_is_valid = tag[1..tag_len - 1] + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_'); + if !tag_is_valid { + return None; + } + let close = tail[tag_len..].find(tag)?; + return Some(start + tag_len + close + tag_len); + } + let (prefix, escapes) = match bytes.get(start..start + 2)? { + [b'u' | b'U', b'&'] => (2, false), + [b'e' | b'E', b'\''] => (1, true), + [b'b' | b'B' | b'x' | b'X' | b'n' | b'N', b'\''] => (1, false), + [b'\'', _] => (0, false), + _ => return None, + }; + let mut cursor = start + prefix; + loop { + // The opening quote. + cursor += 1; + loop { + match bytes.get(cursor)? { + b'\\' if escapes => cursor += 2, + b'\'' if bytes.get(cursor + 1) == Some(&b'\'') => cursor += 2, + b'\'' => break, + _ => cursor += 1, + } + } + cursor += 1; + let mut next = cursor; + let mut newline = false; + while let Some(byte) = bytes.get(next).filter(|byte| byte.is_ascii_whitespace()) { + newline |= *byte == b'\n'; + next += 1; + } + if !(newline && bytes.get(next) == Some(&b'\'')) { + return Some(cursor); + } + cursor = next; + } +} + pub(super) fn operator_end(sql: &str, start: usize) -> Option { let tail = sql.get(start..)?; if tail @@ -156,7 +207,7 @@ pub(super) fn operator_end(sql: &str, start: usize) -> Option { #[cfg(test)] mod tests { - use super::{cast_end, operator_end, reference_end}; + use super::{cast_end, literal_end, operator_end, reference_end}; fn text(sql: &str, start: usize, end: Option) -> &str { &sql[start..end.expect("span")] @@ -187,4 +238,23 @@ mod tests { let start = sql.find("date '").unwrap(); assert_eq!(text(sql, start, cast_end(sql, start)), "date 'today'"); } + + #[test] + fn literal_spans_cover_the_whole_constant() { + let cases = [ + ("'happy'", "'happy'"), + ("'it''s' = x", "'it''s'"), + ("E'a\\'b' || x", "E'a\\'b'"), + ("U&'d\\0061t' x", "U&'d\\0061t'"), + ("$$a'b$$ x", "$$a'b$$"), + ("$t$a$$b$t$ x", "$t$a$$b$t$"), + ("'a'\n 'b' x", "'a'\n 'b'"), + ("'a' 'b'", "'a'"), + ]; + for (sql, expected) in cases { + assert_eq!(text(sql, 0, literal_end(sql, 0)), expected, "{sql}"); + } + assert_eq!(literal_end("x", 0), None); + assert_eq!(literal_end("'open", 0), None); + } } diff --git a/crates/pgls_lsp/src/capabilities.rs b/crates/pgls_lsp/src/capabilities.rs index 9e941638a..ae1aa9b9a 100644 --- a/crates/pgls_lsp/src/capabilities.rs +++ b/crates/pgls_lsp/src/capabilities.rs @@ -4,8 +4,9 @@ use pgls_workspace::features::code_actions::CommandActionCategory; use strum::IntoEnumIterator; use tower_lsp_server::ls_types::{ ClientCapabilities, CompletionOptions, ExecuteCommandOptions, HoverProviderCapability, - InlayHintServerCapabilities, OneOf, PositionEncodingKind, SaveOptions, ServerCapabilities, - TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions, + InlayHintServerCapabilities, OneOf, PositionEncodingKind, SaveOptions, + SemanticTokensFullOptions, SemanticTokensOptions, SemanticTokensServerCapabilities, + ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind, TextDocumentSyncOptions, TextDocumentSyncSaveOptions, WorkDoneProgressOptions, }; @@ -72,6 +73,14 @@ pub(crate) fn server_capabilities(capabilities: &ClientCapabilities) -> ServerCa inlay_hint_provider: Some(OneOf::Right(InlayHintServerCapabilities::Options( Default::default(), ))), + semantic_tokens_provider: Some(SemanticTokensServerCapabilities::SemanticTokensOptions( + SemanticTokensOptions { + legend: crate::handlers::semantic_tokens::legend(), + range: Some(true), + full: Some(SemanticTokensFullOptions::Bool(true)), + ..Default::default() + }, + )), ..Default::default() } } diff --git a/crates/pgls_lsp/src/handlers.rs b/crates/pgls_lsp/src/handlers.rs index e21e1a917..a5fdf9e2a 100644 --- a/crates/pgls_lsp/src/handlers.rs +++ b/crates/pgls_lsp/src/handlers.rs @@ -4,4 +4,5 @@ pub(crate) mod configuration; pub(crate) mod formatting; pub(crate) mod hover; pub(crate) mod inlay_hints; +pub(crate) mod semantic_tokens; pub(crate) mod text_document; diff --git a/crates/pgls_lsp/src/handlers/semantic_tokens.rs b/crates/pgls_lsp/src/handlers/semantic_tokens.rs new file mode 100644 index 000000000..ff029a94d --- /dev/null +++ b/crates/pgls_lsp/src/handlers/semantic_tokens.rs @@ -0,0 +1,125 @@ +use pgls_workspace::{ + WorkspaceError, + features::semantic_tokens::{PullSemanticTokensParams, SemanticTokenKind}, +}; +use tower_lsp_server::ls_types::{ + self, SemanticTokenModifier, SemanticTokenType, SemanticTokens, SemanticTokensLegend, +}; + +use crate::{ + adapters::{from_lsp, to_lsp}, + diagnostics::LspError, + session::Session, +}; + +/// The token types, in the order of their indices in the legend. +const TOKEN_TYPES: [SemanticTokenType; 7] = [ + SemanticTokenType::NAMESPACE, + SemanticTokenType::CLASS, + SemanticTokenType::PROPERTY, + SemanticTokenType::FUNCTION, + SemanticTokenType::TYPE, + SemanticTokenType::PARAMETER, + SemanticTokenType::ENUM_MEMBER, +]; + +/// The bit of the `defaultLibrary` modifier, the only one in the legend. +const DEFAULT_LIBRARY: u32 = 1; + +pub(crate) fn legend() -> SemanticTokensLegend { + SemanticTokensLegend { + token_types: TOKEN_TYPES.to_vec(), + token_modifiers: vec![SemanticTokenModifier::DEFAULT_LIBRARY], + } +} + +fn token_type(kind: SemanticTokenKind) -> u32 { + match kind { + SemanticTokenKind::Namespace => 0, + SemanticTokenKind::Class => 1, + SemanticTokenKind::Property => 2, + SemanticTokenKind::Function => 3, + SemanticTokenKind::Type => 4, + SemanticTokenKind::Parameter => 5, + SemanticTokenKind::EnumMember => 6, + } +} + +#[tracing::instrument(level = "debug", skip(session), err)] +pub(crate) fn semantic_tokens_full( + session: &Session, + params: ls_types::SemanticTokensParams, +) -> Result, LspError> { + Ok(semantic_tokens(session, ¶ms.text_document.uri, None)? + .map(ls_types::SemanticTokensResult::Tokens)) +} + +#[tracing::instrument(level = "debug", skip(session), err)] +pub(crate) fn semantic_tokens_range( + session: &Session, + params: ls_types::SemanticTokensRangeParams, +) -> Result, LspError> { + Ok( + semantic_tokens(session, ¶ms.text_document.uri, Some(params.range))? + .map(ls_types::SemanticTokensRangeResult::Tokens), + ) +} + +/// The tokens of the document, or of the statements that overlap `range`, encoded relative to +/// each other as the protocol requires. +fn semantic_tokens( + session: &Session, + url: &ls_types::Uri, + range: Option, +) -> Result, LspError> { + let path = session.file_path(url)?; + let doc = session.document(url)?; + let encoding = session.position_encoding(); + let range = range + .map(|range| from_lsp::text_range(&doc.line_index, range, encoding)) + .transpose()?; + + let result = match session + .workspace + .pull_semantic_tokens(PullSemanticTokensParams { path, range }) + { + Ok(result) => result, + Err(e) => match e { + WorkspaceError::DatabaseConnectionError(_) => return Ok(None), + _ => return Err(e.into()), + }, + }; + + let mut data = Vec::with_capacity(result.tokens.len()); + let mut previous = ls_types::Position::default(); + for token in result.tokens { + let start = to_lsp::position(&doc.line_index, token.range.start(), encoding)?; + let end = to_lsp::position(&doc.line_index, token.range.end(), encoding)?; + // Clients only support tokens on a single line by default. + if start.line != end.line { + continue; + } + let delta_line = start.line - previous.line; + data.push(ls_types::SemanticToken { + delta_line, + delta_start: if delta_line == 0 { + start.character - previous.character + } else { + start.character + }, + length: end.character - start.character, + token_type: token_type(token.kind), + token_modifiers_bitset: if token.default_library { + DEFAULT_LIBRARY + } else { + 0 + }, + }); + previous = start; + } + + Ok(Some(SemanticTokens { + result_id: None, + data, + })) +} diff --git a/crates/pgls_lsp/src/server.rs b/crates/pgls_lsp/src/server.rs index 94f473687..3a95410f6 100644 --- a/crates/pgls_lsp/src/server.rs +++ b/crates/pgls_lsp/src/server.rs @@ -283,6 +283,28 @@ impl LanguageServer for LSPServer { } } + #[tracing::instrument(level = "trace", skip_all)] + async fn semantic_tokens_full( + &self, + params: SemanticTokensParams, + ) -> LspResult> { + match handlers::semantic_tokens::semantic_tokens_full(&self.session, params) { + Ok(result) => LspResult::Ok(result), + Err(e) => LspResult::Err(into_lsp_error(e)), + } + } + + #[tracing::instrument(level = "trace", skip_all)] + async fn semantic_tokens_range( + &self, + params: SemanticTokensRangeParams, + ) -> LspResult> { + match handlers::semantic_tokens::semantic_tokens_range(&self.session, params) { + Ok(result) => LspResult::Ok(result), + Err(e) => LspResult::Err(into_lsp_error(e)), + } + } + #[tracing::instrument(level = "trace", skip_all)] async fn completion(&self, params: CompletionParams) -> LspResult> { match handlers::completions::get_completions(&self.session, params) { @@ -519,6 +541,7 @@ impl ServerFactory { workspace_method!(builder, check_database_connection); workspace_method!(builder, get_completions); workspace_method!(builder, pull_inlay_hints); + workspace_method!(builder, pull_semantic_tokens); workspace_method!(builder, register_project_folder); workspace_method!(builder, unregister_project_folder); workspace_method!(builder, invalidate_schema_cache); diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index c58197149..5668f0ff6 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -926,6 +926,101 @@ async fn test_inlay_hints(test_db: PgPool) -> Result<()> { Ok(()) } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_semantic_tokens(test_db: PgPool) -> Result<()> { + let factory = ServerFactory::default(); + let mut fs = MemoryFileSystem::default(); + + test_db + .execute("create table people (id int, name text);") + .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() + }); + fs.insert( + url!("postgres-language-server.jsonc") + .to_file_path() + .unwrap() + .into_owned(), + 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, _) = channel(CHANNEL_BUFFER_SIZE); + let reader = tokio::spawn(client_handler(stream, sink, sender)); + + server.initialize().await?; + server.initialized().await?; + + server.load_configuration().await?; + + server + .open_document("select 1;\nselect name, now() from public.people;\n") + .await?; + + let tokens: Option = server + .request( + "textDocument/semanticTokens/full", + "semantic_tokens", + lsp::SemanticTokensParams { + work_done_progress_params: WorkDoneProgressParams::default(), + partial_result_params: lsp::PartialResultParams::default(), + text_document: TextDocumentIdentifier { + uri: url!("document.sql"), + }, + }, + ) + .await?; + + let Some(lsp::SemanticTokensResult::Tokens(tokens)) = tokens else { + bail!("expected semantic tokens, got {tokens:?}"); + }; + let token = + |delta_line, delta_start, length, token_type, token_modifiers_bitset| lsp::SemanticToken { + delta_line, + delta_start, + length, + token_type, + token_modifiers_bitset, + }; + assert_eq!( + tokens.data, + vec![ + // `name`, a property. + token(1, 7, 4, 2, 0), + // `now`, a function of `pg_catalog`. + token(0, 6, 3, 3, 1), + // `public`, a namespace, and `people`, a class. + token(0, 11, 6, 0, 0), + token(0, 7, 6, 1, 0), + ] + ); + + server.shutdown().await?; + reader.abort(); + + Ok(()) +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_completions_is_incomplete_when_truncated(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); diff --git a/crates/pgls_semantic_tokens/Cargo.toml b/crates/pgls_semantic_tokens/Cargo.toml new file mode 100644 index 000000000..de246a2ed --- /dev/null +++ b/crates/pgls_semantic_tokens/Cargo.toml @@ -0,0 +1,25 @@ +[package] +authors.workspace = true +categories.workspace = true +description = "Semantic tokens for names that resolve to catalog objects" +edition.workspace = true +homepage.workspace = true +keywords.workspace = true +license.workspace = true +name = "pgls_semantic_tokens" +repository.workspace = true +version = "0.0.0" + + +[dependencies] +pgls_catalog = { workspace = true } +pgls_query.workspace = true +pgls_text_size.workspace = true + +[dev-dependencies] +pgls_catalog = { workspace = true, features = ["db"] } +pgls_test_utils.workspace = true +sqlx.workspace = true + +[lib] +doctest = false diff --git a/crates/pgls_semantic_tokens/src/lib.rs b/crates/pgls_semantic_tokens/src/lib.rs new file mode 100644 index 000000000..917503880 --- /dev/null +++ b/crates/pgls_semantic_tokens/src/lib.rs @@ -0,0 +1,239 @@ +//! Semantic tokens for a single statement: what the names in it refer to, as in +//! `public.users` (a schema and a table), `make_point` (a function), or `'happy'` (a member of +//! an enum). +//! +//! Tokens only come from the names that the catalog resolver resolves to exactly one object, so +//! they follow Postgres' name resolution rather than guessing from the name alone. Names that +//! match no object or several, and constructs the resolver doesn't model, get no token and keep +//! the highlighting of the editor's grammar. + +use pgls_catalog::CatalogView; +use pgls_catalog::resolve::{ColumnSource, FunctionContext, Reference, ResolveParams, resolve}; +use pgls_query::NodeEnum; +use pgls_text_size::{TextRange, TextSize}; + +pub struct SemanticTokensParams<'a> { + pub stmt: &'a NodeEnum, + /// The statement text. The AST's locations point into it. + pub sql: &'a str, + pub catalog: &'a dyn CatalogView, + /// The explicit search path of the session. + pub search_path: &'a [String], + /// Set when the statement is the body of a SQL function. + pub function: Option<&'a FunctionContext>, +} + +/// What a token refers to. The names follow the standard token types of the language server +/// protocol, which editor themes color. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SemanticTokenKind { + /// A schema. + Namespace, + /// A table, view, or other relation. + Class, + /// A column. + Property, + Function, + Type, + /// A parameter of the SQL function whose body the statement is. + Parameter, + /// A string literal that is coerced to an enum type. + EnumMember, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SemanticToken { + /// The range in the statement. + pub range: TextRange, + pub kind: SemanticTokenKind, + /// Whether the object is built into Postgres: it is in `pg_catalog` or `information_schema`. + pub default_library: bool, +} + +/// The semantic tokens of a statement, ordered by their start. Tokens don't overlap. +pub fn semantic_tokens(params: SemanticTokensParams<'_>) -> Vec { + let references = resolve(ResolveParams { + stmt: params.stmt, + catalog: params.catalog, + search_path: params.search_path, + function: params.function, + sql: Some(params.sql), + collect_references: true, + }) + .references; + + let mut tokens = Vec::new(); + for reference in &references { + let Some(span) = reference.span() else { + continue; + }; + let Some(text) = params.sql.get(std::ops::Range::::from(span)) else { + continue; + }; + let parts: Vec = name_parts(text) + .into_iter() + .map(|part| part + span.start()) + .collect(); + let mut push = |range: Option<&TextRange>, kind, schema: Option<&str>| { + if let Some(range) = range { + tokens.push(SemanticToken { + range: *range, + kind, + default_library: schema.is_some_and(is_builtin_schema), + }); + } + }; + // The schema qualifier of a name, like `public` in `public.users`. + let qualifier = |parts: &[TextRange]| parts.len().checked_sub(2).map(|i| parts[i]); + match reference { + Reference::Relation { relation, .. } => { + let schema = Some(relation.schema.as_str()); + push( + qualifier(&parts).as_ref(), + SemanticTokenKind::Namespace, + schema, + ); + push(parts.last(), SemanticTokenKind::Class, schema); + } + Reference::Function { function, .. } => { + let schema = Some(function.schema.as_str()); + push( + qualifier(&parts).as_ref(), + SemanticTokenKind::Namespace, + schema, + ); + push(parts.last(), SemanticTokenKind::Function, schema); + } + Reference::Type { type_info, .. } => { + let schema = Some(type_info.schema.as_str()); + push( + qualifier(&parts).as_ref(), + SemanticTokenKind::Namespace, + schema, + ); + push(parts.last(), SemanticTokenKind::Type, schema); + } + Reference::Column { source, .. } => { + let relation = match source { + ColumnSource::Relation { schema, name } => Some((schema.as_str(), name)), + ColumnSource::Cte { .. } | ColumnSource::Subquery { .. } => None, + }; + let schema = relation.map(|(schema, _)| schema); + match parts.as_slice() { + // `schema.table.column` + [schema_part, table, _] => { + push(Some(schema_part), SemanticTokenKind::Namespace, schema); + push(Some(table), SemanticTokenKind::Class, schema); + } + // `table.column`, but not `alias.column`. + [table, _] + if relation + .is_some_and(|(_, name)| identifier(¶ms.sql[*table]) == *name) => + { + push(Some(table), SemanticTokenKind::Class, schema); + } + _ => {} + } + push(parts.last(), SemanticTokenKind::Property, schema); + } + Reference::Parameter { .. } => { + // `function_name.parameter` + if parts.len() == 2 { + push(parts.first(), SemanticTokenKind::Function, None); + } + push(parts.last(), SemanticTokenKind::Parameter, None); + } + Reference::EnumLiteral { type_info, .. } => { + push( + Some(&span), + SemanticTokenKind::EnumMember, + Some(type_info.schema.as_str()), + ); + } + } + } + + tokens.sort_by_key(|token| (token.range.start(), token.range.end())); + let mut end = TextSize::default(); + tokens.retain(|token| { + let keep = !token.range.is_empty() && token.range.start() >= end; + if keep { + end = token.range.end(); + } + keep + }); + tokens +} + +fn is_builtin_schema(schema: &str) -> bool { + matches!(schema, "pg_catalog" | "information_schema") +} + +/// The ranges of the identifiers of a (possibly qualified) name like `public . "My Table"`, +/// without the dots and whitespace between them. +fn name_parts(text: &str) -> Vec { + let mut parts = Vec::new(); + let mut chars = text.char_indices().peekable(); + while let Some(&(start, c)) = chars.peek() { + if c.is_whitespace() || c == '.' { + chars.next(); + continue; + } + let mut end = start; + if c == '"' { + chars.next(); + // A doubled quote is a quote inside the identifier. + while let Some((i, c)) = chars.next() { + if c == '"' && chars.next_if(|&(_, c)| c == '"').is_none() { + end = i + 1; + break; + } + } + } else { + while let Some((i, c)) = chars.next_if(|&(_, c)| !c.is_whitespace() && c != '.') { + end = i + c.len_utf8(); + } + } + if let (Ok(start), Ok(end)) = (TextSize::try_from(start), TextSize::try_from(end)) + && start < end + { + parts.push(TextRange::new(start, end)); + } + } + parts +} + +/// The name an identifier refers to: unquoted identifiers are folded to lower case. +fn identifier(text: &str) -> String { + match text + .strip_prefix('"') + .and_then(|text| text.strip_suffix('"')) + { + Some(quoted) => quoted.replace("\"\"", "\""), + None => text.to_lowercase(), + } +} + +#[cfg(test)] +mod tests { + use super::name_parts; + + fn parts(text: &str) -> Vec<&str> { + name_parts(text) + .into_iter() + .map(|range| &text[range]) + .collect() + } + + #[test] + fn splits_qualified_names() { + assert_eq!(parts("users"), ["users"]); + assert_eq!(parts("public.users"), ["public", "users"]); + assert_eq!(parts("public . users"), ["public", "users"]); + assert_eq!( + parts(r#""My.Schema"."a ""b"" c".id"#), + [r#""My.Schema""#, r#""a ""b"" c""#, "id"] + ); + assert_eq!(parts("$1"), ["$1"]); + } +} diff --git a/crates/pgls_semantic_tokens/tests/semantic_tokens_tests.rs b/crates/pgls_semantic_tokens/tests/semantic_tokens_tests.rs new file mode 100644 index 000000000..fbd5f82ab --- /dev/null +++ b/crates/pgls_semantic_tokens/tests/semantic_tokens_tests.rs @@ -0,0 +1,153 @@ +use std::sync::Arc; + +use pgls_catalog::resolve::{FunctionContext, FunctionParam}; +use pgls_catalog::{Catalog, CatalogBase, Snapshot}; +use pgls_semantic_tokens::{SemanticTokenKind, SemanticTokensParams, semantic_tokens}; +use sqlx::{Executor, PgPool}; + +const SETUP: &str = r#" + create type mood as enum ('happy', 'sad'); + create table users (id int8 primary key, name text, current mood); + create table posts (id int8 primary key, user_id int8, title text); + create view active_users as select id, name from users; + create function make_point(x int, y int) returns int language sql as 'select x + y'; + create function amb(a int, b text) returns int language sql as 'select 1'; + create function amb(c text, d int) returns int language sql as 'select 1'; +"#; + +struct Fixture { + catalog: Catalog, + search_path: Vec, +} + +async fn fixture(test_db: &PgPool) -> Fixture { + test_db + .execute(SETUP) + .await + .expect("Failed to setup test database"); + let snapshot = Snapshot::load(test_db) + .await + .expect("Failed to load snapshot"); + Fixture { + search_path: snapshot.expand_search_path(&["public".to_string()]), + catalog: Catalog::new(Some(Arc::new(CatalogBase::new(Arc::new(snapshot))))), + } +} + +impl Fixture { + /// The statement with each token marked as `[text:kind]`, with a `+` after the kind of + /// built-in objects. + fn highlighted(&self, sql: &str) -> String { + self.highlighted_in(sql, None) + } + + fn highlighted_in(&self, sql: &str, function: Option<&FunctionContext>) -> String { + let stmt = pgls_query::parse(sql) + .expect("Failed to parse") + .into_root() + .expect("Expected a statement"); + + let tokens = semantic_tokens(SemanticTokensParams { + stmt: &stmt, + sql, + catalog: &self.catalog, + search_path: &self.search_path, + function, + }); + + let mut result = sql.to_string(); + for token in tokens.iter().rev() { + let kind = match token.kind { + SemanticTokenKind::Namespace => "namespace", + SemanticTokenKind::Class => "class", + SemanticTokenKind::Property => "property", + SemanticTokenKind::Function => "function", + SemanticTokenKind::Type => "type", + SemanticTokenKind::Parameter => "parameter", + SemanticTokenKind::EnumMember => "enumMember", + }; + let builtin = if token.default_library { "+" } else { "" }; + result.insert_str( + usize::from(token.range.end()), + &format!(":{kind}{builtin}]"), + ); + result.insert(usize::from(token.range.start()), '['); + } + result + } +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn relations_and_columns(test_db: PgPool) { + let f = fixture(&test_db).await; + assert_eq!( + f.highlighted("select u.name, title from public.users u join posts on true"), + "select u.[name:property], [title:property] \ + from [public:namespace].[users:class] u join [posts:class] on true" + ); + assert_eq!( + f.highlighted("select users.id from users"), + "select [users:class].[id:property] from [users:class]" + ); + assert_eq!( + f.highlighted("select public.users.id from public.users"), + "select [public:namespace].[users:class].[id:property] \ + from [public:namespace].[users:class]" + ); + assert_eq!( + f.highlighted(r#"select "name" from "active_users""#), + r#"select ["name":property] from ["active_users":class]"# + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn functions_and_types(test_db: PgPool) { + let f = fixture(&test_db).await; + assert_eq!( + f.highlighted("select make_point(1, 2)::int8, pg_catalog.now(), cast(null as text)"), + "select [make_point:function](1, 2)::[int8:type+], \ + [pg_catalog:namespace+].[now:function+](), cast(null as [text:type+])" + ); + // Ambiguous and unknown calls have no token. + assert_eq!( + f.highlighted("select amb('1', '2'), nope(1)"), + "select amb('1', '2'), nope(1)" + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn enum_members(test_db: PgPool) { + let f = fixture(&test_db).await; + assert_eq!( + f.highlighted("update users set current = 'sad' where current in ('happy', 'sad')"), + "update [users:class] set current = ['sad':enumMember] \ + where [current:property] in (['happy':enumMember], ['sad':enumMember])" + ); + assert_eq!( + f.highlighted("select 'happy'::mood, 'happy'"), + "select ['happy':enumMember]::[mood:type], 'happy'" + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn function_parameters(test_db: PgPool) { + let f = fixture(&test_db).await; + let param = |name: &str| FunctionParam { + name: Some(name.to_owned()), + type_schema: None, + type_name: "int8".to_owned(), + is_array: false, + }; + let function = FunctionContext { + function_name: "lookup".to_owned(), + params: vec![param("user_id"), param("name")], + }; + assert_eq!( + f.highlighted_in( + "select title from posts where user_id = $1 or id = lookup.user_id", + Some(&function) + ), + "select [title:property] from [posts:class] where [user_id:property] = [$1:parameter] \ + or [id:property] = [lookup:function].[user_id:parameter]" + ); +} diff --git a/crates/pgls_workspace/Cargo.toml b/crates/pgls_workspace/Cargo.toml index 280b3733c..cf70219ad 100644 --- a/crates/pgls_workspace/Cargo.toml +++ b/crates/pgls_workspace/Cargo.toml @@ -33,6 +33,7 @@ pgls_matcher = { workspace = true } pgls_pretty_print = { workspace = true } pgls_query = { workspace = true } pgls_query_ext = { workspace = true } +pgls_semantic_tokens = { workspace = true } pgls_statement_splitter = { workspace = true } pgls_suppressions = { workspace = true } pgls_text_edit.workspace = true diff --git a/crates/pgls_workspace/src/features/mod.rs b/crates/pgls_workspace/src/features/mod.rs index 6f1747ec1..69db224dc 100644 --- a/crates/pgls_workspace/src/features/mod.rs +++ b/crates/pgls_workspace/src/features/mod.rs @@ -4,3 +4,4 @@ pub mod diagnostics; pub mod format; pub mod inlay_hints; pub mod on_hover; +pub mod semantic_tokens; diff --git a/crates/pgls_workspace/src/features/semantic_tokens.rs b/crates/pgls_workspace/src/features/semantic_tokens.rs new file mode 100644 index 000000000..c49bd2499 --- /dev/null +++ b/crates/pgls_workspace/src/features/semantic_tokens.rs @@ -0,0 +1,63 @@ +use pgls_fs::PgLSPath; +use pgls_text_size::TextRange; + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct PullSemanticTokensParams { + pub path: PgLSPath, + /// The tokens of the statements that overlap this range are returned. `None` for the whole + /// file. + pub range: Option, +} + +#[derive(Debug, serde::Serialize, serde::Deserialize, Default)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct PullSemanticTokensResult { + /// The tokens, ordered by their start. Tokens don't overlap. + pub tokens: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] +pub struct SemanticToken { + /// The range in the file. + pub range: TextRange, + pub kind: SemanticTokenKind, + /// Whether the object is built into Postgres: it is in `pg_catalog` or `information_schema`. + pub default_library: bool, +} + +/// What a token refers to. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] +pub enum SemanticTokenKind { + /// A schema. + Namespace, + /// A table, view, or other relation. + Class, + /// A column. + Property, + Function, + Type, + /// A parameter of a SQL function, in its body. + Parameter, + /// A string literal that is coerced to an enum type. + EnumMember, +} + +impl From for SemanticTokenKind { + fn from(kind: pgls_semantic_tokens::SemanticTokenKind) -> Self { + use pgls_semantic_tokens::SemanticTokenKind as Kind; + match kind { + Kind::Namespace => Self::Namespace, + Kind::Class => Self::Class, + Kind::Property => Self::Property, + Kind::Function => Self::Function, + Kind::Type => Self::Type, + Kind::Parameter => Self::Parameter, + Kind::EnumMember => Self::EnumMember, + } + } +} diff --git a/crates/pgls_workspace/src/workspace.rs b/crates/pgls_workspace/src/workspace.rs index 0966b1e0f..f9a800d45 100644 --- a/crates/pgls_workspace/src/workspace.rs +++ b/crates/pgls_workspace/src/workspace.rs @@ -24,6 +24,7 @@ use crate::{ format::{PullFileFormattingParams, PullFormattingResult}, inlay_hints::{PullInlayHintsParams, PullInlayHintsResult}, on_hover::{OnHoverParams, OnHoverResult}, + semantic_tokens::{PullSemanticTokensParams, PullSemanticTokensResult}, }, }; @@ -154,6 +155,13 @@ pub trait Workspace: Send + Sync + RefUnwindSafe { params: PullInlayHintsParams, ) -> Result; + /// The semantic tokens of the names that resolve to exactly one object, like tables, + /// columns, and functions, in the statements that overlap a range or the whole file. + fn pull_semantic_tokens( + &self, + params: PullSemanticTokensParams, + ) -> Result; + /// Register a possible workspace project folder. Returns the key of said project. Use this key when you want to switch to different projects. fn register_project_folder( &self, diff --git a/crates/pgls_workspace/src/workspace/client.rs b/crates/pgls_workspace/src/workspace/client.rs index ae5789879..261a665d3 100644 --- a/crates/pgls_workspace/src/workspace/client.rs +++ b/crates/pgls_workspace/src/workspace/client.rs @@ -204,6 +204,13 @@ where self.request("pgls/pull_inlay_hints", params) } + fn pull_semantic_tokens( + &self, + params: crate::features::semantic_tokens::PullSemanticTokensParams, + ) -> Result { + self.request("pgls/pull_semantic_tokens", params) + } + fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError> { self.request("pgls/invalidate_schema_cache", all) } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index b186db9f2..0a92c5029 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -68,6 +68,7 @@ use crate::{ format::{PullFileFormattingParams, PullFormattingResult, StatementFormatResult}, inlay_hints::{InlayHint, PullInlayHintsParams, PullInlayHintsResult}, on_hover::{OnHoverParams, OnHoverResult}, + semantic_tokens::{PullSemanticTokensParams, PullSemanticTokensResult, SemanticToken}, }, settings::{Settings, WorkspaceSettings, WorkspaceSettingsHandle, WorkspaceSettingsHandleMut}, workspace::{AnalyserDiagnosticsMapper, WithCSTandASTMapper}, @@ -1853,6 +1854,81 @@ impl Workspace for WorkspaceServer { hints.sort_by_key(|hint| hint.offset); Ok(PullInlayHintsResult { hints }) } + + #[ignored_path(path=¶ms.path)] + #[tracing::instrument(level = "debug", skip_all, fields( + path = params.path.as_os_str().to_str(), + range = ?params.range + ), err)] + fn pull_semantic_tokens( + &self, + params: PullSemanticTokensParams, + ) -> Result { + let documents = self.documents.read().unwrap(); + let doc = documents + .get(¶ms.path) + .ok_or(WorkspaceError::not_found())?; + + #[cfg(feature = "db")] + let Some(snapshot) = self + .workspaces() + .settings() + .and_then(|settings| { + self.connection + .with_pool(&settings.db, |pool| self.snapshot.load(pool)) + }) + .transpose()? + else { + tracing::debug!("No database connection available. Skipping semantic tokens."); + return Ok(PullSemanticTokensResult::default()); + }; + + #[cfg(not(feature = "db"))] + let Some(snapshot) = self.snapshot.get() else { + tracing::debug!("No schema loaded. Skipping semantic tokens."); + return Ok(PullSemanticTokensResult::default()); + }; + + let range = params + .range + .unwrap_or_else(|| TextRange::up_to(TextSize::of(doc.get_document_content()))); + let mut tokens = Vec::new(); + self.walk_statements(doc, &snapshot, range.end(), |stmt, catalog, session| { + let Some(sql) = stmt.sql.as_deref() else { + return; + }; + if stmt.range.intersect(range).is_none() { + return; + } + tokens.extend( + pgls_semantic_tokens::semantic_tokens(pgls_semantic_tokens::SemanticTokensParams { + stmt: &stmt.root, + sql, + catalog, + search_path: session.search_path(), + function: stmt.function.as_ref(), + }) + .into_iter() + .map(|token| SemanticToken { + range: token.range + stmt.range.start(), + kind: token.kind.into(), + default_library: token.default_library, + }), + ); + }); + + // The body of a SQL function is a statement of its own, inside its `CREATE FUNCTION`. + tokens.sort_by_key(|token| (token.range.start(), token.range.end())); + let mut end = TextSize::default(); + tokens.retain(|token| { + let keep = token.range.start() >= end; + if keep { + end = token.range.end(); + } + keep + }); + Ok(PullSemanticTokensResult { tokens }) + } } /// The migration sections of `doc`, like the up and down migrations of dbmate and goose. diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 8f0e857a0..3c420a1e3 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -27,6 +27,7 @@ use crate::{ features::format::PullFileFormattingParams, features::inlay_hints::{InlayHint, InlayHintKind, PullInlayHintsParams}, features::on_hover::OnHoverParams, + features::semantic_tokens::{PullSemanticTokensParams, SemanticTokenKind}, workspace::{ OpenFileParams, RegisterProjectFolderParams, StatementId, UpdateSettingsParams, server::WorkspaceServer, @@ -2043,6 +2044,88 @@ async fn test_inlay_hints(test_db: PgPool) { ); } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_semantic_tokens(test_db: PgPool) { + test_db + .execute("create table db_items (id int, label text);") + .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 content = "create type file_mood as enum ('up', 'down');\n\ + create table file_items (id int, mood file_mood);\n\ + select id from file_items where mood = 'up';\n\ + create function by_id(item_id int) returns text language sql \ + as 'select label from public.db_items where id = item_id';\n"; + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: content.into(), + version: 1, + }) + .expect("Unable to open test file"); + + let tokens = |range: Option| { + workspace + .pull_semantic_tokens(PullSemanticTokensParams { + path: path.clone(), + range, + }) + .expect("Unable to pull semantic tokens") + .tokens + .into_iter() + .map(|token| (&content[token.range], token.kind)) + .collect::>() + }; + + // The table and the enum created in the file, and the body of the SQL function. + assert_eq!( + tokens(None), + vec![ + ("id", SemanticTokenKind::Property), + ("file_items", SemanticTokenKind::Class), + ("mood", SemanticTokenKind::Property), + ("'up'", SemanticTokenKind::EnumMember), + ("label", SemanticTokenKind::Property), + ("public", SemanticTokenKind::Namespace), + ("db_items", SemanticTokenKind::Class), + ("id", SemanticTokenKind::Property), + ("item_id", SemanticTokenKind::Parameter), + ] + ); + + // Only the statements that overlap the range. + let select = content.find("select id").unwrap(); + assert_eq!( + tokens(Some(TextRange::new( + u32::try_from(select).unwrap().into(), + u32::try_from(select + 1).unwrap().into(), + ))), + vec![ + ("id", SemanticTokenKind::Property), + ("file_items", SemanticTokenKind::Class), + ("mood", SemanticTokenKind::Property), + ("'up'", SemanticTokenKind::EnumMember), + ] + ); +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_hover_shows_resolved_references(test_db: PgPool) { test_db diff --git a/docs/features/editor_features.md b/docs/features/editor_features.md index 2d8c598c0..8997a0573 100644 --- a/docs/features/editor_features.md +++ b/docs/features/editor_features.md @@ -43,3 +43,19 @@ The language server caches schema information on startup. Function calls show the names of the parameters their positional arguments are passed to, as in `make_point(x: 1, y: 2)`. Hints are only shown when the call resolves to exactly one function, including functions created earlier in the same file. Named arguments (`f(a => 1)`), unnamed parameters, and arguments that already carry the parameter's name (`f(t.x)` for the parameter `x`) get no hint. Of the elements passed to a `VARIADIC` parameter, only the first is labelled. Inlay hints require a database connection. Enable them in your editor, e.g. with `editor.inlayHints.enabled` in VS Code. + +## Semantic Highlighting + +The language server provides semantic tokens for the names that resolve to exactly one database object, so your editor can color them by what they refer to rather than by how they look: + +- **Schemas** (`namespace`), like `public` in `public.users` +- **Tables, views, and other relations** (`class`) +- **Columns** (`property`) +- **Functions** (`function`), the overload that the call resolves to +- **Types** (`type`), like `int8` in `x::int8` +- **Parameters** (`parameter`) of a SQL function, in its body, by name or as `$1` +- **Enum values** (`enumMember`): string literals that are coerced to an enum type, like `'active'` in `status = 'active'`, `'active'::status`, or `insert into t (status) values ('active')` + +Objects in `pg_catalog` and `information_schema` carry the `defaultLibrary` modifier. Names follow Postgres' name resolution, including objects created earlier in the same file and statements in the bodies of SQL functions. Names that match no object or several get no token and keep the highlighting of your editor's grammar. + +Semantic highlighting requires a database connection. In VS Code, it is controlled by `editor.semanticHighlighting.enabled`. From edbb4f0c1124c2a84ee2b26c23533104775dd219 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Sat, 24 Oct 2026 06:12:59 +0000 Subject: [PATCH 2/2] feat(lsp): refresh and settings for inlay hints and semantic tokens Ask the client to request inlay hints and semantic tokens again when the configuration changes, the schema cache is invalidated, a statement is executed, or the database becomes reachable after they were requested without it. Add inlayHints.enabled and semanticTokens.enabled, and rewrite the editor features docs. --- README.md | 2 +- crates/pgls_configuration/src/inlay_hints.rs | 20 ++ crates/pgls_configuration/src/lib.rs | 25 +++ .../pgls_configuration/src/semantic_tokens.rs | 21 ++ crates/pgls_lsp/src/configuration.rs | 4 +- crates/pgls_lsp/src/handlers/code_actions.rs | 2 + crates/pgls_lsp/src/handlers/inlay_hints.rs | 5 +- .../pgls_lsp/src/handlers/semantic_tokens.rs | 5 +- crates/pgls_lsp/src/server.rs | 1 + crates/pgls_lsp/src/session.rs | 66 +++++- crates/pgls_lsp/tests/server.rs | 199 +++++++++++++++++- crates/pgls_workspace/src/settings.rs | 46 ++++ crates/pgls_workspace/src/workspace/server.rs | 16 ++ docs/features/editor_features.md | 107 +++++++--- docs/index.md | 2 +- docs/schema.json | 50 +++++ mkdocs.yml | 2 +- .../backend-jsonrpc/src/workspace.ts | 26 +++ .../backend-jsonrpc/src/workspace.ts | 26 +++ 19 files changed, 581 insertions(+), 44 deletions(-) create mode 100644 crates/pgls_configuration/src/inlay_hints.rs create mode 100644 crates/pgls_configuration/src/semantic_tokens.rs diff --git a/README.md b/README.md index c2362815a..2a00c6031 100644 --- a/README.md +++ b/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/) diff --git a/crates/pgls_configuration/src/inlay_hints.rs b/crates/pgls_configuration/src/inlay_hints.rs new file mode 100644 index 000000000..937d4bd5a --- /dev/null +++ b/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 } + } +} diff --git a/crates/pgls_configuration/src/lib.rs b/crates/pgls_configuration/src/lib.rs index 579d73086..aa285af14 100644 --- a/crates/pgls_configuration/src/lib.rs +++ b/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, } } diff --git a/crates/pgls_configuration/src/semantic_tokens.rs b/crates/pgls_configuration/src/semantic_tokens.rs new file mode 100644 index 000000000..93273eca1 --- /dev/null +++ b/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 } + } +} diff --git a/crates/pgls_lsp/src/configuration.rs b/crates/pgls_lsp/src/configuration.rs index f876dc6ef..8bdb6cab3 100644 --- a/crates/pgls_lsp/src/configuration.rs +++ b/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. diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index b3c2f6750..bb7a0d20f 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -176,6 +176,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 diff --git a/crates/pgls_lsp/src/handlers/inlay_hints.rs b/crates/pgls_lsp/src/handlers/inlay_hints.rs index bb41b108a..60ac72b9d 100644 --- a/crates/pgls_lsp/src/handlers/inlay_hints.rs +++ b/crates/pgls_lsp/src/handlers/inlay_hints.rs @@ -27,7 +27,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()), }, }; diff --git a/crates/pgls_lsp/src/handlers/semantic_tokens.rs b/crates/pgls_lsp/src/handlers/semantic_tokens.rs index ff029a94d..307fef448 100644 --- a/crates/pgls_lsp/src/handlers/semantic_tokens.rs +++ b/crates/pgls_lsp/src/handlers/semantic_tokens.rs @@ -85,7 +85,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()), }, }; diff --git a/crates/pgls_lsp/src/server.rs b/crates/pgls_lsp/src/server.rs index 3a95410f6..05122ec99 100644 --- a/crates/pgls_lsp/src/server.rs +++ b/crates/pgls_lsp/src/server.rs @@ -185,6 +185,7 @@ impl LanguageServer for LSPServer { { self.session.reload_workspace_settings().await; self.setup_capabilities().await; + self.session.refresh_editor_features(); // self.session.update_all_diagnostics().await; // for now we are only interested to the configuration file, // so it's OK to exist the loop diff --git a/crates/pgls_lsp/src/session.rs b/crates/pgls_lsp/src/session.rs index 7807a87bb..93ab6c5b0 100644 --- a/crates/pgls_lsp/src/session.rs +++ b/crates/pgls_lsp/src/session.rs @@ -125,6 +125,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>>, + + /// 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 @@ -229,6 +234,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), } } @@ -378,9 +384,62 @@ 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 {}); + 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, @@ -425,7 +484,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 @@ -440,6 +500,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`] diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index 5668f0ff6..253b3936b 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -144,9 +144,17 @@ impl Server { } /// Basic implementation of the `initialize` request for tests + async fn initialize(&mut self) -> Result<()> { + self.initialize_with_capabilities(ClientCapabilities::default()) + .await + } + // The `root_path` field is deprecated, but we still need to specify it #[allow(deprecated)] - async fn initialize(&mut self) -> Result<()> { + async fn initialize_with_capabilities( + &mut self, + capabilities: ClientCapabilities, + ) -> Result<()> { let _res: InitializeResult = self .request( "initialize", @@ -156,7 +164,7 @@ impl Server { root_path: None, root_uri: Some(url!("")), initialization_options: None, - capabilities: ClientCapabilities::default(), + capabilities, trace: None, workspace_folders: None, client_info: None, @@ -383,6 +391,8 @@ const CHANNEL_BUFFER_SIZE: usize = 8; #[derive(Debug, PartialEq, Eq)] enum ServerNotification { PublishDiagnostics(PublishDiagnosticsParams), + /// A `workspace/*/refresh` request, by its method. + Refresh(String), } /// Basic handler for requests and notifications coming from the server for tests @@ -413,6 +423,20 @@ where None => continue, }; + if matches!( + req.method(), + "workspace/inlayHint/refresh" | "workspace/semanticTokens/refresh" + ) { + sink.send(Response::from_ok(id.clone(), Value::Null)) + .await + .ok(); + let notification = ServerNotification::Refresh(req.method().to_string()); + match notify.send(notification).await { + Ok(_) => continue, + Err(_) => break, + } + } + let res = Response::from_error(id.clone(), jsonrpc::Error::method_not_found()); sink.send(res).await.ok(); @@ -664,6 +688,7 @@ async fn client_configuration_overrides_are_sticky_and_clearable(test_db: PgPool return true; } } + Some(ServerNotification::Refresh(_)) => {} None => return false, } } @@ -3178,3 +3203,173 @@ async fn test_crash_on_delete_character() -> Result<()> { Ok(()) } + +/// Waits for the `workspace/*/refresh` requests of inlay hints and semantic tokens. +async fn wait_for_editor_feature_refresh( + receiver: &mut Receiver, +) -> Result<()> { + let mut pending = vec![ + "workspace/inlayHint/refresh", + "workspace/semanticTokens/refresh", + ]; + tokio::time::timeout(Duration::from_secs(5), async { + while !pending.is_empty() { + match receiver.next().await { + Some(ServerNotification::Refresh(method)) => { + pending.retain(|pending| *pending != method); + } + Some(ServerNotification::PublishDiagnostics(_)) => {} + None => bail!("the server closed the connection"), + } + } + Ok(()) + }) + .await + .context("timed out waiting for the inlay hint and semantic token refresh")? +} + +async fn inlay_hint_count(server: &mut Server) -> Result { + let hints: Option> = server + .request( + "textDocument/inlayHint", + "inlay_hints", + lsp::InlayHintParams { + work_done_progress_params: WorkDoneProgressParams::default(), + text_document: TextDocumentIdentifier { + uri: url!("document.sql"), + }, + range: Range { + start: Position::new(0, 0), + end: Position::new(1, 0), + }, + }, + ) + .await?; + Ok(hints.map_or(0, |hints| hints.len())) +} + +async fn semantic_token_count(server: &mut Server) -> Result { + let tokens: Option = server + .request( + "textDocument/semanticTokens/full", + "semantic_tokens", + lsp::SemanticTokensParams { + work_done_progress_params: WorkDoneProgressParams::default(), + partial_result_params: lsp::PartialResultParams::default(), + text_document: TextDocumentIdentifier { + uri: url!("document.sql"), + }, + }, + ) + .await?; + Ok(match tokens { + Some(lsp::SemanticTokensResult::Tokens(tokens)) => tokens.data.len(), + _ => 0, + }) +} + +/// Inlay hints and semantic tokens are refreshed when the schema cache is invalidated and when the +/// configuration changes, and the configuration can turn them off. +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_editor_features_refresh_and_settings(test_db: PgPool) -> Result<()> { + let factory = ServerFactory::default(); + let mut fs = MemoryFileSystem::default(); + + test_db + .execute( + "create table people (id int, name text); + create function make_point(x int, y int) returns int language sql as 'select x + y';", + ) + .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() + }); + fs.insert( + url!("postgres-language-server.jsonc") + .to_file_path() + .unwrap() + .into_owned(), + 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_with_capabilities(ClientCapabilities { + workspace: Some(lsp::WorkspaceClientCapabilities { + inlay_hint: Some(lsp::InlayHintWorkspaceClientCapabilities { + refresh_support: Some(true), + }), + semantic_tokens: Some(lsp::SemanticTokensWorkspaceClientCapabilities { + refresh_support: Some(true), + }), + ..Default::default() + }), + ..Default::default() + }) + .await?; + server.initialized().await?; + server.load_configuration().await?; + wait_for_editor_feature_refresh(&mut receiver).await?; + + server + .open_document("select name, make_point(1, 2) from people;\n") + .await?; + + assert_eq!(inlay_hint_count(&mut server).await?, 2); + assert_eq!(semantic_token_count(&mut server).await?, 3); + + drain_notifications(&mut receiver); + server + .request::>( + "workspace/executeCommand", + "_invalidate", + ExecuteCommandParams { + command: "pgls.invalidateSchemaCache".to_string(), + arguments: vec![], + ..Default::default() + }, + ) + .await?; + wait_for_editor_feature_refresh(&mut receiver).await?; + + drain_notifications(&mut receiver); + server + .notify( + "workspace/didChangeConfiguration", + DidChangeConfigurationParams { + settings: json!({ + "inlayHints": { "enabled": false }, + "semanticTokens": { "enabled": false }, + }), + }, + ) + .await?; + wait_for_editor_feature_refresh(&mut receiver).await?; + + assert_eq!(inlay_hint_count(&mut server).await?, 0); + assert_eq!(semantic_token_count(&mut server).await?, 0); + + server.shutdown().await?; + reader.abort(); + Ok(()) +} diff --git a/crates/pgls_workspace/src/settings.rs b/crates/pgls_workspace/src/settings.rs index e289a03b4..0f32ee6bc 100644 --- a/crates/pgls_workspace/src/settings.rs +++ b/crates/pgls_workspace/src/settings.rs @@ -22,9 +22,11 @@ use pgls_configuration::{ CastStyle, ClauseBodyStyle, CommaStyle, FormatConfiguration, IndentStyle, KeywordCase, Layout, LogicalOperatorPlacement, }, + inlay_hints::InlayHintsConfiguration, migrations::{MigrationsConfiguration, PartialMigrationsConfiguration, SectionMarkers}, pglinter::PglinterConfiguration, plpgsql_check::PlPgSqlCheckConfiguration, + semantic_tokens::SemanticTokensConfiguration, splinter::SplinterConfiguration, }; use pgls_fs::PgLSPath; @@ -238,6 +240,12 @@ pub struct Settings { /// Migrations settings pub migrations: Option, + + /// Inlay hint settings for the workspace + pub inlay_hints: InlayHintsSettings, + + /// Semantic token settings for the workspace + pub semantic_tokens: SemanticTokensSettings, } impl Settings { @@ -300,6 +308,20 @@ impl Settings { to_plpgsql_check_settings(PlPgSqlCheckConfiguration::from(plpgsql_check)); } + // inlay hints part + if let Some(inlay_hints) = configuration.inlay_hints { + self.inlay_hints = InlayHintsSettings { + enabled: InlayHintsConfiguration::from(inlay_hints).enabled, + }; + } + + // semantic tokens part + if let Some(semantic_tokens) = configuration.semantic_tokens { + self.semantic_tokens = SemanticTokensSettings { + enabled: SemanticTokensConfiguration::from(semantic_tokens).enabled, + }; + } + // Migrations settings if let Some(migrations) = configuration.migrations { self.migrations = Some(to_migration_settings( @@ -674,6 +696,30 @@ impl Default for PglinterSettings { } } +/// Inlay hint settings for the entire workspace +#[derive(Debug, Clone)] +pub struct InlayHintsSettings { + pub enabled: bool, +} + +impl Default for InlayHintsSettings { + fn default() -> Self { + Self { enabled: true } + } +} + +/// Semantic token settings for the entire workspace +#[derive(Debug, Clone)] +pub struct SemanticTokensSettings { + pub enabled: bool, +} + +impl Default for SemanticTokensSettings { + fn default() -> Self { + Self { enabled: true } + } +} + /// plpgsql_check settings for the entire workspace #[derive(Debug, Clone)] pub struct PlPgSqlCheckSettings { diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 0a92c5029..127a2ac7c 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -1801,6 +1801,14 @@ impl Workspace for WorkspaceServer { .get(¶ms.path) .ok_or(WorkspaceError::not_found())?; + if !self + .workspaces() + .settings() + .is_none_or(|settings| settings.inlay_hints.enabled) + { + return Ok(PullInlayHintsResult::default()); + } + #[cfg(feature = "db")] let Some(snapshot) = self .workspaces() @@ -1869,6 +1877,14 @@ impl Workspace for WorkspaceServer { .get(¶ms.path) .ok_or(WorkspaceError::not_found())?; + if !self + .workspaces() + .settings() + .is_none_or(|settings| settings.semantic_tokens.enabled) + { + return Ok(PullSemanticTokensResult::default()); + } + #[cfg(feature = "db")] let Some(snapshot) = self .workspaces() diff --git a/docs/features/editor_features.md b/docs/features/editor_features.md index 8997a0573..5948a612c 100644 --- a/docs/features/editor_features.md +++ b/docs/features/editor_features.md @@ -1,61 +1,98 @@ -# Autocompletion & Hover +# Editor Features -The language server provides autocompletion and hover information when connected to a database. +When connected to a database, the language server helps you write SQL with autocompletion, hover information, inlay hints, and semantic highlighting. -## Autocompletion +Hover, inlay hints, and semantic highlighting follow Postgres' name resolution. They see the database schema with the DDL of the current file applied on top of it, so a table created or altered earlier in the file is known, and the search path follows `set search_path` in the file. They also work in the bodies of `LANGUAGE sql` functions, where the function's parameters are in scope. -As you type SQL, the language server suggests relevant database objects based on your current context: +## Requirements -- **Tables**: Available tables from your database schema -- **Columns**: Columns from tables referenced in your query -- **Functions**: Database functions and procedures -- **Schemas**: Available schemas in your database -- **Keywords**: SQL keywords and syntax +All these features need: -The suggestions are context-aware - for example, when typing after `FROM`, you'll see table suggestions, and when typing after `SELECT`, you'll see column suggestions from relevant tables. +- A configured database connection. See the [database configuration guide](../guides/configure_database.md). +- A database user that can read the schema. -## Hover Information +Without a database connection, they are not available. The language server caches the schema when it connects. Run the `pgls.invalidateSchemaCache` command after changing the schema outside the editor. -Hovering over database objects in your SQL shows detailed information: +## Autocompletion -- **Tables**: Schema, column list with data types -- **Columns**: Data type, nullable status, table location -- **Functions**: Return type, parameter information +As you type SQL, the language server suggests relevant database objects based on your current context: -The hover information is pulled from your database schema. +- **Tables**: Available tables from your database schema +- **Columns**: Columns from tables referenced in your query +- **Functions**: Database functions and procedures +- **Schemas**: Available schemas in your database +- **Keywords**: SQL keywords and syntax -## Requirements +The suggestions are context-aware - for example, when typing after `FROM`, you'll see table suggestions, and when typing after `SELECT`, you'll see column suggestions from relevant tables. -Both features require: -- A configured database connection -- The language server must be able to read schema information from your database +## Hover -Without a database connection, these features are not available. +Hovering a name shows the object it refers to: -## Configuration +- **Tables and views**: schema, columns with their types, and defaults. +- **Columns**: type, nullability, and the table they belong to. Columns are followed through aliases and joins, so `o.name` shows the column of the table that `o` stands for. +- **Columns of CTEs and subqueries**: their name and inferred type. +- **Functions**: only the overload that the call resolves to, with its signature and return type. +- **Schemas, types, and roles.** -These features work automatically when you have a database connection configured. See the [database configuration guide](../guides/configure_database.md) for setup instructions. +```sql +create function pick(i integer) returns integer language sql as 'select i'; +create function pick(t text) returns text language sql as 'select t'; +select pick(1); -- hover on `pick`: only `public.pick(i integer) → integer` +``` -The language server caches schema information on startup. +When a name doesn't resolve to exactly one object, hover falls back to looking the name up in the schema and shows every match. ## Inlay Hints -Function calls show the names of the parameters their positional arguments are passed to, as in `make_point(x: 1, y: 2)`. Hints are only shown when the call resolves to exactly one function, including functions created earlier in the same file. Named arguments (`f(a => 1)`), unnamed parameters, and arguments that already carry the parameter's name (`f(t.x)` for the parameter `x`) get no hint. Of the elements passed to a `VARIADIC` parameter, only the first is labelled. +Function calls show the names of the parameters their positional arguments are passed to: -Inlay hints require a database connection. Enable them in your editor, e.g. with `editor.inlayHints.enabled` in VS Code. +```sql +select make_point(1, 2 + 3); -- make_point(x: 1, y: 2 + 3) +select make_point(1, y => 2); -- make_point(x: 1, y => 2) +select sum_all(1, 2, 3); -- sum_all(base: 1, nums: 2, 3) +``` + +Hints are only shown when the call resolves to exactly one function, including functions created earlier in the same file. Named arguments (`f(a => 1)`), unnamed parameters, and arguments that already carry the parameter's name (`f(t.x)` for the parameter `x`) get no hint. Of the elements passed to a `VARIADIC` parameter, only the first is labelled. ## Semantic Highlighting The language server provides semantic tokens for the names that resolve to exactly one database object, so your editor can color them by what they refer to rather than by how they look: -- **Schemas** (`namespace`), like `public` in `public.users` -- **Tables, views, and other relations** (`class`) -- **Columns** (`property`) -- **Functions** (`function`), the overload that the call resolves to -- **Types** (`type`), like `int8` in `x::int8` -- **Parameters** (`parameter`) of a SQL function, in its body, by name or as `$1` -- **Enum values** (`enumMember`): string literals that are coerced to an enum type, like `'active'` in `status = 'active'`, `'active'::status`, or `insert into t (status) values ('active')` +| Token | LSP token type | Example | +| --- | --- | --- | +| Schema | `namespace` | `public` in `public.users` | +| Table, view, or other relation | `class` | `users` in `from users`, `users.id` | +| Column | `property` | `name`, `u.name`, `public.users.id` | +| Function | `function` | the overload that a call resolves to | +| Type | `type` | `int8` in `x::int8`, `cast(x as text)`, `date 'today'` | +| Parameter of a SQL function, in its body | `parameter` | `id`, `fn.id`, `$1` | +| Enum value | `enumMember` | `'active'` in `status = 'active'` | + +Objects in `pg_catalog` and `information_schema` carry the `defaultLibrary` modifier. Names that match no object or several get no token and keep the highlighting of your editor's grammar. + +String literals are enum values only where the enum type is certain: casts (`'active'::status`), comparisons with an enum column or value (`=`, `<>`, `IN`, `BETWEEN`, `IS DISTINCT FROM`, `NULLIF`), arguments of function parameters with the enum type, and values assigned to enum columns in `INSERT ... VALUES`, `UPDATE ... SET`, and `ON CONFLICT DO UPDATE SET`. + +## Turning Features Off + +Inlay hints and semantic highlighting are on by default. Turn them off for a project in `postgres-language-server.jsonc`: + +```jsonc +{ + "inlayHints": { "enabled": false }, + "semanticTokens": { "enabled": false } +} +``` + +Without semantic tokens, your editor keeps the highlighting of its SQL grammar. Most editors can also turn these features off on their side, e.g. `editor.inlayHints.enabled` and `editor.semanticHighlighting.enabled` in VS Code. + +The language server asks the editor to request inlay hints and semantic tokens again when the configuration changes, when the schema cache is invalidated or a statement is executed, and when the database becomes reachable, if the editor supports these refreshes. -Objects in `pg_catalog` and `information_schema` carry the `defaultLibrary` modifier. Names follow Postgres' name resolution, including objects created earlier in the same file and statements in the bodies of SQL functions. Names that match no object or several get no token and keep the highlighting of your editor's grammar. +## Limitations -Semantic highlighting requires a database connection. In VS Code, it is controlled by `editor.semanticHighlighting.enabled`. +- Only names that resolve to exactly one object get a hover from the resolver, an inlay hint, or a token. Ambiguous calls, unknown objects, and constructs the resolver doesn't model are skipped. +- After a `DO` block, `CALL`, `EXECUTE`, or a statement that doesn't parse, the rest of the file gets no function hints, function tokens, or enum tokens, because their effects on the schema are unknown. Relations and columns still resolve. +- `CALL proc(...)` gets no hints, tokens, or resolved hover. +- Statements in `BEGIN ATOMIC` function bodies and in `'...'`-quoted bodies that contain `''` are not resolved. `$$`-quoted bodies are. PL/pgSQL bodies are not covered. +- No tokens for roles, for schemas outside qualified names, for DDL (column types in `create table`, function signatures), for CTE names, or for table aliases. Tokens that span several lines are skipped. +- Each request resolves the file up to the visible range. On very long files, like schema dumps with thousands of statements, hints and tokens take longer to appear. diff --git a/docs/index.md b/docs/index.md index a2b92cd43..2f31ae270 100644 --- a/docs/index.md +++ b/docs/index.md @@ -14,7 +14,7 @@ The language server is built on Postgres' own parser `libpg_query` to ensure 100 The following features are available today: -- [Autocompletion & Hover](features/editor_features.md) +- [Autocompletion, Hover, Inlay Hints & Semantic Highlighting](features/editor_features.md) - [Syntax Diagnostics](features/syntax_diagnostics.md) - [Type Checking](features/type_checking.md) - [Formatting](features/formatting.md) diff --git a/docs/schema.json b/docs/schema.json index 94587fa0a..86eb4200a 100644 --- a/docs/schema.json +++ b/docs/schema.json @@ -55,6 +55,17 @@ } ] }, + "inlayHints": { + "description": "The configuration for inlay hints in editors", + "anyOf": [ + { + "$ref": "#/definitions/InlayHintsConfiguration" + }, + { + "type": "null" + } + ] + }, "linter": { "description": "The configuration for the linter", "anyOf": [ @@ -99,6 +110,17 @@ } ] }, + "semanticTokens": { + "description": "The configuration for semantic highlighting in editors", + "anyOf": [ + { + "$ref": "#/definitions/SemanticTokensConfiguration" + }, + { + "type": "null" + } + ] + }, "splinter": { "description": "The configuration for splinter", "anyOf": [ @@ -684,6 +706,20 @@ } ] }, + "InlayHintsConfiguration": { + "description": "The configuration for inlay hints.", + "type": "object", + "properties": { + "enabled": { + "description": "if `false`, the language server sends no inlay hints. `true` by default", + "type": [ + "boolean", + "null" + ] + } + }, + "additionalProperties": false + }, "KeywordCase": { "description": "Keyword casing style for the formatter.", "oneOf": [ @@ -3036,6 +3072,20 @@ }, "additionalProperties": false }, + "SemanticTokensConfiguration": { + "description": "The configuration for semantic highlighting.", + "type": "object", + "properties": { + "enabled": { + "description": "if `false`, the language server sends no semantic tokens, and editors keep the highlighting of their grammar. `true` by default", + "type": [ + "boolean", + "null" + ] + } + }, + "additionalProperties": false + }, "SplinterConfiguration": { "type": "object", "properties": { diff --git a/mkdocs.yml b/mkdocs.yml index 989266397..0fe0afee9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -44,7 +44,7 @@ nav: - Type Checking: features/type_checking.md - Formatting: features/formatting.md - PL/pgSQL Support: features/plpgsql.md - - Autocompletion & Hover: features/editor_features.md + - Editor Features: features/editor_features.md - Guides: - Use in your IDE: guides/ide_setup.md - Checking Migrations: guides/checking_migrations.md diff --git a/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts b/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts index 5a3b0d711..869725ce2 100644 --- a/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts +++ b/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts @@ -378,6 +378,10 @@ export interface PartialConfiguration { * The configuration for the SQL formatter */ format?: PartialFormatConfiguration; + /** + * The configuration for inlay hints in editors + */ + inlayHints?: PartialInlayHintsConfiguration; /** * The configuration for the linter */ @@ -394,6 +398,10 @@ export interface PartialConfiguration { * The configuration for type checking */ plpgsqlCheck?: PartialPlPgSqlCheckConfiguration; + /** + * The configuration for semantic highlighting in editors + */ + semanticTokens?: PartialSemanticTokensConfiguration; /** * The configuration for splinter */ @@ -536,6 +544,15 @@ export interface PartialFormatConfiguration { */ valuesRowLayout?: Layout; } +/** + * The configuration for inlay hints. + */ +export interface PartialInlayHintsConfiguration { + /** + * if `false`, the language server sends no inlay hints. `true` by default + */ + enabled?: boolean; +} export interface PartialLinterConfiguration { /** * if `false`, it disables the feature and the linter won't be executed. `true` by default @@ -642,6 +659,15 @@ export interface PartialPlPgSqlCheckConfiguration { */ withoutWarnings?: boolean; } +/** + * The configuration for semantic highlighting. + */ +export interface PartialSemanticTokensConfiguration { + /** + * if `false`, the language server sends no semantic tokens, and editors keep the highlighting of their grammar. `true` by default + */ + enabled?: boolean; +} export interface PartialSplinterConfiguration { /** * if `false`, it disables the feature and the linter won't be executed. `true` by default diff --git a/packages/@postgrestools/backend-jsonrpc/src/workspace.ts b/packages/@postgrestools/backend-jsonrpc/src/workspace.ts index 5a3b0d711..869725ce2 100644 --- a/packages/@postgrestools/backend-jsonrpc/src/workspace.ts +++ b/packages/@postgrestools/backend-jsonrpc/src/workspace.ts @@ -378,6 +378,10 @@ export interface PartialConfiguration { * The configuration for the SQL formatter */ format?: PartialFormatConfiguration; + /** + * The configuration for inlay hints in editors + */ + inlayHints?: PartialInlayHintsConfiguration; /** * The configuration for the linter */ @@ -394,6 +398,10 @@ export interface PartialConfiguration { * The configuration for type checking */ plpgsqlCheck?: PartialPlPgSqlCheckConfiguration; + /** + * The configuration for semantic highlighting in editors + */ + semanticTokens?: PartialSemanticTokensConfiguration; /** * The configuration for splinter */ @@ -536,6 +544,15 @@ export interface PartialFormatConfiguration { */ valuesRowLayout?: Layout; } +/** + * The configuration for inlay hints. + */ +export interface PartialInlayHintsConfiguration { + /** + * if `false`, the language server sends no inlay hints. `true` by default + */ + enabled?: boolean; +} export interface PartialLinterConfiguration { /** * if `false`, it disables the feature and the linter won't be executed. `true` by default @@ -642,6 +659,15 @@ export interface PartialPlPgSqlCheckConfiguration { */ withoutWarnings?: boolean; } +/** + * The configuration for semantic highlighting. + */ +export interface PartialSemanticTokensConfiguration { + /** + * if `false`, the language server sends no semantic tokens, and editors keep the highlighting of their grammar. `true` by default + */ + enabled?: boolean; +} export interface PartialSplinterConfiguration { /** * if `false`, it disables the feature and the linter won't be executed. `true` by default