Skip to content

Commit 1ce683b

Browse files
WondertanWondertan
authored
feat: context datastore (#238)
* feat: context datastore * fix: split type to avoid collisions
1 parent c4698f6 commit 1ce683b

File tree

3 files changed

+286
-0
lines changed

3 files changed

+286
-0
lines changed

‎context/context.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package contextds
2+
3+
import (
4+
"context"
5+
6+
"github.com/ipfs/go-datastore"
7+
)
8+
9+
// WithWrite adds a write batch to the context.
10+
func WithWrite(ctx context.Context, batch datastore.Write) context.Context {
11+
return context.WithValue(ctx, writeKey, batch)
12+
}
13+
14+
// GetWrite retrieves the write batch from the context.
15+
func GetWrite(ctx context.Context) (datastore.Write, bool) {
16+
batch, ok := ctx.Value(writeKey).(datastore.Write)
17+
return batch, ok
18+
}
19+
20+
// WithRead adds a read batch to the context.
21+
func WithRead(ctx context.Context, batch datastore.Read) context.Context {
22+
return context.WithValue(ctx, readKey, batch)
23+
}
24+
25+
// GetRead retrieves the read batch from the context.
26+
func GetRead(ctx context.Context) (datastore.Read, bool) {
27+
batch, ok := ctx.Value(readKey).(datastore.Read)
28+
return batch, ok
29+
}
30+
31+
type (
32+
writeKeyTp int
33+
readKeyTp int
34+
)
35+
36+
var (
37+
writeKey writeKeyTp
38+
readKey readKeyTp
39+
)

‎context/context_ds.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package contextds
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"github.com/ipfs/go-datastore"
8+
"github.com/ipfs/go-datastore/query"
9+
)
10+
11+
var _ datastore.Datastore = (*Datastore)(nil)
12+
var _ datastore.Batching = (*Datastore)(nil)
13+
var _ datastore.TxnDatastore = (*Datastore)(nil)
14+
15+
// WrapsDatastore wraps around the given datastore.Datastore making its operations context-aware
16+
// It intercepts datastore operations routing them to the current Write or Read if one exists on the context.
17+
func WrapDatastore(ds datastore.Datastore) datastore.Datastore {
18+
return &Datastore{
19+
inner: ds,
20+
}
21+
}
22+
23+
// Datastore is a wrapper around a datastore.Datastore that provides context-aware operations.
24+
// See [WrapDatastore].
25+
type Datastore struct {
26+
inner datastore.Datastore
27+
}
28+
29+
func (ds *Datastore) Put(ctx context.Context, key datastore.Key, value []byte) error {
30+
if write, ok := GetWrite(ctx); ok {
31+
return write.Put(ctx, key, value)
32+
}
33+
return ds.inner.Put(ctx, key, value)
34+
}
35+
36+
func (ds *Datastore) Delete(ctx context.Context, key datastore.Key) error {
37+
if write, ok := GetWrite(ctx); ok {
38+
return write.Delete(ctx, key)
39+
}
40+
return ds.inner.Delete(ctx, key)
41+
}
42+
43+
func (ds *Datastore) Get(ctx context.Context, key datastore.Key) ([]byte, error) {
44+
if read, ok := GetRead(ctx); ok {
45+
return read.Get(ctx, key)
46+
}
47+
return ds.inner.Get(ctx, key)
48+
}
49+
50+
func (ds *Datastore) Has(ctx context.Context, key datastore.Key) (bool, error) {
51+
if read, ok := GetRead(ctx); ok {
52+
return read.Has(ctx, key)
53+
}
54+
return ds.inner.Has(ctx, key)
55+
}
56+
57+
func (ds *Datastore) GetSize(ctx context.Context, key datastore.Key) (int, error) {
58+
if read, ok := GetRead(ctx); ok {
59+
return read.GetSize(ctx, key)
60+
}
61+
return ds.inner.GetSize(ctx, key)
62+
}
63+
64+
func (ds *Datastore) Query(ctx context.Context, q query.Query) (query.Results, error) {
65+
if read, ok := GetRead(ctx); ok {
66+
return read.Query(ctx, q)
67+
}
68+
return ds.inner.Query(ctx, q)
69+
}
70+
71+
func (ds *Datastore) Close() error {
72+
return ds.inner.Close()
73+
}
74+
75+
func (ds *Datastore) Sync(ctx context.Context, prefix datastore.Key) error {
76+
return ds.inner.Sync(ctx, prefix)
77+
}
78+
79+
func (ds *Datastore) Batch(ctx context.Context) (datastore.Batch, error) {
80+
bds, ok := ds.inner.(datastore.Batching)
81+
if !ok {
82+
return nil, datastore.ErrBatchUnsupported
83+
}
84+
85+
return bds.Batch(ctx)
86+
}
87+
88+
func (ds *Datastore) NewTransaction(ctx context.Context, readOnly bool) (datastore.Txn, error) {
89+
tds, ok := ds.inner.(datastore.TxnDatastore)
90+
if !ok {
91+
return nil, fmt.Errorf("transactions not supported")
92+
}
93+
94+
return tds.NewTransaction(ctx, readOnly)
95+
}

