From 15b6f7d34b3e5c1c66d128dc246f2a8430bca05a Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 19 Oct 2026 06:10:50 +0000 Subject: [PATCH 1/7] feat(workspace): return rows from the execute statement command and explain why it is disabled --- crates/pgls_lsp/src/handlers/code_actions.rs | 5 +- crates/pgls_lsp/src/server.rs | 3 +- crates/pgls_lsp/tests/server.rs | 14 +- .../src/features/code_actions.rs | 23 ++ crates/pgls_workspace/src/settings.rs | 116 ++++++++- crates/pgls_workspace/src/workspace/server.rs | 45 ++-- .../src/workspace/server.tests.rs | 227 +++++++++++++++++- .../workspace/server/statement_execution.rs | 89 +++++++ docs/guides/configure_database.md | 24 ++ 9 files changed, 507 insertions(+), 39 deletions(-) create mode 100644 crates/pgls_workspace/src/workspace/server/statement_execution.rs diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 726e9dd61..9e4e72a23 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -116,10 +116,11 @@ pub async fn execute_command( session .client - .show_message(MessageType::INFO, result.message) + .show_message(MessageType::INFO, &result.message) .await; - Ok(None) + // The rows go back to the client, so editor extensions can render them as a table. + Ok(Some(serde_json::to_value(result)?)) } "pgls.invalidateSchemaCache" => { session.workspace.invalidate_schema_cache(true)?; diff --git a/crates/pgls_lsp/src/server.rs b/crates/pgls_lsp/src/server.rs index 5fb902ea2..967412e1e 100644 --- a/crates/pgls_lsp/src/server.rs +++ b/crates/pgls_lsp/src/server.rs @@ -300,8 +300,7 @@ impl LanguageServer for LSPServer { params: ExecuteCommandParams, ) -> LspResult> { match handlers::code_actions::execute_command(&self.session, params).await { - // we'll inform the client within `code_actions::execute_command` - Ok(_) => LspResult::Ok(None), + Ok(result) => LspResult::Ok(result), Err(err) => LspResult::Err(into_lsp_error(err)), } } diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index c47ce78d6..ec2258ad7 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -1239,6 +1239,7 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { conf.merge_with(PartialConfiguration { db: Some(PartialDatabaseConfiguration { database: Some(database), + allow_statement_executions_against: Some(StringSet::from_iter([format!("{host}/*")])), host: Some(host), ..Default::default() }), @@ -1340,7 +1341,7 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { }) .expect("Did not find executeStatement command!"); - server + let result = server .request::>( "workspace/executeCommand", "_execStmt", @@ -1352,6 +1353,17 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { ) .await?; + assert_eq!( + result, + Some(Some(json!({ + "message": "Successfully executed statement. Rows affected: 0", + "columns": [], + "rows": [], + "truncated": false, + "rowsAffected": 0, + }))) + ); + assert!( users_tbl_exists().await, "Users table did not exists even though it should've been created by the workspace/executeStatement command." diff --git a/crates/pgls_workspace/src/features/code_actions.rs b/crates/pgls_workspace/src/features/code_actions.rs index fb16e49d2..7592b02d6 100644 --- a/crates/pgls_workspace/src/features/code_actions.rs +++ b/crates/pgls_workspace/src/features/code_actions.rs @@ -58,8 +58,31 @@ pub struct ExecuteStatementParams { pub path: PgLSPath, } +/// The most rows an executed statement returns; the rest are dropped and the result is marked +/// as truncated. +pub const EXECUTE_STATEMENT_MAX_ROWS: usize = 100; + +/// The outcome of executing a statement. +/// +/// A statement that returns rows (`select`, `values`, `returning`, `show`, `explain`, ...) fills +/// `columns` and `rows`. Any other statement sets `rows_affected` instead. #[derive(Debug, serde::Serialize, serde::Deserialize, Default, PartialEq, Eq)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] pub struct ExecuteStatementResult { + /// A short summary for the user. pub message: String, + /// The names of the returned columns. + #[serde(default)] + pub columns: Vec, + /// The returned rows, at most [`EXECUTE_STATEMENT_MAX_ROWS`]. Each value is in its Postgres + /// text representation, as `psql` shows it, and `null` for SQL `NULL`. + #[serde(default)] + pub rows: Vec>>, + /// Whether the statement returned more rows than `rows` holds. + #[serde(default)] + pub truncated: bool, + /// The number of rows affected, for a statement that returns no rows. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rows_affected: Option, } diff --git a/crates/pgls_workspace/src/settings.rs b/crates/pgls_workspace/src/settings.rs index e289a03b4..5068ebabc 100644 --- a/crates/pgls_workspace/src/settings.rs +++ b/crates/pgls_workspace/src/settings.rs @@ -10,7 +10,7 @@ use std::{ sync::{RwLock, RwLockReadGuard, RwLockWriteGuard}, time::Duration, }; -use tracing::trace; +use tracing::{trace, warn}; use ignore::gitignore::{Gitignore, GitignoreBuilder}; use pgls_configuration::{ @@ -744,6 +744,33 @@ pub struct DatabaseSettings { pub database: String, pub conn_timeout_secs: Duration, pub allow_statement_executions: bool, + /// The patterns of `db.allowStatementExecutionsAgainst` that are not valid globs. + pub invalid_statement_execution_patterns: Vec, +} + +impl DatabaseSettings { + /// Why statements cannot be executed against this database, or `None` if they can. + pub fn statement_execution_disabled_reason(&self) -> Option { + if self.allow_statement_executions { + return None; + } + + let host = &self.host; + let database = &self.database; + let mut reason = format!( + "Statement execution is disabled for `{host}/{database}`. Add a matching pattern to `db.allowStatementExecutionsAgainst`, e.g. [\"{host}/*\"]." + ); + if !self.invalid_statement_execution_patterns.is_empty() { + let patterns = self + .invalid_statement_execution_patterns + .iter() + .map(|pattern| format!("`{pattern}`")) + .collect::>() + .join(", "); + reason.push_str(&format!(" Ignored invalid patterns: {patterns}.")); + } + Some(reason) + } } impl std::fmt::Debug for DatabaseSettings { @@ -764,6 +791,10 @@ impl std::fmt::Debug for DatabaseSettings { "allow_statement_executions", &self.allow_statement_executions, ) + .field( + "invalid_statement_execution_patterns", + &self.invalid_statement_execution_patterns, + ) .finish() } } @@ -780,6 +811,7 @@ impl Default for DatabaseSettings { database: "postgres".to_string(), conn_timeout_secs: Duration::from_secs(10), allow_statement_executions: true, + invalid_statement_execution_patterns: Vec::new(), } } } @@ -827,18 +859,26 @@ impl From for DatabaseSettings { } } - let allow_statement_executions = value + // An invalid pattern is ignored rather than failing the whole configuration; the reason a + // statement cannot be executed names it. + let target = format!("{host}/{database}"); + let mut allow_statement_executions = false; + let mut invalid_statement_execution_patterns = Vec::new(); + for pattern in value .allow_statement_executions_against - .map(|stringset| { - stringset.iter().any(|pattern| { - let glob = Glob::new(pattern) - .unwrap_or_else(|_| panic!("Invalid pattern: {pattern}")) - .compile_matcher(); - - glob.is_match(format!("{host}/{database}")) - }) - }) - .unwrap_or(false); + .iter() + .flat_map(|set| set.iter()) + { + match Glob::new(pattern) { + Ok(glob) => allow_statement_executions |= glob.compile_matcher().is_match(&target), + Err(err) => { + warn!( + "Ignoring invalid pattern `{pattern}` in `db.allowStatementExecutionsAgainst`: {err}" + ); + invalid_statement_execution_patterns.push(pattern.clone()); + } + } + } Self { enable_connection, @@ -857,6 +897,7 @@ impl From for DatabaseSettings { .unwrap_or(d.conn_timeout_secs), allow_statement_executions, + invalid_statement_execution_patterns, } } } @@ -974,7 +1015,56 @@ mod tests { let config = DatabaseSettings::from(partial_config); - assert!(!config.allow_statement_executions) + assert!(!config.allow_statement_executions); + assert_eq!( + config.statement_execution_disabled_reason().as_deref(), + Some( + "Statement execution is disabled for `production/test-db`. Add a matching pattern to `db.allowStatementExecutionsAgainst`, e.g. [\"production/*\"]." + ) + ); + } + + #[test] + fn ignores_invalid_statement_execution_patterns() { + let partial_config = PartialDatabaseConfiguration { + allow_statement_executions_against: Some(StringSet::from_iter(vec![ + String::from("local[host/*"), + String::from("localhost/*"), + ])), + host: Some("localhost".into()), + database: Some("test-db".into()), + ..Default::default() + }; + + let config = DatabaseSettings::from(partial_config); + + assert!(config.allow_statement_executions); + assert_eq!( + config.invalid_statement_execution_patterns, + ["local[host/*"] + ); + } + + #[test] + fn names_invalid_statement_execution_patterns_in_disabled_reason() { + let partial_config = PartialDatabaseConfiguration { + allow_statement_executions_against: Some(StringSet::from_iter(vec![String::from( + "local[host/*", + )])), + host: Some("localhost".into()), + database: Some("test-db".into()), + ..Default::default() + }; + + let config = DatabaseSettings::from(partial_config); + + assert!(!config.allow_statement_executions); + assert!( + config + .statement_execution_disabled_reason() + .unwrap() + .ends_with("Ignored invalid patterns: `local[host/*`.") + ); } #[cfg(feature = "db")] diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 33946c974..cef3436ac 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -37,8 +37,6 @@ use pgls_text_size::{TextRange, TextSize}; use pgls_typecheck::TypecheckParams; use pgls_workspace_macros::ignored_path; use snapshot_manager::SnapshotManager; -#[cfg(feature = "db")] -use sqlx::Executor; use tracing::{debug, info}; use crate::{ @@ -84,6 +82,8 @@ mod pg_query; mod pgpass; mod snapshot_manager; mod sql_function; +#[cfg(feature = "db")] +mod statement_execution; mod statement_identifier; mod tree_sitter; @@ -654,8 +654,7 @@ impl Workspace for WorkspaceServer { let settings = settings.settings(); let disabled_reason = match settings { - Some(settings) if settings.db.allow_statement_executions => None, - Some(_) => Some("Statement execution is disabled in the settings.".into()), + Some(settings) => settings.db.statement_execution_disabled_reason(), None => Some("Statement execution not allowed against database.".into()), }; @@ -724,6 +723,7 @@ impl Workspace for WorkspaceServer { if stmt.is_none() { return Ok(ExecuteStatementResult { message: "Statement was not found in document.".into(), + ..Default::default() }); }; @@ -732,29 +732,35 @@ impl Workspace for WorkspaceServer { if ast.is_none() { return Ok(ExecuteStatementResult { message: "Statement is invalid.".into(), + ..Default::default() }); }; let settings = self.workspaces(); - let Some(result) = settings.settings().and_then(|settings| { - self.connection.with_pool(&settings.db, |pool| { - let pool = pool.clone(); - run_async(async move { pool.execute(sqlx::query(&content)).await })? - .map_err(WorkspaceError::from) - }) - }) else { + let no_connection = ExecuteStatementResult { + message: "No database connection available.".into(), + ..Default::default() + }; + let Some(settings) = settings.settings() else { + return Ok(no_connection); + }; + + if let Some(reason) = settings.db.statement_execution_disabled_reason() { return Ok(ExecuteStatementResult { - message: "No database connection available.".into(), + message: reason, + ..Default::default() }); + } + + let Some(result) = self.connection.with_pool(&settings.db, |pool| { + let pool = pool.clone(); + run_async(async move { statement_execution::execute_statement(&pool, &content).await })? + .map_err(WorkspaceError::from) + }) else { + return Ok(no_connection); }; - let result = result?; - Ok(ExecuteStatementResult { - message: format!( - "Successfully executed statement. Rows affected: {}", - result.rows_affected() - ), - }) + result } #[cfg(not(feature = "db"))] @@ -766,6 +772,7 @@ impl Workspace for WorkspaceServer { ) -> Result { Ok(ExecuteStatementResult { message: "Statement execution requires database support.".into(), + ..Default::default() }) } diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 3cdcf78c7..cadb826cb 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -16,12 +16,15 @@ use pgls_configuration::{ use pgls_configuration::plpgsql_check::PartialPlPgSqlCheckConfiguration; use pgls_diagnostics::Diagnostic; use pgls_fs::PgLSPath; -use pgls_text_size::TextRange; +use pgls_text_size::{TextRange, TextSize}; use sqlx::{Executor, PgPool}; use crate::{ Workspace, WorkspaceError, - features::code_actions::ExecuteStatementResult, + features::code_actions::{ + CodeActionKind, CodeActionsParams, CommandAction, CommandActionCategory, + EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementParams, ExecuteStatementResult, + }, features::format::PullFileFormattingParams, workspace::{ OpenFileParams, RegisterProjectFolderParams, StatementId, UpdateSettingsParams, @@ -1814,3 +1817,223 @@ fn is_typecheck(diagnostic: &pgls_diagnostics::serde::Diagnostic) -> bool { }) }) } + +fn get_statement_execution_workspace( + test_db: &PgPool, + allow_statement_executions_against: &[String], +) -> WorkspaceServer { + let mut conf = PartialConfiguration::init(); + conf.merge_with(PartialConfiguration { + db: Some(PartialDatabaseConfiguration { + host: Some(test_db.connect_options().get_host().to_string()), + database: Some( + test_db + .connect_options() + .get_database() + .unwrap() + .to_string(), + ), + allow_statement_executions_against: Some(StringSet::from_iter( + allow_statement_executions_against.iter().cloned(), + )), + ..Default::default() + }), + ..Default::default() + }); + + get_test_workspace(Some(conf)).expect("Unable to create test workspace") +} + +/// Executes the single statement in `sql` through its code action, returning the action's +/// disabled reason and the result. +fn execute_sql(workspace: &WorkspaceServer, sql: &str) -> (Option, ExecuteStatementResult) { + let path = PgLSPath::new("execute.sql"); + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: sql.into(), + version: 1, + }) + .expect("Unable to open test file"); + + let action = workspace + .pull_code_actions(CodeActionsParams { + path: path.clone(), + cursor_position: TextSize::new(0), + only: vec![], + skip: vec![], + }) + .expect("Unable to pull code actions") + .actions + .into_iter() + .find_map(|action| match action.kind { + CodeActionKind::Command(CommandAction { + category: CommandActionCategory::ExecuteStatement(statement_id), + }) => Some((action.disabled_reason, statement_id)), + _ => None, + }); + let (disabled_reason, statement_id) = + action.expect("Did not find the execute statement action"); + + let result = workspace + .execute_statement(ExecuteStatementParams { + path: path.clone(), + statement_id, + }) + .expect("Unable to execute statement"); + + workspace + .close_file(crate::workspace::CloseFileParams { path }) + .expect("Unable to close test file"); + + (disabled_reason, result) +} + +fn allow_test_host(test_db: &PgPool) -> Vec { + vec![format!("{}/*", test_db.connect_options().get_host())] +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_is_disabled_without_matching_pattern(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &[]); + + let (disabled_reason, result) = execute_sql(&workspace, "create table not_created (id int);"); + + let host = test_db.connect_options().get_host().to_string(); + let database = test_db + .connect_options() + .get_database() + .unwrap() + .to_string(); + let reason = format!( + "Statement execution is disabled for `{host}/{database}`. Add a matching pattern to `db.allowStatementExecutionsAgainst`, e.g. [\"{host}/*\"]." + ); + assert_eq!(disabled_reason.as_deref(), Some(reason.as_str())); + assert_eq!( + result, + ExecuteStatementResult { + message: reason, + ..Default::default() + } + ); + + let table_exists: bool = sqlx::query_scalar( + "select exists (select 1 from pg_catalog.pg_tables where tablename = 'not_created')", + ) + .fetch_one(&test_db) + .await + .unwrap(); + assert!(!table_exists, "A disabled statement must not be executed"); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_ignores_invalid_patterns(test_db: PgPool) { + let mut patterns = vec!["local[host/*".to_string()]; + patterns.extend(allow_test_host(&test_db)); + let workspace = get_statement_execution_workspace(&test_db, &patterns); + + let (disabled_reason, result) = execute_sql(&workspace, "select 1 as one;"); + + assert_eq!(disabled_reason, None); + assert_eq!(result.rows, [[Some("1".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_returns_rows(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql( + &workspace, + "select 1 as id, null::text as name, array[1, 2] as numbers, interval '1 day' as span, '{\"a\": 1}'::jsonb as data + union all + select 2, 'two', null, null, null;", + ); + + assert_eq!( + result, + ExecuteStatementResult { + message: "2 rows".into(), + columns: ["id", "name", "numbers", "span", "data"] + .map(String::from) + .into(), + rows: vec![ + vec![ + Some("1".into()), + None, + Some("{1,2}".into()), + Some("1 day".into()), + Some("{\"a\": 1}".into()), + ], + vec![Some("2".into()), Some("two".into()), None, None, None], + ], + truncated: false, + rows_affected: None, + } + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_returns_columns_without_rows(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql(&workspace, "select 1 as id where false;"); + + assert_eq!( + result, + ExecuteStatementResult { + message: "0 rows".into(), + columns: vec!["id".into()], + ..Default::default() + } + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_truncates_rows(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql(&workspace, "select generate_series(1, 150) as n;"); + + assert!(result.truncated); + assert_eq!(result.rows.len(), EXECUTE_STATEMENT_MAX_ROWS); + assert_eq!(result.rows.last(), Some(&vec![Some("100".to_string())])); + assert_eq!(result.rows_affected, None); + assert_eq!(result.message, "Showing the first 100 rows (truncated)."); + + // The connection is still usable after the dropped rows. + let (_, result) = execute_sql(&workspace, "select 'next' as value;"); + assert_eq!(result.rows, [[Some("next".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_reports_rows_affected(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql(&workspace, "create table executed (id int);"); + assert_eq!( + result, + ExecuteStatementResult { + message: "Successfully executed statement. Rows affected: 0".into(), + rows_affected: Some(0), + ..Default::default() + } + ); + + let (_, result) = execute_sql(&workspace, "insert into executed values (1), (2);"); + assert_eq!( + result, + ExecuteStatementResult { + message: "Successfully executed statement. Rows affected: 2".into(), + rows_affected: Some(2), + ..Default::default() + } + ); + + let (_, result) = execute_sql( + &workspace, + "delete from executed where id = 1 returning id;", + ); + assert_eq!(result.columns, ["id"]); + assert_eq!(result.rows, [[Some("1".to_string())]]); + assert_eq!(result.rows_affected, None); +} diff --git a/crates/pgls_workspace/src/workspace/server/statement_execution.rs b/crates/pgls_workspace/src/workspace/server/statement_execution.rs new file mode 100644 index 000000000..b348da8c0 --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/statement_execution.rs @@ -0,0 +1,89 @@ +use futures::TryStreamExt; +use sqlx::{ + Column, Either, Executor, PgPool, Row, Statement, ValueRef, + postgres::{PgColumn, PgRow}, +}; + +use crate::features::code_actions::{EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementResult}; + +/// Executes `sql` and collects the rows it returns, if any. +/// +/// The statement runs over the simple query protocol, so Postgres sends every value in its text +/// representation and no type needs decoding. +pub(super) async fn execute_statement( + pool: &PgPool, + sql: &str, +) -> Result { + let mut columns = None; + let mut rows = Vec::new(); + let mut rows_affected = 0; + let mut truncated = false; + + let mut results = sqlx::raw_sql(sql).fetch_many(pool); + while let Some(result) = results.try_next().await? { + match result { + Either::Left(result) => rows_affected += result.rows_affected(), + Either::Right(row) => { + if rows.len() == EXECUTE_STATEMENT_MAX_ROWS { + truncated = true; + break; + } + columns.get_or_insert_with(|| column_names(row.columns())); + rows.push(row_values(&row)); + } + } + } + drop(results); + + // Postgres sends no columns along with zero rows, so a query without rows looks like a command. + // Describing the statement tells them apart, unless it affected rows and is thus a command. + if columns.is_none() && rows_affected == 0 { + columns = pool + .prepare(sql) + .await + .ok() + .map(|statement| column_names(statement.columns())) + .filter(|columns| !columns.is_empty()); + } + + Ok(match columns { + Some(columns) => ExecuteStatementResult { + message: if truncated { + format!("Showing the first {EXECUTE_STATEMENT_MAX_ROWS} rows (truncated).") + } else if rows.len() == 1 { + "1 row".to_string() + } else { + format!("{} rows", rows.len()) + }, + columns, + rows, + truncated, + rows_affected: None, + }, + None => ExecuteStatementResult { + message: format!("Successfully executed statement. Rows affected: {rows_affected}"), + rows_affected: Some(rows_affected), + ..Default::default() + }, + }) +} + +fn column_names(columns: &[PgColumn]) -> Vec { + columns + .iter() + .map(|column| column.name().to_string()) + .collect() +} + +fn row_values(row: &PgRow) -> Vec> { + (0..row.len()) + .map(|index| { + let value = row.try_get_raw(index).ok()?; + if value.is_null() { + return None; + } + let bytes = value.as_bytes().unwrap_or_default(); + Some(String::from_utf8_lossy(bytes).into_owned()) + }) + .collect() +} diff --git a/docs/guides/configure_database.md b/docs/guides/configure_database.md index 8504104a7..2f850b943 100644 --- a/docs/guides/configure_database.md +++ b/docs/guides/configure_database.md @@ -113,6 +113,30 @@ Code actions that execute the statement under the cursor are disabled by default The patterns do not refer to schemas. To allow a single database only, name it, e.g. `"localhost/my_dev_db"`. +When no pattern matches, the code action is disabled and its reason names the setting. An invalid pattern is ignored with a warning in the log, and the reason lists it. + +### Statement Execution Results +The language server runs the statement with the LSP command `pgls.executeStatement`, shows a short summary with `window/showMessage`, and returns the result of `workspace/executeCommand`, so an editor extension can render the rows as a table: + +```json +{ + "message": "2 rows", + "columns": ["id", "name"], + "rows": [ + ["1", "Alice"], + ["2", null] + ], + "truncated": false +} +``` + +- `columns` and `rows` are filled for statements that return rows, such as `select`, `values`, `show`, `explain`, and statements with `returning`. +- Each value is a string in its Postgres text representation, as `psql` shows it, or `null` for SQL `NULL`. +- At most 100 rows are returned. `truncated` is `true` if the statement returned more. +- `rowsAffected` is set instead of `columns` and `rows` for statements that return no rows, such as `insert` without `returning` or `create table`. For example: `{ "message": "Successfully executed statement. Rows affected: 2", "columns": [], "rows": [], "truncated": false, "rowsAffected": 2 }`. +- If the statement is not executed, for example because statement execution is disabled, only `message` is meaningful: `columns` and `rows` are empty and `rowsAffected` is not set. +- If Postgres rejects the statement, the command fails with a JSON-RPC error carrying the Postgres error message. + ## Disabling Database Features If you prefer to work without a database connection, you can disable all database-related features: From 54744f4102f7a6d5da0a09ca56953ac58d7e067d Mon Sep 17 00:00:00 2001 From: psteinroe Date: Tue, 20 Oct 2026 07:35:57 +0000 Subject: [PATCH 2/7] fix(workspace): disable the database connection on an invalid connection string instead of panicking --- crates/pgls_cli/tests/assert_check.rs | 30 +++++ ...ck_fails_on_invalid_connection_string.snap | 19 +++ crates/pgls_lsp/src/session.rs | 15 +++ crates/pgls_lsp/tests/server.rs | 112 +++++++++++++++- crates/pgls_workspace/src/settings.rs | 120 +++++++++++++++--- crates/pgls_workspace/src/workspace/server.rs | 6 + docs/guides/configure_database.md | 2 + 7 files changed, 281 insertions(+), 23 deletions(-) create mode 100644 crates/pgls_cli/tests/snapshots/assert_check__check_fails_on_invalid_connection_string.snap diff --git a/crates/pgls_cli/tests/assert_check.rs b/crates/pgls_cli/tests/assert_check.rs index 53c48a6d3..7818863c1 100644 --- a/crates/pgls_cli/tests/assert_check.rs +++ b/crates/pgls_cli/tests/assert_check.rs @@ -270,6 +270,36 @@ fn check_fails_when_database_is_unreachable() { assert!(output.contains("database/connection"), "{output}"); } +#[cfg(target_os = "linux")] +#[test] +fn check_fails_on_invalid_connection_string() { + let output = run_check_without_db_env( + &[ + "--connection-string", + "postgres://postgres:hunter2@localhost:not-a-port/postgres", + UNKNOWN_COLUMN_PATH, + ], + &[], + ); + + assert!(!output.contains("hunter2"), "{output}"); + assert_snapshot!(output); +} + +#[cfg(target_os = "linux")] +#[test] +fn check_fails_on_invalid_database_url() { + let output = run_check_without_db_env( + &[UNKNOWN_COLUMN_PATH], + &[("DATABASE_URL", "host=localhost password=hunter2")], + ); + + assert!(output.starts_with("status: failure"), "{output}"); + assert!(output.contains("database/connection"), "{output}"); + assert!(output.contains("Invalid `db.connectionString`"), "{output}"); + assert!(!output.contains("hunter2"), "{output}"); +} + #[cfg(target_os = "linux")] #[test] fn check_without_database_succeeds_with_disable_db() { diff --git a/crates/pgls_cli/tests/snapshots/assert_check__check_fails_on_invalid_connection_string.snap b/crates/pgls_cli/tests/snapshots/assert_check__check_fails_on_invalid_connection_string.snap new file mode 100644 index 000000000..a78684a9f --- /dev/null +++ b/crates/pgls_cli/tests/snapshots/assert_check__check_fails_on_invalid_connection_string.snap @@ -0,0 +1,19 @@ +--- +source: crates/pgls_cli/tests/assert_check.rs +expression: output +--- +status: failure +stdout: +Checked 1 file in . No fixes applied. +Found 1 error. +stderr: +database/connection ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + + × Invalid `db.connectionString`: invalid port number. Expected a URL like `postgres://user:password@host:5432/database`. + + i The checks that need the database were skipped. Fix the connection settings, or pass --disable-db or set db.disableConnection to check without a database. + + +check ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + + × Some errors were emitted while running checks. diff --git a/crates/pgls_lsp/src/session.rs b/crates/pgls_lsp/src/session.rs index 7807a87bb..2b367418e 100644 --- a/crates/pgls_lsp/src/session.rs +++ b/crates/pgls_lsp/src/session.rs @@ -663,6 +663,21 @@ impl Session { return ConfigurationStatus::Error; } + // An invalid connection string only disables the database features, + // so the user would not notice it otherwise. + if let Some(error) = configuration + .db + .as_ref() + .and_then(pgls_workspace::settings::connection_string_error) + { + self.client + .show_message( + MessageType::WARNING, + format!("{error} Database features are disabled."), + ) + .await; + } + let result = self.workspace.update_settings(UpdateSettingsParams { workspace_directory: self.fs.working_directory(), configuration, diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index ec2258ad7..2e360636d 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -48,8 +48,8 @@ use tower_lsp_server::ls_types::{ ClientCapabilities, DidChangeConfigurationParams, DidChangeTextDocumentParams, DidChangeWatchedFilesParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams, FileChangeType, FileEvent, InitializeResult, InitializedParams, PublishDiagnosticsParams, - TextDocumentContentChangeEvent, TextDocumentIdentifier, TextDocumentItem, Uri, - VersionedTextDocumentIdentifier, + ShowMessageParams, TextDocumentContentChangeEvent, TextDocumentIdentifier, TextDocumentItem, + Uri, VersionedTextDocumentIdentifier, }; use tower_lsp_server::{jsonrpc::Request, ls_types::InitializeParams}; @@ -383,6 +383,7 @@ const CHANNEL_BUFFER_SIZE: usize = 8; #[derive(Debug, PartialEq, Eq)] enum ServerNotification { PublishDiagnostics(PublishDiagnosticsParams), + ShowMessage(ShowMessageParams), } /// Basic handler for requests and notifications coming from the server for tests @@ -408,6 +409,15 @@ where } } + if req.method() == "window/showMessage" { + let params = req.params().expect("invalid request"); + let message = from_value(params.clone()).expect("invalid params"); + match notify.send(ServerNotification::ShowMessage(message)).await { + Ok(_) => continue, + Err(_) => break, + } + } + let id = match req.id() { Some(id) => id, None => continue, @@ -580,11 +590,14 @@ async fn wait_for_diagnostics( ) -> Result { tokio::time::timeout(Duration::from_secs(5), async { loop { - let Some(ServerNotification::PublishDiagnostics(params)) = receiver.next().await else { - bail!("the server closed the connection"); - }; - if params.version == Some(version) { - return Ok(params); + match receiver.next().await { + Some(ServerNotification::PublishDiagnostics(params)) + if params.version == Some(version) => + { + return Ok(params); + } + Some(_) => {} + None => bail!("the server closed the connection"), } } }) @@ -664,6 +677,7 @@ async fn client_configuration_overrides_are_sticky_and_clearable(test_db: PgPool return true; } } + Some(ServerNotification::ShowMessage(_)) => {} None => return false, } } @@ -1223,6 +1237,90 @@ async fn test_issue_271(test_db: PgPool) -> Result<()> { Ok(()) } +/// An invalid connection string disables the database features with a warning instead of crashing +/// the server, on every configuration transition. +#[tokio::test] +async fn invalid_connection_string_warns_instead_of_crashing() -> Result<()> { + let factory = ServerFactory::default(); + let mut fs = MemoryFileSystem::default(); + fs.insert( + url!("postgres-language-server.jsonc") + .to_file_path() + .unwrap() + .into_owned(), + "{}".to_string(), + ); + + let (service, client) = factory + .create_with_fs(None, DynRef::Owned(Box::new(fs))) + .into_inner(); + + let (stream, sink) = client.split(); + let mut server = Server::new(service); + + let (sender, mut receiver) = channel(CHANNEL_BUFFER_SIZE); + let reader = tokio::spawn(client_handler(stream, sink, sender)); + + server.initialize().await?; + server.initialized().await?; + server.open_document("select 1 from;").await?; + drain_notifications(&mut receiver); + + async fn wait_for_warning(receiver: &mut Receiver) -> Option { + tokio::time::timeout(Duration::from_secs(5), async { + loop { + match receiver.next().await { + Some(ServerNotification::ShowMessage(params)) + if params.typ == lsp::MessageType::WARNING => + { + return Some(params.message); + } + Some(_) => {} + None => return None, + } + } + }) + .await + .ok() + .flatten() + } + + server + .set_configuration_overrides(json!({ + "overrides": { + "db": { "connectionString": "postgres://postgres:hunter2@localhost:not-a-port/postgres" } + } + })) + .await?; + + let warning = wait_for_warning(&mut receiver) + .await + .expect("expected a warning about the connection string"); + assert_eq!( + warning, + "Invalid `db.connectionString`: invalid port number. Expected a URL like `postgres://user:password@host:5432/database`. Database features are disabled." + ); + + // A reload with the invalid connection string warns again, and the server keeps working. + server.change_watched_configuration_file().await?; + assert!(wait_for_warning(&mut receiver).await.is_some()); + + server.replace_document(1, "select 1 from;").await?; + let diagnostics = wait_for_diagnostics(&mut receiver, 1).await?; + assert!( + diagnostics + .diagnostics + .iter() + .any(|diagnostic| diagnostic.message.contains("syntax error")), + "{diagnostics:?}" + ); + + server.shutdown().await?; + reader.abort(); + + Ok(()) +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_execute_statement(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); diff --git a/crates/pgls_workspace/src/settings.rs b/crates/pgls_workspace/src/settings.rs index 5068ebabc..e346b815b 100644 --- a/crates/pgls_workspace/src/settings.rs +++ b/crates/pgls_workspace/src/settings.rs @@ -746,6 +746,8 @@ pub struct DatabaseSettings { pub allow_statement_executions: bool, /// The patterns of `db.allowStatementExecutionsAgainst` that are not valid globs. pub invalid_statement_execution_patterns: Vec, + /// Why `db.connectionString` is invalid. The connection is disabled then. + pub connection_string_error: Option, } impl DatabaseSettings { @@ -795,6 +797,7 @@ impl std::fmt::Debug for DatabaseSettings { "invalid_statement_execution_patterns", &self.invalid_statement_execution_patterns, ) + .field("connection_string_error", &self.connection_string_error) .finish() } } @@ -812,6 +815,7 @@ impl Default for DatabaseSettings { conn_timeout_secs: Duration::from_secs(10), allow_statement_executions: true, invalid_statement_execution_patterns: Vec::new(), + connection_string_error: None, } } } @@ -829,12 +833,7 @@ impl From for DatabaseSettings { } }); - // "host" OR "connectionString" is the minimum required setting for database features - // to be enabled. - let disable_connection = value.disable_connection.is_some_and(|disabled| disabled); - let enable_connection = - (connection_string.is_some() || value.host.as_ref().is_some()) && !disable_connection; - + let has_host = value.host.is_some(); #[cfg_attr(not(feature = "db"), allow(unused_mut))] let mut database = value.database.unwrap_or(d.database); #[cfg_attr(not(feature = "db"), allow(unused_mut))] @@ -844,21 +843,38 @@ impl From for DatabaseSettings { #[cfg_attr(not(feature = "db"), allow(unused_mut))] let mut username = value.username.unwrap_or(d.username); - // Parse connection string to extract host/port/username/database (DB mode only) + let disable_connection = value.disable_connection.is_some_and(|disabled| disabled); + + // Parse connection string to extract host/port/username/database (DB mode only). An + // invalid one disables the connection instead of failing the whole configuration. + #[cfg_attr(not(feature = "db"), allow(unused_mut))] + let mut connection_string_error: Option = None; #[cfg(feature = "db")] if let Some(uri) = connection_string.as_ref() { - let opts = PgConnectOptions::from_str(uri) - .unwrap_or_else(|err| panic!("Invalid db.connectionString provided: {err}")); - - // we only extract the values we need for statement execution checks and connection key - host = opts.get_host().to_string(); - port = opts.get_port(); - username = opts.get_username().to_string(); - if let Some(db) = opts.get_database() { - database = db.to_string(); + match parse_connection_string(uri) { + Ok(opts) => { + // we only extract the values we need for statement execution checks and connection key + host = opts.get_host().to_string(); + port = opts.get_port(); + username = opts.get_username().to_string(); + if let Some(db) = opts.get_database() { + database = db.to_string(); + } + } + Err(error) if !disable_connection => { + warn!("{error} The database connection is disabled."); + connection_string_error = Some(error); + } + Err(_) => {} } } + // "host" OR "connectionString" is the minimum required setting for database features + // to be enabled. + let enable_connection = (connection_string.is_some() || has_host) + && !disable_connection + && connection_string_error.is_none(); + // An invalid pattern is ignored rather than failing the whole configuration; the reason a // statement cannot be executed names it. let target = format!("{host}/{database}"); @@ -898,10 +914,41 @@ impl From for DatabaseSettings { allow_statement_executions, invalid_statement_execution_patterns, + connection_string_error, } } } +/// Parses `db.connectionString`. The error says why it is invalid without repeating it, as it may +/// contain a password. +#[cfg(feature = "db")] +fn parse_connection_string(uri: &str) -> Result { + PgConnectOptions::from_str(uri).map_err(|err| { + let reason = match err { + sqlx::Error::Configuration(source) => source.to_string(), + err => err.to_string(), + }; + format!( + "Invalid `db.connectionString`: {reason}. Expected a URL like `postgres://user:password@host:5432/database`." + ) + }) +} + +/// Why the `db.connectionString` of `configuration` is invalid, if it is. The database connection +/// is disabled while it is. +pub fn connection_string_error(configuration: &PartialDatabaseConfiguration) -> Option { + if configuration.disable_connection == Some(true) { + return None; + } + #[cfg(feature = "db")] + if let Some(uri) = configuration.connection_string.as_deref().map(str::trim) + && !uri.is_empty() + { + return parse_connection_string(uri).err(); + } + None +} + /// Filesystem settings for the entire workspace #[derive(Debug)] pub struct FilesSettings { @@ -1103,4 +1150,45 @@ mod tests { assert!(!config.allow_statement_executions) } + + #[cfg(feature = "db")] + #[test] + fn invalid_connection_string_disables_the_connection() { + let partial_config = PartialDatabaseConfiguration { + connection_string: Some( + "postgres://postgres:hunter2@localhost:not-a-port/test-db".to_string(), + ), + allow_statement_executions_against: Some(StringSet::from_iter(vec![String::from( + "*/*", + )])), + ..Default::default() + }; + + let error = super::connection_string_error(&partial_config); + let config = DatabaseSettings::from(partial_config); + + assert!(!config.enable_connection); + assert_eq!(config.connection_string_error, error); + assert_eq!( + error.as_deref(), + Some( + "Invalid `db.connectionString`: invalid port number. Expected a URL like `postgres://user:password@host:5432/database`." + ) + ); + } + + #[cfg(feature = "db")] + #[test] + fn disabled_connection_ignores_invalid_connection_string() { + let partial_config = PartialDatabaseConfiguration { + connection_string: Some("host=localhost".to_string()), + disable_connection: Some(true), + ..Default::default() + }; + + assert_eq!(super::connection_string_error(&partial_config), None); + let config = DatabaseSettings::from(partial_config); + assert!(!config.enable_connection); + assert_eq!(config.connection_string_error, None); + } } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index cef3436ac..63b9d0f0a 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -1301,6 +1301,12 @@ impl Workspace for WorkspaceServer { return Ok(CheckDatabaseConnectionResult::default()); }; + if let Some(error) = &settings.db.connection_string_error { + return Ok(CheckDatabaseConnectionResult { + error: Some(error.clone()), + }); + } + let result = self .connection .try_with_pool(&settings.db, |pool| self.snapshot.load(pool).map(|_| ())); diff --git a/docs/guides/configure_database.md b/docs/guides/configure_database.md index 2f850b943..9a6bc7273 100644 --- a/docs/guides/configure_database.md +++ b/docs/guides/configure_database.md @@ -62,6 +62,8 @@ To check without a database, disable the connection with `--disable-db` or `"dis The language server keeps working without the database when the connection fails, and retries it in the background. +An invalid `connectionString` (or `DATABASE_URL`), such as one with a non-numeric port, disables the database connection instead of stopping the language server, which shows a warning. `postgres-language-server check` reports it as a `database/connection` error, like a failed connection. + ## Setting the connection from the editor An editor extension can set the connection for a running language server session with the LSP request `pgls/set_configuration_overrides`, which takes any configuration and applies it on top of the configuration file: From 50397cf156f2f01fc26757afe893e3fe380ac1bf Mon Sep 17 00:00:00 2001 From: psteinroe Date: Fri, 23 Oct 2026 14:15:06 +0000 Subject: [PATCH 3/7] fix(workspace): only connection failures put the database into backoff Every sqlx error became a DatabaseConnectionError, and the connection manager put every such error into the reconnect backoff, so an executed statement that divides by zero disabled completions, hover and typecheck for 5 to 60 seconds. Only errors that mean the database could not be reached or refused the connection start the backoff now. An error Postgres raises for an executed statement is part of the result, with its SQLSTATE, detail, hint and position in the document. --- crates/pgls_lsp/src/handlers/code_actions.rs | 22 ++- crates/pgls_lsp/tests/server.rs | 161 +++++++++++++----- crates/pgls_workspace/src/diagnostics.rs | 34 ++++ .../src/features/code_actions.rs | 21 +++ crates/pgls_workspace/src/workspace/server.rs | 12 +- .../src/workspace/server.tests.rs | 72 +++++++- .../workspace/server/connection_manager.rs | 26 ++- .../workspace/server/statement_execution.rs | 53 +++++- docs/guides/configure_database.md | 4 +- 9 files changed, 344 insertions(+), 61 deletions(-) diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 9e4e72a23..525c37ad5 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -1,4 +1,7 @@ -use crate::{adapters::get_cursor_position, session::Session}; +use crate::{ + adapters::{get_cursor_position, to_lsp}, + session::Session, +}; use anyhow::{Result, anyhow}; use tower_lsp_server::ls_types::{ self, CodeAction, CodeActionDisabled, CodeActionOrCommand, Command, ExecuteCommandParams, @@ -114,13 +117,26 @@ pub async fn execute_command( */ session.update_all_diagnostics().await; + let message_type = if result.error.is_some() { + MessageType::ERROR + } else { + MessageType::INFO + }; session .client - .show_message(MessageType::INFO, &result.message) + .show_message(message_type, &result.message) .await; // The rows go back to the client, so editor extensions can render them as a table. - Ok(Some(serde_json::to_value(result)?)) + let mut value = serde_json::to_value(&result)?; + // The workspace locates an error by its offset, the client by its line and character. + if let Some(offset) = result.error.and_then(|error| error.position) { + let position = session.document(&doc_url).ok().and_then(|doc| { + to_lsp::position(&doc.line_index, offset, session.position_encoding()).ok() + }); + value["error"]["position"] = serde_json::to_value(position)?; + } + Ok(Some(value)) } "pgls.invalidateSchemaCache" => { session.workspace.invalidate_schema_cache(true)?; diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index 2e360636d..4c40f47fb 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -1321,8 +1321,8 @@ async fn invalid_connection_string_warns_instead_of_crashing() -> Result<()> { Ok(()) } -#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] -async fn test_execute_statement(test_db: PgPool) -> Result<()> { +/// A server whose configuration allows executing statements against `test_db`. +fn statement_execution_server(test_db: &PgPool) -> Server { let factory = ServerFactory::default(); let mut fs = MemoryFileSystem::default(); @@ -1357,50 +1357,20 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { .into_inner(); let (stream, sink) = client.split(); - let mut server = Server::new(service); + let 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?; - - let users_tbl_exists = async || { - sqlx::query_scalar::<_, bool>( - r#" - select exists ( - select 1 - from pg_catalog.pg_tables - where tablename = 'users' - ) - "#, - ) - .fetch_one(&test_db.clone()) - .await - .unwrap() - }; - - assert!( - !(users_tbl_exists().await), - "The user table shouldn't exist at this point." - ); - - let doc_content = r#" - create table users ( - id serial primary key, - name text, - email text - ); - "#; - - let doc_url = url!("test.sql"); + tokio::spawn(client_handler(stream, sink, sender)); server - .open_named_document(doc_content.to_string(), doc_url.clone(), "sql") - .await?; +} +/// Runs the `pgls.executeStatement` command of the statement at `position`. +async fn execute_statement_at( + server: &mut Server, + doc_url: &Uri, + position: Position, +) -> Result> { let code_actions_response = server .request::( "textDocument/codeAction", @@ -1410,9 +1380,9 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { uri: doc_url.clone(), }, range: Range { - start: Position::new(3, 7), - end: Position::new(3, 7), - }, // just somewhere within the statement. + start: position, + end: position, + }, context: CodeActionContext::default(), partial_result_params: PartialResultParams::default(), work_done_progress_params: WorkDoneProgressParams::default(), @@ -1451,15 +1421,64 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { ) .await?; + Ok(result.flatten()) +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn test_execute_statement(test_db: PgPool) -> Result<()> { + let mut server = statement_execution_server(&test_db); + + server.initialize().await?; + server.initialized().await?; + + server.load_configuration().await?; + + let users_tbl_exists = async || { + sqlx::query_scalar::<_, bool>( + r#" + select exists ( + select 1 + from pg_catalog.pg_tables + where tablename = 'users' + ) + "#, + ) + .fetch_one(&test_db.clone()) + .await + .unwrap() + }; + + assert!( + !(users_tbl_exists().await), + "The user table shouldn't exist at this point." + ); + + let doc_content = r#" + create table users ( + id serial primary key, + name text, + email text + ); + "#; + + let doc_url = url!("test.sql"); + + server + .open_named_document(doc_content.to_string(), doc_url.clone(), "sql") + .await?; + + // Just somewhere within the statement. + let result = execute_statement_at(&mut server, &doc_url, Position::new(3, 7)).await?; + assert_eq!( result, - Some(Some(json!({ + Some(json!({ "message": "Successfully executed statement. Rows affected: 0", "columns": [], "rows": [], "truncated": false, "rowsAffected": 0, - }))) + })) ); assert!( @@ -1468,7 +1487,55 @@ async fn test_execute_statement(test_db: PgPool) -> Result<()> { ); server.shutdown().await?; - reader.abort(); + + Ok(()) +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_returns_the_postgres_error_at_its_position( + test_db: PgPool, +) -> Result<()> { + let mut server = statement_execution_server(&test_db); + + server.initialize().await?; + server.initialized().await?; + server.load_configuration().await?; + + let doc_url = url!("test.sql"); + server + .open_named_document( + "select 1;\nselect 'ä' as a, missing_column from pg_class;\n", + doc_url.clone(), + "sql", + ) + .await?; + + let result = execute_statement_at(&mut server, &doc_url, Position::new(1, 3)).await?; + + // The position counts UTF-16 code units, the default position encoding. + assert_eq!( + result, + Some(json!({ + "message": "Statement failed: column \"missing_column\" does not exist", + "columns": [], + "rows": [], + "truncated": false, + "error": { + "message": "column \"missing_column\" does not exist", + "code": "42703", + "position": { "line": 1, "character": 17 }, + }, + })) + ); + + // The error does not take the database offline. + let result = execute_statement_at(&mut server, &doc_url, Position::new(0, 3)).await?; + assert_eq!( + result.as_ref().and_then(|result| result.get("rows")), + Some(&json!([["1"]])) + ); + + server.shutdown().await?; Ok(()) } diff --git a/crates/pgls_workspace/src/diagnostics.rs b/crates/pgls_workspace/src/diagnostics.rs index 7db9e96c3..5a1a791ab 100644 --- a/crates/pgls_workspace/src/diagnostics.rs +++ b/crates/pgls_workspace/src/diagnostics.rs @@ -253,26 +253,60 @@ impl From for WorkspaceError { pub struct DatabaseConnectionError { pub(crate) message: String, pub(crate) code: Option, + /// Whether the database could not be reached or refused the connection, as opposed to a + /// query failing on a working connection. Only a connection failure puts the connection into + /// the reconnect backoff. + #[serde(default)] + pub(crate) connection_failure: bool, } #[cfg(feature = "db")] impl From for WorkspaceError { fn from(err: sqlx::Error) -> Self { + let connection_failure = is_connection_failure(&err); let db_err = err.as_database_error(); if let Some(db_err) = db_err { Self::DatabaseConnectionError(DatabaseConnectionError { message: db_err.message().to_string(), code: db_err.code().map(|c| c.to_string()), + connection_failure, }) } else { Self::DatabaseConnectionError(DatabaseConnectionError { message: err.to_string(), code: None, + connection_failure, }) } } } +/// Whether `err` means that the database could not be reached or refused the connection. +/// +/// Any other error, such as a syntax error, a division by zero, or a missing permission, comes +/// from a query on a working connection. +#[cfg(feature = "db")] +pub(crate) fn is_connection_failure(err: &sqlx::Error) -> bool { + match err { + sqlx::Error::Configuration(_) + | sqlx::Error::Io(_) + | sqlx::Error::Tls(_) + | sqlx::Error::Protocol(_) + | sqlx::Error::PoolTimedOut + | sqlx::Error::PoolClosed + | sqlx::Error::WorkerCrashed => true, + sqlx::Error::Database(err) => err.code().is_some_and(|code| { + // Class 08 is a connection exception and class 28 a failed authentication. 3D000: + // the database does not exist, 53300: too many connections, 57P03: the server is + // starting up or shutting down. + code.starts_with("08") + || code.starts_with("28") + || matches!(code.as_ref(), "3D000" | "53300" | "57P03") + }), + _ => false, + } +} + #[derive(Debug, Serialize, Deserialize, Diagnostic)] #[diagnostic( category = "internalError/fs", diff --git a/crates/pgls_workspace/src/features/code_actions.rs b/crates/pgls_workspace/src/features/code_actions.rs index 7592b02d6..0b9bb73c5 100644 --- a/crates/pgls_workspace/src/features/code_actions.rs +++ b/crates/pgls_workspace/src/features/code_actions.rs @@ -85,4 +85,25 @@ pub struct ExecuteStatementResult { /// The number of rows affected, for a statement that returns no rows. #[serde(default, skip_serializing_if = "Option::is_none")] pub rows_affected: Option, + /// Why Postgres rejected the statement, if it did. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +/// An error that Postgres raised while executing a statement. +#[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "camelCase")] +pub struct ExecuteStatementError { + /// The primary error message, e.g. `division by zero`. + pub message: String, + /// The SQLSTATE error code, e.g. `22012`. + pub code: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub detail: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub hint: Option, + /// The offset of the error in the document, if Postgres reported a position. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub position: Option, } diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 63b9d0f0a..94c47adb8 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -727,7 +727,7 @@ impl Workspace for WorkspaceServer { }); }; - let (_id, _range, content, ast) = stmt.unwrap(); + let (_id, range, content, ast) = stmt.unwrap(); if ast.is_none() { return Ok(ExecuteStatementResult { @@ -760,7 +760,15 @@ impl Workspace for WorkspaceServer { return Ok(no_connection); }; - result + let mut result = result?; + if let Some(position) = result + .error + .as_mut() + .and_then(|error| error.position.as_mut()) + { + *position += range.start(); + } + Ok(result) } #[cfg(not(feature = "db"))] diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index cadb826cb..4ff9d54bc 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -23,7 +23,8 @@ use crate::{ Workspace, WorkspaceError, features::code_actions::{ CodeActionKind, CodeActionsParams, CommandAction, CommandActionCategory, - EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementParams, ExecuteStatementResult, + EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementError, ExecuteStatementParams, + ExecuteStatementResult, }, features::format::PullFileFormattingParams, workspace::{ @@ -1847,6 +1848,15 @@ fn get_statement_execution_workspace( /// Executes the single statement in `sql` through its code action, returning the action's /// disabled reason and the result. fn execute_sql(workspace: &WorkspaceServer, sql: &str) -> (Option, ExecuteStatementResult) { + execute_sql_at(workspace, sql, TextSize::new(0)) +} + +/// Like [execute_sql], for the statement at `cursor_position`. +fn execute_sql_at( + workspace: &WorkspaceServer, + sql: &str, + cursor_position: TextSize, +) -> (Option, ExecuteStatementResult) { let path = PgLSPath::new("execute.sql"); workspace .open_file(OpenFileParams { @@ -1859,7 +1869,7 @@ fn execute_sql(workspace: &WorkspaceServer, sql: &str) -> (Option, Execu let action = workspace .pull_code_actions(CodeActionsParams { path: path.clone(), - cursor_position: TextSize::new(0), + cursor_position, only: vec![], skip: vec![], }) @@ -1968,6 +1978,7 @@ async fn execute_statement_returns_rows(test_db: PgPool) { ], truncated: false, rows_affected: None, + error: None, } ); } @@ -2037,3 +2048,60 @@ async fn execute_statement_reports_rows_affected(test_db: PgPool) { assert_eq!(result.rows, [[Some("1".to_string())]]); assert_eq!(result.rows_affected, None); } + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_returns_sql_errors_without_disabling_the_database(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql(&workspace, "select 1 / 0;"); + assert_eq!( + result, + ExecuteStatementResult { + message: "Statement failed: division by zero".into(), + error: Some(ExecuteStatementError { + message: "division by zero".into(), + code: "22012".into(), + detail: None, + hint: None, + position: None, + }), + ..Default::default() + } + ); + + // A failed statement says nothing about the connection, so the database stays available. + let (_, result) = execute_sql(&workspace, "select 1 as one;"); + assert_eq!(result.rows, [[Some("1".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_locates_the_error_in_the_document(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + // Postgres counts characters, so the multibyte character before the error shifts it. + let sql = "select 1;\nselect 'ä' as a, missing_column from pg_class;"; + let second_statement = TextSize::new(sql.find("select 'ä'").unwrap() as u32); + let (_, result) = execute_sql_at(&workspace, sql, second_statement); + + let error = result.error.expect("expected an error"); + assert_eq!(error.code, "42703"); + assert_eq!( + error.position, + Some(TextSize::new(sql.find("missing_column").unwrap() as u32)) + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_returns_detail_and_hint(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql( + &workspace, + "do $$ begin raise exception 'failed' using detail = 'the detail', hint = 'the hint'; end $$;", + ); + + let error = result.error.expect("expected an error"); + assert_eq!(error.code, "P0001"); + assert_eq!(error.detail.as_deref(), Some("the detail")); + assert_eq!(error.hint.as_deref(), Some("the hint")); +} diff --git a/crates/pgls_workspace/src/workspace/server/connection_manager.rs b/crates/pgls_workspace/src/workspace/server/connection_manager.rs index 0f3767d51..f51863526 100644 --- a/crates/pgls_workspace/src/workspace/server/connection_manager.rs +++ b/crates/pgls_workspace/src/workspace/server/connection_manager.rs @@ -235,9 +235,12 @@ impl ConnectionManager { ) { match result { Ok(_) => self.clear_failure(key), - Err(err @ WorkspaceError::DatabaseConnectionError(error)) => { + Err(err @ WorkspaceError::DatabaseConnectionError(error)) + if error.connection_failure => + { self.record_failure(key, settings_fingerprint, error, &err.to_string()); } + // A query that fails, e.g. on a division by zero, says nothing about the connection. Err(_) => {} } } @@ -398,6 +401,7 @@ mod tests { DatabaseConnectionError { message: "password authentication failed".to_string(), code: Some("28P01".to_string()), + connection_failure: true, } } @@ -421,4 +425,24 @@ mod tests { if error.message == "password authentication failed" )); } + + #[test] + fn only_connection_failures_start_the_backoff() { + let manager = ConnectionManager::new(); + let settings = settings(); + + let query_error = manager.with_pool(&settings, |_| { + Err::<(), _>(WorkspaceError::from(sqlx::Error::ColumnNotFound( + "id".to_string(), + ))) + }); + assert!(matches!(query_error, Some(Err(_)))); + assert!(manager.with_pool(&settings, |_| Ok(())).is_some()); + + let connection_failure = manager.with_pool(&settings, |_| { + Err::<(), _>(WorkspaceError::from(sqlx::Error::PoolTimedOut)) + }); + assert!(matches!(connection_failure, Some(Err(_)))); + assert!(manager.with_pool(&settings, |_| Ok(())).is_none()); + } } diff --git a/crates/pgls_workspace/src/workspace/server/statement_execution.rs b/crates/pgls_workspace/src/workspace/server/statement_execution.rs index b348da8c0..ca7b0d091 100644 --- a/crates/pgls_workspace/src/workspace/server/statement_execution.rs +++ b/crates/pgls_workspace/src/workspace/server/statement_execution.rs @@ -1,19 +1,37 @@ use futures::TryStreamExt; +use pgls_text_size::TextSize; use sqlx::{ Column, Either, Executor, PgPool, Row, Statement, ValueRef, - postgres::{PgColumn, PgRow}, + postgres::{PgColumn, PgDatabaseError, PgErrorPosition, PgRow}, }; -use crate::features::code_actions::{EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementResult}; +use crate::diagnostics::is_connection_failure; +use crate::features::code_actions::{ + EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementError, ExecuteStatementResult, +}; /// Executes `sql` and collects the rows it returns, if any. /// /// The statement runs over the simple query protocol, so Postgres sends every value in its text -/// representation and no type needs decoding. +/// representation and no type needs decoding. An error that Postgres raises for the statement is +/// part of the result, with its position relative to `sql`. Only a failed connection is an error. pub(super) async fn execute_statement( pool: &PgPool, sql: &str, ) -> Result { + match collect_result(pool, sql).await { + Err(err) if !is_connection_failure(&err) => match err + .as_database_error() + .and_then(|err| err.try_downcast_ref::()) + { + Some(err) => Ok(rejected(err, sql)), + None => Err(err), + }, + result => result, + } +} + +async fn collect_result(pool: &PgPool, sql: &str) -> Result { let mut columns = None; let mut rows = Vec::new(); let mut rows_affected = 0; @@ -58,7 +76,7 @@ pub(super) async fn execute_statement( columns, rows, truncated, - rows_affected: None, + ..Default::default() }, None => ExecuteStatementResult { message: format!("Successfully executed statement. Rows affected: {rows_affected}"), @@ -68,6 +86,33 @@ pub(super) async fn execute_statement( }) } +fn rejected(err: &PgDatabaseError, sql: &str) -> ExecuteStatementResult { + // Postgres counts the characters from 1, and points one past the end for an error at the end + // of the input. + let position = match err.position() { + Some(PgErrorPosition::Original(position)) => position.checked_sub(1).and_then(|index| { + sql.char_indices() + .map(|(offset, _)| offset) + .chain([sql.len()]) + .nth(index) + .and_then(|offset| TextSize::try_from(offset).ok()) + }), + _ => None, + }; + + ExecuteStatementResult { + message: format!("Statement failed: {}", err.message()), + error: Some(ExecuteStatementError { + message: err.message().to_string(), + code: err.code().to_string(), + detail: err.detail().map(str::to_string), + hint: err.hint().map(str::to_string), + position, + }), + ..Default::default() + } +} + fn column_names(columns: &[PgColumn]) -> Vec { columns .iter() diff --git a/docs/guides/configure_database.md b/docs/guides/configure_database.md index 9a6bc7273..9944ff451 100644 --- a/docs/guides/configure_database.md +++ b/docs/guides/configure_database.md @@ -60,7 +60,7 @@ When a database is configured (`host` or `connectionString`, or `PGHOST` or `DAT To check without a database, disable the connection with `--disable-db` or `"disableConnection": true` (see [Disabling Database Features](#disabling-database-features)). -The language server keeps working without the database when the connection fails, and retries it in the background. +The language server keeps working without the database when the connection fails, and retries it in the background. A query that fails, such as an executed statement that divides by zero, does not count as a failed connection. An invalid `connectionString` (or `DATABASE_URL`), such as one with a non-numeric port, disables the database connection instead of stopping the language server, which shows a warning. `postgres-language-server check` reports it as a `database/connection` error, like a failed connection. @@ -137,7 +137,7 @@ The language server runs the statement with the LSP command `pgls.executeStateme - At most 100 rows are returned. `truncated` is `true` if the statement returned more. - `rowsAffected` is set instead of `columns` and `rows` for statements that return no rows, such as `insert` without `returning` or `create table`. For example: `{ "message": "Successfully executed statement. Rows affected: 2", "columns": [], "rows": [], "truncated": false, "rowsAffected": 2 }`. - If the statement is not executed, for example because statement execution is disabled, only `message` is meaningful: `columns` and `rows` are empty and `rowsAffected` is not set. -- If Postgres rejects the statement, the command fails with a JSON-RPC error carrying the Postgres error message. +- If Postgres rejects the statement, `error` holds the Postgres error: its `message`, the SQLSTATE `code`, and `detail` and `hint` if Postgres sent them. `position` is the LSP position of the error in the document, if Postgres reported one. For example: `{ "message": "Statement failed: column \"nme\" does not exist", "columns": [], "rows": [], "truncated": false, "error": { "message": "column \"nme\" does not exist", "code": "42703", "position": { "line": 3, "character": 7 } } }`. Only a failed connection makes the command fail with a JSON-RPC error. ## Disabling Database Features From 5621c827db39fa61d1b8fe6a0cd2f5ccbea553c8 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Fri, 23 Oct 2026 14:21:24 +0000 Subject: [PATCH 4/7] fix(workspace): don't leak transactions or session settings from executed statements Executed statements ran on the pool that the language server uses for its own queries, and sqlx returned the connection to the pool as the statement left it. After an executed `begin`, later statements and the server's own snapshot and typecheck queries ran inside that transaction, and a `set search_path` applied to them as well. Executed statements now run on a pool of their own. After each one, the session is reset with `discard all`, after rolling back a transaction the statement left open. The pool caches no prepared statements, which `discard all` would drop. --- crates/pgls_workspace/src/workspace/server.rs | 4 +- .../src/workspace/server.tests.rs | 64 ++++++++++++++++++- .../workspace/server/connection_manager.rs | 64 ++++++++++++++----- .../workspace/server/statement_execution.rs | 55 ++++++++++++++-- docs/guides/configure_database.md | 1 + 5 files changed, 164 insertions(+), 24 deletions(-) diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 94c47adb8..8772f3f45 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -752,8 +752,8 @@ impl Workspace for WorkspaceServer { }); } - let Some(result) = self.connection.with_pool(&settings.db, |pool| { - let pool = pool.clone(); + let Some(result) = self.connection.with_pools(&settings.db, |pools| { + let pool = pools.statement_pool.clone(); run_async(async move { statement_execution::execute_statement(&pool, &content).await })? .map_err(WorkspaceError::from) }) else { diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 4ff9d54bc..61e29208a 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -17,7 +17,7 @@ use pgls_configuration::plpgsql_check::PartialPlPgSqlCheckConfiguration; use pgls_diagnostics::Diagnostic; use pgls_fs::PgLSPath; use pgls_text_size::{TextRange, TextSize}; -use sqlx::{Executor, PgPool}; +use sqlx::{Executor, PgPool, Row}; use crate::{ Workspace, WorkspaceError, @@ -2105,3 +2105,65 @@ async fn execute_statement_returns_detail_and_hint(test_db: PgPool) { assert_eq!(error.detail.as_deref(), Some("the detail")); assert_eq!(error.hint.as_deref(), Some("the hint")); } + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_rolls_back_the_transaction_it_opens(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql(&workspace, "begin;"); + assert_eq!( + result.message, + "Successfully executed statement. Rows affected: 0. The open transaction was rolled back, as executed statements don't share a session." + ); + + // In autocommit mode, `now()` is the start of this very statement. + let (_, result) = execute_sql( + &workspace, + "select now() = statement_timestamp() as autocommit;", + ); + assert_eq!(result.rows, [[Some("t".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_resets_session_settings(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql(&workspace, "set search_path to pg_catalog;"); + assert_eq!(result.error, None); + let (_, result) = execute_sql(&workspace, "show search_path;"); + assert_eq!(result.rows, [[Some("\"$user\", public".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn executed_statements_do_not_affect_the_server_connections(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + let settings = workspace.workspaces(); + let pool = workspace + .connection + .get_pool(&settings.settings().unwrap().db) + .unwrap(); + + execute_sql(&workspace, "begin;"); + execute_sql(&workspace, "set search_path to pg_catalog;"); + + // The session reset after a statement drops the prepared statements that sqlx caches on the + // connection, so the language server's own queries must not share it. + assert_eq!( + pool.size(), + 0, + "the statements ran on a connection of the language server's own pool" + ); + // The language server's own queries see neither the transaction nor the setting. + workspace.snapshot.load(&pool).expect("snapshot"); + // `now()` is the start of the statement in autocommit mode. Over the extended query + // protocol, `statement_timestamp()` moves on with every message, so use the simple one. + let row = sqlx::raw_sql("select now() = statement_timestamp(), current_setting('search_path')") + .fetch_one(&pool) + .await + .unwrap(); + assert!( + row.get::(0), + "the connection is inside a transaction" + ); + assert_eq!(row.get::(1), "\"$user\", public"); +} diff --git a/crates/pgls_workspace/src/workspace/server/connection_manager.rs b/crates/pgls_workspace/src/workspace/server/connection_manager.rs index f51863526..107688fe1 100644 --- a/crates/pgls_workspace/src/workspace/server/connection_manager.rs +++ b/crates/pgls_workspace/src/workspace/server/connection_manager.rs @@ -71,9 +71,20 @@ fn connection_string_options(uri: &str) -> Result }) } +/// The connection pools of one database. +#[derive(Clone)] +pub(crate) struct DatabasePools { + /// The pool for the language server's own queries. + pub(crate) pool: PgPool, + /// The pool for the statements a user executes. Its connections never run the language + /// server's own queries, so the session state that a statement leaves behind can't leak into + /// them. + pub(crate) statement_pool: PgPool, +} + /// Cached connection pool with last access time struct CachedPool { - pool: PgPool, + pools: DatabasePools, settings_fingerprint: u64, last_accessed: Instant, idle_timeout: Duration, @@ -101,11 +112,16 @@ impl ConnectionManager { } } - /// Get a connection pool for the given database settings. - /// If a pool already exists for these settings, it will be returned. - /// If not, a new pool will be created if connections are enabled. - /// Will also clean up idle connections that haven't been accessed for a while. + #[cfg(test)] pub(crate) fn get_pool(&self, settings: &DatabaseSettings) -> Option { + self.get_pools(settings).map(|pools| pools.pool) + } + + /// Get the connection pools for the given database settings. + /// If pools already exist for these settings, they will be returned. + /// If not, new pools will be created if connections are enabled. + /// Will also clean up idle connections that haven't been accessed for a while. + pub(crate) fn get_pools(&self, settings: &DatabaseSettings) -> Option { let key = ConnectionKey::from(settings); if !settings.enable_connection { @@ -123,7 +139,7 @@ impl ConnectionManager { && let Some(cached_pool) = pools.get(&key) && cached_pool.settings_fingerprint == settings_fingerprint { - return Some(cached_pool.pool.clone()); + return Some(cached_pool.pools.clone()); } } @@ -140,7 +156,7 @@ impl ConnectionManager { // Double-check after acquiring write lock if let Some(cached_pool) = pools.get_mut(&key) { cached_pool.last_accessed = Instant::now(); - return Some(cached_pool.pool.clone()); + return Some(cached_pool.pools.clone()); } // Clean up idle connections before creating new ones to avoid unbounded growth @@ -178,13 +194,20 @@ impl ConnectionManager { let timeout = settings.conn_timeout_secs; - let pool = PoolOptions::::new() - .acquire_timeout(timeout) - .acquire_slow_threshold(Duration::from_secs(2)) - .connect_lazy_with(config); + let pool_options = || { + PoolOptions::::new() + .acquire_timeout(timeout) + .acquire_slow_threshold(Duration::from_secs(2)) + }; + let database_pools = DatabasePools { + pool: pool_options().connect_lazy_with(config.clone()), + // Without cached prepared statements, `discard all` can reset the whole session after + // a statement. + statement_pool: pool_options().connect_lazy_with(config.statement_cache_capacity(0)), + }; let cached_pool = CachedPool { - pool: pool.clone(), + pools: database_pools.clone(), settings_fingerprint, last_accessed: Instant::now(), // TODO: add this to the db settings, for now default to five minutes @@ -193,7 +216,7 @@ impl ConnectionManager { pools.insert(key, cached_pool); - Some(pool) + Some(database_pools) } /// Like [`Self::with_pool`], but returns the error of a connection that failed recently @@ -220,9 +243,18 @@ impl ConnectionManager { settings: &DatabaseSettings, operation: impl FnOnce(&PgPool) -> Result, ) -> Option> { - let pool = self.get_pool(settings)?; + self.with_pools(settings, |pools| operation(&pools.pool)) + } + + /// Like [`Self::with_pool`], but passes the pool for executed statements as well. + pub(crate) fn with_pools( + &self, + settings: &DatabaseSettings, + operation: impl FnOnce(&DatabasePools) -> Result, + ) -> Option> { + let pools = self.get_pools(settings)?; let key = ConnectionKey::from(settings); - let result = operation(&pool); + let result = operation(&pools); self.record_result(&key, settings_fingerprint(settings), &result); Some(result) } @@ -367,7 +399,7 @@ mod tests { assert_ne!(first_fingerprint, changed_fingerprint); assert_eq!(cached.settings_fingerprint, changed_fingerprint); assert_eq!( - cached.pool.connect_options().get_host(), + cached.pools.pool.connect_options().get_host(), pool.connect_options().get_host() ); assert_eq!(pools.len(), 1); diff --git a/crates/pgls_workspace/src/workspace/server/statement_execution.rs b/crates/pgls_workspace/src/workspace/server/statement_execution.rs index ca7b0d091..18c81288d 100644 --- a/crates/pgls_workspace/src/workspace/server/statement_execution.rs +++ b/crates/pgls_workspace/src/workspace/server/statement_execution.rs @@ -1,7 +1,7 @@ use futures::TryStreamExt; use pgls_text_size::TextSize; use sqlx::{ - Column, Either, Executor, PgPool, Row, Statement, ValueRef, + Column, Either, Executor, PgConnection, PgPool, Row, Statement, ValueRef, postgres::{PgColumn, PgDatabaseError, PgErrorPosition, PgRow}, }; @@ -15,11 +15,16 @@ use crate::features::code_actions::{ /// The statement runs over the simple query protocol, so Postgres sends every value in its text /// representation and no type needs decoding. An error that Postgres raises for the statement is /// part of the result, with its position relative to `sql`. Only a failed connection is an error. +/// +/// Each statement runs on its own: afterwards, the session is reset, so a `begin` or `set` does +/// not carry over to the next statement on the connection. pub(super) async fn execute_statement( pool: &PgPool, sql: &str, ) -> Result { - match collect_result(pool, sql).await { + let mut conn = pool.acquire().await?; + + let result = match collect_result(&mut conn, sql).await { Err(err) if !is_connection_failure(&err) => match err .as_database_error() .and_then(|err| err.try_downcast_ref::()) @@ -28,16 +33,56 @@ pub(super) async fn execute_statement( None => Err(err), }, result => result, + }; + + match reset_session(&mut conn).await { + Ok(rolled_back) => result.map(|mut result| { + if rolled_back { + if !result.message.ends_with('.') { + result.message.push('.'); + } + result.message.push_str( + " The open transaction was rolled back, as executed statements don't share a session.", + ); + } + result + }), + Err(err) => { + tracing::warn!("Closing the connection, as its session could not be reset: {err}"); + let _ = conn.close().await; + result + } } } -async fn collect_result(pool: &PgPool, sql: &str) -> Result { +/// Rolls back the transaction that a statement left open and resets the session, like +/// `pgbouncer`'s `server_reset_query`. Returns whether a transaction was open. +async fn reset_session(conn: &mut PgConnection) -> Result { + match conn.execute(sqlx::raw_sql("discard all")).await { + Ok(_) => Ok(false), + // `discard all` can't run inside a transaction block (25001), and a failed transaction + // only accepts a rollback (25P02). + Err(sqlx::Error::Database(err)) + if matches!(err.code().as_deref(), Some("25001" | "25P02")) => + { + conn.execute(sqlx::raw_sql("rollback")).await?; + conn.execute(sqlx::raw_sql("discard all")).await?; + Ok(true) + } + Err(err) => Err(err), + } +} + +async fn collect_result( + conn: &mut PgConnection, + sql: &str, +) -> Result { let mut columns = None; let mut rows = Vec::new(); let mut rows_affected = 0; let mut truncated = false; - let mut results = sqlx::raw_sql(sql).fetch_many(pool); + let mut results = sqlx::raw_sql(sql).fetch_many(&mut *conn); while let Some(result) = results.try_next().await? { match result { Either::Left(result) => rows_affected += result.rows_affected(), @@ -56,7 +101,7 @@ async fn collect_result(pool: &PgPool, sql: &str) -> Result Date: Sat, 24 Oct 2026 06:44:52 +0000 Subject: [PATCH 5/7] feat(workspace): guard statement execution Executed statements had no timeout and could not be cancelled, sent values of any size to the editor, made Postgres produce the whole result of a truncated query, failed on COPY with a protocol error, dropped the notices, and left the schema snapshot stale after DDL. - A statement runs for at most db.statementTimeoutSecs (60 by default). - $/cancelRequest cancels it in the database. The language server runs it on a blocking thread, so it keeps handling requests meanwhile. - A value is cut to 10,000 bytes and the result to 1 MB, both marked. - A select or values runs in a cursor that fetches 101 rows. Other statements still run to completion. - COPY FROM STDIN and TO STDOUT are refused with a clear message. - Notices and warnings are returned and appended to the message. - After a statement that may change the schema, the snapshot is reloaded. --- crates/pgls_configuration/src/database.rs | 8 + crates/pgls_configuration/src/lib.rs | 1 + crates/pgls_lsp/src/handlers/code_actions.rs | 35 ++- crates/pgls_lsp/tests/server.rs | 106 +++++++- .../src/features/code_actions.rs | 17 +- crates/pgls_workspace/src/settings.rs | 11 + crates/pgls_workspace/src/workspace.rs | 7 + crates/pgls_workspace/src/workspace/client.rs | 7 + crates/pgls_workspace/src/workspace/server.rs | 115 +++++++- .../src/workspace/server.tests.rs | 240 ++++++++++++++++- .../workspace/server/statement_execution.rs | 252 +++++++++++++++--- .../server/statement_execution/notices.rs | 128 +++++++++ docs/guides/configure_database.md | 7 +- docs/schema.json | 9 + .../backend-jsonrpc/src/workspace.ts | 4 + .../backend-jsonrpc/src/workspace.ts | 4 + 16 files changed, 881 insertions(+), 70 deletions(-) create mode 100644 crates/pgls_workspace/src/workspace/server/statement_execution/notices.rs diff --git a/crates/pgls_configuration/src/database.rs b/crates/pgls_configuration/src/database.rs index 2e650031d..7185c0f1c 100644 --- a/crates/pgls_configuration/src/database.rs +++ b/crates/pgls_configuration/src/database.rs @@ -50,6 +50,11 @@ pub struct DatabaseConfiguration { #[partial(bpaf(long("conn_timeout_secs"), fallback(Some(10)), debug_fallback))] pub conn_timeout_secs: u16, + /// How long a statement executed from the editor may run, in seconds, before Postgres cancels + /// it. `0` keeps the `statement_timeout` of the database. Default: `60`. + #[partial(bpaf(hide))] + pub statement_timeout_secs: u16, + /// Actively disable all database-related features. #[partial(bpaf(long("disable-db"), req_flag(true)))] #[partial(cfg_attr(feature = "schema", schemars(skip)))] @@ -73,6 +78,7 @@ impl std::fmt::Debug for DatabaseConfiguration { &self.allow_statement_executions_against, ) .field("conn_timeout_secs", &self.conn_timeout_secs) + .field("statement_timeout_secs", &self.statement_timeout_secs) .field("disable_connection", &self.disable_connection) .finish() } @@ -95,6 +101,7 @@ impl std::fmt::Debug for PartialDatabaseConfiguration { &self.allow_statement_executions_against, ) .field("conn_timeout_secs", &self.conn_timeout_secs) + .field("statement_timeout_secs", &self.statement_timeout_secs) .field("disable_connection", &self.disable_connection) .finish() } @@ -112,6 +119,7 @@ impl Default for DatabaseConfiguration { database: "postgres".to_string(), allow_statement_executions_against: Default::default(), conn_timeout_secs: 10, + statement_timeout_secs: 60, } } } diff --git a/crates/pgls_configuration/src/lib.rs b/crates/pgls_configuration/src/lib.rs index 579d73086..0f8c1debb 100644 --- a/crates/pgls_configuration/src/lib.rs +++ b/crates/pgls_configuration/src/lib.rs @@ -195,6 +195,7 @@ impl PartialConfiguration { database: Some("postgres".to_string()), allow_statement_executions_against: Default::default(), conn_timeout_secs: Some(10), + statement_timeout_secs: None, disable_connection: Some(false), }), } diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 525c37ad5..873cd739a 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -8,9 +8,11 @@ use tower_lsp_server::ls_types::{ MessageType, }; +use pgls_workspace::Workspace; use pgls_workspace::features::code_actions::{ CodeActionKind, CodeActionsParams, CommandActionCategory, ExecuteStatementParams, }; +use std::sync::Arc; #[tracing::instrument(level = "debug", skip(session), err)] pub fn get_actions( @@ -107,9 +109,15 @@ pub async fn execute_command( let path = session.file_path(&doc_url)?; - let result = session - .workspace - .execute_statement(ExecuteStatementParams { statement_id, path })?; + let params = ExecuteStatementParams { statement_id, path }; + // The statement runs on a blocking thread, so that the server keeps handling requests, + // `$/cancelRequest` included. That drops this future, and the guard then cancels the + // statement in the database. + let guard = CancelStatementOnDrop(Some((session.workspace.clone(), params.clone()))); + let workspace = session.workspace.clone(); + let result = + tokio::task::spawn_blocking(move || workspace.execute_statement(params)).await??; + guard.disarm(); /* * Updating all diagnostics: the changes caused by the statement execution @@ -151,3 +159,24 @@ pub async fn execute_command( any => Err(anyhow!(format!("Unknown command: {any}"))), } } + +/// Cancels the execution of a statement when dropped, unless disarmed. +struct CancelStatementOnDrop(Option<(Arc, ExecuteStatementParams)>); + +impl CancelStatementOnDrop { + fn disarm(mut self) { + self.0 = None; + } +} + +impl Drop for CancelStatementOnDrop { + fn drop(&mut self) { + if let Some((workspace, params)) = self.0.take() { + tokio::task::spawn_blocking(move || { + if let Err(err) = workspace.cancel_statement_execution(params) { + tracing::warn!("Could not cancel the statement execution: {err}"); + } + }); + } + } +} diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index 4c40f47fb..487a903fc 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -1371,6 +1371,24 @@ async fn execute_statement_at( doc_url: &Uri, position: Position, ) -> Result> { + let params = execute_statement_command(server, doc_url, position).await?; + let result = server + .request::>( + "workspace/executeCommand", + "_execStmt", + params, + ) + .await?; + + Ok(result.flatten()) +} + +/// The `pgls.executeStatement` command of the statement at `position`. +async fn execute_statement_command( + server: &mut Server, + doc_url: &Uri, + position: Position, +) -> Result { let code_actions_response = server .request::( "textDocument/codeAction", @@ -1409,19 +1427,11 @@ async fn execute_statement_at( }) .expect("Did not find executeStatement command!"); - let result = server - .request::>( - "workspace/executeCommand", - "_execStmt", - ExecuteCommandParams { - command: exec_statement_command.0, - arguments: exec_statement_command.1, - ..Default::default() - }, - ) - .await?; - - Ok(result.flatten()) + Ok(ExecuteCommandParams { + command: exec_statement_command.0, + arguments: exec_statement_command.1, + ..Default::default() + }) } #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] @@ -1540,6 +1550,76 @@ async fn execute_statement_returns_the_postgres_error_at_its_position( Ok(()) } +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn cancel_request_cancels_the_executed_statement(test_db: PgPool) -> Result<()> { + let mut server = statement_execution_server(&test_db); + + server.initialize().await?; + server.initialized().await?; + server.load_configuration().await?; + + let doc_url = url!("test.sql"); + server + .open_named_document("select pg_sleep(30);", doc_url.clone(), "sql") + .await?; + let params = execute_statement_command(&mut server, &doc_url, Position::new(0, 3)).await?; + + let sleeping = async || { + sqlx::query_scalar::<_, i64>( + "select count(*) from pg_stat_activity where state = 'active' and query like '%pg_sleep(30)%' and pid <> pg_backend_pid()", + ) + .fetch_one(&test_db) + .await + .unwrap() + }; + + // The executed statement runs on a blocking thread, so the server handles the cancel + // request while it waits for the database. + let execution = tokio::spawn( + server.service.ready().await.map_err(Error::msg)?.call( + Request::build("workspace/executeCommand") + .id("exec") + .params(to_value(params)?) + .finish(), + ), + ); + tokio::time::timeout(Duration::from_secs(5), async { + while sleeping().await == 0 { + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .context("the statement did not start")?; + + server + .notify( + "$/cancelRequest", + lsp::CancelParams { + id: lsp::NumberOrString::String("exec".into()), + }, + ) + .await?; + + let response = execution.await?.map_err(Error::msg)?.unwrap(); + assert_eq!( + response.error().map(|error| error.code), + Some(tower_lsp_server::jsonrpc::ErrorCode::RequestCancelled) + ); + + // The database cancels the statement as well. + tokio::time::timeout(Duration::from_secs(5), async { + while sleeping().await > 0 { + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .context("the statement still runs")?; + + server.shutdown().await?; + + Ok(()) +} + #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] async fn test_invalidate_schema_cache(test_db: PgPool) -> Result<()> { let factory = ServerFactory::default(); diff --git a/crates/pgls_workspace/src/features/code_actions.rs b/crates/pgls_workspace/src/features/code_actions.rs index 0b9bb73c5..af98c601a 100644 --- a/crates/pgls_workspace/src/features/code_actions.rs +++ b/crates/pgls_workspace/src/features/code_actions.rs @@ -51,7 +51,7 @@ pub enum CommandActionCategory { InvalidateSchemaCache, } -#[derive(Debug, serde::Serialize, serde::Deserialize)] +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct ExecuteStatementParams { pub statement_id: StatementId, @@ -62,6 +62,14 @@ pub struct ExecuteStatementParams { /// as truncated. pub const EXECUTE_STATEMENT_MAX_ROWS: usize = 100; +/// The most bytes of a value that an executed statement returns. A longer value is cut and ends +/// with `…`, and the result is marked with `truncated_values`. +pub const EXECUTE_STATEMENT_MAX_VALUE_SIZE: usize = 10_000; + +/// The most bytes of values that an executed statement returns. The rows that would exceed it are +/// dropped, and the result is marked as truncated. +pub const EXECUTE_STATEMENT_MAX_RESULT_SIZE: usize = 1_000_000; + /// The outcome of executing a statement. /// /// A statement that returns rows (`select`, `values`, `returning`, `show`, `explain`, ...) fills @@ -82,12 +90,19 @@ pub struct ExecuteStatementResult { /// Whether the statement returned more rows than `rows` holds. #[serde(default)] pub truncated: bool, + /// Whether a value in `rows` was cut to [`EXECUTE_STATEMENT_MAX_VALUE_SIZE`] bytes. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub truncated_values: bool, /// The number of rows affected, for a statement that returns no rows. #[serde(default, skip_serializing_if = "Option::is_none")] pub rows_affected: Option, /// Why Postgres rejected the statement, if it did. #[serde(default, skip_serializing_if = "Option::is_none")] pub error: Option, + /// The notices and warnings Postgres sent, e.g. for `raise notice`, as `NOTICE: message`. + /// `message` ends with them as well. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub notices: Vec, } /// An error that Postgres raised while executing a statement. diff --git a/crates/pgls_workspace/src/settings.rs b/crates/pgls_workspace/src/settings.rs index e346b815b..dd849c081 100644 --- a/crates/pgls_workspace/src/settings.rs +++ b/crates/pgls_workspace/src/settings.rs @@ -732,6 +732,7 @@ impl Default for TypecheckSettings { pub const DEFAULT_DATABASE_PASSWORD: &str = "postgres"; /// Database settings for the entire workspace +#[derive(Clone)] pub struct DatabaseSettings { pub enable_connection: bool, pub connection_string: Option, @@ -743,6 +744,9 @@ pub struct DatabaseSettings { pub password: Option, pub database: String, pub conn_timeout_secs: Duration, + /// How long an executed statement may run. `None` keeps the `statement_timeout` of the + /// database. + pub statement_timeout: Option, pub allow_statement_executions: bool, /// The patterns of `db.allowStatementExecutionsAgainst` that are not valid globs. pub invalid_statement_execution_patterns: Vec, @@ -789,6 +793,7 @@ impl std::fmt::Debug for DatabaseSettings { .field("password", &"[redacted]") .field("database", &self.database) .field("conn_timeout_secs", &self.conn_timeout_secs) + .field("statement_timeout", &self.statement_timeout) .field( "allow_statement_executions", &self.allow_statement_executions, @@ -813,6 +818,7 @@ impl Default for DatabaseSettings { password: None, database: "postgres".to_string(), conn_timeout_secs: Duration::from_secs(10), + statement_timeout: Some(Duration::from_secs(60)), allow_statement_executions: true, invalid_statement_execution_patterns: Vec::new(), connection_string_error: None, @@ -911,6 +917,11 @@ impl From for DatabaseSettings { .conn_timeout_secs .map(|s| Duration::from_secs(s.into())) .unwrap_or(d.conn_timeout_secs), + statement_timeout: match value.statement_timeout_secs { + Some(0) => None, + Some(secs) => Some(Duration::from_secs(secs.into())), + None => d.statement_timeout, + }, allow_statement_executions, invalid_statement_execution_patterns, diff --git a/crates/pgls_workspace/src/workspace.rs b/crates/pgls_workspace/src/workspace.rs index 0ae2afed3..9e0b07572 100644 --- a/crates/pgls_workspace/src/workspace.rs +++ b/crates/pgls_workspace/src/workspace.rs @@ -184,6 +184,13 @@ pub trait Workspace: Send + Sync + RefUnwindSafe { params: ExecuteStatementParams, ) -> Result; + /// Cancels the execution of the statement, if it is running. The execution then returns the + /// error that Postgres raises for the cancelled statement. + fn cancel_statement_execution( + &self, + params: ExecuteStatementParams, + ) -> Result<(), WorkspaceError>; + /// Invalidate the schema cache. /// /// # Arguments diff --git a/crates/pgls_workspace/src/workspace/client.rs b/crates/pgls_workspace/src/workspace/client.rs index 803f05cf3..9cd0c608f 100644 --- a/crates/pgls_workspace/src/workspace/client.rs +++ b/crates/pgls_workspace/src/workspace/client.rs @@ -106,6 +106,13 @@ where self.request("pgls/execute_statement", params) } + fn cancel_statement_execution( + &self, + params: crate::features::code_actions::ExecuteStatementParams, + ) -> Result<(), WorkspaceError> { + self.request("pgls/cancel_statement_execution", params) + } + fn register_project_folder( &self, params: RegisterProjectFolderParams, diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index 8772f3f45..0d0dd8d33 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -104,6 +104,16 @@ pub struct WorkspaceServer { #[cfg(feature = "db")] connection: ConnectionManager, + + /// The statements that are being executed, to cancel them + #[cfg(feature = "db")] + running_statements: Mutex>, +} + +#[cfg(feature = "db")] +struct RunningStatement { + params: ExecuteStatementParams, + cancel: futures::channel::oneshot::Sender<()>, } /// The `Workspace` object is long-lived, so we want it to be able to cross @@ -131,6 +141,7 @@ impl WorkspaceServer { catalog_base: Mutex::default(), snapshot_before: Mutex::default(), connection: ConnectionManager::new(), + running_statements: Mutex::default(), } } @@ -718,7 +729,9 @@ impl Workspace for WorkspaceServer { .get(¶ms.path) .ok_or(WorkspaceError::not_found())?; - let stmt = parser.find(params.statement_id, ExecuteStatementMapper); + let stmt = parser.find(params.statement_id.clone(), ExecuteStatementMapper); + // Files can change while the statement runs. + drop(documents); if stmt.is_none() { return Ok(ExecuteStatementResult { @@ -729,34 +742,89 @@ impl Workspace for WorkspaceServer { let (_id, range, content, ast) = stmt.unwrap(); - if ast.is_none() { + let Some(ast) = ast else { return Ok(ExecuteStatementResult { message: "Statement is invalid.".into(), ..Default::default() }); }; + // The language server has no data to send, and nowhere to put the data it receives. + if let pgls_query::NodeEnum::CopyStmt(copy) = &ast + && copy.filename.is_empty() + && !copy.is_program + { + let copy = if copy.is_from { + "COPY ... FROM STDIN" + } else { + "COPY ... TO STDOUT" + }; + return Ok(ExecuteStatementResult { + message: format!( + "`{copy}` can't be executed from the editor. Use `\\copy` in psql, or a file on the database server." + ), + ..Default::default() + }); + } + let settings = self.workspaces(); let no_connection = ExecuteStatementResult { message: "No database connection available.".into(), ..Default::default() }; - let Some(settings) = settings.settings() else { + let Some(db) = settings.settings().map(|settings| settings.db.clone()) else { return Ok(no_connection); }; + // The settings can change while the statement runs. + drop(settings); - if let Some(reason) = settings.db.statement_execution_disabled_reason() { + if let Some(reason) = db.statement_execution_disabled_reason() { return Ok(ExecuteStatementResult { message: reason, ..Default::default() }); } - let Some(result) = self.connection.with_pools(&settings.db, |pools| { - let pool = pools.statement_pool.clone(); - run_async(async move { statement_execution::execute_statement(&pool, &content).await })? - .map_err(WorkspaceError::from) - }) else { + let options = statement_execution::ExecuteOptions { + is_query: statement_execution::is_query(&ast), + timeout: db.statement_timeout, + }; + let (cancel, cancelled) = futures::channel::oneshot::channel(); + self.running_statements + .lock() + .unwrap() + .push(RunningStatement { + params: params.clone(), + cancel, + }); + + let result = self.connection.with_pools(&db, |pools| { + let statement_pools = pools.clone(); + let result = run_async(async move { + statement_execution::execute_statement( + &statement_pools, + &content, + options, + cancelled, + ) + .await + })? + .map_err(WorkspaceError::from)?; + + // Diagnostics, completions and hover should see the objects the statement created. + if result.error.is_none() && statement_execution::may_change_schema(&ast) { + self.snapshot.clear_connection(&pools.pool); + } + Ok(result) + }); + + // The execution is over once it dropped the receiver. + self.running_statements + .lock() + .unwrap() + .retain(|running| !running.cancel.is_canceled()); + + let Some(result) = result else { return Ok(no_connection); }; @@ -784,6 +852,35 @@ impl Workspace for WorkspaceServer { }) } + #[cfg(feature = "db")] + fn cancel_statement_execution( + &self, + params: ExecuteStatementParams, + ) -> Result<(), WorkspaceError> { + let cancelled = self + .running_statements + .lock() + .unwrap() + .extract_if(.., |running| { + running.params.path == params.path + && running.params.statement_id == params.statement_id + }) + .collect::>(); + for running in cancelled { + let _ = running.cancel.send(()); + } + Ok(()) + } + + #[cfg(not(feature = "db"))] + #[allow(unused_variables)] + fn cancel_statement_execution( + &self, + params: ExecuteStatementParams, + ) -> Result<(), WorkspaceError> { + Ok(()) + } + #[cfg(feature = "db")] fn invalidate_schema_cache(&self, all: bool) -> Result<(), WorkspaceError> { if all { diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 61e29208a..d6b07daf3 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -23,7 +23,8 @@ use crate::{ Workspace, WorkspaceError, features::code_actions::{ CodeActionKind, CodeActionsParams, CommandAction, CommandActionCategory, - EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementError, ExecuteStatementParams, + EXECUTE_STATEMENT_MAX_RESULT_SIZE, EXECUTE_STATEMENT_MAX_ROWS, + EXECUTE_STATEMENT_MAX_VALUE_SIZE, ExecuteStatementError, ExecuteStatementParams, ExecuteStatementResult, }, features::format::PullFileFormattingParams, @@ -1822,6 +1823,18 @@ fn is_typecheck(diagnostic: &pgls_diagnostics::serde::Diagnostic) -> bool { fn get_statement_execution_workspace( test_db: &PgPool, allow_statement_executions_against: &[String], +) -> WorkspaceServer { + get_statement_execution_workspace_with_timeout( + test_db, + allow_statement_executions_against, + None, + ) +} + +fn get_statement_execution_workspace_with_timeout( + test_db: &PgPool, + allow_statement_executions_against: &[String], + statement_timeout_secs: Option, ) -> WorkspaceServer { let mut conf = PartialConfiguration::init(); conf.merge_with(PartialConfiguration { @@ -1837,6 +1850,7 @@ fn get_statement_execution_workspace( allow_statement_executions_against: Some(StringSet::from_iter( allow_statement_executions_against.iter().cloned(), )), + statement_timeout_secs, ..Default::default() }), ..Default::default() @@ -1977,8 +1991,7 @@ async fn execute_statement_returns_rows(test_db: PgPool) { vec![Some("2".into()), Some("two".into()), None, None, None], ], truncated: false, - rows_affected: None, - error: None, + ..Default::default() } ); } @@ -2167,3 +2180,224 @@ async fn executed_statements_do_not_affect_the_server_connections(test_db: PgPoo ); assert_eq!(row.get::(1), "\"$user\", public"); } + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_times_out(test_db: PgPool) { + let workspace = get_statement_execution_workspace_with_timeout( + &test_db, + &allow_test_host(&test_db), + Some(1), + ); + + let start = Instant::now(); + let (_, result) = execute_sql(&workspace, "select pg_sleep(10);"); + + assert_eq!( + result.error.map(|error| error.code).as_deref(), + Some("57014") + ); + assert!( + start.elapsed() < Duration::from_secs(5), + "{:?}", + start.elapsed() + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_can_be_cancelled(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + let path = PgLSPath::new("execute.sql"); + workspace + .open_file(OpenFileParams { + path: path.clone(), + content: "select pg_sleep(30);".into(), + version: 1, + }) + .unwrap(); + let statement_id = workspace + .pull_code_actions(CodeActionsParams { + path: path.clone(), + cursor_position: TextSize::new(0), + only: vec![], + skip: vec![], + }) + .unwrap() + .actions + .into_iter() + .find_map(|action| match action.kind { + CodeActionKind::Command(CommandAction { + category: CommandActionCategory::ExecuteStatement(statement_id), + }) => Some(statement_id), + _ => None, + }) + .unwrap(); + let params = ExecuteStatementParams { path, statement_id }; + + let start = Instant::now(); + let result = std::thread::scope(|scope| { + let execution = scope.spawn(|| workspace.execute_statement(params.clone())); + std::thread::sleep(Duration::from_millis(500)); + workspace + .cancel_statement_execution(params.clone()) + .unwrap(); + execution.join().unwrap().unwrap() + }); + + assert_eq!( + result.error.map(|error| error.code).as_deref(), + Some("57014") + ); + assert!( + start.elapsed() < Duration::from_secs(10), + "{:?}", + start.elapsed() + ); + assert!(workspace.running_statements.lock().unwrap().is_empty()); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_stops_a_truncated_query_early(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + // Without the cursor, Postgres would produce and send all 50 million rows. + let start = Instant::now(); + let (_, result) = execute_sql(&workspace, "select generate_series(1, 50000000) as n"); + let (_, next) = execute_sql(&workspace, "select 1 as one;"); + + assert!(result.truncated); + assert_eq!(result.rows.len(), EXECUTE_STATEMENT_MAX_ROWS); + assert_eq!(next.rows, [[Some("1".to_string())]]); + assert!( + start.elapsed() < Duration::from_secs(3), + "{:?}", + start.elapsed() + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_completes_a_truncated_command(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + execute_sql(&workspace, "create table inserted (id int);"); + let (_, result) = execute_sql( + &workspace, + "insert into inserted select generate_series(1, 1000) returning id;", + ); + assert!(result.truncated); + assert_eq!(result.rows.len(), EXECUTE_STATEMENT_MAX_ROWS); + + let (_, result) = execute_sql(&workspace, "select count(*) from inserted;"); + assert_eq!(result.rows, [[Some("1000".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_cuts_long_values(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + // The cut falls into the middle of an `ä`. + let (_, result) = execute_sql( + &workspace, + "select repeat('x', 20000000) as big, 'x' || repeat('ä', 10000) as umlauts, 'short' as short;", + ); + + assert!(result.truncated_values); + let [big, umlauts, short] = result.rows[0].clone().try_into().unwrap(); + let big = big.unwrap(); + assert_eq!(big.len(), EXECUTE_STATEMENT_MAX_VALUE_SIZE + '…'.len_utf8()); + assert!(big.ends_with('…')); + let umlauts = umlauts.unwrap(); + assert_eq!( + umlauts.len(), + EXECUTE_STATEMENT_MAX_VALUE_SIZE - 1 + '…'.len_utf8() + ); + assert!(umlauts.ends_with("ää…")); + assert_eq!(short.as_deref(), Some("short")); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_caps_the_result_size(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql( + &workspace, + "select repeat('x', 10000) as a, repeat('y', 10000) as b from generate_series(1, 150);", + ); + + assert!(result.truncated); + assert!(!result.truncated_values); + assert_eq!( + result.rows.len(), + EXECUTE_STATEMENT_MAX_RESULT_SIZE / (2 * EXECUTE_STATEMENT_MAX_VALUE_SIZE) + ); + assert_eq!(result.message, "Showing the first 50 rows (truncated)."); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_refuses_copy_with_the_client(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + execute_sql(&workspace, "create table copy_target (id int);"); + + let (_, result) = execute_sql(&workspace, "copy (select 1) to stdout;"); + assert_eq!( + result.message, + "`COPY ... TO STDOUT` can't be executed from the editor. Use `\\copy` in psql, or a file on the database server." + ); + let (_, result) = execute_sql(&workspace, "copy copy_target from stdin;"); + assert!( + result + .message + .starts_with("`COPY ... FROM STDIN` can't be executed") + ); + + let (_, result) = execute_sql(&workspace, "select 1 as one;"); + assert_eq!(result.rows, [[Some("1".to_string())]]); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn execute_statement_returns_notices(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql( + &workspace, + "do $$ begin raise notice 'hello'; raise warning 'careful'; end $$;", + ); + + assert_eq!(result.notices, ["NOTICE: hello", "WARNING: careful"]); + assert_eq!( + result.message, + "Successfully executed statement. Rows affected: 0\nNOTICE: hello\nWARNING: careful" + ); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +async fn executed_ddl_refreshes_the_snapshot(test_db: PgPool) { + let workspace = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + let other = PgLSPath::new("other.sql"); + workspace + .open_file(OpenFileParams { + path: other.clone(), + content: "select id from created_by_execute;".into(), + version: 1, + }) + .unwrap(); + let typecheck_findings = || { + workspace + .pull_file_diagnostics(crate::workspace::PullFileDiagnosticsParams { + path: other.clone(), + categories: RuleCategories::all(), + max_diagnostics: 100, + only: vec![], + skip: vec![], + }) + .expect("Unable to pull diagnostics") + .diagnostics + .iter() + .filter(|diagnostic| is_typecheck(diagnostic)) + .count() + }; + assert!(typecheck_findings() > 0, "the table should be unknown"); + + execute_sql(&workspace, "create table created_by_execute (id int);"); + + assert_eq!(typecheck_findings(), 0); +} diff --git a/crates/pgls_workspace/src/workspace/server/statement_execution.rs b/crates/pgls_workspace/src/workspace/server/statement_execution.rs index 18c81288d..15f2e1899 100644 --- a/crates/pgls_workspace/src/workspace/server/statement_execution.rs +++ b/crates/pgls_workspace/src/workspace/server/statement_execution.rs @@ -1,15 +1,70 @@ -use futures::TryStreamExt; +use std::time::Duration; + +use futures::{FutureExt, TryStreamExt, channel::oneshot, future}; +use pgls_query::NodeEnum; use pgls_text_size::TextSize; use sqlx::{ - Column, Either, Executor, PgConnection, PgPool, Row, Statement, ValueRef, + Column, Either, Executor, PgConnection, Row, Statement, ValueRef, postgres::{PgColumn, PgDatabaseError, PgErrorPosition, PgRow}, }; +use tracing::instrument::WithSubscriber; use crate::diagnostics::is_connection_failure; use crate::features::code_actions::{ - EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementError, ExecuteStatementResult, + EXECUTE_STATEMENT_MAX_RESULT_SIZE, EXECUTE_STATEMENT_MAX_ROWS, + EXECUTE_STATEMENT_MAX_VALUE_SIZE, ExecuteStatementError, ExecuteStatementResult, }; +use super::connection_manager::DatabasePools; + +mod notices; + +/// The cursor that a query runs in, so that Postgres stops once it produced the rows the result +/// holds. +const CURSOR: &str = "pgls_statement"; + +/// How to execute a statement. +pub(super) struct ExecuteOptions { + /// Whether the statement is a query that a cursor can run: a `select` or `values` without + /// `into` and without a data-modifying `with`. + pub(super) is_query: bool, + /// How long the statement may run, or `None` to keep the session's `statement_timeout`. + pub(super) timeout: Option, +} + +/// Whether a cursor can run `ast`: `declare` takes a `select` or `values` without `into` and +/// without a data-modifying `with`. +pub(super) fn is_query(ast: &NodeEnum) -> bool { + let NodeEnum::SelectStmt(select) = ast else { + return false; + }; + select.into_clause.is_none() + && select.with_clause.as_ref().is_none_or(|with| { + with.ctes.iter().all(|cte| match &cte.node { + Some(NodeEnum::CommonTableExpr(cte)) => matches!( + cte.ctequery + .as_deref() + .and_then(|query| query.node.as_ref()), + Some(NodeEnum::SelectStmt(_)) + ), + _ => false, + }) + }) +} + +/// Whether executing `ast` may change the schema, i.e. it does more than reading or writing rows. +pub(super) fn may_change_schema(ast: &NodeEnum) -> bool { + match ast { + NodeEnum::SelectStmt(select) => select.into_clause.is_some(), + NodeEnum::InsertStmt(_) + | NodeEnum::UpdateStmt(_) + | NodeEnum::DeleteStmt(_) + | NodeEnum::MergeStmt(_) + | NodeEnum::VariableShowStmt(_) => false, + _ => true, + } +} + /// Executes `sql` and collects the rows it returns, if any. /// /// The statement runs over the simple query protocol, so Postgres sends every value in its text @@ -18,43 +73,94 @@ use crate::features::code_actions::{ /// /// Each statement runs on its own: afterwards, the session is reset, so a `begin` or `set` does /// not carry over to the next statement on the connection. +/// +/// When `cancel` fires, Postgres cancels the statement. pub(super) async fn execute_statement( - pool: &PgPool, + pools: &DatabasePools, sql: &str, + options: ExecuteOptions, + cancel: oneshot::Receiver<()>, ) -> Result { - let mut conn = pool.acquire().await?; + let notices = notices::NoticeCollector::default(); + let result = execute_on_own_connection(pools, sql, options, cancel) + .with_subscriber(notices.dispatch()) + .await; - let result = match collect_result(&mut conn, sql).await { - Err(err) if !is_connection_failure(&err) => match err - .as_database_error() - .and_then(|err| err.try_downcast_ref::()) - { - Some(err) => Ok(rejected(err, sql)), - None => Err(err), - }, - result => result, - }; + result.map(|mut result| { + result.notices = notices.take(); + for notice in &result.notices { + result.message.push('\n'); + result.message.push_str(notice); + } + result + }) +} - match reset_session(&mut conn).await { - Ok(rolled_back) => result.map(|mut result| { - if rolled_back { - if !result.message.ends_with('.') { - result.message.push('.'); +async fn execute_on_own_connection( + pools: &DatabasePools, + sql: &str, + options: ExecuteOptions, + cancel: oneshot::Receiver<()>, +) -> Result { + let mut pool_conn = pools.statement_pool.acquire().await?; + let conn: &mut PgConnection = &mut pool_conn; + let backend_pid = prepare_session(conn, options.timeout).await?; + + // Boxing the future pins down the lifetime of the borrowed connection, which the `Send` + // check of `run_async` can't infer otherwise. + let result = match future::select(run(conn, sql, options.is_query).boxed(), cancel).await { + future::Either::Left((result, _)) => result, + future::Either::Right((cancelled, run)) => { + if cancelled.is_ok() { + // A connection can't cancel its own statement while it waits for the result. + if let Err(err) = sqlx::query("select pg_cancel_backend($1)") + .bind(backend_pid) + .execute(&pools.pool) + .await + { + tracing::warn!("Could not cancel the executed statement: {err}"); } - result.message.push_str( - " The open transaction was rolled back, as executed statements don't share a session.", - ); } + run.await + } + }; + + match reset_session(conn).await { + // The cursor of a query has a transaction of its own, which a failed query leaves open. + Ok(rolled_back) if rolled_back && !options.is_query => result.map(|mut result| { + if !result.message.ends_with('.') { + result.message.push('.'); + } + result.message.push_str( + " The open transaction was rolled back, as executed statements don't share a session.", + ); result }), + Ok(_) => result, Err(err) => { tracing::warn!("Closing the connection, as its session could not be reset: {err}"); - let _ = conn.close().await; + let _ = pool_conn.close().await; result } } } +/// Applies the statement timeout and returns the process ID of the connection's backend, which +/// `pg_cancel_backend` takes. +async fn prepare_session( + conn: &mut PgConnection, + timeout: Option, +) -> Result { + let query = match timeout { + Some(timeout) => format!( + "select pg_backend_pid(), set_config('statement_timeout', '{}', false)", + timeout.as_millis() + ), + None => "select pg_backend_pid()".to_string(), + }; + conn.fetch_one(sqlx::raw_sql(&query)).await?.try_get(0) +} + /// Rolls back the transaction that a statement left open and resets the session, like /// `pgbouncer`'s `server_reset_query`. Returns whether a transaction was open. async fn reset_session(conn: &mut PgConnection) -> Result { @@ -73,26 +179,74 @@ async fn reset_session(conn: &mut PgConnection) -> Result { } } +/// Runs `sql`, a query in a cursor that fetches one row more than the result holds. +async fn run( + conn: &mut PgConnection, + sql: &str, + is_query: bool, +) -> Result { + let prefix = format!("begin; declare {CURSOR} no scroll cursor for\n"); + let query = if is_query { + let statement = sql.trim_end(); + let statement = statement.strip_suffix(';').unwrap_or(statement); + // The newline ends a trailing line comment. + format!( + "{prefix}{statement}\n; fetch {} from {CURSOR}; commit", + EXECUTE_STATEMENT_MAX_ROWS + 1 + ) + } else { + sql.to_string() + }; + + match collect_result(conn, &query, sql).await { + Err(err) if !is_connection_failure(&err) => match err + .as_database_error() + .and_then(|err| err.try_downcast_ref::()) + { + Some(err) => { + let offset = if is_query { prefix.len() } else { 0 }; + Ok(rejected(err, sql, offset)) + } + None => Err(err), + }, + result => result, + } +} + +/// Runs `query` and collects its rows. `sql` is the statement as written, to describe it. async fn collect_result( conn: &mut PgConnection, + query: &str, sql: &str, ) -> Result { let mut columns = None; let mut rows = Vec::new(); + let mut size = 0; let mut rows_affected = 0; let mut truncated = false; + let mut truncated_values = false; - let mut results = sqlx::raw_sql(sql).fetch_many(&mut *conn); + // Read the whole result even when it is truncated, so that the statement runs to completion + // before the session is reset. A query in a cursor stops early instead. + let mut results = sqlx::raw_sql(query).fetch_many(&mut *conn); while let Some(result) = results.try_next().await? { match result { Either::Left(result) => rows_affected += result.rows_affected(), Either::Right(row) => { - if rows.len() == EXECUTE_STATEMENT_MAX_ROWS { + columns.get_or_insert_with(|| column_names(row.columns())); + if truncated { + continue; + } + let values = row_values(&row, &mut truncated_values); + let row_size = values.iter().flatten().map(String::len).sum::(); + if rows.len() == EXECUTE_STATEMENT_MAX_ROWS + || size + row_size > EXECUTE_STATEMENT_MAX_RESULT_SIZE + { truncated = true; - break; + continue; } - columns.get_or_insert_with(|| column_names(row.columns())); - rows.push(row_values(&row)); + size += row_size; + rows.push(values); } } } @@ -112,7 +266,7 @@ async fn collect_result( Ok(match columns { Some(columns) => ExecuteStatementResult { message: if truncated { - format!("Showing the first {EXECUTE_STATEMENT_MAX_ROWS} rows (truncated).") + format!("Showing the first {} rows (truncated).", rows.len()) } else if rows.len() == 1 { "1 row".to_string() } else { @@ -121,6 +275,7 @@ async fn collect_result( columns, rows, truncated, + truncated_values, ..Default::default() }, None => ExecuteStatementResult { @@ -131,17 +286,21 @@ async fn collect_result( }) } -fn rejected(err: &PgDatabaseError, sql: &str) -> ExecuteStatementResult { +/// The result for an error that Postgres raised. Its position counts the characters of the query, +/// which has `offset` bytes of ASCII before `sql`. +fn rejected(err: &PgDatabaseError, sql: &str, offset: usize) -> ExecuteStatementResult { // Postgres counts the characters from 1, and points one past the end for an error at the end // of the input. let position = match err.position() { - Some(PgErrorPosition::Original(position)) => position.checked_sub(1).and_then(|index| { - sql.char_indices() - .map(|(offset, _)| offset) - .chain([sql.len()]) - .nth(index) - .and_then(|offset| TextSize::try_from(offset).ok()) - }), + Some(PgErrorPosition::Original(position)) => position + .checked_sub(1 + offset) + .and_then(|index| { + sql.char_indices() + .map(|(offset, _)| offset) + .chain([sql.len()]) + .nth(index) + }) + .and_then(|offset| TextSize::try_from(offset).ok()), _ => None, }; @@ -165,7 +324,8 @@ fn column_names(columns: &[PgColumn]) -> Vec { .collect() } -fn row_values(row: &PgRow) -> Vec> { +/// The values of `row`, each cut to [`EXECUTE_STATEMENT_MAX_VALUE_SIZE`] bytes. +fn row_values(row: &PgRow, truncated_values: &mut bool) -> Vec> { (0..row.len()) .map(|index| { let value = row.try_get_raw(index).ok()?; @@ -173,7 +333,19 @@ fn row_values(row: &PgRow) -> Vec> { return None; } let bytes = value.as_bytes().unwrap_or_default(); - Some(String::from_utf8_lossy(bytes).into_owned()) + if bytes.len() <= EXECUTE_STATEMENT_MAX_VALUE_SIZE { + return Some(String::from_utf8_lossy(bytes).into_owned()); + } + + *truncated_values = true; + // Cut before a UTF-8 continuation byte, so that no character is split. + let mut end = EXECUTE_STATEMENT_MAX_VALUE_SIZE; + while end > 0 && bytes[end] & 0b1100_0000 == 0b1000_0000 { + end -= 1; + } + let mut value = String::from_utf8_lossy(&bytes[..end]).into_owned(); + value.push('…'); + Some(value) }) .collect() } diff --git a/crates/pgls_workspace/src/workspace/server/statement_execution/notices.rs b/crates/pgls_workspace/src/workspace/server/statement_execution/notices.rs new file mode 100644 index 000000000..83ca3fc7d --- /dev/null +++ b/crates/pgls_workspace/src/workspace/server/statement_execution/notices.rs @@ -0,0 +1,128 @@ +//! sqlx has no handler for the notices Postgres sends, such as `raise notice` or a warning. It +//! only reports them as tracing events with the target `sqlx::postgres::notice`, so the statement +//! runs with a subscriber that collects these events and passes everything else on. + +use std::fmt; +use std::sync::{Arc, Mutex}; + +use tracing::field::{Field, Visit}; +use tracing::span::{Attributes, Id, Record}; +use tracing::subscriber::Interest; +use tracing::{Dispatch, Event, Level, Metadata, Subscriber}; + +const NOTICE_TARGET: &str = "sqlx::postgres::notice"; + +/// Collects the notices of the code that runs with [`Self::dispatch`] as its subscriber. +#[derive(Clone, Default)] +pub(super) struct NoticeCollector { + notices: Arc>>, +} + +impl NoticeCollector { + /// A subscriber that collects the notices and passes every other event and span on to the + /// current one. + pub(super) fn dispatch(&self) -> Dispatch { + Dispatch::new(Collector { + inner: tracing::dispatcher::get_default(Dispatch::clone), + notices: Arc::clone(&self.notices), + }) + } + + /// The collected notices, each as `SEVERITY: message`. + pub(super) fn take(&self) -> Vec { + std::mem::take(&mut self.notices.lock().unwrap()) + } +} + +struct Collector { + inner: Dispatch, + notices: Arc>>, +} + +fn is_notice(metadata: &Metadata<'_>) -> bool { + metadata.target() == NOTICE_TARGET +} + +impl Subscriber for Collector { + fn register_callsite(&self, metadata: &'static Metadata<'static>) -> Interest { + if is_notice(metadata) { + Interest::always() + } else { + self.inner.register_callsite(metadata) + } + } + + fn enabled(&self, metadata: &Metadata<'_>) -> bool { + is_notice(metadata) || self.inner.enabled(metadata) + } + + fn event(&self, event: &Event<'_>) { + let metadata = event.metadata(); + if is_notice(metadata) { + let mut visitor = MessageVisitor(None); + event.record(&mut visitor); + if let Some(message) = visitor.0 { + // sqlx maps the severity to a level; INFO and LOG both become TRACE. + let severity = match *metadata.level() { + Level::ERROR => "ERROR", + Level::WARN => "WARNING", + Level::INFO => "NOTICE", + Level::DEBUG => "DEBUG", + Level::TRACE => "INFO", + }; + self.notices + .lock() + .unwrap() + .push(format!("{severity}: {message}")); + } + if !self.inner.enabled(metadata) { + return; + } + } + self.inner.event(event); + } + + fn new_span(&self, span: &Attributes<'_>) -> Id { + self.inner.new_span(span) + } + + fn record(&self, span: &Id, values: &Record<'_>) { + self.inner.record(span, values); + } + + fn record_follows_from(&self, span: &Id, follows: &Id) { + self.inner.record_follows_from(span, follows); + } + + fn enter(&self, span: &Id) { + self.inner.enter(span); + } + + fn exit(&self, span: &Id) { + self.inner.exit(span); + } + + fn clone_span(&self, id: &Id) -> Id { + self.inner.clone_span(id) + } + + fn try_close(&self, id: Id) -> bool { + self.inner.try_close(id) + } +} + +struct MessageVisitor(Option); + +impl Visit for MessageVisitor { + fn record_str(&mut self, field: &Field, value: &str) { + if field.name() == "message" { + self.0 = Some(value.to_string()); + } + } + + fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) { + if field.name() == "message" { + self.0 = Some(format!("{value:?}")); + } + } +} diff --git a/docs/guides/configure_database.md b/docs/guides/configure_database.md index 7580b75b6..06e13a40a 100644 --- a/docs/guides/configure_database.md +++ b/docs/guides/configure_database.md @@ -134,9 +134,14 @@ The language server runs the statement with the LSP command `pgls.executeStateme - `columns` and `rows` are filled for statements that return rows, such as `select`, `values`, `show`, `explain`, and statements with `returning`. - Each value is a string in its Postgres text representation, as `psql` shows it, or `null` for SQL `NULL`. -- At most 100 rows are returned. `truncated` is `true` if the statement returned more. +- At most 100 rows and 1 MB of values are returned. `truncated` is `true` if the statement returned more. A `select` or `values` runs in a cursor that fetches 101 rows, so Postgres stops once it produced them. Any other statement runs to completion, and the rows after the first 100 are read and dropped: an `insert ... returning` inserts every row. +- A value longer than 10,000 bytes is cut and ends with `…`. `truncatedValues` is `true` then. +- `notices` lists the notices and warnings Postgres sent, e.g. for `raise notice`, as `"NOTICE: message"`. `message` ends with them as well. - `rowsAffected` is set instead of `columns` and `rows` for statements that return no rows, such as `insert` without `returning` or `create table`. For example: `{ "message": "Successfully executed statement. Rows affected: 2", "columns": [], "rows": [], "truncated": false, "rowsAffected": 2 }`. - Every statement runs on its own. Afterwards, the language server rolls back a transaction the statement left open and resets the session with `discard all`, so a `begin` or `set` does not carry over to the next executed statement. Executed statements never share a connection with the language server's own queries either. +- A statement runs for at most `db.statementTimeoutSecs` seconds, 60 by default, before Postgres cancels it with `canceling statement due to statement timeout`. `0` keeps the `statement_timeout` of the database. Cancelling the `workspace/executeCommand` request with `$/cancelRequest` cancels the statement in the database too. +- After a statement that may change the schema, i.e. anything but a query, `insert`, `update`, `delete`, `merge` or `show`, the language server reloads the schema, so diagnostics, completions and hover see the change. +- `COPY ... FROM STDIN` and `COPY ... TO STDOUT` are not executed. Use `\copy` in `psql`, or a file on the database server. - If the statement is not executed, for example because statement execution is disabled, only `message` is meaningful: `columns` and `rows` are empty and `rowsAffected` is not set. - If Postgres rejects the statement, `error` holds the Postgres error: its `message`, the SQLSTATE `code`, and `detail` and `hint` if Postgres sent them. `position` is the LSP position of the error in the document, if Postgres reported one. For example: `{ "message": "Statement failed: column \"nme\" does not exist", "columns": [], "rows": [], "truncated": false, "error": { "message": "column \"nme\" does not exist", "code": "42703", "position": { "line": 3, "character": 7 } } }`. Only a failed connection makes the command fail with a JSON-RPC error. diff --git a/docs/schema.json b/docs/schema.json index 68ca485ad..aacd7e18c 100644 --- a/docs/schema.json +++ b/docs/schema.json @@ -427,6 +427,15 @@ "format": "uint16", "minimum": 0.0 }, + "statementTimeoutSecs": { + "description": "How long a statement executed from the editor may run, in seconds, before Postgres cancels it. `0` keeps the `statement_timeout` of the database. Default: `60`.", + "type": [ + "integer", + "null" + ], + "format": "uint16", + "minimum": 0.0 + }, "username": { "description": "The username to connect to the database. Can also be set via the `PGUSER` environment variable.", "type": [ diff --git a/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts b/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts index 90572e043..d7d54043a 100644 --- a/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts +++ b/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts @@ -417,6 +417,10 @@ export interface PartialDatabaseConfiguration { * The port of the database. Can also be set via the `PGPORT` environment variable. */ port?: number; + /** + * How long a statement executed from the editor may run, in seconds, before Postgres cancels it. `0` keeps the `statement_timeout` of the database. Default: `60`. + */ + statementTimeoutSecs?: number; /** * The username to connect to the database. Can also be set via the `PGUSER` environment variable. */ diff --git a/packages/@postgrestools/backend-jsonrpc/src/workspace.ts b/packages/@postgrestools/backend-jsonrpc/src/workspace.ts index 90572e043..d7d54043a 100644 --- a/packages/@postgrestools/backend-jsonrpc/src/workspace.ts +++ b/packages/@postgrestools/backend-jsonrpc/src/workspace.ts @@ -417,6 +417,10 @@ export interface PartialDatabaseConfiguration { * The port of the database. Can also be set via the `PGPORT` environment variable. */ port?: number; + /** + * How long a statement executed from the editor may run, in seconds, before Postgres cancels it. `0` keeps the `statement_timeout` of the database. Default: `60`. + */ + statementTimeoutSecs?: number; /** * The username to connect to the database. Can also be set via the `PGUSER` environment variable. */ From 9324158cd1b03f150d667448ceda562372d27d60 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 26 Oct 2026 14:57:14 +0000 Subject: [PATCH 6/7] refactor(workspace): encapsulate the running statements in their own type The registry of executed statements was a Mutex> on the WorkspaceServer, and execute_statement and cancel_statement_execution locked and searched it themselves. RunningStatements in features/code_actions/execute_statement/running.rs now owns it: start registers an execution and returns a guard and the cancel receiver, and the guard removes the execution when dropped, also if the execution panics. cancel cancels the executions of a statement. The check for COPY with the client moves next to the other statement checks in execution.rs. --- .../code_actions/execute_statement.rs | 2 + .../execute_statement/execution.rs | 20 +++ .../code_actions/execute_statement/running.rs | 142 ++++++++++++++++++ crates/pgls_workspace/src/workspace/server.rs | 55 +------ .../src/workspace/server.tests.rs | 2 +- 5 files changed, 172 insertions(+), 49 deletions(-) create mode 100644 crates/pgls_workspace/src/features/code_actions/execute_statement/running.rs diff --git a/crates/pgls_workspace/src/features/code_actions/execute_statement.rs b/crates/pgls_workspace/src/features/code_actions/execute_statement.rs index 81062afea..2bbb1d529 100644 --- a/crates/pgls_workspace/src/features/code_actions/execute_statement.rs +++ b/crates/pgls_workspace/src/features/code_actions/execute_statement.rs @@ -10,6 +10,8 @@ use crate::workspace::{DefaultMapper, Document, StatementId}; pub(crate) mod execution; #[cfg(feature = "db")] mod notices; +#[cfg(feature = "db")] +pub(crate) mod running; #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] diff --git a/crates/pgls_workspace/src/features/code_actions/execute_statement/execution.rs b/crates/pgls_workspace/src/features/code_actions/execute_statement/execution.rs index a67ed8415..4f17a05c6 100644 --- a/crates/pgls_workspace/src/features/code_actions/execute_statement/execution.rs +++ b/crates/pgls_workspace/src/features/code_actions/execute_statement/execution.rs @@ -49,6 +49,26 @@ pub(crate) fn is_query(ast: &NodeEnum) -> bool { }) } +/// Why `ast` isn't executed, if it isn't: `COPY ... FROM STDIN` and `COPY ... TO STDOUT` exchange +/// data with the client, but the language server has no data to send, and nowhere to put the data +/// it receives. +pub(crate) fn refusal(ast: &NodeEnum) -> Option { + let NodeEnum::CopyStmt(copy) = ast else { + return None; + }; + if !copy.filename.is_empty() || copy.is_program { + return None; + } + let copy = if copy.is_from { + "COPY ... FROM STDIN" + } else { + "COPY ... TO STDOUT" + }; + Some(format!( + "`{copy}` can't be executed from the editor. Use `\\copy` in psql, or a file on the database server." + )) +} + /// Whether executing `ast` may change the schema, i.e. it does more than reading or writing rows. pub(crate) fn may_change_schema(ast: &NodeEnum) -> bool { match ast { diff --git a/crates/pgls_workspace/src/features/code_actions/execute_statement/running.rs b/crates/pgls_workspace/src/features/code_actions/execute_statement/running.rs new file mode 100644 index 000000000..b51c233c6 --- /dev/null +++ b/crates/pgls_workspace/src/features/code_actions/execute_statement/running.rs @@ -0,0 +1,142 @@ +//! The statements that are being executed, so that a request can cancel them. + +use std::sync::{Mutex, MutexGuard, PoisonError}; + +use futures::channel::oneshot; + +use super::ExecuteStatementParams; + +/// The statements that are being executed. An execution registers with [`Self::start`] for as +/// long as it runs, and [`Self::cancel`] cancels it. +#[derive(Default)] +pub(crate) struct RunningStatements { + registry: Mutex, +} + +#[derive(Default)] +struct Registry { + next_id: u64, + running: Vec, +} + +struct RunningStatement { + /// Tells this execution apart from other executions of the same statement. + id: u64, + params: ExecuteStatementParams, + cancel: oneshot::Sender<()>, +} + +impl RunningStatements { + /// Registers an execution of the statement of `params` until the returned guard is dropped. + /// The receiver fires when [`Self::cancel`] cancels the execution. + pub(crate) fn start( + &self, + params: &ExecuteStatementParams, + ) -> (RunningStatementGuard<'_>, oneshot::Receiver<()>) { + let (cancel, cancelled) = oneshot::channel(); + let mut registry = self.lock(); + let id = registry.next_id; + registry.next_id += 1; + registry.running.push(RunningStatement { + id, + params: params.clone(), + cancel, + }); + ( + RunningStatementGuard { + statements: self, + id, + }, + cancelled, + ) + } + + /// Cancels the executions of the statement of `params`, if any are running. + pub(crate) fn cancel(&self, params: &ExecuteStatementParams) { + let cancelled: Vec = self + .lock() + .running + .extract_if(.., |running| { + running.params.path == params.path + && running.params.statement_id == params.statement_id + }) + .collect(); + for running in cancelled { + // The execution may have finished in the meantime and dropped the receiver. + let _ = running.cancel.send(()); + } + } + + /// Whether no statement is being executed. + #[cfg(test)] + pub(crate) fn is_empty(&self) -> bool { + self.lock().running.is_empty() + } + + fn lock(&self) -> MutexGuard<'_, Registry> { + // No code panics while it holds the lock, and the guard must not panic in `drop`. + self.registry.lock().unwrap_or_else(PoisonError::into_inner) + } +} + +/// Removes an execution from the [`RunningStatements`] when dropped, also if the execution +/// panicked. +pub(crate) struct RunningStatementGuard<'a> { + statements: &'a RunningStatements, + id: u64, +} + +impl Drop for RunningStatementGuard<'_> { + fn drop(&mut self) { + self.statements + .lock() + .running + .retain(|running| running.id != self.id); + } +} + +#[cfg(test)] +mod tests { + use pgls_fs::PgLSPath; + + use super::RunningStatements; + use crate::features::code_actions::ExecuteStatementParams; + use crate::workspace::{ProjectKey, StatementId}; + + fn params(statement: &str) -> ExecuteStatementParams { + ExecuteStatementParams { + project_key: ProjectKey::default(), + statement_id: StatementId::new(statement), + path: PgLSPath::new("test.sql"), + } + } + + #[test] + fn cancels_the_executions_of_the_statement() { + let statements = RunningStatements::default(); + let (_first, mut first) = statements.start(¶ms("select 1")); + let (_second, mut second) = statements.start(¶ms("select 1")); + let (_other, mut other) = statements.start(¶ms("select 2")); + + statements.cancel(¶ms("select 1")); + + assert_eq!(first.try_recv(), Ok(Some(()))); + assert_eq!(second.try_recv(), Ok(Some(()))); + assert_eq!(other.try_recv(), Ok(None)); + } + + #[test] + fn the_guard_removes_the_execution() { + let statements = RunningStatements::default(); + let (first, _) = statements.start(¶ms("select 1")); + let (second, _) = statements.start(¶ms("select 1")); + + drop(first); + assert!(!statements.is_empty()); + drop(second); + assert!(statements.is_empty()); + + // Cancelling a statement that doesn't run does nothing. + statements.cancel(¶ms("select 1")); + } +} diff --git a/crates/pgls_workspace/src/workspace/server.rs b/crates/pgls_workspace/src/workspace/server.rs index bf57a72c2..70db5f404 100644 --- a/crates/pgls_workspace/src/workspace/server.rs +++ b/crates/pgls_workspace/src/workspace/server.rs @@ -8,7 +8,7 @@ use std::{ }; #[cfg(feature = "db")] -use crate::features::code_actions::execute_statement::execution; +use crate::features::code_actions::execute_statement::{execution, running::RunningStatements}; use analyser::AnalyserVisitorBuilder; #[cfg(feature = "db")] use async_helper::run_async; @@ -115,13 +115,7 @@ pub struct WorkspaceServer { /// The statements that are being executed, to cancel them #[cfg(feature = "db")] - running_statements: Mutex>, -} - -#[cfg(feature = "db")] -struct RunningStatement { - params: ExecuteStatementParams, - cancel: futures::channel::oneshot::Sender<()>, + running_statements: RunningStatements, } /// A database snapshot with the changes of the statements before a statement. @@ -158,7 +152,7 @@ impl WorkspaceServer { catalog_base: Mutex::default(), snapshot_before: Mutex::default(), connection: ConnectionManager::new(), - running_statements: Mutex::default(), + running_statements: RunningStatements::default(), } } @@ -1153,20 +1147,9 @@ impl Workspace for WorkspaceServer { }); }; - // The language server has no data to send, and nowhere to put the data it receives. - if let pgls_query::NodeEnum::CopyStmt(copy) = &ast - && copy.filename.is_empty() - && !copy.is_program - { - let copy = if copy.is_from { - "COPY ... FROM STDIN" - } else { - "COPY ... TO STDOUT" - }; + if let Some(message) = execution::refusal(&ast) { return Ok(ExecuteStatementResult { - message: format!( - "`{copy}` can't be executed from the editor. Use `\\copy` in psql, or a file on the database server." - ), + message, ..Default::default() }); } @@ -1196,14 +1179,7 @@ impl Workspace for WorkspaceServer { is_query: execution::is_query(&ast), timeout: db.statement_timeout, }; - let (cancel, cancelled) = futures::channel::oneshot::channel(); - self.running_statements - .lock() - .unwrap() - .push(RunningStatement { - params: params.clone(), - cancel, - }); + let (_running, cancelled) = self.running_statements.start(¶ms); let result = self.connection.with_pools(&db, |pools| { let statement_pools = pools.clone(); @@ -1219,12 +1195,6 @@ impl Workspace for WorkspaceServer { Ok(result) }); - // The execution is over once it dropped the receiver. - self.running_statements - .lock() - .unwrap() - .retain(|running| !running.cancel.is_canceled()); - let Some(result) = result else { return Ok(no_connection); }; @@ -1258,18 +1228,7 @@ impl Workspace for WorkspaceServer { &self, params: ExecuteStatementParams, ) -> Result<(), WorkspaceError> { - let cancelled = self - .running_statements - .lock() - .unwrap() - .extract_if(.., |running| { - running.params.path == params.path - && running.params.statement_id == params.statement_id - }) - .collect::>(); - for running in cancelled { - let _ = running.cancel.send(()); - } + self.running_statements.cancel(¶ms); Ok(()) } diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 1661b5d82..057154793 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2460,7 +2460,7 @@ async fn execute_statement_can_be_cancelled(test_db: PgPool) { "{:?}", start.elapsed() ); - assert!(workspace.running_statements.lock().unwrap().is_empty()); + assert!(workspace.running_statements.is_empty()); } #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] From b55d5f52665ad668b578574adbad07e611dd4510 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Mon, 26 Oct 2026 15:02:21 +0000 Subject: [PATCH 7/7] test(workspace): don't time the first execution in the truncated query test Since the code actions run the linter, the first execution of a workspace also connects and loads the schema, which can take more than the 3 s that the test allows for the query. --- crates/pgls_workspace/src/workspace/server.tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/pgls_workspace/src/workspace/server.tests.rs b/crates/pgls_workspace/src/workspace/server.tests.rs index 057154793..8d531a06e 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -2468,6 +2468,9 @@ async fn execute_statement_stops_a_truncated_query_early(test_db: PgPool) { let (workspace, project_key) = get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + // The first execution connects and loads the schema for the code actions, which isn't timed. + execute_sql(&workspace, project_key, "select 1 as one;"); + // Without the cursor, Postgres would produce and send all 50 million rows. let start = Instant::now(); let (_, result) = execute_sql(