Multiple pools support in query stake snapshot

cardano-api/src/Cardano/Api/Orphans.hs

@@ -19,7 +19,7 @@ module Cardano.Api.Orphans () where
 
 import           Prelude
 
-import           Data.Aeson (FromJSON (..), ToJSON (..), object, (.=))
+import           Data.Aeson (FromJSON (..), ToJSON (..), object, pairs, (.=))
 import qualified Data.Aeson as Aeson
 import           Data.Aeson.Types (ToJSONKey (..), toJSONKeyText)
 import           Data.BiMap (BiMap (..), Bimap)
@@ -70,6 +70,7 @@ import qualified Cardano.Ledger.Shelley.Rewards as Shelley
 import qualified Cardano.Ledger.Shelley.RewardUpdate as Shelley
 import           Cardano.Ledger.Val (Val)
 import qualified Ouroboros.Consensus.Shelley.Eras as Consensus
+import qualified Ouroboros.Consensus.Shelley.Ledger.Query as Consensus
 
 -- Orphan instances involved in the JSON output of the API queries.
 -- We will remove/replace these as we provide more API wrapper types
@@ -687,3 +688,37 @@ instance Crypto.Crypto crypto => ToJSON (VMap VB VB (Shelley.KeyHash    'Shelley
 instance Crypto.Crypto crypto => ToJSON (VMap VB VP (Shelley.Credential 'Shelley.Staking   crypto) (Shelley.CompactForm Shelley.Coin)) where
   toJSON = toJSON . fmap fromCompact . VMap.toMap
   toEncoding = toEncoding . fmap fromCompact . VMap.toMap
+
+instance Crypto.Crypto crypto => ToJSON (Consensus.StakeSnapshots crypto) where
+  toJSON = object . stakeSnapshotsToPair
+  toEncoding = pairs . mconcat . stakeSnapshotsToPair
+
+stakeSnapshotsToPair :: (Aeson.KeyValue a, Crypto.Crypto crypto) => Consensus.StakeSnapshots crypto -> [a]
+stakeSnapshotsToPair Consensus.StakeSnapshots
+    { Consensus.ssStakeSnapshots
+    , Consensus.ssMarkTotal
+    , Consensus.ssSetTotal
+    , Consensus.ssGoTotal
+    } =
+    [ "pools" .= ssStakeSnapshots
+    , "total" .= object
+      [ "stakeMark" .= ssMarkTotal
+      , "stakeSet" .= ssSetTotal
+      , "stakeGo" .= ssGoTotal
+      ]
+    ]
+
+instance ToJSON (Consensus.StakeSnapshot crypto) where
+  toJSON = object . stakeSnapshotToPair
+  toEncoding = pairs . mconcat . stakeSnapshotToPair
+
+stakeSnapshotToPair :: Aeson.KeyValue a => Consensus.StakeSnapshot crypto -> [a]
+stakeSnapshotToPair Consensus.StakeSnapshot
+    { Consensus.ssMarkPool
+    , Consensus.ssSetPool
+    , Consensus.ssGoPool
+    } =
+    [ "stakeMark" .= ssMarkPool
+    , "stakeSet" .= ssSetPool
+    , "stakeGo" .= ssGoPool
+    ]

cardano-api/src/Cardano/Api/Query.hs

@@ -260,7 +260,7 @@ data QueryInShelleyBasedEra era result where
     -> QueryInShelleyBasedEra era (SerialisedPoolDistribution era)
 
   QueryStakeSnapshot
-    :: PoolId
+    :: Maybe (Set PoolId)
     -> QueryInShelleyBasedEra era (SerialisedStakeSnapshots era)
 
 deriving instance Show (QueryInShelleyBasedEra era result)
@@ -631,8 +631,8 @@ toConsensusQueryShelleyBased erainmode QueryCurrentEpochState =
 toConsensusQueryShelleyBased erainmode (QueryPoolState poolIds) =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolState (Set.map unStakePoolKeyHash <$> poolIds))))
 
-toConsensusQueryShelleyBased erainmode (QueryStakeSnapshot poolId) =
-    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetStakeSnapshots (Just (Set.singleton (unStakePoolKeyHash poolId))))))
+toConsensusQueryShelleyBased erainmode (QueryStakeSnapshot mPoolIds) =
+    Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetStakeSnapshots (fmap (Set.map unStakePoolKeyHash) mPoolIds))))
 
 toConsensusQueryShelleyBased erainmode (QueryPoolDistribution poolIds) =
     Some (consensusQueryInEraInMode erainmode (Consensus.GetCBOR (Consensus.GetPoolDistr (getPoolIds <$> poolIds))))

