Skip to content

Commit

Permalink
Support watch for get
Browse files Browse the repository at this point in the history
Signed-off-by: Shiming Zhang <[email protected]>
  • Loading branch information
wzshiming committed Aug 13, 2024
1 parent 9d42b38 commit a2a05ca
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 6 deletions.
34 changes: 28 additions & 6 deletions augerctl/command/get_command.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ type getFlagpole struct {
Output string
ChunkSize int64
Prefix string

Watch bool
WatchOnly bool
}

var (
Expand Down Expand Up @@ -89,6 +92,8 @@ func newCtlGetCommand(f *flagpole) *cobra.Command {
cmd.Flags().Int64Var(&flags.ChunkSize, "chunk-size", 500, "chunk size of the list pager")
cmd.Flags().StringVar(&flags.Prefix, "prefix", "/registry", "prefix to prepend to the resource")

cmd.Flags().BoolVarP(&flags.Watch, "watch", "w", false, "after listing/getting the requested object, watch for changes")
cmd.Flags().BoolVar(&flags.WatchOnly, "watch-only", false, "watch for changes to the requested object(s), without listing/getting first")
return cmd
}

Expand Down Expand Up @@ -124,13 +129,30 @@ func getCommand(ctx context.Context, etcdclient client.Client, flags *getFlagpol
client.WithResponse(printer.Print),
}

// TODO: Support watch
if flags.Watch {
if !flags.WatchOnly {
rev, err := etcdclient.Get(ctx, flags.Prefix,
opOpts...,
)
if err != nil {
return err
}
opOpts = append(opOpts, client.WithRevision(rev))
}

_, err := etcdclient.Get(ctx, flags.Prefix,
opOpts...,
)
if err != nil {
return err
err := etcdclient.Watch(ctx, flags.Prefix,
opOpts...,
)
if err != nil {
return err
}
} else {
_, err := etcdclient.Get(ctx, flags.Prefix,
opOpts...,
)
if err != nil {
return err
}
}

return nil
Expand Down
3 changes: 3 additions & 0 deletions augerctl/command/printer.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@ func NewPrinter(w io.Writer, printerType string) Printer {

func formatResponse(w io.Writer, outMediaType string, kv *client.KeyValue) error {
value := kv.Value
if kv.PrevValue != nil {
value = kv.PrevValue
}
inMediaType, _, err := encoding.DetectAndExtract(value)
if err != nil {
_, err0 := fmt.Fprintf(w, "---\n# %s | raw | %v\n# %s\n", kv.Key, err, value)
Expand Down
13 changes: 13 additions & 0 deletions pkg/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ type Client interface {
// Get is a method that retrieves a key-value pair from the etcd server.
// It returns the revision of the key-value pair
Get(ctx context.Context, prefix string, opOpts ...OpOption) (rev int64, err error)

// Watch is a method that watches for changes to a key-value pair on the etcd server.
Watch(ctx context.Context, prefix string, opOpts ...OpOption) error
}

// client is the etcd client.
Expand Down Expand Up @@ -89,6 +92,13 @@ func WithChunkSize(chunkSize int64) OpOption {
}
}

// WithRevision sets the revision for the target.
func WithRevision(revision int64) OpOption {
return func(o *op) {
o.revision = revision
}
}

func opOption(opts []OpOption) op {
var opt op
for _, o := range opts {
Expand All @@ -101,6 +111,9 @@ func opOption(opts []OpOption) op {
type KeyValue struct {
Key []byte
Value []byte

// For delete event
PrevValue []byte
}

func iterateList(kvs []*mvccpb.KeyValue, callback func(kv *KeyValue) error) error {
Expand Down
66 changes: 66 additions & 0 deletions pkg/client/client_watch.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
Copyright 2024 The Kubernetes Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package client

import (
"context"
"fmt"

clientv3 "go.etcd.io/etcd/client/v3"
)

func (c *client) Watch(ctx context.Context, prefix string, opOpts ...OpOption) error {
opt := opOption(opOpts)
if opt.response == nil {
return fmt.Errorf("response is required")
}

path, single, err := getPrefix(prefix, opt.gr, opt.name, opt.namespace)
if err != nil {
return err
}

opts := make([]clientv3.OpOption, 0, 3)

if !single {
opts = append(opts, clientv3.WithPrefix())
}

if opt.revision != 0 {
opts = append(opts, clientv3.WithRev(opt.revision))
}

opts = append(opts, clientv3.WithPrevKV())

watchChan := c.client.Watch(ctx, path, opts...)
for watchResp := range watchChan {
for _, event := range watchResp.Events {
r := &KeyValue{
Key: event.Kv.Key,
Value: event.Kv.Value,
}
if event.PrevKv != nil {
r.PrevValue = event.PrevKv.Value
}
err := opt.response(r)
if err != nil {
return err
}
}
}
return nil
}

0 comments on commit a2a05ca

Please sign in to comment.