From 74229ac48174471552a40fe9d2aac3bf715235e1 Mon Sep 17 00:00:00 2001
From: Mehul Arora <mehul_arora@live.com>
Date: Wed, 6 Nov 2024 22:52:16 -0500
Subject: [PATCH] feat: allow S2_CELL to be set

---
 src/config.rs | 11 ++++++-----
 src/main.rs   | 43 ++++++++++++++++++++++++++++++-------------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 8fcc43e..168fb62 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -7,9 +7,12 @@ use thiserror::Error;
 
 use crate::error::S2CliError;
 
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Debug, Deserialize, Serialize, Default)]
 pub struct S2Config {
     pub auth_token: String,
+    pub cell: String,
+    #[serde(skip_serializing)]
+    pub basin_zone: Option<String>,
 }
 
 #[cfg(target_os = "windows")]
@@ -41,14 +44,12 @@ pub fn load_config(path: &Path) -> Result<S2Config, S2ConfigError> {
     Ok(builder.build()?.try_deserialize::<S2Config>()?)
 }
 
-pub fn create_config(config_path: &PathBuf, auth_token: String) -> Result<(), S2ConfigError> {
-    let cfg = S2Config { auth_token };
-
+pub fn create_config(config_path: &PathBuf, config: S2Config) -> Result<(), S2ConfigError> {
     if let Some(parent) = config_path.parent() {
         std::fs::create_dir_all(parent).map_err(S2ConfigError::WriteError)?;
     }
 
-    let toml = toml::to_string(&cfg).unwrap();
+    let toml = toml::to_string(&config).unwrap();
     std::fs::write(config_path, toml).map_err(S2ConfigError::WriteError)?;
 
     Ok(())
diff --git a/src/main.rs b/src/main.rs
index c3d05e8..701240e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,9 +2,9 @@ use std::path::PathBuf;
 
 use account::AccountService;
 use basin::BasinService;
-use clap::{builder::styling, Parser, Subcommand};
+use clap::{builder::styling, ArgGroup, Parser, Subcommand};
 use colored::*;
-use config::{config_path, create_config};
+use config::{config_path, create_config, load_config, S2Config};
 use error::S2CliError;
 use stream::{RecordStream, StreamService, StreamServiceError};
 use streamstore::{
@@ -85,11 +85,22 @@ enum Commands {
 
 #[derive(Subcommand, Debug)]
 enum ConfigActions {
-    /// Set the authentication token to be reused in subsequent commands.
-    /// Alternatively, use the S2_AUTH_TOKEN environment variable.
+    /// Set s2 configuration options.
+    #[command(group(
+        ArgGroup::new("config")
+            .required(true)
+            .args(["auth_token", "cell"]),
+    ))]
     Set {
-        #[arg(short, long)]
-        auth_token: String,
+        /// Set the authentication token to be reused in subsequent commands.
+        /// Alternatively, use the S2_AUTH_TOKEN environment variable.
+        #[arg(short, long, group = "config")]
+        auth_token: Option<String>,
+
+        /// Set the cell endpoint to be reused in subsequent commands.
+        /// Alternatively, use the S2_CELL environment variable.
+        #[arg(short, long, group = "config")]
+        cell: Option<String>,
     },
 }
 
@@ -296,13 +307,19 @@ async fn run() -> Result<(), S2CliError> {
 
     match commands.command {
         Commands::Config { action } => match action {
-            ConfigActions::Set { auth_token } => {
-                create_config(&config_path, auth_token)?;
-                eprintln!("{}", "✓ Token set successfully".green().bold());
-                eprintln!(
-                    "  Configuration saved to: {}",
-                    config_path.display().to_string().cyan()
-                );
+            ConfigActions::Set { auth_token, cell } => {
+                let existing_config =
+                    load_config(&config_path).unwrap_or_else(|_| S2Config::default());
+
+                let new_config = S2Config {
+                    auth_token: auth_token.unwrap_or(existing_config.auth_token),
+                    cell: cell.unwrap_or(existing_config.cell),
+                    basin_zone: existing_config.basin_zone,
+                };
+
+                create_config(&config_path, new_config)?;
+                eprintln!("{}", "✓ Confg set successfully".green().bold());
+                eprintln!("  Saved to: {}", config_path.display().to_string().cyan());
             }
         },