database transaction support in typescript (#1847)

Adds support for transactions in runtimes/core and runtimes/js

We should probably merge
encoredev/rust-postgres#1 and update Cargo.toml
again. (suggested that change upstream as well
sfackler/rust-postgres#1230)

Author: fredr

Cargo.lock

@@ -9,8 +9,8 @@ insta.opt-level = 3
 lto = true
 
 [patch.crates-io]
-tokio-postgres = { git = "https://github.com/encoredev/rust-postgres", branch = "proxy" }
-postgres-protocol = { git = "https://github.com/encoredev/rust-postgres", branch = "proxy" }
+tokio-postgres = { git = "https://github.com/encoredev/rust-postgres", branch = "encore-patches" }
+postgres-protocol = { git = "https://github.com/encoredev/rust-postgres", branch = "encore-patches" }
 swc_ecma_parser = { git = "https://github.com/encoredev/swc", branch = "node-resolve-exports" }
 swc_ecma_ast = { git = "https://github.com/encoredev/swc", branch = "node-resolve-exports" }
 swc_ecma_transforms_base = { git = "https://github.com/encoredev/swc", branch = "node-resolve-exports" }

Cargo.toml

@@ -171,6 +171,31 @@ await db.rawExec(
 );
 ```
 
+### Transactions
+
+Transactions allow you to group multiple database operations into a single unit of work. If any operation within the transaction fails, the entire transaction is rolled back, ensuring data consistency.
+
+The transaction type implements `AsyncDisposable`, which automatically rolls back the transaction if it is not explicitly committed or rolled back. This ensures that no open transactions are left accidentally.
+
+For example:
+
+```ts
+await using tx = await db.begin();
+
+await db.exec`
+  INSERT INTO todo_item (title, done)
+  VALUES (${title1}, false)
+`;
+
+await db.exec`
+  INSERT INTO todo_item (title, done)
+  VALUES (${title2}, false)
+`;
+
+await tx.commit();
+```
+
+
 ## Connecting to databases
 
 It's often useful to be able to connect to the database from outside the backend application. For example for scripts, ad-hoc querying, or dumping data for analysis.

docs/ts/primitives/databases.md

@@ -12,6 +12,8 @@ use crate::sqldb::val::RowValue;
 use crate::trace::{protocol, Tracer};
 use crate::{model, sqldb};
 
+use super::transaction::Transaction;
+
 type Mgr = PostgresConnectionManager<postgres_native_tls::MakeTlsConnector>;
 
 pub struct Pool {
@@ -99,6 +101,14 @@ impl Pool {
             tracer: self.tracer.clone(),
         })
     }
+
+    pub async fn begin(&self, source: Option<&model::Request>) -> Result<Transaction, Error> {
+        let conn = self.pool.get_owned().await.map_err(|e| match e {
+            RunError::User(err) => err,
+            RunError::TimedOut => tokio_postgres::Error::__private_api_timeout(),
+        })?;
+        Transaction::begin(conn, self.tracer.clone(), source).await
+    }
 }
 
 pub struct Cursor {
@@ -135,37 +145,14 @@ impl Row {
     }
 }
 
-type PooledConn =
+pub(crate) type PooledConn =
     PooledConnection<'static, PostgresConnectionManager<postgres_native_tls::MakeTlsConnector>>;
 
 pub struct Connection {
     conn: tokio::sync::RwLock<Option<PooledConn>>,
     tracer: QueryTracer,
 }
 
-#[derive(Debug)]
-pub enum Error {
-    DB(tokio_postgres::Error),
-    Closed,
-    ConnectTimeout,
-}
-
-impl std::fmt::Display for Error {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Error::DB(err) => <tokio_postgres::Error as std::fmt::Display>::fmt(err, f),
-            Error::Closed => f.write_str("connection_closed"),
-            Error::ConnectTimeout => f.write_str("timeout establishing connection"),
-        }
-    }
-}
-
-impl From<tokio_postgres::Error> for Error {
-    fn from(err: tokio_postgres::Error) -> Self {
-        Error::DB(err)
-    }
-}
-
 impl Connection {
     pub async fn close(&self) {
         let mut guard = self.conn.write().await;
@@ -197,11 +184,34 @@ impl Connection {
     }
 }
 
+#[derive(Debug)]
+pub enum Error {
+    DB(tokio_postgres::Error),
+    Closed,
+    ConnectTimeout,
+}
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Error::DB(err) => <tokio_postgres::Error as std::fmt::Display>::fmt(err, f),
+            Error::Closed => f.write_str("connection_closed"),
+            Error::ConnectTimeout => f.write_str("timeout establishing connection"),
+        }
+    }
+}
+
+impl From<tokio_postgres::Error> for Error {
+    fn from(err: tokio_postgres::Error) -> Self {
+        Error::DB(err)
+    }
+}
+
 #[derive(Debug, Clone)]
-struct QueryTracer(Tracer);
+pub(crate) struct QueryTracer(Tracer);
 
 impl QueryTracer {
-    async fn trace<F, Fut>(
+    pub(crate) async fn trace<F, Fut>(
         &self,
         source: Option<&model::Request>,
         query: &str,
@@ -235,4 +245,36 @@ impl QueryTracer {
             stream: Box::pin(stream),
         })
     }
+
+    pub(crate) async fn trace_batch_execute<F, Fut>(
+        &self,
+        source: Option<&model::Request>,
+        query: &str,
+        exec: F,
+    ) -> Result<(), Error>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = Result<(), Error>>,
+    {
+        let start_id = if let Some(source) = source {
+            let id = self
+                .0
+                .db_query_start(protocol::DBQueryStartData { source, query });
+            Some(id)
+        } else {
+            None
+        };
+
+        let result = exec().await;
+
+        if let Some(start_id) = start_id {
+            self.0.db_query_end(protocol::DBQueryEndData {
+                start_id,
+                source: source.unwrap(),
+                error: result.as_ref().err(),
+            });
+        }
+
+        result
+    }
 }

runtimes/core/src/sqldb/client.rs

@@ -1,7 +1,9 @@
 mod client;
 mod manager;
+mod transaction;
 mod val;
 
 pub use client::{Connection, Cursor, Pool, Row};
 pub use manager::{Database, DatabaseImpl, Manager, ManagerConfig};
+pub use transaction::Transaction;
 pub use val::RowValue;

runtimes/core/src/sqldb/mod.rs

@@ -0,0 +1,122 @@
+use tokio_postgres::types::BorrowToSql;
+
+use crate::model;
+
+use super::{
+    client::{Error, PooledConn, QueryTracer},
+    Cursor,
+};
+
+// Heavily inspired by rust-postgres, but where the transaction doesnt have a lifetime, so it can
+// be shared via napi-rs.
+//
+// https://github.com/sfackler/rust-postgres/blob/720ffe83216714bf9716a03122c547a2e8e9bfd9/tokio-postgres/src/transaction.rs
+
+pub struct Transaction {
+    conn: PooledConn,
+    tracer: QueryTracer,
+    done: bool,
+}
+
+impl Transaction {
+    pub(crate) async fn begin(
+        conn: PooledConn,
+        tracer: QueryTracer,
+        source: Option<&model::Request>,
+    ) -> Result<Self, Error> {
+        struct RollbackIfNotDone<'me> {
+            client: &'me tokio_postgres::Client,
+            done: bool,
+        }
+
+        impl Drop for RollbackIfNotDone<'_> {
+            fn drop(&mut self) {
+                if self.done {
+                    return;
+                }
+
+                self.client.__private_api_rollback(None);
+            }
+        }
+
+        // This is done, as `Future` created by this method can be dropped after
+        // `RequestMessages` is synchronously send to the `Connection` by
+        // `batch_execute()`, but before `Responses` is asynchronously polled to
+        // completion. In that case `Transaction` won't be created and thus
+        // won't be rolled back.
+        {
+            let mut cleaner = RollbackIfNotDone {
+                client: &conn,
+                done: false,
+            };
+
+            tracer
+                .trace_batch_execute(source, "BEGIN", || async {
+                    conn.batch_execute("BEGIN").await.map_err(Error::from)
+                })
+                .await?;
+
+            cleaner.done = true;
+        }
+
+        Ok(Transaction {
+            conn,
+            tracer,
+            done: false,
+        })
+    }
+
+    pub async fn commit(mut self, source: Option<&model::Request>) -> Result<(), Error> {
+        self.done = true;
+        self.batch_execute("COMMIT", source).await
+    }
+
+    pub async fn rollback(mut self, source: Option<&model::Request>) -> Result<(), Error> {
+        self.done = true;
+        self.batch_execute("ROLLBACK", source).await
+    }
+
+    async fn batch_execute(
+        &self,
+        query: &str,
+        source: Option<&model::Request>,
+    ) -> Result<(), Error> {
+        self.tracer
+            .trace_batch_execute(source, query, || async {
+                self.conn.batch_execute(query).await.map_err(Error::from)
+            })
+            .await
+    }
+
+    pub async fn query_raw<P, I>(
+        &self,
+        query: &str,
+        params: I,
+        source: Option<&model::Request>,
+    ) -> Result<Cursor, Error>
+    where
+        P: BorrowToSql,
+        I: IntoIterator<Item = P>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        self.tracer
+            .trace(source, query, || async {
+                self.conn
+                    .query_raw(query, params)
+                    .await
+                    .map_err(Error::from)
+            })
+            .await
+    }
+}
+
+impl Drop for Transaction {
+    fn drop(&mut self) {
+        if self.done {
+            return;
+        }
+
+        log::warn!("transaction not completed, forcing rollback");
+        self.conn.__private_api_rollback(None);
+    }
+}