waymaker-client/go/kv.go

package waymaker

// KV subsystem — thin RPC binding over the server's Kv* RPCs.
// All conventions (subject patterns, tombstone marker, TTL header) live
// server-side in waymaker_streams::wire_conventions. This client just
// calls the typed RPCs — no subject-level knowledge required.
//
// Entry points on *Client:
//   - client.CreateKV(ctx, KVConfig{…})
//   - client.GetOrCreateKV(ctx, KVConfig{…})
//   - client.KV(name) — handle without creation
//   - client.DeleteKV(ctx, name)

import (
	"context"
	"io"
	"time"

	pb "git.awesomike.com/pub/waymaker-client/go/genpb/kv"
)

// KVConfig is the bucket creation config.
type KVConfig struct {
	Name         string
	MaxBytes     *uint64
	MaxValueSize *uint64
	MaxAge       *time.Duration
	Ephemeral    bool
	// MaxRevisions caps per-key revision count. 0 = unbounded.
	MaxRevisions uint64
}

// Bucket is a reference to a KV bucket. Cheap to copy.
type Bucket struct {
	client *Client
	Name   string
}

func newBucket(c *Client, name string) *Bucket {
	return &Bucket{client: c, Name: name}
}

// Put stores value under key. Latest-write-wins. Returns the new revision.
func (b *Bucket) Put(ctx context.Context, key string, value []byte) (uint64, error) {
	return b.putInternal(ctx, key, value, 0)
}

// PutWithTTL stores value under key with a per-key TTL.
func (b *Bucket) PutWithTTL(ctx context.Context, key string, value []byte, ttl time.Duration) (uint64, error) {
	return b.putInternal(ctx, key, value, uint64(ttl.Milliseconds()))
}

func (b *Bucket) putInternal(ctx context.Context, key string, value []byte, ttlMs uint64) (uint64, error) {
	c := b.client.kvClient()
	r, err := c.Put(ctx, &pb.KvPutRequest{
		Bucket: b.Name,
		Key:    key,
		Value:  value,
		TtlMs:  ttlMs,
	})
	if err != nil {
		return 0, rpcErr(err)
	}
	if !r.GetSuccess() {
		return 0, serverErr(r.GetResultCode(), r.GetMessage())
	}
	return r.GetRevision(), nil
}

// Create is an atomic create — fails with code "wrong_revision" if the key
// already exists.
func (b *Bucket) Create(ctx context.Context, key string, value []byte) (uint64, error) {
	c := b.client.kvClient()
	r, err := c.Create(ctx, &pb.KvCreateRequest{
		Bucket: b.Name,
		Key:    key,
		Value:  value,
		TtlMs:  0,
	})
	if err != nil {
		return 0, rpcErr(err)
	}
	if !r.GetSuccess() {
		return 0, serverErr(r.GetResultCode(), r.GetMessage())
	}
	return r.GetRevision(), nil
}

// Update is a CAS update — succeeds only if current revision matches
// expectedRevision.
func (b *Bucket) Update(ctx context.Context, key string, value []byte, expectedRevision uint64) (uint64, error) {
	c := b.client.kvClient()
	r, err := c.Update(ctx, &pb.KvUpdateRequest{
		Bucket:           b.Name,
		Key:              key,
		Value:            value,
		ExpectedRevision: expectedRevision,
		TtlMs:            0,
	})
	if err != nil {
		return 0, rpcErr(err)
	}
	if !r.GetSuccess() {
		return 0, serverErr(r.GetResultCode(), r.GetMessage())
	}
	return r.GetRevision(), nil
}

// Get returns the latest value. Returns (nil, nil) when absent or
// tombstoned.
func (b *Bucket) Get(ctx context.Context, key string) ([]byte, error) {
	v, _, err := b.GetWithRevision(ctx, key)
	return v, err
}

// GetWithRevision returns the latest value + revision (for chaining CAS).
// Returns (nil, 0, nil) when absent or tombstoned.
func (b *Bucket) GetWithRevision(ctx context.Context, key string) ([]byte, uint64, error) {
	c := b.client.kvClient()
	r, err := c.Get(ctx, &pb.KvGetRequest{Bucket: b.Name, Key: key})
	if err != nil {
		return nil, 0, rpcErr(err)
	}
	if !r.GetSuccess() {
		return nil, 0, serverErr(r.GetResultCode(), r.GetMessage())
	}
	e := r.GetEntry()
	if e == nil {
		return nil, 0, nil
	}
	return e.GetValue(), e.GetRevision(), nil
}

// Delete tombstones key.
func (b *Bucket) Delete(ctx context.Context, key string) error {
	c := b.client.kvClient()
	r, err := c.Delete(ctx, &pb.KvDeleteRequest{Bucket: b.Name, Key: key})
	if err != nil {
		return rpcErr(err)
	}
	if !r.GetSuccess() {
		return serverErr(r.GetResultCode(), r.GetMessage())
	}
	return nil
}

// Touch extends the TTL on key without changing its value.
func (b *Bucket) Touch(ctx context.Context, key string, ttl time.Duration) (uint64, error) {
	c := b.client.kvClient()
	r, err := c.Touch(ctx, &pb.KvTouchRequest{
		Bucket: b.Name,
		Key:    key,
		TtlMs:  uint64(ttl.Milliseconds()),
	})
	if err != nil {
		return 0, rpcErr(err)
	}
	if !r.GetSuccess() {
		return 0, serverErr(r.GetResultCode(), r.GetMessage())
	}
	return r.GetRevision(), nil
}