cardano-cli/src/Cardano/CLI/Shelley/Commands.hs

@@ -361,7 +361,11 @@ data QueryCmd =
   | QueryUTxO' AnyConsensusModeParams QueryUTxOFilter NetworkId (Maybe OutputFile)
   | QueryDebugLedgerState' AnyConsensusModeParams NetworkId (Maybe OutputFile)
   | QueryProtocolState' AnyConsensusModeParams NetworkId (Maybe OutputFile)
-  | QueryStakeSnapshot' AnyConsensusModeParams NetworkId (Hash StakePoolKey)
+  | QueryStakeSnapshot'
+      AnyConsensusModeParams
+      NetworkId
+      (AllOrOnly [Hash StakePoolKey])
+      (Maybe OutputFile)
   | QueryKesPeriodInfo
       AnyConsensusModeParams
       NetworkId
@@ -510,8 +514,9 @@ data MetadataFile = MetadataFileJSON FilePath
 
   deriving Show
 
-newtype OutputFile
-  = OutputFile FilePath
+newtype OutputFile = OutputFile
+  { unOutputFile :: FilePath
+  }
   deriving Show
 
 newtype PoolMetadataFile = PoolMetadataFile

cardano-cli/src/Cardano/CLI/Shelley/Parsers.hs

@@ -990,11 +990,22 @@ pQueryCmd =
                             <*> pNetworkId
                             <*> pMaybeOutputFile
 
+    pAllStakePoolsOrOnly :: Parser (AllOrOnly [Hash StakePoolKey])
+    pAllStakePoolsOrOnly = pAll <|> pOnly
+      where pAll :: Parser (AllOrOnly [Hash StakePoolKey])
+            pAll = Opt.flag' All
+              (  Opt.long "all-stake-pools"
+              <> Opt.help "Query for all stake pools"
+              )
+            pOnly :: Parser (AllOrOnly [Hash StakePoolKey])
+            pOnly = Only <$> many pStakePoolVerificationKeyHash
+
     pQueryStakeSnapshot :: Parser QueryCmd
     pQueryStakeSnapshot = QueryStakeSnapshot'
       <$> pConsensusModeParams
       <*> pNetworkId
-      <*> pStakePoolVerificationKeyHash
+      <*> pAllStakePoolsOrOnly
+      <*> pMaybeOutputFile
 
     pQueryPoolState :: Parser QueryCmd
     pQueryPoolState = QueryPoolState'
@@ -2572,10 +2583,12 @@ pStakePoolVerificationKeyHash :: Parser (Hash StakePoolKey)
 pStakePoolVerificationKeyHash =
     Opt.option
       (pBech32StakePoolId <|> pHexStakePoolId)
-        (  Opt.long "stake-pool-id"
-        <> Opt.metavar "STAKE-POOL-ID"
-        <> Opt.help "Stake pool ID/verification key hash (either \
-                    \Bech32-encoded or hex-encoded)."
+        (   Opt.long "stake-pool-id"
+        <>  Opt.metavar "STAKE_POOL_ID"
+        <>  Opt.help
+            (   "Stake pool ID/verification key hash (either Bech32-encoded or hex-encoded).  "
+            <>  "Zero or more occurences of this option is allowed."
+            )
         )
   where
     pHexStakePoolId :: ReadM (Hash StakePoolKey)

cardano-cli/src/Cardano/CLI/Shelley/Run/Query.hs

@@ -25,7 +25,7 @@ module Cardano.CLI.Shelley.Run.Query
   , executeQuery
   ) where
 
