feat(workspace): guard statement execution

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,
         }
     }
 }

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,

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<dyn Workspace>, 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};

crates/pgls_lsp/tests/server.rs

@@ -1631,6 +1631,24 @@ async fn execute_statement_at(
     doc_url: &Uri,
     position: Position,
 ) -> Result<Option<Value>> {
+    let params = execute_statement_command(server, doc_url, position).await?;
+    let result = server
+        .request::<ExecuteCommandParams, Option<Value>>(
+            "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<ExecuteCommandParams> {
     let code_actions_response = server
         .request::<CodeActionParams, CodeActionResponse>(
             "textDocument/codeAction",
@@ -1669,19 +1687,11 @@ async fn execute_statement_at(
         })
         .expect("Did not find executeStatement command!");
 
-    let result = server
-        .request::<ExecuteCommandParams, Option<Value>>(
-            "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();

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};

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<u64>,
     /// Why Postgres rejected the statement, if it did.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub error: Option<ExecuteStatementError>,
+    /// 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<String>,
 }
 
 /// An error that Postgres raised while executing a statement.