new on_reset.sql

Author: lovasoa

CHANGELOG.md

@@ -2,7 +2,7 @@
 
 ## 0.32.0 (unreleased)
 
-- Rollback any open transactions when returning a connection to the pool.
+- Rollback any open transactions when an error occurs in a SQL file.
   - Previously, if an error occurred in the middle of a transaction, the transaction would be left open, and the connection would be returned to the pool. The next request could get a connection with an open half-completed transaction, which could lead to hard to debug issues.
   - This allows safely using features that require a transaction, like
     - ```sql
@@ -21,6 +21,8 @@
     ```
 - Fix `error returned from database: 1295 (HY000): This command is not supported in the prepared statement protocol yet` when trying to use transactions with MySQL. `START TRANSACTION` now works as expected in MySQL.
 - Fix a bug where a multi-select dropdown would unexpectedly open when the form was reset.
+- Add a new optional `sqlpage/on_reset.sql` file that can be used to execute some SQL code after the end of each page execution.
+   - Useful to reset a connection to the database after each request.
 
 ## 0.31.0 (2024-11-24)
 

configuration.md

@@ -106,7 +106,9 @@ For instance, if you want to create a custom `my_component` component, that disp
 
 [See the full custom component documentation](https://sql-page.com/custom_components.sql).
 
-## Connection initialization scripts
+## Connection management
+
+### Connection initialization scripts
 
 SQLPage allows you to run a SQL script when a new database connection is opened,
 by simply creating a `sqlpage/on_connect.sql` file.
@@ -135,6 +137,52 @@ CREATE TEMPORARY TABLE my_temporary_table(
 );
 ```
 
+### Connection cleanup scripts: `on_reset.sql`
+
+SQLPage allows you to run a SQL script after a request has been processed,
+by simply creating a `sqlpage/on_reset.sql` file.
+
+This can be useful to clean up temporary tables,
+rollback transactions that were left open,
+or other resources that were created during the request.
+
+You can also use this script to close database connections that are 
+in an undesirable state, such as being in a transaction that was left open.
+To close a connection, write a select statement that returns a single row
+with a single boolean column named `is_healthy`, and set it to false.
+
+#### Rollback transactions
+
+You can automatically rollback any open transactions
+when a connection is returned to the pool,
+so that a new request is never executed in the context of an open transaction from a previous request.
+
+For this to work, you need to create a `sqlpage/on_reset.sql` containing the following line:
+
+```sql
+ROLLBACK;
+```
+
+#### Cleaning up all connection state
+
+Some databases allow you to clean up all the state associatPed with a connection.
+
+##### PostgreSQL
+
+By creating a `sqlpage/on_disconnect.sql` file containing a [`DISCARD ALL`](https://www.postgresql.org/docs/current/sql-discard.html) statement.
+
+```sql
+DISCARD ALL;
+```
+
+##### SQL Server
+
+By creating a `sqlpage/on_disconnect.sql` file containing a call to the [`sp_reset_connection`](https://learn.microsoft.com/en-us/sql/relational-databases/system-stored-procedures/system-stored-procedures-transact-sql?view=sql-server-ver16#api-system-stored-procedures) stored procedure.
+
+```sql
+EXEC sp_reset_connection;
+```
+
 ## Migrations
 
 SQLPage allows you to run SQL scripts when the database schema changes, by creating a `sqlpage/migrations` directory.

src/lib.rs

@@ -94,6 +94,7 @@ use webserver::Database;
 pub const TEMPLATES_DIR: &str = "sqlpage/templates/";
 pub const MIGRATIONS_DIR: &str = "migrations";
 pub const ON_CONNECT_FILE: &str = "on_connect.sql";
+pub const ON_RESET_FILE: &str = "on_reset.sql";
 
 pub struct AppState {
     pub db: Database,

src/webserver/database/connect.rs

@@ -1,7 +1,7 @@
 use std::{mem::take, time::Duration};
 
 use super::Database;
-use crate::{app_config::AppConfig, ON_CONNECT_FILE};
+use crate::{app_config::AppConfig, ON_CONNECT_FILE, ON_RESET_FILE};
 use anyhow::Context;
 use futures_util::future::BoxFuture;
 use sqlx::{
@@ -93,37 +93,50 @@ impl Database {
             )
             .acquire_timeout(Duration::from_secs_f64(
                 config.database_connection_acquire_timeout_seconds,
-            ))
-            .after_release(on_return_to_pool);
+            ));
+        pool_options = add_on_return_to_pool(config, pool_options);
         pool_options = add_on_connection_handler(config, pool_options);
         pool_options
     }
 }
 
+fn add_on_return_to_pool(config: &AppConfig, pool_options: PoolOptions<Any>) -> PoolOptions<Any> {
+    let on_disconnect_file = config.configuration_directory.join(ON_RESET_FILE);
+    if !on_disconnect_file.exists() {
+        log::debug!("Not creating a custom SQL connection cleanup handler because {on_disconnect_file:?} does not exist");
+        return pool_options;
+    }
+    log::info!("Creating a custom SQL connection cleanup handler from {on_disconnect_file:?}");
+    let sql = match std::fs::read_to_string(&on_disconnect_file) {
+        Ok(sql) => std::sync::Arc::new(sql),
+        Err(e) => {
+            log::error!("Unable to read the file {on_disconnect_file:?}: {e}");
+            return pool_options;
+        }
+    };
+    log::trace!("The custom SQL connection cleanup handler is:\n{sql}");
+    let sql = sql.clone();
+    pool_options
+        .after_release(move |conn, meta| on_return_to_pool(conn, meta, std::sync::Arc::clone(&sql)))
+}
+
 fn on_return_to_pool(
     conn: &mut sqlx::AnyConnection,
     meta: sqlx::pool::PoolConnectionMetadata,
+    sql: std::sync::Arc<String>,
 ) -> BoxFuture<'_, Result<bool, sqlx::Error>> {
+    use sqlx::Row;
     Box::pin(async move {
-        match conn.execute("ROLLBACK").await {
-            Ok(query_result) => {
-                if query_result.rows_affected() > 0 {
-                    log::warn!(
-                        "Rolled back a transaction, because it was left open after a page was rendered, and it affected {} rows",
-                        query_result.rows_affected()
-                    );
-                } else {
-                    log::trace!(
-                        "Rolled back a transaction before returning a connection to the pool"
-                    );
-                }
-            }
-            Err(e) => log::trace!(
-                "Failed to rollback before returning a connection to the pool. There was probably no transaction left open: {e:?}"
-            ),
+        log::trace!("Running the custom SQL connection cleanup handler. {meta:?}");
+        let query_result = conn.fetch_optional(sql.as_str()).await?;
+        if let Some(query_result) = query_result {
+            let is_healthy = query_result.try_get::<bool, _>(0);
+            log::debug!("Is the connection healthy? {is_healthy:?}");
+            is_healthy
+        } else {
+            log::debug!("No result from the custom SQL connection cleanup handler");
+            Ok(true)
         }
-        log::trace!("Releasing connection: {meta:#?}");
-        Ok(true)
     })
 }
 

src/webserver/database/execute_queries.rs

@@ -21,7 +21,9 @@ use super::syntax_tree::{extract_req_param, StmtParam};
 use super::{error_highlighting::display_db_error, Database, DbItem};
 use sqlx::any::{AnyArguments, AnyQueryResult, AnyRow, AnyStatement, AnyTypeInfo};
 use sqlx::pool::PoolConnection;
-use sqlx::{Any, Arguments, Column, Either, Executor, Row as _, Statement, ValueRef};
+use sqlx::{
+    Any, AnyConnection, Arguments, Column, Either, Executor, Row as _, Statement, ValueRef,
+};
 
 pub type DbConn = Option<PoolConnection<sqlx::Any>>;
 
@@ -58,18 +60,24 @@ pub fn stream_query_results_with_conn<'a>(
                     let connection = take_connection(&request.app_state.db, db_connection).await?;
                     log::trace!("Executing query {:?}", query.sql);
                     let mut stream = connection.fetch_many(query);
+                    let mut error = None;
                     while let Some(elem) = stream.next().await {
-                        let is_err = elem.is_err();
                         let mut query_result = parse_single_sql_result(source_file, &stmt.query, elem);
+                        if let DbItem::Error(e) = query_result {
+                            error = Some(e);
+                            break;
+                        }
                         apply_json_columns(&mut query_result, &stmt.json_columns);
                         apply_delayed_functions(request, &stmt.delayed_functions, &mut query_result).await?;
-                        for i in parse_dynamic_rows(query_result) {
-                            yield i;
-                        }
-                        if is_err {
-                            break;
+                        for db_item in parse_dynamic_rows(query_result) {
+                            yield db_item;
                         }
                     }
+                    drop(stream);
+                    if let Some(error) = error {
+                        try_rollback_transaction(connection).await;
+                        yield DbItem::Error(error);
+                    }
                 },
                 ParsedStatement::SetVariable { variable, value} => {
                     execute_set_variable_query(db_connection, request, variable, value, source_file).await
@@ -131,6 +139,16 @@ async fn exec_static_simple_select(
     Ok(serde_json::Value::Object(map))
 }
 
+async fn try_rollback_transaction(db_connection: &mut AnyConnection) {
+    log::debug!("Attempting to rollback transaction");
+    match db_connection.execute("ROLLBACK").await {
+        Ok(_) => log::debug!("Rolled back transaction"),
+        Err(e) => {
+            log::debug!("There was probably no transaction in progress when this happened: {e:?}");
+        }
+    }
+}
+
 /// Extracts the value of a parameter from the request.
 /// Returns `Ok(None)` when NULL should be used as the parameter value.
 async fn extract_req_param_as_json(
@@ -177,6 +195,7 @@ async fn execute_set_variable_query<'a>(
         Ok(Some(row)) => row_to_string(&row),
         Ok(None) => None,
         Err(e) => {
+            try_rollback_transaction(connection).await;
             let err = display_db_error(source_file, &statement.query, e);
             return Err(err);
         }