feat: context datastore

context/context.go

@@ -0,0 +1,39 @@
+package contextds
+
+import (
+	"context"
+
+	"github.com/ipfs/go-datastore"
+)
+
+// WithWrite adds a write batch to the context.
+func WithWrite(ctx context.Context, batch datastore.Write) context.Context {
+	return context.WithValue(ctx, writeKey, batch)
+}
+
+// GetWrite retrieves the write batch from the context.
+func GetWrite(ctx context.Context) (datastore.Write, bool) {
+	batch, ok := ctx.Value(writeKey).(datastore.Write)
+	return batch, ok
+}
+
+// WithRead adds a read batch to the context.
+func WithRead(ctx context.Context, batch datastore.Read) context.Context {
+	return context.WithValue(ctx, readKey, batch)
+}
+
+// GetRead retrieves the read batch from the context.
+func GetRead(ctx context.Context) (datastore.Read, bool) {
+	batch, ok := ctx.Value(readKey).(datastore.Read)
+	return batch, ok
+}
+
+type (
+	writeKeyTp int
+	readKeyTp  int
+)
+
+var (
+	writeKey writeKeyTp
+	readKey  readKeyTp
+)

context/context_ds.go

@@ -0,0 +1,95 @@
+package contextds
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/ipfs/go-datastore"
+	"github.com/ipfs/go-datastore/query"
+)
+
+var _ datastore.Datastore = (*Datastore)(nil)
+var _ datastore.Batching = (*Datastore)(nil)
+var _ datastore.TxnDatastore = (*Datastore)(nil)
+
+// WrapsDatastore wraps around the given datastore.Datastore making its operations context-aware
+// It intercepts datastore operations routing them to the current Write or Read if one exists on the context.
+func WrapDatastore(ds datastore.Datastore) datastore.Datastore {
+	return &Datastore{
+		inner: ds,
+	}
+}
+
+// Datastore is a wrapper around a datastore.Datastore that provides context-aware operations.
+// See [WrapDatastore].
+type Datastore struct {
+	inner datastore.Datastore
+}
+
+func (ds *Datastore) Put(ctx context.Context, key datastore.Key, value []byte) error {
+	if write, ok := GetWrite(ctx); ok {
+		return write.Put(ctx, key, value)
+	}
+	return ds.inner.Put(ctx, key, value)
+}
+
+func (ds *Datastore) Delete(ctx context.Context, key datastore.Key) error {
+	if write, ok := GetWrite(ctx); ok {
+		return write.Delete(ctx, key)
+	}
+	return ds.inner.Delete(ctx, key)
+}
+
+func (ds *Datastore) Get(ctx context.Context, key datastore.Key) ([]byte, error) {
+	if read, ok := GetRead(ctx); ok {
+		return read.Get(ctx, key)
+	}
+	return ds.inner.Get(ctx, key)
+}
+
+func (ds *Datastore) Has(ctx context.Context, key datastore.Key) (bool, error) {
+	if read, ok := GetRead(ctx); ok {
+		return read.Has(ctx, key)
+	}
+	return ds.inner.Has(ctx, key)
+}
+
+func (ds *Datastore) GetSize(ctx context.Context, key datastore.Key) (int, error) {
+	if read, ok := GetRead(ctx); ok {
+		return read.GetSize(ctx, key)
+	}
+	return ds.inner.GetSize(ctx, key)
+}
+
+func (ds *Datastore) Query(ctx context.Context, q query.Query) (query.Results, error) {
+	if read, ok := GetRead(ctx); ok {
+		return read.Query(ctx, q)
+	}
+	return ds.inner.Query(ctx, q)
+}
+
+func (ds *Datastore) Close() error {
+	return ds.inner.Close()
+}
+
+func (ds *Datastore) Sync(ctx context.Context, prefix datastore.Key) error {
+	return ds.inner.Sync(ctx, prefix)
+}
+
+func (ds *Datastore) Batch(ctx context.Context) (datastore.Batch, error) {
+	bds, ok := ds.inner.(datastore.Batching)
+	if !ok {
+		return nil, datastore.ErrBatchUnsupported
+	}
+
+	return bds.Batch(ctx)
+}
+
+func (ds *Datastore) NewTransaction(ctx context.Context, readOnly bool) (datastore.Txn, error) {
+	tds, ok := ds.inner.(datastore.TxnDatastore)
+	if !ok {
+		return nil, fmt.Errorf("transactions not supported")
+	}
+
+	return tds.NewTransaction(ctx, readOnly)
+}

context/context_ds_test.go

