From 8a25d54835ee6d57912c440650208dfda73d5d5a Mon Sep 17 00:00:00 2001
From: Jack Westwood <jack.westwood@couchbase.com>
Date: Wed, 14 Aug 2024 16:31:57 +0100
Subject: [PATCH] Allow id to be an int in replace, upsert, insert and import
 cmds

---
 src/cli/doc_import.rs | 11 +++++------
 src/cli/doc_remove.rs |  4 ++--
 src/cli/doc_upsert.rs | 17 ++++++++++++++++-
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/cli/doc_import.rs b/src/cli/doc_import.rs
index b20d954c..7203349b 100644
--- a/src/cli/doc_import.rs
+++ b/src/cli/doc_import.rs
@@ -1,17 +1,16 @@
-use crate::cli::doc_upsert::run_kv_mutations;
+use crate::cli::doc_upsert::{id_from_value, run_kv_mutations};
+use crate::cli::error::serialize_error;
+use crate::cli::util::convert_nu_value_to_json_value;
 use crate::client::KeyValueRequest;
 use crate::state::State;
 use nu_command::Open;
 use nu_engine::CallExt;
-use std::sync::{Arc, Mutex};
-
-use crate::cli::error::serialize_error;
-use crate::cli::util::convert_nu_value_to_json_value;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
 use nu_protocol::{
     Category, IntoPipelineData, PipelineData, ShellError, Signature, SyntaxShape, Value,
 };
+use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 pub struct DocImport {
@@ -120,7 +119,7 @@ fn run_import(
                 let mut content = serde_json::Map::new();
                 for (k, v) in val.iter() {
                     if k.clone() == id_column {
-                        id = v.as_str().ok();
+                        id = id_from_value(v, span);
                     }
 
                     content.insert(k.clone(), convert_nu_value_to_json_value(v, span).ok()?);
diff --git a/src/cli/doc_remove.rs b/src/cli/doc_remove.rs
index 116f7701..1071a7cd 100644
--- a/src/cli/doc_remove.rs
+++ b/src/cli/doc_remove.rs
@@ -85,11 +85,11 @@ impl Command for DocRemove {
         call: &Call,
         input: PipelineData,
     ) -> Result<PipelineData, ShellError> {
-        run_get(self.state.clone(), engine_state, stack, call, input)
+        run_remove(self.state.clone(), engine_state, stack, call, input)
     }
 }
 
-fn run_get(
+fn run_remove(
     state: Arc<Mutex<State>>,
     engine_state: &EngineState,
     stack: &mut Stack,
diff --git a/src/cli/doc_upsert.rs b/src/cli/doc_upsert.rs
index 714433ac..c17fe1a2 100644
--- a/src/cli/doc_upsert.rs
+++ b/src/cli/doc_upsert.rs
@@ -10,6 +10,7 @@ use crate::state::State;
 use crate::RemoteCluster;
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
+use log::info;
 use nu_engine::CallExt;
 use nu_protocol::ast::Call;
 use nu_protocol::engine::{Command, EngineState, Stack};
@@ -164,7 +165,7 @@ pub(crate) fn run_kv_store_ops(
             let mut content = None;
             for (k, v) in val.iter() {
                 if k.clone() == id_column {
-                    id = v.as_str().ok();
+                    id = id_from_value(v, span);
                 }
                 if k.clone() == content_column {
                     content = convert_nu_value_to_json_value(v, span).ok();
@@ -198,6 +199,20 @@ pub(crate) fn run_kv_store_ops(
     )
 }
 
+pub fn id_from_value(v: &Value, span: Span) -> Option<String> {
+    match v {
+        Value::String { val, .. } => Some(val.clone()),
+        Value::Int { val, .. } => Some(val.to_string()),
+        _ => {
+            info!(
+                "Skipping doc with id '{}' as id is not an int or string",
+                convert_nu_value_to_json_value(v, span).unwrap_or("error parsing id".into())
+            );
+            None
+        }
+    }
+}
+
 pub fn run_kv_mutations(
     state: Arc<Mutex<State>>,
     engine_state: &EngineState,