‎context/context_ds_test.go‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
package contextds_test
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/ipfs/go-datastore"
8+
contextds "github.com/ipfs/go-datastore/context"
9+
"github.com/ipfs/go-datastore/query"
10+
)
11+
12+
// AI generated tests and mocks
13+
14+
type mockWrite struct {
15+
putCalled bool
16+
putKey datastore.Key
17+
putValue []byte
18+
putErr error
19+
deleteCalled bool
20+
deleteKey datastore.Key
21+
deleteErr error
22+
}
23+
24+
func (m *mockWrite) Put(ctx context.Context, key datastore.Key, value []byte) error {
25+
m.putCalled = true
26+
m.putKey = key
27+
m.putValue = value
28+
return m.putErr
29+
}
30+
31+
func (m *mockWrite) Delete(ctx context.Context, key datastore.Key) error {
32+
m.deleteCalled = true
33+
m.deleteKey = key
34+
return m.deleteErr
35+
}
36+
37+
type mockRead struct {
38+
getCalled bool
39+
getKey datastore.Key
40+
getValue []byte
41+
getErr error
42+
hasCalled bool
43+
hasKey datastore.Key
44+
hasValue bool
45+
hasErr error
46+
47+
getSizeCalled bool
48+
getSizeKey datastore.Key
49+
getSizeValue int
50+
getSizeErr error
51+
52+
queryCalled bool
53+
queryQ query.Query
54+
queryResults query.Results
55+
queryErr error
56+
}
57+
58+
func (m *mockRead) Get(ctx context.Context, key datastore.Key) ([]byte, error) {
59+
m.getCalled = true
60+
m.getKey = key
61+
return m.getValue, m.getErr
62+
}
63+
64+
func (m *mockRead) Has(ctx context.Context, key datastore.Key) (bool, error) {
65+
m.hasCalled = true
66+
m.hasKey = key
67+
return m.hasValue, m.hasErr
68+
}
69+
70+
func (m *mockRead) GetSize(ctx context.Context, key datastore.Key) (int, error) {
71+
m.getSizeCalled = true
72+
m.getSizeKey = key
73+
return m.getSizeValue, m.getSizeErr
74+
}
75+
76+
func (m *mockRead) Query(ctx context.Context, q query.Query) (query.Results, error) {
77+
m.queryCalled = true
78+
m.queryQ = q
79+
return m.queryResults, m.queryErr
80+
}
81+
82+
func TestDatastore_WithWriteContext(t *testing.T) {
83+
inner := datastore.NewMapDatastore()
84+
mock := &mockWrite{}
85+
ds := contextds.WrapDatastore(inner)
86+
ctx := contextds.WithWrite(context.Background(), mock)
87+
88+
key := datastore.NewKey("foo")
89+
value := []byte("bar")
90+
91+
err := ds.Put(ctx, key, value)
92+
if err != nil {
93+
t.Fatalf("Put failed: %v", err)
94+
}
95+
if !mock.putCalled || mock.putKey != key || string(mock.putValue) != string(value) {
96+
t.Errorf("Put did not delegate to mockWrite correctly")
97+
}
98+
99+
err = ds.Delete(ctx, key)
100+
if err != nil {
101+
t.Fatalf("Delete failed: %v", err)
102+
}
103+
if !mock.deleteCalled || mock.deleteKey != key {
104+
t.Errorf("Delete did not delegate to mockWrite correctly")
105+
}
106+
}
107+
108+
func TestDatastore_WithReadContext(t *testing.T) {
109+
inner := datastore.NewMapDatastore()
110+
mock := &mockRead{
111+
getValue: []byte("baz"),
112+
hasValue: true,
113+
getSizeValue: 3,
114+
}
115+
ds := contextds.WrapDatastore(inner)
116+
ctx := contextds.WithRead(context.Background(), mock)
117+
118+
key := datastore.NewKey("foo")
119+
120+
val, err := ds.Get(ctx, key)
121+
if err != nil {
122+
t.Fatalf("Get failed: %v", err)
123+
}
124+
if !mock.getCalled || mock.getKey != key || string(val) != "baz" {
125+
t.Errorf("Get did not delegate to mockRead correctly")
126+
}
127+
128+
has, err := ds.Has(ctx, key)
129+
if err != nil {
130+
t.Fatalf("Has failed: %v", err)
131+
}
132+
if !mock.hasCalled || mock.hasKey != key || !has {
133+
t.Errorf("Has did not delegate to mockRead correctly")
134+
}
135+
136+
sz, err := ds.GetSize(ctx, key)
137+
if err != nil {
138+
t.Fatalf("GetSize failed: %v", err)
139+
}
140+
if !mock.getSizeCalled || mock.getSizeKey != key || sz != 3 {
141+
t.Errorf("GetSize did not delegate to mockRead correctly")
142+
}
143+
144+
q := query.Query{}
145+
_, err = ds.Query(ctx, q)
146+
if err != nil {
147+
t.Fatalf("Query failed: %v", err)
148+
}
149+
if !mock.queryCalled {
150+
t.Errorf("Query did not delegate to mockRead correctly")
151+
}
152+
}

0 commit comments

Comments
 (0)