feat(batch): add batch statements support (#30)

Signed-off-by: Daniel Boll <danielboll.academico@gmail.com>

Author: Daniel-Boll

examples/auth.mts

@@ -12,7 +12,7 @@ const cluster = new Cluster({
 
 const session = await cluster.connect();
 
-await session.execute("CREATE KEYSPACE IF NOT EXISTS examples_ks WITH REPLICATION = {'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}");
-await session.execute("DROP TABLE IF EXISTS examples_ks.auth");
+await session.execute("CREATE KEYSPACE IF NOT EXISTS auth WITH REPLICATION = {'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}");
+await session.execute("DROP TABLE IF EXISTS auth.auth");
 
 console.log("Ok.");

examples/batch-statements.mts

@@ -0,0 +1,22 @@
+import { Cluster, BatchStatement, Query, Uuid } from "../index.js"
+
+const nodes = process.env.CLUSTER_NODES?.split(",") ?? ["127.0.0.1:9042"];
+
+const cluster = new Cluster({ nodes });
+const session = await cluster.connect();
+
+const batch = new BatchStatement();
+
+await session.execute("CREATE KEYSPACE IF NOT EXISTS batch_statements WITH REPLICATION = { 'class' : 'SimpleStrategy', 'replication_factor' : 1 }");
+await session.useKeyspace("batch_statements");
+await session.execute("CREATE TABLE IF NOT EXISTS users (id UUID PRIMARY KEY, name TEXT)");
+
+const simpleStatement = new Query("INSERT INTO users (id, name) VALUES (?, ?)");
+const preparedStatement = await session.prepare("INSERT INTO users (id, name) VALUES (?, ?)");
+
+batch.appendStatement(simpleStatement);
+batch.appendStatement(preparedStatement);
+
+await session.batch(batch, [[Uuid.randomV4(), "Alice"], [Uuid.randomV4(), "Bob"]]);
+
+console.log(await session.execute("SELECT * FROM users"));

examples/lwt.mts

@@ -6,10 +6,10 @@ const cluster = new Cluster({ nodes });
 
 const session = await cluster.connect();
 
-await session.execute("CREATE KEYSPACE IF NOT EXISTS examples_ks WITH REPLICATION = {'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}");
-await session.execute("CREATE TABLE IF NOT EXISTS examples_ks.tab (a int PRIMARY KEY)");
+await session.execute("CREATE KEYSPACE IF NOT EXISTS lwt WITH REPLICATION = {'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}");
+await session.execute("CREATE TABLE IF NOT EXISTS lwt.tab (a int PRIMARY KEY)");
 
-const query = new Query("INSERT INTO examples_ks.tab (a) VALUES(?) IF NOT EXISTS");
+const query = new Query("INSERT INTO lwt.tab (a) VALUES(?) IF NOT EXISTS");
 query.setConsistency(Consistency.One);
 query.setSerialConsistency(SerialConsistency.Serial);
 

index.d.ts

@@ -71,16 +71,33 @@ export class Cluster {
   /** Connect to the cluster */
   connect(keyspaceOrOptions?: string | ConnectionOptions | undefined | null, options?: ConnectionOptions | undefined | null): Promise<ScyllaSession>
 }
-export type ScyllaQuery = Query
-export class Query {
-  constructor(query: string)
+export type ScyllaBatchStatement = BatchStatement
+/**
+ * Batch statements
+ *
+ * A batch statement allows to execute many data-modifying statements at once.
+ * These statements can be simple or prepared.
+ * Only INSERT, UPDATE and DELETE statements are allowed.
+ */
+export class BatchStatement {
+  constructor()
+  /**
+   * Appends a statement to the batch.
+   *
+   * _Warning_
+   * Using simple statements with bind markers in batches is strongly discouraged. For each simple statement with a non-empty list of values in the batch, the driver will send a prepare request, and it will be done sequentially. Results of preparation are not cached between `session.batch` calls. Consider preparing the statements before putting them into the batch.
+   */
+  appendStatement(statement: Query | PreparedStatement): void
+}
+export class PreparedStatement {
   setConsistency(consistency: Consistency): void
   setSerialConsistency(serialConsistency: SerialConsistency): void
-  setPageSize(pageSize: number): void
 }
-export class ScyllaPreparedStatement {
+export class Query {
+  constructor(query: string)
   setConsistency(consistency: Consistency): void
   setSerialConsistency(serialConsistency: SerialConsistency): void
+  setPageSize(pageSize: number): void
 }
 export class Metrics {
   /** Returns counter for nonpaged queries */
@@ -104,9 +121,10 @@ export class Metrics {
 }
 export class ScyllaSession {
   metrics(): Metrics
-  execute(query: string | Query | ScyllaPreparedStatement, parameters?: Array<number | string | Uuid> | undefined | null): Promise<any>
+  execute(query: string | Query | PreparedStatement, parameters?: Array<number | string | Uuid> | undefined | null): Promise<any>
   query(scyllaQuery: Query, parameters?: Array<number | string | Uuid> | undefined | null): Promise<any>
-  prepare(query: string): Promise<ScyllaPreparedStatement>
+  prepare(query: string): Promise<PreparedStatement>
+  batch(batch: BatchStatement, parameters: Array<Array<number | string | Uuid> | undefined | null>): Promise<any>
   /**
    * Sends `USE <keyspace_name>` request on all connections\
    * This allows to write `SELECT * FROM table` instead of `SELECT * FROM keyspace.table`\

index.js

@@ -310,15 +310,16 @@ if (!nativeBinding) {
   throw new Error(`Failed to load native binding`)
 }
 
-const { Compression, Consistency, SerialConsistency, Cluster, VerifyMode, Query, ScyllaPreparedStatement, Metrics, ScyllaSession, Uuid } = nativeBinding
+const { Compression, Consistency, SerialConsistency, Cluster, VerifyMode, BatchStatement, PreparedStatement, Query, Metrics, ScyllaSession, Uuid } = nativeBinding
 
 module.exports.Compression = Compression
 module.exports.Consistency = Consistency
 module.exports.SerialConsistency = SerialConsistency
 module.exports.Cluster = Cluster
 module.exports.VerifyMode = VerifyMode
+module.exports.BatchStatement = BatchStatement
+module.exports.PreparedStatement = PreparedStatement
 module.exports.Query = Query
-module.exports.ScyllaPreparedStatement = ScyllaPreparedStatement
 module.exports.Metrics = Metrics
 module.exports.ScyllaSession = ScyllaSession
 module.exports.Uuid = Uuid

src/query/batch_statement.rs

@@ -0,0 +1,64 @@
+use std::fmt::Display;
+
+use napi::Either;
+use scylla::batch::Batch;
+
+use super::{scylla_prepared_statement::PreparedStatement, scylla_query::Query};
+
+/// Batch statements
+///
+/// A batch statement allows to execute many data-modifying statements at once.
+/// These statements can be simple or prepared.
+/// Only INSERT, UPDATE and DELETE statements are allowed.
+#[napi(js_name = "BatchStatement")]
+pub struct ScyllaBatchStatement {
+  pub(crate) batch: Batch,
+}
+
+impl Display for ScyllaBatchStatement {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    write!(
+      f,
+      "ScyllaBatchStatement: {:?}",
+      self
+        .batch
+        .statements
+        .iter()
+        .map(|s| match s {
+          scylla::batch::BatchStatement::Query(q) => q.contents.clone(),
+          scylla::batch::BatchStatement::PreparedStatement(p) => p.get_statement().to_string(),
+        })
+        .collect::<Vec<_>>()
+    )
+  }
+}
+
+#[napi]
+impl ScyllaBatchStatement {
+  #[napi(constructor)]
+  pub fn new() -> Self {
+    Self {
+      batch: Default::default(),
+    }
+  }
+
+  /// Appends a statement to the batch.
+  ///
+  /// _Warning_
+  /// Using simple statements with bind markers in batches is strongly discouraged. For each simple statement with a non-empty list of values in the batch, the driver will send a prepare request, and it will be done sequentially. Results of preparation are not cached between `session.batch` calls. Consider preparing the statements before putting them into the batch.
+  #[napi]
+  pub fn append_statement(&mut self, statement: Either<&Query, &PreparedStatement>) {
+    match statement {
+      Either::A(simple_query) => self.batch.append_statement(simple_query.query.clone()),
+      Either::B(prepared_statement) => self
+        .batch
+        .append_statement(prepared_statement.prepared.clone()),
+    }
+  }
+}
+
+impl Default for ScyllaBatchStatement {
+  fn default() -> Self {
+    Self::new()
+  }
+}

src/query/mod.rs

@@ -1,2 +1,3 @@
-pub mod scylla_query;
+pub mod batch_statement;
 pub mod scylla_prepared_statement;
+pub mod scylla_query;

src/query/scylla_prepared_statement.rs

@@ -1,17 +1,17 @@
-use scylla::prepared_statement::PreparedStatement;
+use scylla::prepared_statement;
 
 use crate::cluster::execution_profile::{
   consistency::Consistency, serial_consistency::SerialConsistency,
 };
 
 #[napi]
-pub struct ScyllaPreparedStatement {
-  pub(crate) prepared: PreparedStatement,
+pub struct PreparedStatement {
+  pub(crate) prepared: prepared_statement::PreparedStatement,
 }
 
 #[napi]
-impl ScyllaPreparedStatement {
-  pub fn new(prepared: PreparedStatement) -> Self {
+impl PreparedStatement {
+  pub fn new(prepared: prepared_statement::PreparedStatement) -> Self {
     Self { prepared }
   }
 

src/query/scylla_query.rs

@@ -3,25 +3,25 @@ use std::fmt::Display;
 use crate::cluster::execution_profile::{
   consistency::Consistency, serial_consistency::SerialConsistency,
 };
-use scylla::query::Query;
+use scylla::query;
 
-#[napi(js_name = "Query")]
-pub struct ScyllaQuery {
-  pub(crate) query: Query,
+#[napi]
+pub struct Query {
+  pub(crate) query: query::Query,
 }
 
-impl Display for ScyllaQuery {
+impl Display for Query {
   fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     write!(f, "ScyllaQuery: {}", self.query.contents)
   }
 }
 
 #[napi]
-impl ScyllaQuery {
+impl Query {
   #[napi(constructor)]
   pub fn new(query: String) -> Self {
     Self {
-      query: Query::new(query),
+      query: query::Query::new(query),
     }
   }
 

src/session/scylla_session.rs

@@ -1,7 +1,8 @@
 use crate::helpers::query_parameter::QueryParameter;
 use crate::helpers::query_results::QueryResult;
-use crate::query::scylla_prepared_statement::ScyllaPreparedStatement;
-use crate::query::scylla_query::ScyllaQuery;
+use crate::query::batch_statement::ScyllaBatchStatement;
+use crate::query::scylla_prepared_statement::PreparedStatement;
+use crate::query::scylla_query::Query;
 use crate::types::uuid::Uuid;
 use napi::bindgen_prelude::Either3;
 
@@ -26,7 +27,7 @@ impl ScyllaSession {
   #[napi]
   pub async fn execute(
     &self,
-    query: Either3<String, &ScyllaQuery, &ScyllaPreparedStatement>,
+    query: Either3<String, &Query, &PreparedStatement>,
     parameters: Option<Vec<Either3<u32, String, &Uuid>>>,
   ) -> napi::Result<serde_json::Value> {
     let values = QueryParameter::parser(parameters.clone()).ok_or(napi::Error::new(
@@ -58,7 +59,7 @@ impl ScyllaSession {
   #[napi]
   pub async fn query(
     &self,
-    scylla_query: &ScyllaQuery,
+    scylla_query: &Query,
     parameters: Option<Vec<Either3<u32, String, &Uuid>>>,
   ) -> napi::Result<serde_json::Value> {
     let values = QueryParameter::parser(parameters.clone()).ok_or(napi::Error::new(
@@ -81,15 +82,46 @@ impl ScyllaSession {
   }
 
   #[napi]
-  pub async fn prepare(&self, query: String) -> napi::Result<ScyllaPreparedStatement> {
+  pub async fn prepare(&self, query: String) -> napi::Result<PreparedStatement> {
     let prepared = self.session.prepare(query.clone()).await.map_err(|_| {
       napi::Error::new(
         napi::Status::InvalidArg,
         format!("Something went wrong with your prepared statement. - [{query}]"),
       )
     })?;
 
-    Ok(ScyllaPreparedStatement::new(prepared))
+    Ok(PreparedStatement::new(prepared))
+  }
+
+  #[napi]
+  #[allow(clippy::type_complexity)]
+  pub async fn batch(
+    &self,
+    batch: &ScyllaBatchStatement,
+    parameters: Vec<Option<Vec<Either3<u32, String, &Uuid>>>>,
+  ) -> napi::Result<serde_json::Value> {
+    let values = parameters
+      .iter()
+      .map(|params| {
+        QueryParameter::parser(params.clone()).ok_or(napi::Error::new(
+          napi::Status::InvalidArg,
+          format!("Something went wrong with your batch parameters. {parameters:?}"),
+        ))
+      })
+      .collect::<napi::Result<Vec<_>>>()?;
+
+    let query_result = self
+      .session
+      .batch(&batch.batch, values)
+      .await
+      .map_err(|e| {
+        napi::Error::new(
+          napi::Status::InvalidArg,
+          format!("Something went wrong with your batch. - [{batch}] - {parameters:?}\n{e}"),
+        )
+      })?;
+
+    Ok(QueryResult::parser(query_result))
   }
 
   /// Sends `USE <keyspace_name>` request on all connections\