zed manage: Add -pool flag (#5164)

Add a -pool flag to zed manage that specifies a pool to run Zed manage
on.

cmd/zed/internal/lakemanage/manage.go

@@ -3,11 +3,15 @@ package lakemanage
 import (
 	"context"
 	"errors"
+	"fmt"
 	"syscall"
 	"time"
 
 	"github.com/brimdata/zed/api/client"
 	lakeapi "github.com/brimdata/zed/lake/api"
+	"github.com/brimdata/zed/lake/pools"
+	"github.com/brimdata/zed/lakeparse"
+	"github.com/segmentio/ksuid"
 	"go.uber.org/zap"
 	"golang.org/x/sync/errgroup"
 )
@@ -66,7 +70,7 @@ func monitor(ctx context.Context, lk lakeapi.Interface, conf Config, logger *zap
 }
 
 func getBranches(ctx context.Context, conf Config, lk lakeapi.Interface, logger *zap.Logger) ([]*branch, error) {
-	pools, err := lakeapi.GetPools(ctx, lk)
+	pools, err := getPools(ctx, conf, lk)
 	if err != nil {
 		return nil, err
 	}
@@ -78,3 +82,37 @@ func getBranches(ctx context.Context, conf Config, lk lakeapi.Interface, logger
 	}
 	return branches, nil
 }
+
+func getPools(ctx context.Context, conf Config, lk lakeapi.Interface) ([]*pools.Config, error) {
+	pls, err := lakeapi.GetPools(ctx, lk)
+	if err != nil {
+		return nil, err
+	}
+	if len(conf.Pools) == 0 {
+		return pls, nil
+	}
+	m := map[ksuid.KSUID]struct{}{}
+	var out []*pools.Config
+	for _, c := range conf.Pools {
+		p := selectPool(c, pls)
+		if p == nil {
+			return nil, fmt.Errorf("pool %q not found", c.Pool)
+		}
+		if _, ok := m[p.ID]; ok {
+			return nil, fmt.Errorf("duplicate pool in configuration: %q", c.Pool)
+		}
+		m[p.ID] = struct{}{}
+		out = append(out, p)
+	}
+	return out, nil
+}
+
+func selectPool(c PoolConfig, pools []*pools.Config) *pools.Config {
+	id, _ := lakeparse.ParseID(c.Pool)
+	for _, p := range pools {
+		if id == p.ID || c.Pool == p.Name {
+			return p
+		}
+	}
+	return nil
+}

cmd/zed/manage/command.go

@@ -64,6 +64,10 @@ func New(parent charm.Command, f *flag.FlagSet) (charm.Command, error) {
 		d.KnownFields(true) // returns error for unknown fields
 		return d.Decode(&c.config)
 	})
+	f.Func("pool", "pool to manage (all if unset, can be specified multiple times)", func(s string) error {
+		c.config.Pools = append(c.config.Pools, lakemanage.PoolConfig{Pool: s, Branch: "main"})
+		return nil
+	})
 	c.config.Interval = f.Duration("interval", lakemanage.DefaultInterval, "interval between updates (only applicable with -monitor")
 	f.BoolVar(&c.monitor, "monitor", false, "continuously monitor the lake for updates")
 	f.BoolVar(&c.config.Vectors, "vectors", false, "create vectors for objects")

cmd/zed/manage/ztests/config.yaml

@@ -28,10 +28,6 @@ outputs:
           name: "test2",
           branch: "live"
       }
-      {
-          name: "test3",
-          branch: "main"
-      }
   - name: stderr
     data: ""
   - name: stdout

cmd/zed/manage/ztests/pool-flag.yaml

@@ -0,0 +1,15 @@
+script: |
+  export ZED_LAKE=test
+  zed init -q
+  zed create -q test1
+  zed create -q test2
+  zed create -q test3
+  zed create -q test4
+  zed manage -pool test2 -pool test4 -log.path=manage.log
+  zq -z 'msg == "updating pool" | yield name' manage.log
+
+outputs:
+  - name: stdout
+    data: |
+      "test2"
+      "test4"