From e3061f93074c0f5adae02777ce2688802a951379 Mon Sep 17 00:00:00 2001
From: Robsdedude <rouven.bauer@neo4j.com>
Date: Sat, 13 Apr 2024 19:28:06 +0200
Subject: [PATCH] TestKit: add executeQuery sessionAuth config

---
 testkit_backend/src/testkit_backend/driver_holder.rs | 5 +++++
 testkit_backend/src/testkit_backend/requests.rs      | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/testkit_backend/src/testkit_backend/driver_holder.rs b/testkit_backend/src/testkit_backend/driver_holder.rs
index 3026346..f840cfd 100644
--- a/testkit_backend/src/testkit_backend/driver_holder.rs
+++ b/testkit_backend/src/testkit_backend/driver_holder.rs
@@ -655,6 +655,7 @@ impl DriverHolderRunner {
                     bookmark_manager,
                     tx_meta,
                     timeout,
+                    auth,
                 }) => {
                     let mut builder = self.driver.execute_query(query);
                     if let Some(params) = params {
@@ -684,6 +685,9 @@ impl DriverHolderRunner {
                     if let Some(timeout) = timeout {
                         builder = builder.with_transaction_timeout(timeout);
                     }
+                    if let Some(auth) = auth {
+                        builder = builder.with_session_auth(auth);
+                    }
                     let result = builder
                         .run_with_retry(ExponentialBackoff::default())
                         .map_err(Into::into);
@@ -1051,6 +1055,7 @@ pub(super) struct ExecuteQuery {
     pub(crate) bookmark_manager: ExecuteQueryBookmarkManager,
     pub(super) tx_meta: Option<HashMap<String, ValueSend>>,
     pub(super) timeout: Option<TransactionTimeout>,
+    pub(super) auth: Option<Arc<AuthToken>>,
 }
 
 impl From<ExecuteQuery> for Command {
diff --git a/testkit_backend/src/testkit_backend/requests.rs b/testkit_backend/src/testkit_backend/requests.rs
index 6d1c2e3..e147f79 100644
--- a/testkit_backend/src/testkit_backend/requests.rs
+++ b/testkit_backend/src/testkit_backend/requests.rs
@@ -599,6 +599,8 @@ pub(super) struct ExecuteQueryConfig {
     bookmark_manager_id: Option<ExecuteQueryBmmId>,
     tx_meta: Option<HashMap<String, CypherValue>>,
     timeout: Option<i64>,
+    #[serde(rename = "authorizationToken")]
+    auth: MaybeTestKitAuth,
 }
 
 #[derive(Deserialize, Debug)]
@@ -1466,6 +1468,7 @@ impl Request {
             bookmark_manager_id,
             tx_meta,
             timeout,
+            auth,
         } = config;
         let data = backend.data.borrow_mut();
         let driver_holder = get_driver(&data, &driver_id)?;
@@ -1483,6 +1486,7 @@ impl Request {
             .map(cypher_value_map_to_value_send_map)
             .transpose()?;
         let timeout = read_transaction_timeout(timeout)?;
+        let auth = auth.0.map(|auth| Arc::new(auth.into()));
         let result = driver_holder
             .execute_query(ExecuteQuery {
                 query,
@@ -1493,6 +1497,7 @@ impl Request {
                 bookmark_manager,
                 tx_meta,
                 timeout,
+                auth,
             })
             .result?;
         let response: Response = result.try_into()?;