// Keys lists every non-tombstoned key in the bucket.
func (b *Bucket) Keys(ctx context.Context) ([]string, error) {
	c := b.client.kvClient()
	r, err := c.Keys(ctx, &pb.KvKeysRequest{Bucket: b.Name})
	if err != nil {
		return nil, rpcErr(err)
	}
	if !r.GetSuccess() {
		return nil, serverErr(r.GetResultCode(), r.GetMessage())
	}
	var out []string
	for _, e := range r.GetEntries() {
		if !e.GetDeleted() {
			out = append(out, e.GetKey())
		}
	}
	return out, nil
}

// HistoryEntry is one revision of a key.
type HistoryEntry struct {
	Value     []byte
	Revision  uint64
	TsMs      int64
	Tombstone bool
}

// History returns historical values at key in publish order.
func (b *Bucket) History(ctx context.Context, key string) ([]HistoryEntry, error) {
	c := b.client.kvClient()
	r, err := c.History(ctx, &pb.KvHistoryRequest{
		Bucket:       b.Name,
		Key:          key,
		FromRevision: 0,
		Limit:        0,
	})
	if err != nil {
		return nil, rpcErr(err)
	}
	if !r.GetSuccess() {
		return nil, serverErr(r.GetResultCode(), r.GetMessage())
	}
	out := make([]HistoryEntry, len(r.GetEntries()))
	for i, e := range r.GetEntries() {
		out[i] = HistoryEntry{
			Value:     e.GetValue(),
			Revision:  e.GetRevision(),
			TsMs:      e.GetTsMs(),
			Tombstone: e.GetDeleted(),
		}
	}
	return out, nil
}

// KVEvent is a live change event delivered by Watch.
type KVEvent struct {
	// Put is non-nil for a put event.
	Put    *KVPutEvent
	// Delete is non-nil for a delete event.
	Delete *KVDeleteEvent
}

// KVPutEvent carries a put notification.
type KVPutEvent struct {
	Key      string
	Value    []byte
	Revision uint64
	TsMs     int64
}

// KVDeleteEvent carries a delete/tombstone notification.
type KVDeleteEvent struct {
	Key      string
	Revision uint64
	TsMs     int64
}

// WatchStream is a live watch on a KV bucket.
type WatchStream struct {
	inner pb.WaymakerKvService_WatchClient
}

// Next blocks until the next event arrives or an error occurs.
// Returns (zero, nil) on normal end-of-stream.
func (w *WatchStream) Next() (KVEvent, error) {
	for {
		ev, err := w.inner.Recv()
		if err != nil {
			if err == io.EOF {
				return KVEvent{}, nil
			}
			return KVEvent{}, rpcErr(err)
		}
		switch e := ev.GetEvent().(type) {
		case *pb.KvWatchEvent_Put:
			return KVEvent{Put: &KVPutEvent{
				Key:      e.Put.GetKey(),
				Value:    e.Put.GetValue(),
				Revision: e.Put.GetRevision(),
				TsMs:     e.Put.GetTsMs(),
			}}, nil
		case *pb.KvWatchEvent_Delete:
			return KVEvent{Delete: &KVDeleteEvent{
				Key:      e.Delete.GetKey(),
				Revision: e.Delete.GetRevision(),
				TsMs:     e.Delete.GetTsMs(),
			}}, nil
		}
	}
}

// Watch opens a live watch on key.
func (b *Bucket) Watch(ctx context.Context, key string) (*WatchStream, error) {
	return b.watchInner(ctx, key)
}

// WatchAll opens a live watch on every key in the bucket.
func (b *Bucket) WatchAll(ctx context.Context) (*WatchStream, error) {
	return b.watchInner(ctx, "")
}

func (b *Bucket) watchInner(ctx context.Context, key string) (*WatchStream, error) {
	c := b.client.kvClient()
	stream, err := c.Watch(ctx, &pb.KvWatchRequest{Bucket: b.Name, Key: key})
	if err != nil {
		return nil, rpcErr(err)
	}
	return &WatchStream{inner: stream}, nil
}

// --- Client entry points ---

// CreateKV creates a new KV bucket.
func (c *Client) CreateKV(ctx context.Context, config KVConfig) (*Bucket, error) {
	kvc := c.kvClient()
	var maxAgeMs uint64
	if config.MaxAge != nil {
		maxAgeMs = uint64(config.MaxAge.Milliseconds())
	}
	r, err := kvc.CreateBucket(ctx, &pb.KvCreateBucketRequest{
		Bucket:       config.Name,
		MaxBytes:     uint64OrZero(config.MaxBytes),
		MaxValueSize: uint64OrZero(config.MaxValueSize),
		MaxAgeMs:     maxAgeMs,
		Ephemeral:    config.Ephemeral,
		MaxRevisions: config.MaxRevisions,
	})
	if err != nil {
		return nil, rpcErr(err)
	}
	if !r.GetSuccess() {
		return nil, serverErr(r.GetResultCode(), r.GetMessage())
	}
	return newBucket(c, config.Name), nil
}

// GetOrCreateKV is the idempotent create-or-get.
func (c *Client) GetOrCreateKV(ctx context.Context, config KVConfig) (*Bucket, error) {
	b, err := c.CreateKV(ctx, config)
	if err == nil {
		return b, nil
	}
	if IsServerCode(err, "already_exists") {
		return newBucket(c, config.Name), nil
	}
	return nil, err
}

// KV returns a bucket handle without verifying existence.
func (c *Client) KV(name string) *Bucket {
	return newBucket(c, name)
}

// DeleteKV deletes the KV bucket.
func (c *Client) DeleteKV(ctx context.Context, name string) error {
	kvc := c.kvClient()
	r, err := kvc.DeleteBucket(ctx, &pb.KvDeleteBucketRequest{Bucket: name})
	if err != nil {
		return rpcErr(err)
	}
	if !r.GetSuccess() {
		return serverErr(r.GetResultCode(), r.GetMessage())
	}
	return nil
}