Merge pull request #557 from IntersectMBO/ch/query-drep-stake-conveni…

…ence Convenience flag to query stake distribution with `drep-state` query
IntersectMBO · Feb 13, 2024 · 61cf8da · 61cf8da
2 parents 07180c3 + 3f44c29
commit 61cf8da
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 8 deletions.
diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Commands/Query.hs
@@ -25,6 +25,7 @@ module Cardano.CLI.EraBased.Commands.Query
   , QueryDRepStateCmdArgs(..)
   , QueryDRepStakeDistributionCmdArgs(..)
   , renderQueryCmds
+  , IncludeStake (..)
   ) where
 
 import           Cardano.Api.Shelley hiding (QueryInShelleyBasedEra (..))
@@ -203,10 +204,16 @@ data QueryDRepStateCmdArgs era = QueryDRepStateCmdArgs
   , consensusModeParams :: !ConsensusModeParams
   , networkId           :: !NetworkId
   , drepKeys            :: !(AllOrOnly (VerificationKeyOrHashOrFile DRepKey))
+  , includeStake        :: !IncludeStake
   , target              :: !(Consensus.Target ChainPoint)
   , mOutFile            :: !(Maybe (File () Out))
   } deriving Show
 
+-- | Whether to include the stake, as queried by drep-stake-distribution, in
+-- the output of drep-state. This is (computationally) expensive, but sometimes
+-- convenient.
+data IncludeStake = WithStake | NoStake deriving Show
+
 data QueryDRepStakeDistributionCmdArgs era = QueryDRepStakeDistributionCmdArgs
   { eon                 :: !(ConwayEraOnwards era)
   , nodeSocketPath      :: !SocketPath

diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Options/Query.hs
@@ -335,6 +335,16 @@ pQueryDRepStateCmd era envCli = do
         <*> pConsensusModeParams
         <*> pNetworkId envCli
         <*> pAllOrOnlyDRepVerificationKeyOrHashOrFile
+        <*> Opt.flag WithStake NoStake (mconcat
+              [ Opt.long "include-stake"
+              , Opt.help $ mconcat
+                 [ "Also return the stake associated with each DRep. "
+                 , "The result is the same as with \"drep-stake-distribution\"; "
+                 , "this is a convenience option to obtain all information concerning a DRep at once. "
+                 , "This is a potentially expensive query."
+                 ]
+              ]
+            )
         <*> pTarget era
         <*> optional pOutputFile
 

diff --git a/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs b/cardano-cli/src/Cardano/CLI/EraBased/Run/Query.hs
@@ -8,6 +8,7 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
 
@@ -1434,6 +1435,7 @@ runQueryDRepState
       , Cmd.consensusModeParams
       , Cmd.networkId
       , Cmd.drepKeys = drepKeys'
+      , Cmd.includeStake
       , Cmd.target
       , Cmd.mOutFile
       } = conwayEraOnwardsConstraints eon $ do
@@ -1442,17 +1444,32 @@ runQueryDRepState
   let drepKeys = case drepKeys' of
                    All -> []
                    Only l -> l
-  drepCreds <- Set.fromList <$> mapM (firstExceptT QueryCmdDRepKeyError . getDRepCredentialFromVerKeyHashOrFile) drepKeys
+  drepCreds <- mapM (firstExceptT QueryCmdDRepKeyError . getDRepCredentialFromVerKeyHashOrFile) drepKeys
+
+  drepState <- runQuery localNodeConnInfo target $ queryDRepState eon $ Set.fromList drepCreds
+
+  drepStakeDistribution <-
+    case includeStake of
+      Cmd.WithStake -> runQuery localNodeConnInfo target $
+        queryDRepStakeDistribution eon (Set.fromList $ Ledger.DRepCredential <$> drepCreds)
+      Cmd.NoStake -> return mempty
 
-  drepState <- runQuery localNodeConnInfo target $ queryDRepState eon drepCreds
   writeOutput mOutFile $
-    second drepStateToJson <$> Map.assocs drepState
+    drepStateToJson drepStakeDistribution <$> Map.assocs drepState
   where
-    drepStateToJson ds = A.object
-      [ "expiry" .= (ds ^. Ledger.drepExpiryL)
-      , "anchor" .= (ds ^. Ledger.drepAnchorL)
-      , "deposit" .= (ds ^. Ledger.drepDepositL)
-      ]
+    drepStateToJson stakeDistr (cred, ds) = (cred,) . A.object $
+      if Map.null stakeDistr
+      then
+        [ "expiry" .= (ds ^. Ledger.drepExpiryL)
+        , "anchor" .= (ds ^. Ledger.drepAnchorL)
+        , "deposit" .= (ds ^. Ledger.drepDepositL)
+        ]
+      else
+        [ "expiry" .= (ds ^. Ledger.drepExpiryL)
+        , "anchor" .= (ds ^. Ledger.drepAnchorL)
+        , "deposit" .= (ds ^. Ledger.drepDepositL)
+        , "stake" .= Map.lookup (Ledger.DRepCredential cred) stakeDistr
+        ]
 
 runQueryDRepStakeDistribution
   :: Cmd.QueryDRepStakeDistributionCmdArgs era

diff --git a/cardano-cli/test/cardano-cli-golden/files/golden/help.cli b/cardano-cli/test/cardano-cli-golden/files/golden/help.cli
@@ -6743,6 +6743,7 @@ Usage: cardano-cli conway query drep-state --socket-path SOCKET_PATH
                                              | --drep-key-hash HASH
                                              )
                                              )
+                                             [--include-stake]
                                              [--volatile-tip | --immutable-tip]
                                              [--out-file FILE]
 

diff --git a/cardano-cli/test/cardano-cli-golden/files/golden/help/conway_query_drep-state.cli b/cardano-cli/test/cardano-cli-golden/files/golden/help/conway_query_drep-state.cli
@@ -11,6 +11,7 @@ Usage: cardano-cli conway query drep-state --socket-path SOCKET_PATH
                                              | --drep-key-hash HASH
                                              )
                                              )
+                                             [--include-stake]
                                              [--volatile-tip | --immutable-tip]
                                              [--out-file FILE]
 
@@ -37,6 +38,11 @@ Available options:
                            Filepath of the DRep verification key.
   --drep-key-hash HASH     DRep verification key hash (either Bech32-encoded or
                            hex-encoded).
+  --include-stake          Also return the stake associated with each DRep. The
+                           result is the same as with "drep-stake-distribution";
+                           this is a convenience option to obtain all
+                           information concerning a DRep at once. This is a
+                           potentially expensive query.
   --volatile-tip           Use the volatile tip as a target. (This is the
                            default)
   --immutable-tip          Use the immutable tip as a target.