From 1597fcd4fa908f9e718104df8895d52ad74ca999 Mon Sep 17 00:00:00 2001 From: Wondertan Date: Fri, 25 Jul 2025 17:13:43 +0200 Subject: [PATCH 1/2] feat: context datastore --- context/context.go | 36 +++++++++ context/context_ds.go | 95 +++++++++++++++++++++++ context/context_ds_test.go | 152 +++++++++++++++++++++++++++++++++++++ 3 files changed, 283 insertions(+) create mode 100644 context/context.go create mode 100644 context/context_ds.go create mode 100644 context/context_ds_test.go diff --git a/context/context.go b/context/context.go new file mode 100644 index 0000000..b6ee2e7 --- /dev/null +++ b/context/context.go @@ -0,0 +1,36 @@ +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 key int + +var ( + writeKey key + readKey key +) diff --git a/context/context_ds.go b/context/context_ds.go new file mode 100644 index 0000000..b927cde --- /dev/null +++ b/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) +} diff --git a/context/context_ds_test.go b/context/context_ds_test.go new file mode 100644 index 0000000..3dbeaaf --- /dev/null +++ b/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") + } +} From d6c676753344ea6c53ae9814cf0f4e805f4a490c Mon Sep 17 00:00:00 2001 From: Wondertan Date: Fri, 1 Aug 2025 15:21:59 +0200 Subject: [PATCH 2/2] fix: split type to avoid collisions --- context/context.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/context/context.go b/context/context.go index b6ee2e7..ee6d12a 100644 --- a/context/context.go +++ b/context/context.go @@ -28,9 +28,12 @@ func GetRead(ctx context.Context) (datastore.Read, bool) { return batch, ok } -type key int +type ( + writeKeyTp int + readKeyTp int +) var ( - writeKey key - readKey key + writeKey writeKeyTp + readKey readKeyTp )