-import           Cardano.Prelude
+import           Cardano.Prelude hiding (All)
 import           Prelude (String, id)
 
 import           Cardano.Api
@@ -95,7 +95,6 @@ import           Ouroboros.Network.Block (Serialised (..))
 import qualified Ouroboros.Consensus.HardFork.History as Consensus
 import qualified Ouroboros.Consensus.Protocol.Abstract as Consensus
 import qualified Ouroboros.Consensus.Protocol.Praos.Common as Consensus
-import qualified Ouroboros.Consensus.Shelley.Ledger.Query as Consensus
 
 import qualified Ouroboros.Consensus.HardFork.History.Qry as Qry
 import qualified Ouroboros.Network.Protocol.LocalStateQuery.Type as LocalStateQuery
@@ -187,8 +186,8 @@ runQueryCmd cmd =
       runQueryStakeAddressInfo consensusModeParams addr network mOutFile
     QueryDebugLedgerState' consensusModeParams network mOutFile ->
       runQueryLedgerState consensusModeParams network mOutFile
-    QueryStakeSnapshot' consensusModeParams network poolid ->
-      runQueryStakeSnapshot consensusModeParams network poolid
+    QueryStakeSnapshot' consensusModeParams network allOrOnlyPoolIds mOutFile ->
+      runQueryStakeSnapshot consensusModeParams network allOrOnlyPoolIds mOutFile
     QueryProtocolState' consensusModeParams network mOutFile ->
       runQueryProtocolState consensusModeParams network mOutFile
     QueryUTxO' consensusModeParams qFilter networkId mOutFile ->
@@ -671,9 +670,10 @@ runQueryTxMempool (AnyConsensusModeParams cModeParams) network query mOutFile =
 runQueryStakeSnapshot
   :: AnyConsensusModeParams
   -> NetworkId
-  -> Hash StakePoolKey
+  -> AllOrOnly [Hash StakePoolKey]
+  -> Maybe OutputFile
   -> ExceptT ShelleyQueryCmdError IO ()
-runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolId = do
+runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network allOrOnlyPoolIds mOutFile = do
   SocketPath sockPath <- firstExceptT ShelleyQueryCmdEnvVarSocketErr $ newExceptT readEnvSocketPath
   let localNodeConnInfo = LocalNodeConnectInfo cModeParams network sockPath
 
@@ -687,9 +687,12 @@ runQueryStakeSnapshot (AnyConsensusModeParams cModeParams) network poolId = do
   eInMode <- toEraInMode era cMode
     & hoistMaybe (ShelleyQueryCmdEraConsensusModeMismatch (AnyConsensusMode cMode) anyE)
 
-  let qInMode = QueryInEra eInMode . QueryInShelleyBasedEra sbe $ QueryStakeSnapshot poolId
+  let qInMode = QueryInEra eInMode . QueryInShelleyBasedEra sbe $ QueryStakeSnapshot $ case allOrOnlyPoolIds of
+        All -> Nothing
+        Only poolIds -> Just $ Set.fromList poolIds
+
   result <- executeQuery era cModeParams localNodeConnInfo qInMode
-  obtainLedgerEraClassConstraints sbe writeStakeSnapshot result
+  obtainLedgerEraClassConstraints sbe (writeStakeSnapshots mOutFile) result
 
 
 runQueryLedgerState
@@ -850,30 +853,19 @@ writeLedgerState mOutFile qState@(SerialisedDebugLedgerState serLedgerState) =
       handleIOExceptT (ShelleyQueryCmdWriteFileError . FileIOError fpath)
         $ LBS.writeFile fpath $ unSerialised serLedgerState
 
-writeStakeSnapshot :: forall era ledgerera. ()
+writeStakeSnapshots :: forall era ledgerera. ()
   => ShelleyLedgerEra era ~ ledgerera
   => Era.Crypto ledgerera ~ StandardCrypto
