Allow overriding of the token through a arg (#261)

* Allow the token to be passed along through a argument (or envvar) * Update changelog
fiberplane · Dec 11, 2023 · a97f34e · a97f34e
1 parent dd6482e
commit a97f34e
Show file tree

Hide file tree

Showing 20 changed files with 422 additions and 121 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format of this file is based on [Keep a Changelog](https://keepachangelog.co
 ## [Unreleased]
 
 - Update all dependencies, except `hyper` (#260)
+- Allow overriding the token through an argument or envvar (#261)
 
 ## [2.11.0]
 

diff --git a/src/auth.rs b/src/auth.rs
@@ -103,22 +103,35 @@ pub async fn handle_login_command(args: Arguments) -> Result<(), Error> {
     Ok(())
 }
 
-/// Logout from Fiberplane and delete the API Token from the config file
+/// Invalidate the API token.
+///
+/// If a token is set using the `--token` flag, then that will be used,
+/// otherwise it will use the token from the config file. The token will also be
+/// removed from the config file if that was used.
 pub async fn handle_logout_command(args: Arguments) -> Result<(), Error> {
-    let mut config = Config::load(args.config).await?;
+    if let Some(token) = args.token {
+        let api_config = api_client_configuration_from_token(&token, args.base_url)?;
 
-    match config.api_token {
-        Some(token) => {
-            let api_config = api_client_configuration_from_token(&token, args.base_url)?;
-            logout(&api_config).await?;
+        logout(&api_config).await?;
 
-            config.api_token = None;
-            config.save().await?;
+        info!("You are logged out");
+    } else {
+        let mut config = Config::load(args.config).await?;
 
-            info!("You are logged out");
-        }
-        None => {
-            warn!("You are already logged out");
+        match config.api_token {
+            Some(token) => {
+                let api_config = api_client_configuration_from_token(&token, args.base_url)?;
+
+                logout(&api_config).await?;
+
+                config.api_token = None;
+                config.save().await?;
+
+                info!("You are logged out");
+            }
+            None => {
+                warn!("You are already logged out");
+            }
         }
     }
 

diff --git a/src/config.rs b/src/config.rs
@@ -78,18 +78,34 @@ fn default_config_file_path() -> PathBuf {
         .join("config.toml")
 }
 
+/// Returns a ApiClient with the token set. If `token` is `Some` then that will
+/// be used, otherwise the config will be used for the token.
 pub(crate) async fn api_client_configuration(
+    token: Option<String>,
     config_path: Option<PathBuf>,
     base_url: Url,
 ) -> Result<ApiClient> {
-    let token = Config::load(config_path).await?.api_token.ok_or_else(|| {
-        anyhow!("Must be logged in to run this command. Please run `fp login` first.")
-    })?;
+    let token = {
+        if let Some(token) = token {
+            // an override has been specified, use that
+            token
+        } else {
+            // no override, so try loading the config
+            let config = Config::load(config_path).await?;
+            if let Some(token) = config.api_token {
+                token
+            } else {
+                return Err(anyhow!(
+                    "Must be logged in to run this command. Please run `fp login` first."
+                ));
+            }
+        }
+    };
 
     api_client_configuration_from_token(&token, base_url)
 }
 
-pub(crate) fn api_client_configuration_from_token(token: &str, base_url: Url) -> Result<ApiClient> {
+pub fn api_client_configuration_from_token(token: &str, base_url: Url) -> Result<ApiClient> {
     let mut headers = HeaderMap::new();
     headers.insert(
         "Authorization",

diff --git a/src/daemons.rs b/src/daemons.rs
@@ -62,6 +62,9 @@ pub struct CreateArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -79,6 +82,9 @@ pub struct ListArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -96,6 +102,9 @@ pub struct DataSourcesArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -116,6 +125,9 @@ pub struct GetArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -132,6 +144,9 @@ pub struct DeleteArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 /// A generic output for daemon related commands.
@@ -158,7 +173,7 @@ pub async fn handle_command(args: Arguments) -> Result<()> {
 async fn handle_create_command(args: CreateArgs) -> Result<()> {
     let default_name = Name::new(petname(2, "-")).expect("petname should be valid name");
     let name = name_req("Daemon name", args.name, Some(default_name))?;
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
 
     let mut new_proxy = NewProxy::builder().name(name).build();
@@ -189,7 +204,7 @@ struct ProxySummaryWithConnectedDataSources {
 }
 
 async fn handle_list_command(args: ListArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
     let proxies = proxy_list(&client, workspace_id).await?;
     let data_sources = data_source_list(&client, workspace_id).await?;
@@ -249,7 +264,7 @@ async fn handle_list_command(args: ListArgs) -> Result<()> {
 }
 
 async fn handle_get_command(args: GetArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
     let proxy_name = interactive::proxy_picker(
         &client,
@@ -270,7 +285,7 @@ async fn handle_get_command(args: GetArgs) -> Result<()> {
 }
 
 async fn handle_data_sources_command(args: DataSourcesArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
     let data_sources = data_source_list(&client, workspace_id).await?;
 
@@ -286,7 +301,7 @@ async fn handle_data_sources_command(args: DataSourcesArgs) -> Result<()> {
 }
 
 async fn handle_delete_command(args: DeleteArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
     let proxy_name = interactive::proxy_picker(
         &client,

diff --git a/src/data_sources.rs b/src/data_sources.rs
@@ -96,6 +96,9 @@ struct CreateArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -117,6 +120,9 @@ struct GetArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -134,6 +140,9 @@ struct DeleteArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -163,6 +172,9 @@ struct UpdateArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -180,6 +192,9 @@ struct ListArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 pub async fn handle_command(args: Arguments) -> Result<()> {
@@ -196,7 +211,7 @@ pub async fn handle_command(args: Arguments) -> Result<()> {
 }
 
 async fn handle_create(args: CreateArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
 
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
     let name = name_req("Data source name", args.name, None)?;
@@ -240,7 +255,7 @@ async fn handle_create(args: CreateArgs) -> Result<()> {
 }
 
 async fn handle_delete(args: DeleteArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
 
     let data_source = data_source_picker(&client, Some(workspace_id), args.name).await?;
@@ -251,7 +266,7 @@ async fn handle_delete(args: DeleteArgs) -> Result<()> {
 }
 
 async fn handle_get(args: GetArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
 
     let data_source = data_source_picker(&client, Some(workspace_id), args.name).await?;
@@ -266,7 +281,7 @@ async fn handle_get(args: GetArgs) -> Result<()> {
 }
 
 async fn handle_update(args: UpdateArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
 
     let data_source = data_source_picker(&client, Some(workspace_id), args.name).await?;
@@ -287,7 +302,7 @@ async fn handle_update(args: UpdateArgs) -> Result<()> {
 }
 
 async fn handle_list(args: ListArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
 
     let data_sources = data_source_list(&client, workspace_id).await?;

diff --git a/src/events.rs b/src/events.rs
@@ -79,6 +79,9 @@ struct CreateArguments {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -124,10 +127,13 @@ pub struct SearchArguments {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 async fn handle_event_create_command(args: CreateArguments) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
 
     let key_values: HashMap<_, _> = args
         .labels
@@ -159,7 +165,7 @@ async fn handle_event_create_command(args: CreateArguments) -> Result<()> {
 }
 
 async fn handle_event_search_command(args: SearchArguments) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
 
     let workspace_id = workspace_picker(&client, args.workspace_id).await?;
 
@@ -196,10 +202,13 @@ pub struct DeleteArguments {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 async fn handle_event_delete_command(args: DeleteArguments) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
 
     event_delete(&client, args.id).await?;
 

diff --git a/src/experiments.rs b/src/experiments.rs
@@ -68,12 +68,15 @@ struct MessageArgs {
     #[clap(from_global)]
     base_url: Url,
 
-    #[clap(from_global)]
-    config: Option<PathBuf>,
-
     /// Output type to display
     #[clap(long, short, default_value = "table", value_enum)]
     output: MessageOutput,
+
+    #[clap(from_global)]
+    config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -92,6 +95,9 @@ struct CrawlArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(Parser)]
@@ -115,6 +121,9 @@ struct PrometheusGraphToNotebookArgs {
 
     #[clap(from_global)]
     config: Option<PathBuf>,
+
+    #[clap(from_global)]
+    token: Option<String>,
 }
 
 #[derive(ValueEnum, Clone)]
@@ -138,7 +147,7 @@ pub async fn handle_command(args: Arguments) -> Result<()> {
 }
 
 async fn handle_message_command(args: MessageArgs) -> Result<()> {
-    let client = api_client_configuration(args.config, args.base_url).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url).await?;
     let notebook_id = interactive::notebook_picker(&client, args.notebook_id, None).await?;
     let mut cache = Cache::load().await?;
 
@@ -214,7 +223,7 @@ async fn handle_crawl_command(args: CrawlArgs) -> Result<()> {
     let mut notebook_titles: HashMap<String, usize> = HashMap::new();
     let mut notebooks_to_crawl = VecDeque::new();
 
-    let client = api_client_configuration(args.config, args.base_url.clone()).await?;
+    let client = api_client_configuration(args.token, args.config, args.base_url.clone()).await?;
     let starting_notebook_id =
         interactive::notebook_picker(&client, args.notebook_id, None).await?;
 
@@ -384,7 +393,7 @@ fn cache_file_path() -> PathBuf {
 }
 
 async fn handle_prometheus_redirect_command(args: PrometheusGraphToNotebookArgs) -> Result<()> {
-    let client = Arc::new(api_client_configuration(args.config, args.base_url).await?);
+    let client = Arc::new(api_client_configuration(args.token, args.config, args.base_url).await?);
     let workspace_id = interactive::workspace_picker(&client, args.workspace_id).await?;
     let notebook_id =
         interactive::notebook_picker(&client, args.notebook_id, Some(workspace_id)).await?;