commit 38731258aced7d53bf87cbe4acaf35f7550fc69f Author: Milas Bowman Date: Tue Feb 25 20:24:09 2025 -0500 feat: initial prototype diff --git a/README.md b/README.md new file mode 100644 index 0000000..fc3a09f --- /dev/null +++ b/README.md @@ -0,0 +1,44 @@ +# garage-k2v-go + +Go client for [K2V][k2v-about], an experimental small object key-value storage engine built as part of [Garage][garage-about]. + +Because the [K2V API][k2v-api] is not stable, breaking changes in this Go module should be expected until then. + +## Import +```go +import k2v "code.notaphish.fyi/milas/garage-k2v-go" +``` + +## Create API client +```go +endpoint := "http://localhost:3904" + +// Read K2V_KEY_ID and K2V_SECRET_KEY from OS environment variables. +key := k2v.KeyFromEnv() +// key := k2v.Key{ID: "GK49e661847883e4813993a0db", Secret: "..."} + +client := k2v.NewClient(endpoint, key) +``` + +## Operations +```go +type Client + func NewClient(endpoint string, key Key, opts ...ClientOption) *Client + func (c *Client) Clone(opts ...ClientOption) *Client + func (c *Client) Close() + func (c *Client) DeleteItem(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken) error + func (c *Client) InsertBatch(ctx context.Context, b Bucket, items []BatchInsertItem) error + func (c *Client) InsertItem(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken, item []byte) error + func (c *Client) PollItem(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken, timeout time.Duration) (Item, CausalityToken, error) + func (c *Client) ReadBatch(ctx context.Context, b Bucket, q []ReadBatchSearch) ([]BatchSearchResult, error) + func (c *Client) ReadIndex(ctx context.Context, b Bucket, q ReadIndexQuery) (*ReadIndexResponse, error) + func (c *Client) ReadItemMulti(ctx context.Context, b Bucket, pk string, sk string) ([]Item, CausalityToken, error) + func (c *Client) ReadItemSingle(ctx context.Context, b Bucket, pk string, sk string) (Item, CausalityToken, error) +``` + +## Usage +Review the [K2V API spec][k2v-api] and the integration tests in this module for complete examples. + +[garage-about]: https://garagehq.deuxfleurs.fr/ +[k2v-about]: https://garagehq.deuxfleurs.fr/documentation/reference-manual/k2v/ +[k2v-api]: https://git.deuxfleurs.fr/Deuxfleurs/garage/src/branch/main/doc/drafts/k2v-spec.md diff --git a/auth.go b/auth.go new file mode 100644 index 0000000..488c258 --- /dev/null +++ b/auth.go @@ -0,0 +1,13 @@ +package k2v + +import "os" + +const EnvVarKeyID = "K2V_KEY_ID" +const EnvVarKeySecret = "K2V_KEY_SECRET" + +func KeyFromEnv() Key { + return Key{ + ID: os.Getenv(EnvVarKeyID), + Secret: os.Getenv(EnvVarKeySecret), + } +} diff --git a/client.go b/client.go new file mode 100644 index 0000000..8b63909 --- /dev/null +++ b/client.go @@ -0,0 +1,384 @@ +package k2v + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "strings" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4" +) + +const CausalityTokenHeader = "X-Garage-Causality-Token" + +var TombstoneItemErr = errors.New("item is a tombstone") +var NoSuchItemErr = errors.New("item does not exist") +var ConcurrentItemsErr = errors.New("item has multiple concurrent values") + +var awsSigner = v4.NewSigner() + +type Bucket string + +type CausalityToken string + +type Item []byte + +func (i Item) GoString() string { + return string(i) +} + +type Key struct { + ID string + Secret string +} + +type Client struct { + key Key + endpoint string + httpClient *http.Client + middleware []RequestMiddleware +} + +type ClientOption func(*Client) + +type RequestMiddleware func(*http.Request) error + +func WithHTTPClient(httpClient *http.Client) ClientOption { + return func(c *Client) { + c.httpClient = httpClient + } +} + +func WithRequestMiddleware(middleware ...RequestMiddleware) ClientOption { + return func(c *Client) { + c.middleware = append(c.middleware, middleware...) + } +} + +func NewClient(endpoint string, key Key, opts ...ClientOption) *Client { + cli := &Client{ + endpoint: endpoint, + key: key, + httpClient: http.DefaultClient, + } + + for _, opt := range opts { + opt(cli) + } + + return cli +} + +func (c *Client) Clone(opts ...ClientOption) *Client { + cli := *c + for _, opt := range opts { + opt(&cli) + } + return &cli +} + +func (c *Client) Close() { + c.httpClient.CloseIdleConnections() +} + +func (c *Client) executeRequest(req *http.Request) (*http.Response, error) { + if err := c.signRequest(req); err != nil { + return nil, err + } + + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + // caller is responsible for closing body + return resp, nil +} + +func (c *Client) signRequest(req *http.Request) error { + creds := aws.Credentials{ + AccessKeyID: c.key.ID, + SecretAccessKey: c.key.Secret, + } + const noBody = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + req.Header.Set("X-Amz-Content-Sha256", noBody) + + err := awsSigner.SignHTTP(req.Context(), creds, req, noBody, "k2v", "garage", time.Now()) + if err != nil { + return err + } + return nil +} + +type ReadIndexQuery struct { + // Prefix restricts listing to partition keys that start with this value. + Prefix string + + // Start is the first partition key to list, in lexicographical order. + Start string + + // End is the last partition key to list (excluded). + End string + + // Limit for maximum number of partition keys to list. + Limit int + + // Reverse iterates in reverse lexicographical order. + Reverse bool +} + +type ReadIndexResponsePartitionKey struct { + PK string `json:"pk"` + Entries int `json:"entries"` + Conflicts int `json:"conflicts"` + Values int `json:"values"` + Bytes int `json:"bytes"` +} + +type ReadIndexResponse struct { + Prefix any `json:"prefix"` + Start any `json:"start"` + End any `json:"end"` + Limit any `json:"limit"` + Reverse bool `json:"reverse"` + PartitionKeys []ReadIndexResponsePartitionKey `json:"partitionKeys"` + More bool `json:"more"` + NextStart any `json:"nextStart"` +} + +func (c *Client) ReadIndex(ctx context.Context, b Bucket, q ReadIndexQuery) (*ReadIndexResponse, error) { + u, err := url.Parse(c.endpoint) + if err != nil { + return nil, err + } + u.Path = string(b) + + query := make(url.Values) + if q.Prefix != "" { + query.Set("prefix", q.Prefix) + } + if q.Start != "" { + query.Set("start", q.Start) + } + if q.End != "" { + query.Set("end", q.End) + } + if q.Limit > 0 { + query.Set("limit", strconv.Itoa(q.Limit)) + } + if q.Reverse { + query.Set("reverse", "true") + } + u.RawQuery = query.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) + if err != nil { + return nil, err + } + + resp, err := c.executeRequest(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + var ret ReadIndexResponse + if err := json.Unmarshal(body, &ret); err != nil { + return nil, err + } + return &ret, nil +} + +func (c *Client) ReadItemMulti(ctx context.Context, b Bucket, pk string, sk string) ([]Item, CausalityToken, error) { + u, err := url.Parse(c.endpoint) + if err != nil { + return nil, "", err + } + u.Path = string(b) + "/" + pk + query := make(url.Values) + query.Set("sort_key", sk) + u.RawQuery = query.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) + if err != nil { + return nil, "", err + } + req.Header.Set("Accept", strings.Join([]string{"application/octet-stream", "application/json"}, ",")) + + resp, err := c.executeRequest(req) + if err != nil { + return nil, "", err + } + defer resp.Body.Close() + + ct := CausalityToken(resp.Header.Get("X-Garage-Causality-Token")) + + switch resp.StatusCode { + case http.StatusOK: + break + case http.StatusNoContent: + return nil, ct, TombstoneItemErr + case http.StatusNotFound: + return nil, "", NoSuchItemErr + default: + return nil, "", fmt.Errorf("http status code %d", resp.StatusCode) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, "", err + } + + switch resp.Header.Get("Content-Type") { + case "application/octet-stream": + // single item, return as-is + return []Item{body}, ct, nil + case "application/json": + var items []Item + if err != nil { + return nil, "", err + } + if err := json.Unmarshal(body, &items); err != nil { + return nil, "", err + } + return items, ct, nil + default: + return nil, "", fmt.Errorf("unsupported content-type: %s", resp.Header.Get("Content-Type")) + } +} + +func (c *Client) ReadItemSingle(ctx context.Context, b Bucket, pk string, sk string) (Item, CausalityToken, error) { + return c.readItemSingle(ctx, b, pk, sk, "", 0) +} + +func (c *Client) readItemSingle(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken, timeout time.Duration) (Item, CausalityToken, error) { + u, err := url.Parse(c.endpoint) + if err != nil { + return nil, "", err + } + u.Path = string(b) + "/" + pk + query := make(url.Values) + query.Set("sort_key", sk) + if ct != "" && timeout > 0 { + query.Set("causality_token", string(ct)) + query.Set("timeout", strconv.Itoa(int(timeout.Seconds()))) + } + u.RawQuery = query.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil) + if err != nil { + return nil, "", err + } + req.Header.Set("Accept", "application/octet-stream") + + resp, err := c.executeRequest(req) + if err != nil { + return nil, "", err + } + defer resp.Body.Close() + + ct = CausalityToken(resp.Header.Get("X-Garage-Causality-Token")) + + switch resp.StatusCode { + case http.StatusOK: + break + case http.StatusNoContent: + return nil, ct, TombstoneItemErr + case http.StatusNotFound: + return nil, "", NoSuchItemErr + case http.StatusConflict: + return nil, ct, ConcurrentItemsErr + default: + return nil, "", fmt.Errorf("http status code %d", resp.StatusCode) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, "", err + } + return body, ct, nil +} + +func (c *Client) DeleteItem(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken) error { + if ct == "" { + return errors.New("continuity token is required for delete") + } + + u, err := url.Parse(c.endpoint) + if err != nil { + return err + } + u.Path = string(b) + "/" + pk + query := make(url.Values) + query.Set("sort_key", sk) + u.RawQuery = query.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodDelete, u.String(), nil) + if err != nil { + return err + } + req.Header.Set(CausalityTokenHeader, string(ct)) + + resp, err := c.executeRequest(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNoContent { + body, err := io.ReadAll(resp.Body) + if err != nil { + return err + } + return fmt.Errorf("http status code %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (c *Client) InsertItem(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken, item []byte) error { + u, err := url.Parse(c.endpoint) + if err != nil { + return err + } + u.Path = string(b) + "/" + pk + query := make(url.Values) + query.Set("sort_key", sk) + u.RawQuery = query.Encode() + + req, err := http.NewRequestWithContext(ctx, http.MethodPut, u.String(), bytes.NewReader(item)) + if err != nil { + return err + } + if ct != "" { + req.Header.Set("X-Garage-Causality-Token", string(ct)) + } + + resp, err := c.executeRequest(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNoContent { + body, err := io.ReadAll(resp.Body) + if err != nil { + return err + } + return fmt.Errorf("http status code %d: %s", resp.StatusCode, string(body)) + } + + return nil +} diff --git a/client_test.go b/client_test.go new file mode 100644 index 0000000..c3d84a3 --- /dev/null +++ b/client_test.go @@ -0,0 +1,202 @@ +package k2v_test + +import ( + "context" + "math/rand/v2" + "net/http/httptrace" + "strconv" + "testing" + "time" + + k2v "code.notaphish.fyi/milas/garage-k2v-go" + "github.com/stretchr/testify/require" + "go.uber.org/goleak" +) + +const endpoint = "http://127.0.0.1:3904" + +type fixture struct { + t testing.TB + ctx context.Context + cli *k2v.Client + bucket k2v.Bucket +} + +func newFixture(t testing.TB) (*fixture, context.Context) { + t.Helper() + + t.Cleanup(func() { + goleak.VerifyNone(t) + }) + + ctx := testContext(t) + + cli := k2v.NewClient(endpoint, k2v.KeyFromEnv()) + t.Cleanup(cli.Close) + + f := &fixture{ + t: t, + ctx: ctx, + cli: cli, + bucket: k2v.Bucket("k2v-test"), + } + + return f, ctx +} + +func testContext(t testing.TB) context.Context { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + return ctx +} + +func randomKey() string { + return "key-" + strconv.Itoa(rand.IntN(1000000)) +} + +func TestClient_InsertItem(t *testing.T) { + f, ctx := newFixture(t) + err := f.cli.InsertItem(ctx, f.bucket, randomKey(), randomKey(), "", []byte("hello")) + require.NoError(t, err) +} + +func TestClient_ReadItemNotExist(t *testing.T) { + f, ctx := newFixture(t) + + pk := randomKey() + sk := randomKey() + + t.Run("Single", func(t *testing.T) { + item, ct, err := f.cli.ReadItemSingle(ctx, f.bucket, pk, sk) + require.ErrorIs(t, err, k2v.NoSuchItemErr) + require.Nil(t, item) + require.Empty(t, ct) + }) + + t.Run("Multi", func(t *testing.T) { + items, ct, err := f.cli.ReadItemMulti(ctx, f.bucket, pk, sk) + require.ErrorIs(t, err, k2v.NoSuchItemErr) + require.Empty(t, items) + require.Empty(t, ct) + }) +} + +func TestClient_ReadItemTombstone(t *testing.T) { + f, ctx := newFixture(t) + + pk := randomKey() + sk := randomKey() + + t.Logf("Creating item: PK=%s, SK=%s", pk, sk) + + err := f.cli.InsertItem(ctx, f.bucket, pk, sk, "", []byte("hello")) + require.NoError(t, err) + _, ct, err := f.cli.ReadItemSingle(ctx, f.bucket, pk, sk) + require.NoError(t, err) + err = f.cli.DeleteItem(ctx, f.bucket, pk, sk, ct) + require.NoError(t, err) + + t.Run("Single", func(t *testing.T) { + item, ct, err := f.cli.ReadItemSingle(ctx, f.bucket, pk, sk) + require.ErrorIs(t, err, k2v.TombstoneItemErr) + require.Nil(t, item) + require.NotEmpty(t, ct) + }) + + t.Run("Multi", func(t *testing.T) { + items, ct, err := f.cli.ReadItemMulti(ctx, f.bucket, pk, sk) + require.ErrorIs(t, err, k2v.TombstoneItemErr) + require.Empty(t, items) + require.NotEmpty(t, ct) + }) +} + +func TestClient_ReadItemSingleRevision(t *testing.T) { + f, ctx := newFixture(t) + + pk := randomKey() + sk := randomKey() + + err := f.cli.InsertItem(ctx, f.bucket, pk, sk, "", []byte("hello")) + require.NoError(t, err) + + t.Run("Single", func(t *testing.T) { + item, ct, err := f.cli.ReadItemSingle(ctx, f.bucket, pk, sk) + require.NoError(t, err) + require.Equal(t, "hello", string(item)) + require.NotEmpty(t, ct) + }) + + t.Run("Multi", func(t *testing.T) { + items, ct, err := f.cli.ReadItemMulti(ctx, f.bucket, pk, sk) + require.NoError(t, err) + require.Len(t, items, 1) + require.Equal(t, "hello", string(items[0])) + require.NotEmpty(t, ct) + }) +} + +func TestClient_ReadItemMultipleRevisions(t *testing.T) { + f, ctx := newFixture(t) + + pk := randomKey() + sk := randomKey() + + err := f.cli.InsertItem(ctx, f.bucket, pk, sk, "", []byte("hello1")) + require.NoError(t, err) + + // don't use a continuation token to intentionally create 2x concurrent revisions + err = f.cli.InsertItem(ctx, f.bucket, pk, sk, "", []byte("hello2")) + require.NoError(t, err) + + t.Run("Single", func(t *testing.T) { + item, ct, err := f.cli.ReadItemSingle(ctx, f.bucket, pk, sk) + require.ErrorIs(t, err, k2v.ConcurrentItemsErr) + require.Nil(t, item) + require.NotEmpty(t, ct) + }) + + t.Run("Multi", func(t *testing.T) { + items, ct, err := f.cli.ReadItemMulti(ctx, f.bucket, pk, sk) + require.NoError(t, err) + require.Len(t, items, 2) + require.Equal(t, "hello1", string(items[0])) + require.Equal(t, "hello2", string(items[1])) + require.NotEmpty(t, ct) + }) +} + +func TestClient_PollItem(t *testing.T) { + f, ctx := newFixture(t) + + pk := randomKey() + sk := randomKey() + + err := f.cli.InsertItem(ctx, f.bucket, pk, sk, "", []byte("hello1")) + require.NoError(t, err) + + _, ct, err := f.cli.ReadItemSingle(ctx, f.bucket, pk, sk) + + pollReadyCh := make(chan struct{}) + go func(ct k2v.CausalityToken) { + select { + case <-ctx.Done(): + t.Errorf("Context canceled: %v", ctx.Err()) + return + case <-pollReadyCh: + t.Logf("PollItem connected") + } + err = f.cli.InsertItem(ctx, f.bucket, pk, sk, ct, []byte("hello2")) + require.NoError(t, err) + }(ct) + + trace := &httptrace.ClientTrace{ + WroteRequest: func(_ httptrace.WroteRequestInfo) { + pollReadyCh <- struct{}{} + }, + } + item, ct, err := f.cli.PollItem(httptrace.WithClientTrace(ctx, trace), f.bucket, pk, sk, ct, 5*time.Second) + require.NoError(t, err) + require.Equal(t, "hello2", string(item)) + require.NotEmpty(t, ct) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..9c37efd --- /dev/null +++ b/go.mod @@ -0,0 +1,17 @@ +module code.notaphish.fyi/milas/garage-k2v-go + +go 1.23.1 + +require ( + github.com/aws/aws-sdk-go-v2 v1.32.2 + github.com/davecgh/go-spew v1.1.1 + github.com/stretchr/testify v1.8.0 + go.uber.org/goleak v1.3.0 +) + +require ( + github.com/aws/smithy-go v1.22.0 // indirect + github.com/kr/text v0.2.0 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..052ea4d --- /dev/null +++ b/go.sum @@ -0,0 +1,27 @@ +github.com/aws/aws-sdk-go-v2 v1.32.2 h1:AkNLZEyYMLnx/Q/mSKkcMqwNFXMAvFto9bNsHqcTduI= +github.com/aws/aws-sdk-go-v2 v1.32.2/go.mod h1:2SK5n0a2karNTv5tbP1SjsX0uhttou00v/HpXKM1ZUo= +github.com/aws/smithy-go v1.22.0 h1:uunKnWlcoL3zO7q+gG2Pk53joueEOsnNB28QdMsmiMM= +github.com/aws/smithy-go v1.22.0/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/insert_batch.go b/insert_batch.go new file mode 100644 index 0000000..bb236c9 --- /dev/null +++ b/insert_batch.go @@ -0,0 +1,54 @@ +package k2v + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" +) + +type BatchInsertItem struct { + PartitionKey string `json:"pk"` + SortKey string `json:"sk"` + CausalityToken *string `json:"ct"` + Value Item `json:"v"` +} + +func (c *Client) InsertBatch(ctx context.Context, b Bucket, items []BatchInsertItem) error { + u, err := url.Parse(c.endpoint) + if err != nil { + return err + } + u.Path = string(b) + + reqBody, err := json.Marshal(items) + if err != nil { + return err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u.String(), bytes.NewReader(reqBody)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + + resp, err := c.executeRequest(req) + if err != nil { + return err + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return err + } + + if resp.StatusCode != http.StatusNoContent { + return fmt.Errorf("http status code %d: %s", resp.StatusCode, string(body)) + } + + return nil +} diff --git a/poll_single.go b/poll_single.go new file mode 100644 index 0000000..4e854a4 --- /dev/null +++ b/poll_single.go @@ -0,0 +1,10 @@ +package k2v + +import ( + "context" + "time" +) + +func (c *Client) PollItem(ctx context.Context, b Bucket, pk string, sk string, ct CausalityToken, timeout time.Duration) (Item, CausalityToken, error) { + return c.readItemSingle(ctx, b, pk, sk, ct, timeout) +} diff --git a/read_batch.go b/read_batch.go new file mode 100644 index 0000000..876c95c --- /dev/null +++ b/read_batch.go @@ -0,0 +1,138 @@ +package k2v + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" +) + +type ReadBatchSearch struct { + PartitionKey string `json:"partitionKey"` + + // Prefix restricts listing to partition keys that start with this value. + Prefix string `json:"prefix,omitempty"` + + // Start is the first partition key to list, in lexicographical order. + Start string `json:"start,omitempty"` + + // End is the last partition key to list (excluded). + End string `json:"end,omitempty"` + + // Limit for maximum number of partition keys to list. + Limit int `json:"limit,omitempty"` + + // Reverse iterates in reverse lexicographical order. + Reverse bool `json:"reverse,omitempty"` + + // SingleItem determines whether to return only the item with sort key start. + SingleItem bool `json:"singleItem,omitempty"` + + // ConflictsOnly determines whether to return only items that have several concurrent values. + ConflictsOnly bool `json:"conflictsOnly,omitempty"` + + // Tombstones determines whether or not to return tombstone lines to indicate the presence of old deleted items. + Tombstones bool `json:"tombstones,omitempty"` +} + +type BatchSearchResult struct { + PartitionKey string `json:"partitionKey"` + Prefix *string `json:"prefix"` + Start *string `json:"start"` + End *string `json:"end"` + Limit *int `json:"limit"` + Reverse bool `json:"reverse"` + SingleItem bool `json:"singleItem"` + ConflictsOnly bool `json:"conflictsOnly"` + Tombstones bool `json:"tombstones"` + Items []SearchResultItem `json:"items"` + More bool `json:"more"` + NextStart *string `json:"nextStart"` +} + +type SearchResultItem struct { + SortKey string `json:"sk"` + CausalityToken string `json:"ct"` + Values []Item `json:"v"` +} + +func (c *Client) ReadBatch(ctx context.Context, b Bucket, q []ReadBatchSearch) ([]BatchSearchResult, error) { + u, err := url.Parse(c.endpoint) + if err != nil { + return nil, err + } + u.Path = string(b) + + reqBody, err := json.Marshal(q) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, "SEARCH", u.String(), bytes.NewReader(reqBody)) + if err != nil { + return nil, err + } + + resp, err := c.executeRequest(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("http status code %d: %s", resp.StatusCode, body) + } + + var items []BatchSearchResult + if err := json.Unmarshal(body, &items); err != nil { + return nil, err + } + + return items, err +} + +type ItemKey struct { + PartitionKey string + SortKey string +} + +type BulkGetItem struct { + PartitionKey string + SortKey string + CausalityToken CausalityToken + Values []Item +} + +func BulkGet(ctx context.Context, cli *Client, b Bucket, keys []ItemKey) ([]BulkGetItem, error) { + q := make([]ReadBatchSearch, len(keys)) + for i := range keys { + q[i] = ReadBatchSearch{ + PartitionKey: keys[i].PartitionKey, + Start: keys[i].SortKey, + SingleItem: true, + Tombstones: true, + } + } + results, err := cli.ReadBatch(ctx, b, q) + if err != nil { + return nil, err + } + ret := make([]BulkGetItem, len(results)) + for i := range results { + ret[i] = BulkGetItem{ + PartitionKey: results[i].PartitionKey, + SortKey: results[i].Items[0].SortKey, + CausalityToken: CausalityToken(results[i].Items[0].CausalityToken), + Values: results[i].Items[0].Values, + } + } + return ret, nil +} diff --git a/read_batch_test.go b/read_batch_test.go new file mode 100644 index 0000000..47bd07b --- /dev/null +++ b/read_batch_test.go @@ -0,0 +1,79 @@ +package k2v_test + +import ( + k2v "code.notaphish.fyi/milas/garage-k2v-go" + "github.com/davecgh/go-spew/spew" + "github.com/stretchr/testify/require" + "math/rand/v2" + "strconv" + "testing" +) + +func TestClient_ReadBatch(t *testing.T) { + f, ctx := newFixture(t) + + pk1 := randomKey() + sk1 := randomKey() + + require.NoError(t, f.cli.InsertItem(ctx, f.bucket, pk1, sk1, "", []byte("hello"))) + + pk2 := randomKey() + for i := range 5 { + sk := randomKey() + require.NoError(t, f.cli.InsertItem(ctx, f.bucket, pk2, sk, "", []byte("hello-"+strconv.Itoa(i)))) + } + + pk3 := randomKey() + sk3 := randomKey() + for i := range 5 { + require.NoError(t, f.cli.InsertItem(ctx, f.bucket, pk3, sk3, "", []byte("hello-"+strconv.Itoa(i)))) + } + + q := []k2v.ReadBatchSearch{ + { + PartitionKey: pk1, + }, + { + PartitionKey: pk2, + }, + { + PartitionKey: pk3, + SingleItem: true, + Start: sk3, + }, + } + + items, err := f.cli.ReadBatch(ctx, f.bucket, q) + require.NoError(t, err) + require.NotEmpty(t, items) + + spew.Dump(items) +} + +func TestBulkGet(t *testing.T) { + f, ctx := newFixture(t) + + keys := make([]k2v.ItemKey, 500) + for i := range keys { + keys[i] = k2v.ItemKey{ + PartitionKey: randomKey(), + SortKey: randomKey(), + } + require.NoError(t, f.cli.InsertItem(ctx, f.bucket, keys[i].PartitionKey, keys[i].SortKey, "", []byte("hello"+strconv.Itoa(i)))) + } + + rand.Shuffle(len(keys), func(i, j int) { + keys[i], keys[j] = keys[j], keys[i] + }) + + items, err := k2v.BulkGet(ctx, f.cli, f.bucket, keys) + require.NoError(t, err) + require.NotEmpty(t, items) + + require.Equal(t, len(keys), len(items)) + for i := range keys { + require.Equal(t, keys[i].SortKey, items[i].SortKey) + require.Len(t, items[i].Values, 1) + require.Contains(t, string(items[i].Values[0]), "hello") + } +}