From a04d2b196e0824ed2512f6a33b426affaeb4310c Mon Sep 17 00:00:00 2001
From: Chris Marslender <chrismarslender@gmail.com>
Date: Thu, 3 Oct 2024 15:53:02 -0500
Subject: [PATCH] Add a bulk subscribe option to chia-data

---
 cmd/datalayer/bulksub.go | 54 ++++++++++++++++++++++++++++++++++++++++
 cmd/network/switch.go    |  1 -
 cmd/root.go              |  4 +--
 go.mod                   |  2 +-
 go.sum                   |  4 +--
 5 files changed, 59 insertions(+), 6 deletions(-)
 create mode 100644 cmd/datalayer/bulksub.go

diff --git a/cmd/datalayer/bulksub.go b/cmd/datalayer/bulksub.go
new file mode 100644
index 0000000..69da0d9
--- /dev/null
+++ b/cmd/datalayer/bulksub.go
@@ -0,0 +1,54 @@
+package datalayer
+
+import (
+	"encoding/json"
+	"os"
+
+	"github.com/chia-network/go-chia-libs/pkg/rpc"
+	"github.com/chia-network/go-modules/pkg/slogs"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+// bulkSubCmd Subscribes to multiple datastores at once using the output of chia data get_subscriptions
+var bulkSubCmd = &cobra.Command{
+	Use:     "bulk-subscribe",
+	Short:   "Subscribes to multiple datastores at once using the output of chia data get_subscriptions",
+	Example: "chia-tools data bulk-subscribe -f subscriptions.json",
+	//Args:    cobra.ExactArgs(1),
+	Run: func(cmd *cobra.Command, args []string) {
+		client, err := rpc.NewClient(rpc.ConnectionModeHTTP, rpc.WithAutoConfig())
+		if err != nil {
+			slogs.Logr.Fatal("error creating chia RPC client", "error", err)
+		}
+
+		jsonFile := viper.GetString("bulksub-file")
+		content, err := os.ReadFile(jsonFile)
+		if err != nil {
+			slogs.Logr.Fatal("Unable to read input file", "error", err)
+		}
+
+		subs := &rpc.DatalayerGetSubscriptionsResponse{}
+		err = json.Unmarshal(content, subs)
+		if err != nil {
+			slogs.Logr.Fatal("Could not parse the subscriptions json file", "error", err)
+		}
+
+		for _, id := range subs.StoreIDs {
+			slogs.Logr.Info("Subscription to store", "id", id)
+			_, _, err = client.DataLayerService.Subscribe(&rpc.DatalayerSubscribeOptions{
+				ID: id,
+			})
+			if err != nil {
+				slogs.Logr.Error("Error subscribing to datastore", "id", id, "error", err)
+			}
+		}
+	},
+}
+
+func init() {
+	bulkSubCmd.PersistentFlags().StringP("file", "f", "", "The file containing the json of subscriptions to add")
+	cobra.CheckErr(viper.BindPFlag("bulksub-file", bulkSubCmd.PersistentFlags().Lookup("file")))
+
+	datalayerCmd.AddCommand(bulkSubCmd)
+}
diff --git a/cmd/network/switch.go b/cmd/network/switch.go
index 815044d..8aaf105 100644
--- a/cmd/network/switch.go
+++ b/cmd/network/switch.go
@@ -51,7 +51,6 @@ var switchCmd = &cobra.Command{
 			slogs.Logr.Fatal("selected network does not exist in config's network override config", "network", networkName)
 		}
 
-
 		// Ensure a folder to store the current network's sub-epoch-summaries and height-to-hash files exists
 		cacheFileDirOldNetwork := path.Join(chiaRoot, "db", currentNetwork)
 		cacheFileDirNewNetwork := path.Join(chiaRoot, "db", networkName)
diff --git a/cmd/root.go b/cmd/root.go
index 2bac90d..4eef402 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -19,8 +19,8 @@ var (
 
 // RootCmd represents the base command when called without any subcommands
 var RootCmd = &cobra.Command{
-	Use:   "chia-tools",
-	Short: "Collection of CLI tools for working with Chia Blockchain",
+	Use:     "chia-tools",
+	Short:   "Collection of CLI tools for working with Chia Blockchain",
 	Version: fmt.Sprintf("%s (%s)", gitVersion, buildTime),
 }
 
diff --git a/go.mod b/go.mod
index 6f8fcad..3d56b70 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/chia-network/chia-tools
 go 1.22.4
 
 require (
-	github.com/chia-network/go-chia-libs v0.15.5
+	github.com/chia-network/go-chia-libs v0.15.6
 	github.com/chia-network/go-modules v0.0.7
 	github.com/spf13/cast v1.7.0
 	github.com/spf13/cobra v1.8.1
diff --git a/go.sum b/go.sum
index 981aa8c..641927d 100644
--- a/go.sum
+++ b/go.sum
@@ -1,5 +1,5 @@
-github.com/chia-network/go-chia-libs v0.15.5 h1:HsBhOnxHuriOLVqJiQQu+8egirQxjA2Qt2Iz6nS7K4A=
-github.com/chia-network/go-chia-libs v0.15.5/go.mod h1:npTqaFSjTdMxE7hc0LOmWJmWGqcs+IERarK5fDxXk/I=
+github.com/chia-network/go-chia-libs v0.15.6 h1:8XxvR9EQKFSnIZtGzCsSlsy4BpWSIJsEHxut/wZR1WU=
+github.com/chia-network/go-chia-libs v0.15.6/go.mod h1:npTqaFSjTdMxE7hc0LOmWJmWGqcs+IERarK5fDxXk/I=
 github.com/chia-network/go-modules v0.0.7 h1:l3EV5rGrE2Rvoxyhbxfk4cZy2Pc/l2hNQ54pmB/+NE4=
 github.com/chia-network/go-modules v0.0.7/go.mod h1:ZoHvOfnjNj3ydtJxXMxBxM0fBXMWiPlj13IlYgpSj8I=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=