Merge pull request #2 from marktoda/add-key-override

feat: add key override

src/chain.rs

@@ -14,6 +14,7 @@ use colored::*;
 use dialoguer::{FuzzySelect, Input};
 use serde::{Deserialize, Serialize};
 use std::fmt::{Debug, Display};
+use std::sync::Arc;
 
 pub const DEFAULT_KEY_NAME: &str = "default";
 
@@ -28,13 +29,30 @@ pub struct ChainDefinition {
     pub key_name: String,
 }
 
+#[derive(Clone)]
 pub struct ChainInstance {
     pub definition: ChainDefinition,
-    pub provider: Box<dyn Provider<BoxTransport>>,
+    pub provider: Arc<dyn Provider<BoxTransport>>,
     pub rpc_url: String,
     pub key: Key,
 }
 
+impl ChainInstance {
+    pub fn new(definition: ChainDefinition, provider: Box<dyn Provider<BoxTransport>>, rpc_url: String, key: Key) -> Self {
+        Self {
+            definition,
+            provider: Arc::from(provider),
+            rpc_url,
+            key,
+        }
+    }
+
+    pub fn with_key(mut self, key: Key) -> Self {
+        self.key = key;
+        self
+    }
+}
+
 pub struct Rpc {
     pub rpc_url: String,
     pub provider: Box<dyn Provider<BoxTransport>>,

src/config.rs

@@ -43,27 +43,31 @@ impl Chainz {
         })
     }
 
-    pub async fn get_chain(&mut self, name_or_id: &str) -> Result<&ChainInstance> {
+    pub async fn get_chain(&mut self, name_or_id: &str) -> Result<ChainInstance> {
         let definition = self.config.get_chain(name_or_id)?;
         let name = definition.name.clone();
 
         if !self.active_chains.contains_key(&name) {
             let instance = self.instantiate_chain(&definition).await?;
             self.active_chains.insert(name.clone(), instance);
         }
-        Ok(&self.active_chains[&name])
+        Ok(self
+            .active_chains
+            .get(&name)
+            .expect("Chain was just inserted")
+            .clone())
     }
 
     async fn instantiate_chain(&self, def: &ChainDefinition) -> Result<ChainInstance> {
         let rpc = def.get_rpc(&self.config.globals).await?;
-        let key = self.get_key(&def.key_name.clone())?;
+        let key = self.get_key(&def.key_name)?;
 
-        Ok(ChainInstance {
-            definition: def.clone(),
-            provider: rpc.provider,
-            rpc_url: rpc.rpc_url,
+        Ok(ChainInstance::new(
+            def.clone(),
+            rpc.provider,
+            rpc.rpc_url,
             key,
-        })
+        ))
     }
 
     // Key management methods

src/main.rs

@@ -8,10 +8,12 @@ pub mod config;
 pub mod init;
 pub mod key;
 pub mod opt;
+pub mod util;
 pub mod variables;
 
 use config::Chainz;
 use opt::Opt;
+use util::Pipe;
 use variables::ChainVariables;
 
 #[tokio::main]
@@ -48,10 +50,19 @@ async fn main() -> Result<()> {
             name_or_id,
             print,
             export,
+            key,
         } => {
-            let chain = chainz.get_chain(&name_or_id).await?;
+            let chain = chainz
+                .get_chain(&name_or_id)
+                .await?
+                .pipe(|chain| -> Result<_> {
+                    Ok(match key {
+                        Some(key_name) => chain.with_key(chainz.get_key(&key_name)?),
+                        None => chain,
+                    })
+                })?;
             eprintln!("{}", chain);
-            let variables = ChainVariables::new(chain)?;
+            let variables = ChainVariables::new(&chain)?;
             if export {
                 print!("{}", variables.as_exports());
             } else {
@@ -64,12 +75,21 @@ async fn main() -> Result<()> {
         opt::Command::Exec {
             name_or_id,
             command,
+            key,
         } => {
             if command.is_empty() {
                 anyhow::bail!("No command specified");
             }
-            let chain = chainz.get_chain(&name_or_id).await?;
-            let variables = ChainVariables::new(chain)?;
+            let chain = chainz
+                .get_chain(&name_or_id)
+                .await?
+                .pipe(|chain| -> Result<_> {
+                    Ok(match key {
+                        Some(key_name) => chain.with_key(chainz.get_key(&key_name)?),
+                        None => chain,
+                    })
+                })?;
+            let variables = ChainVariables::new(&chain)?;
             let expanded_command = variables.expand(command);
 
             let status = ProcessCommand::new(&expanded_command[0])

src/opt.rs

@@ -44,6 +44,7 @@ pub enum Command {
     /// Flags:
     ///     -p, --print  : Print variables to stdout instead of writing to .env
     ///     -e, --export : Include 'export' prefix in output
+    ///     -k, --key    : Override the key to use for this command
     ///
     /// Example: chainz use ethereum --print
     Use {
@@ -55,6 +56,9 @@ pub enum Command {
         /// Include 'export' prefix in output
         #[structopt(short, long)]
         export: bool,
+        /// Override the key to use for this command
+        #[structopt(short, long)]
+        key: Option<String>,
     },
 
     /// List all configured chains
@@ -84,6 +88,9 @@ pub enum Command {
         /// Command to execute (after --)
         #[structopt(last = true)]
         command: Vec<String>,
+        /// Override the key to use for this command
+        #[structopt(short, long)]
+        key: Option<String>,
     },
 
     /// Manage private keys

src/util.rs

@@ -0,0 +1,10 @@
+pub trait Pipe: Sized {
+    fn pipe<B, F>(self, f: F) -> B
+    where
+        F: FnOnce(Self) -> B,
+    {
+        f(self)
+    }
+}
+
+impl<T: Sized> Pipe for T {}