-  => SerialisedStakeSnapshots era
+  => Maybe OutputFile
+  -> SerialisedStakeSnapshots era
   -> ExceptT ShelleyQueryCmdError IO ()
-writeStakeSnapshot qState =
+writeStakeSnapshots mOutFile qState =
   case decodeStakeSnapshot qState of
     Left err -> left (ShelleyQueryCmdStakeSnapshotDecodeError err)
 
     Right (StakeSnapshot snapshot) -> do
       -- Calculate the three pool and active stake values for the given pool
-      liftIO . LBS.putStrLn $ encodePretty $ Aeson.object $
-        [ "activeStakeMark" .= Consensus.ssMarkTotal snapshot
-        , "activeStakeSet"  .= Consensus.ssSetTotal snapshot
-        , "activeStakeGo"   .= Consensus.ssGoTotal snapshot
-        ] <> poolFields snapshot
-  where poolFields :: Consensus.StakeSnapshots (Ledger.Crypto (ShelleyLedgerEra era)) -> [Aeson.Pair]
-        poolFields snapshot = case Map.elems (Consensus.ssStakeSnapshots snapshot) of
-          [pool] -> 
-            [ "poolStakeMark"   .= Consensus.ssMarkPool pool
-            , "poolStakeSet"    .= Consensus.ssSetPool pool
-            , "poolStakeGo"     .= Consensus.ssGoPool pool
-            ]
-          _ -> []
+      liftIO . (maybe LBS.putStrLn (LBS.writeFile . unOutputFile) mOutFile) $ encodePretty snapshot
 
 -- | This function obtains the pool parameters, equivalent to the following jq query on the output of query ledger-state
 --   .nesEs.esLState._delegationState._pstate._pParams.<pool_id>

cardano-cli/src/Cardano/CLI/Types.hs

@@ -38,6 +38,7 @@ module Cardano.CLI.Types
   , VerificationKeyFile (..)
   , Params (..)
   , RequiredSigner (..)
+  , AllOrOnly(..)
   ) where
 
 import           Cardano.Prelude hiding (Word64)
@@ -181,6 +182,8 @@ data OutputFormat
   | OutputFormatBech32
   deriving (Eq, Show)
 
+data AllOrOnly a = All | Only a deriving (Eq, Show)
+
 -- | This data structure is used to allow nicely formatted output in the query pool-params command.
 -- params are the current pool parameter settings, futureparams are new parameters, retiringEpoch is the
 -- epoch that has been set for pool retirement.  Any of these may be Nothing.

doc/stake-pool-operations/10_query_stakepool.md

@@ -16,21 +16,27 @@ $ cardano-cli query stake-snapshot \
     --stake-pool-id 00beef0a9be2f6d897ed24a613cf547bb20cd282a04edfc53d477114 \
     --mainnet
 {
-    "poolStakeGo": 40278547538358,
-    "activeStakeGo": 22753958467474959,
-    "poolStakeMark": 40424218559492,
-    "activeStakeMark": 22670949084364797,
-    "poolStakeSet": 39898761956772,
-    "activeStakeSet": 22488877070796904
+  "pools": {
+    "00beef0a9be2f6d897ed24a613cf547bb20cd282a04edfc53d477114": {
+      "stakeGo": 40278547538358,
+      "stakeMark": 40424218559492,
+      "stakeSet": 39898761956772
+    }
+  },
+  "total": {
+    "stakeGo": 22753958467474959,
+    "stakeMark": 22670949084364797,
+    "stakeSet": 22488877070796904
+  }
 }
 ```
 
 Each snapshot is taken at the end of a different era.  The `go` snapshot is the current one and
 was taken two epochs earlier, `set` was taken one epoch ago, and `mark` was taken immediately
 before the start of the current epoch.
 
-This command is for debugging purposes only and may fail when used in a memory constrained
-environment due to the size of the ledger state.
+The command accepts zero or more occurences of the `--stake-pool-id` option.  Alternatively, to query
+all pools supply the `--all-stake-pools` option instead.
 
 # Querying for pool parameters