diff --git a/crates/pgls_configuration/src/database.rs b/crates/pgls_configuration/src/database.rs index de87a2296..4a62911cc 100644 --- a/crates/pgls_configuration/src/database.rs +++ b/crates/pgls_configuration/src/database.rs @@ -53,6 +53,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)))] @@ -76,6 +81,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() } @@ -98,6 +104,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() } @@ -115,6 +122,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 aa285af14..9d1e8d2e9 100644 --- a/crates/pgls_configuration/src/lib.rs +++ b/crates/pgls_configuration/src/lib.rs @@ -218,6 +218,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), }), inlay_hints: None, diff --git a/crates/pgls_lsp/src/handlers/code_actions.rs b/crates/pgls_lsp/src/handlers/code_actions.rs index 981b389dc..6d23888f9 100644 --- a/crates/pgls_lsp/src/handlers/code_actions.rs +++ b/crates/pgls_lsp/src/handlers/code_actions.rs @@ -13,11 +13,13 @@ use tower_lsp_server::ls_types::{ }; use pgls_text_size::{TextRange, TextSize}; +use pgls_workspace::Workspace; use pgls_workspace::features::code_actions::{ CodeActionKind, CodeActionsParams, CommandActionCategory, ExecuteStatementParams, FixFileMode, FixFileParams, }; use pgls_workspace::workspace::GetFileContentParams; +use std::sync::Arc; #[tracing::instrument(level = "debug", skip(session), err)] pub fn get_actions( @@ -310,13 +312,19 @@ pub async fn execute_command( let path = session.file_path(&doc_url)?; let project_key = session.project_key(&path); - let result = session - .workspace - .execute_statement(ExecuteStatementParams { - project_key, - statement_id, - path, - })?; + let params = ExecuteStatementParams { + project_key, + 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 @@ -366,6 +374,27 @@ pub async fn execute_command( } } +/// 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}"); + } + }); + } + } +} + #[cfg(test)] mod tests { use super::{FIX_ALL_KIND, changed_range, covers}; diff --git a/crates/pgls_lsp/tests/server.rs b/crates/pgls_lsp/tests/server.rs index 2cfd60141..ca4a8e47f 100644 --- a/crates/pgls_lsp/tests/server.rs +++ b/crates/pgls_lsp/tests/server.rs @@ -1631,6 +1631,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", @@ -1669,19 +1687,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")] @@ -2036,6 +2046,76 @@ async fn test_expand_star_code_action(test_db: PgPool) -> Result<()> { 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 b5162cb1a..a1a120868 100644 --- a/crates/pgls_workspace/src/features/code_actions.rs +++ b/crates/pgls_workspace/src/features/code_actions.rs @@ -14,7 +14,8 @@ pub(crate) mod lint; pub(crate) mod schema_cache; pub use execute_statement::{ - EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementError, ExecuteStatementParams, + EXECUTE_STATEMENT_MAX_RESULT_SIZE, EXECUTE_STATEMENT_MAX_ROWS, + EXECUTE_STATEMENT_MAX_VALUE_SIZE, ExecuteStatementError, ExecuteStatementParams, ExecuteStatementResult, }; pub use lint::{FixFileMode, FixFileParams, FixFileResult}; 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 6a6974a57..2bbb1d529 100644 --- a/crates/pgls_workspace/src/features/code_actions/execute_statement.rs +++ b/crates/pgls_workspace/src/features/code_actions/execute_statement.rs @@ -8,8 +8,12 @@ use crate::workspace::{DefaultMapper, Document, StatementId}; #[cfg(feature = "db")] pub(crate) mod execution; +#[cfg(feature = "db")] +mod notices; +#[cfg(feature = "db")] +pub(crate) mod running; -#[derive(Debug, serde::Serialize, serde::Deserialize)] +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] pub struct ExecuteStatementParams { pub project_key: crate::workspace::ProjectKey, @@ -21,6 +25,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 @@ -41,12 +53,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/features/code_actions/execute_statement/execution.rs b/crates/pgls_workspace/src/features/code_actions/execute_statement/execution.rs index 2e5460385..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 @@ -1,12 +1,86 @@ -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 super::{EXECUTE_STATEMENT_MAX_ROWS, ExecuteStatementError, ExecuteStatementResult}; +use super::{ + EXECUTE_STATEMENT_MAX_RESULT_SIZE, EXECUTE_STATEMENT_MAX_ROWS, + EXECUTE_STATEMENT_MAX_VALUE_SIZE, ExecuteStatementError, ExecuteStatementResult, notices, +}; use crate::diagnostics::is_connection_failure; +use crate::workspace::server::DatabasePools; + +/// 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(crate) 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(crate) is_query: bool, + /// How long the statement may run, or `None` to keep the session's `statement_timeout`. + pub(crate) timeout: Option, +} + +/// Whether a cursor can run `ast`: `declare` takes a `select` or `values` without `into` and +/// without a data-modifying `with`. +pub(crate) 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, + }) + }) +} + +/// 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 { + 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. /// @@ -16,43 +90,94 @@ use crate::diagnostics::is_connection_failure; /// /// 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(crate) 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 + }) +} + +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?; - match reset_session(&mut conn).await { - Ok(rolled_back) => result.map(|mut result| { - if rolled_back { - if !result.message.ends_with('.') { - result.message.push('.'); + // 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 { @@ -71,26 +196,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); } } } @@ -110,7 +283,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 { @@ -119,6 +292,7 @@ async fn collect_result( columns, rows, truncated, + truncated_values, ..Default::default() }, None => ExecuteStatementResult { @@ -129,17 +303,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, }; @@ -163,7 +341,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()?; @@ -171,7 +350,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/features/code_actions/execute_statement/notices.rs b/crates/pgls_workspace/src/features/code_actions/execute_statement/notices.rs new file mode 100644 index 000000000..83ca3fc7d --- /dev/null +++ b/crates/pgls_workspace/src/features/code_actions/execute_statement/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/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/settings.rs b/crates/pgls_workspace/src/settings.rs index 36fa0fd1b..383a1d287 100644 --- a/crates/pgls_workspace/src/settings.rs +++ b/crates/pgls_workspace/src/settings.rs @@ -759,6 +759,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, @@ -770,6 +771,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, @@ -826,6 +830,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, @@ -850,6 +855,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: false, invalid_statement_execution_patterns: Vec::new(), connection_string_error: None, @@ -957,6 +963,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 ad204d74a..fb38a6cfc 100644 --- a/crates/pgls_workspace/src/workspace.rs +++ b/crates/pgls_workspace/src/workspace.rs @@ -261,6 +261,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 45f4dd825..21530530e 100644 --- a/crates/pgls_workspace/src/workspace/client.rs +++ b/crates/pgls_workspace/src/workspace/client.rs @@ -113,6 +113,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 d4c4df773..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; @@ -83,6 +83,8 @@ mod async_helper; mod connection_key; #[cfg(feature = "db")] mod connection_manager; +#[cfg(feature = "db")] +pub(crate) use connection_manager::DatabasePools; pub(crate) mod document; mod migration; mod pg_query; @@ -110,6 +112,10 @@ pub struct WorkspaceServer { #[cfg(feature = "db")] connection: ConnectionManager, + + /// The statements that are being executed, to cancel them + #[cfg(feature = "db")] + running_statements: RunningStatements, } /// A database snapshot with the changes of the statements before a statement. @@ -146,6 +152,7 @@ impl WorkspaceServer { catalog_base: Mutex::default(), snapshot_before: Mutex::default(), connection: ConnectionManager::new(), + running_statements: RunningStatements::default(), } } @@ -1120,7 +1127,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 { @@ -1131,35 +1140,62 @@ 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() }); }; + if let Some(message) = execution::refusal(&ast) { + return Ok(ExecuteStatementResult { + message, + ..Default::default() + }); + } + let settings = self.workspaces(); let no_connection = ExecuteStatementResult { message: "No database connection available.".into(), ..Default::default() }; - let Some(settings) = settings.settings_for(params.project_key, params.path.as_path()) + let Some(db) = settings + .settings_for(params.project_key, params.path.as_path()) + .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 { execution::execute_statement(&pool, &content).await })? - .map_err(WorkspaceError::from) - }) else { + let options = execution::ExecuteOptions { + is_query: execution::is_query(&ast), + timeout: db.statement_timeout, + }; + let (_running, cancelled) = self.running_statements.start(¶ms); + + let result = self.connection.with_pools(&db, |pools| { + let statement_pools = pools.clone(); + let result = run_async(async move { + 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() && execution::may_change_schema(&ast) { + self.snapshot.clear_connection(&pools.pool); + } + Ok(result) + }); + + let Some(result) = result else { return Ok(no_connection); }; @@ -1187,6 +1223,24 @@ impl Workspace for WorkspaceServer { }) } + #[cfg(feature = "db")] + fn cancel_statement_execution( + &self, + params: ExecuteStatementParams, + ) -> Result<(), WorkspaceError> { + self.running_statements.cancel(¶ms); + 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 d635618a7..8d531a06e 100644 --- a/crates/pgls_workspace/src/workspace/server.tests.rs +++ b/crates/pgls_workspace/src/workspace/server.tests.rs @@ -30,7 +30,8 @@ use crate::{ }, 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, FixFileMode, FixFileParams, TextEdit, }, features::format::PullFileFormattingParams, @@ -1988,6 +1989,18 @@ fn is_typecheck(diagnostic: &pgls_diagnostics::serde::Diagnostic) -> bool { fn get_statement_execution_workspace( test_db: &PgPool, allow_statement_executions_against: &[String], +) -> (WorkspaceServer, ProjectKey) { + 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, ProjectKey) { let mut conf = PartialConfiguration::init(); conf.merge_with(PartialConfiguration { @@ -2003,6 +2016,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() @@ -2157,8 +2171,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() } ); } @@ -2369,6 +2382,262 @@ 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, project_key) = get_statement_execution_workspace_with_timeout( + &test_db, + &allow_test_host(&test_db), + Some(1), + ); + + let start = Instant::now(); + let (_, result) = execute_sql(&workspace, project_key, "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + let path = PgLSPath::new("execute.sql"); + workspace + .open_file(OpenFileParams { + project_key, + path: path.clone(), + content: "select pg_sleep(30);".into(), + version: 1, + }) + .unwrap(); + let statement_id = workspace + .pull_code_actions(CodeActionsParams { + project_key, + 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 { + project_key, + 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.is_empty()); +} + +#[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")] +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( + &workspace, + project_key, + "select generate_series(1, 50000000) as n", + ); + let (_, next) = execute_sql(&workspace, project_key, "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + execute_sql(&workspace, project_key, "create table inserted (id int);"); + let (_, result) = execute_sql( + &workspace, + project_key, + "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, project_key, "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + // The cut falls into the middle of an `ä`. + let (_, result) = execute_sql( + &workspace, + project_key, + "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql( + &workspace, + project_key, + "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + execute_sql( + &workspace, + project_key, + "create table copy_target (id int);", + ); + + let (_, result) = execute_sql(&workspace, project_key, "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, project_key, "copy copy_target from stdin;"); + assert!( + result + .message + .starts_with("`COPY ... FROM STDIN` can't be executed") + ); + + let (_, result) = execute_sql(&workspace, project_key, "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + + let (_, result) = execute_sql( + &workspace, + project_key, + "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, project_key) = + get_statement_execution_workspace(&test_db, &allow_test_host(&test_db)); + let other = PgLSPath::new("other.sql"); + workspace + .open_file(OpenFileParams { + project_key, + path: other.clone(), + content: "select id from created_by_execute;".into(), + version: 1, + }) + .unwrap(); + let typecheck_findings = || { + workspace + .pull_file_diagnostics(crate::workspace::PullFileDiagnosticsParams { + project_key, + 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, + project_key, + "create table created_by_execute (id int);", + ); + + assert_eq!(typecheck_findings(), 0); +} + /// A workspace that reads the section markers of a migration tool. fn section_workspace(section_markers: SectionMarkers) -> (WorkspaceServer, ProjectKey) { let mut conf = PartialConfiguration::init(); diff --git a/crates/pgls_workspace/src/workspace/server/snapshot_manager.rs b/crates/pgls_workspace/src/workspace/server/snapshot_manager.rs index 37727cfff..93d8d20a3 100644 --- a/crates/pgls_workspace/src/workspace/server/snapshot_manager.rs +++ b/crates/pgls_workspace/src/workspace/server/snapshot_manager.rs @@ -102,6 +102,14 @@ impl SnapshotManager { Ok(snapshot) } + /// Clear the snapshot for a specific connection. + #[cfg(feature = "db")] + pub fn clear_connection(&self, pool: &PgPool) { + let key: ConnectionKey = pool.into(); + let mut schemas = self.db_schemas.write().unwrap(); + schemas.remove(&key); + } + /// Clear all connection-based snapshots. #[cfg(feature = "db")] pub fn clear_all_connections(&self) { diff --git a/docs/guides/configure_database.md b/docs/guides/configure_database.md index e46693f00..b26d8e591 100644 --- a/docs/guides/configure_database.md +++ b/docs/guides/configure_database.md @@ -140,9 +140,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 ea7e8ec5f..e7e2ce29b 100644 --- a/docs/schema.json +++ b/docs/schema.json @@ -450,6 +450,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 9c0f9fa32..86483a50e 100644 --- a/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts +++ b/packages/@postgres-language-server/backend-jsonrpc/src/workspace.ts @@ -561,6 +561,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 9c0f9fa32..86483a50e 100644 --- a/packages/@postgrestools/backend-jsonrpc/src/workspace.ts +++ b/packages/@postgrestools/backend-jsonrpc/src/workspace.ts @@ -561,6 +561,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. */