@@ -0,0 +1,152 @@
+package contextds_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ipfs/go-datastore"
+	contextds "github.com/ipfs/go-datastore/context"
+	"github.com/ipfs/go-datastore/query"
+)
+
+// AI generated tests and mocks
+
+type mockWrite struct {
+	putCalled    bool
+	putKey       datastore.Key
+	putValue     []byte
+	putErr       error
+	deleteCalled bool
+	deleteKey    datastore.Key
+	deleteErr    error
+}
+
+func (m *mockWrite) Put(ctx context.Context, key datastore.Key, value []byte) error {
+	m.putCalled = true
+	m.putKey = key
+	m.putValue = value
+	return m.putErr
+}
+
+func (m *mockWrite) Delete(ctx context.Context, key datastore.Key) error {
+	m.deleteCalled = true
+	m.deleteKey = key
+	return m.deleteErr
+}
+
+type mockRead struct {
+	getCalled bool
+	getKey    datastore.Key
+	getValue  []byte
+	getErr    error
+	hasCalled bool
+	hasKey    datastore.Key
+	hasValue  bool
+	hasErr    error
+
+	getSizeCalled bool
+	getSizeKey    datastore.Key
+	getSizeValue  int
+	getSizeErr    error
+
+	queryCalled  bool
+	queryQ       query.Query
+	queryResults query.Results
+	queryErr     error
+}
+
+func (m *mockRead) Get(ctx context.Context, key datastore.Key) ([]byte, error) {
+	m.getCalled = true
+	m.getKey = key
+	return m.getValue, m.getErr
+}
+
+func (m *mockRead) Has(ctx context.Context, key datastore.Key) (bool, error) {
+	m.hasCalled = true
+	m.hasKey = key
+	return m.hasValue, m.hasErr
+}
+
+func (m *mockRead) GetSize(ctx context.Context, key datastore.Key) (int, error) {
+	m.getSizeCalled = true
+	m.getSizeKey = key
+	return m.getSizeValue, m.getSizeErr
+}
+
+func (m *mockRead) Query(ctx context.Context, q query.Query) (query.Results, error) {
+	m.queryCalled = true
+	m.queryQ = q
+	return m.queryResults, m.queryErr
+}
+
+func TestDatastore_WithWriteContext(t *testing.T) {
+	inner := datastore.NewMapDatastore()
+	mock := &mockWrite{}
+	ds := contextds.WrapDatastore(inner)
+	ctx := contextds.WithWrite(context.Background(), mock)
+
+	key := datastore.NewKey("foo")
+	value := []byte("bar")
+
+	err := ds.Put(ctx, key, value)
+	if err != nil {
+		t.Fatalf("Put failed: %v", err)
+	}
+	if !mock.putCalled || mock.putKey != key || string(mock.putValue) != string(value) {
+		t.Errorf("Put did not delegate to mockWrite correctly")
+	}
+
+	err = ds.Delete(ctx, key)
+	if err != nil {
+		t.Fatalf("Delete failed: %v", err)
+	}
+	if !mock.deleteCalled || mock.deleteKey != key {
+		t.Errorf("Delete did not delegate to mockWrite correctly")
+	}
+}
+
+func TestDatastore_WithReadContext(t *testing.T) {
+	inner := datastore.NewMapDatastore()
+	mock := &mockRead{
+		getValue:     []byte("baz"),
+		hasValue:     true,
+		getSizeValue: 3,
+	}
+	ds := contextds.WrapDatastore(inner)
+	ctx := contextds.WithRead(context.Background(), mock)
+
+	key := datastore.NewKey("foo")
+
+	val, err := ds.Get(ctx, key)
+	if err != nil {
+		t.Fatalf("Get failed: %v", err)
+	}
+	if !mock.getCalled || mock.getKey != key || string(val) != "baz" {
+		t.Errorf("Get did not delegate to mockRead correctly")
+	}
+
+	has, err := ds.Has(ctx, key)
+	if err != nil {
+		t.Fatalf("Has failed: %v", err)
+	}
+	if !mock.hasCalled || mock.hasKey != key || !has {
+		t.Errorf("Has did not delegate to mockRead correctly")
+	}
+
+	sz, err := ds.GetSize(ctx, key)
+	if err != nil {
+		t.Fatalf("GetSize failed: %v", err)
+	}
+	if !mock.getSizeCalled || mock.getSizeKey != key || sz != 3 {
+		t.Errorf("GetSize did not delegate to mockRead correctly")
+	}
+
+	q := query.Query{}
+	_, err = ds.Query(ctx, q)
+	if err != nil {
+		t.Fatalf("Query failed: %v", err)
+	}
+	if !mock.queryCalled {
+		t.Errorf("Query did not delegate to mockRead correctly")
+	}
+}