From 3848141b26a295d7ebb58f7032d8b5915cc07573 Mon Sep 17 00:00:00 2001
From: Daniel Boll <danielboll.academico@gmail.com>
Date: Sat, 14 Oct 2023 19:14:40 -0300
Subject: [PATCH] fix(index.d.ts): add ConnectionOptions and Auth interfaces to
 provide more flexibility in connecting to the cluster
 feat(scylla_cluster.rs): add support for ConnectionOptions and Auth structs
 to allow connecting to the cluster with configurable keyspace and
 authentication credentials

---
 index.d.ts                    | 10 ++++++-
 src/cluster/scylla_cluster.rs | 51 ++++++++++++++++++++++++++++++++---
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/index.d.ts b/index.d.ts
index 3a36650..f327ba5 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -24,6 +24,14 @@ export interface ExecutionProfile {
   consistency?: Consistency
   requestTimeout?: number
 }
+export interface ConnectionOptions {
+  keyspace?: string
+  auth?: Auth
+}
+export interface Auth {
+  username: string
+  password: string
+}
 export type ScyllaCluster = Cluster
 export class Cluster {
   /**
@@ -34,7 +42,7 @@ export class Cluster {
    */
   constructor(clusterConfig: ClusterConfig)
   /** Connect to the cluster */
-  connect(keyspace?: string | undefined | null): Promise<ScyllaSession>
+  connect(keyspaceOrOptions?: string | ConnectionOptions | undefined | null, options?: ConnectionOptions | undefined | null): Promise<ScyllaSession>
 }
 export class ScyllaSession {
   execute(query: string, parameters?: Array<number | string | Uuid> | undefined | null): Promise<any>
diff --git a/src/cluster/scylla_cluster.rs b/src/cluster/scylla_cluster.rs
index cc31e54..230eef5 100644
--- a/src/cluster/scylla_cluster.rs
+++ b/src/cluster/scylla_cluster.rs
@@ -1,3 +1,5 @@
+use napi::Either;
+
 use crate::{
   cluster::{cluster_config::ClusterConfig, execution_profile::ExecutionProfile},
   session::scylla_session::ScyllaSession,
@@ -9,6 +11,18 @@ struct ScyllaCluster {
   default_execution_profile: Option<ExecutionProfile>,
 }
 
+#[napi(object)]
+struct ConnectionOptions {
+  pub keyspace: Option<String>,
+  pub auth: Option<Auth>,
+}
+
+#[napi(object)]
+struct Auth {
+  pub username: String,
+  pub password: String,
+}
+
 #[napi]
 impl ScyllaCluster {
   /// Object config is in the format:
@@ -32,17 +46,48 @@ impl ScyllaCluster {
 
   /// Connect to the cluster
   #[napi]
-  pub async fn connect(&self, keyspace: Option<String>) -> ScyllaSession {
+  pub async fn connect(
+    &self,
+    keyspace_or_options: Option<Either<String, ConnectionOptions>>,
+    options: Option<ConnectionOptions>,
+  ) -> napi::Result<ScyllaSession> {
     let mut builder = scylla::SessionBuilder::new().known_node(self.uri.as_str());
 
-    if let Some(keyspace) = keyspace {
+    let keyspace = match (&keyspace_or_options, &options) {
+      (Some(Either::A(keyspace)), _) => Ok(Some(keyspace.clone())),
+      (Some(Either::B(_)), Some(_)) => Err(napi::Error::new(
+        napi::Status::InvalidArg,
+        "Options cannot be provided twice",
+      )),
+      (Some(Either::B(options)), _) => Ok(options.keyspace.clone()),
+      (None, Some(options)) => Ok(options.keyspace.clone()),
+      (None, None) => Ok(None),
+    };
+
+    let auth = match (keyspace_or_options, options) {
+      (Some(Either::A(_)), Some(options)) => Ok(options.auth),
+      (Some(Either::B(_)), Some(_)) => Err(napi::Error::new(
+        napi::Status::InvalidArg,
+        "Options cannot be provided twice",
+      )),
+      (Some(Either::B(options)), None) => Ok(options.auth),
+      (None, Some(options)) => Ok(options.auth),
+      (None, None) => Ok(None),
+      (Some(Either::A(_)), None) => Ok(None),
+    };
+
+    if let Some(keyspace) = keyspace? {
       builder = builder.use_keyspace(keyspace, false);
     }
 
+    if let Some(auth) = auth? {
+      builder = builder.user(auth.username, auth.password);
+    }
+
     if let Some(default_execution_profile) = &self.default_execution_profile {
       builder = builder.default_execution_profile_handle(default_execution_profile.into_handle());
     }
 
-    ScyllaSession::new(builder.build().await.unwrap())
+    Ok(ScyllaSession::new(builder.build().await.unwrap()))
   }
 }