testutils: Add script-based test utilities

any_table.go

@@ -0,0 +1,63 @@
+package statedb
+
+import (
+	"iter"
+)
+
+// AnyTable allows any-typed access to a StateDB table. This is intended
+// for building generic tooling for accessing the table and should be
+// avoided if possible.
+type AnyTable struct {
+	Meta TableMeta
+}
+
+func (t AnyTable) All(txn ReadTxn) iter.Seq2[any, Revision] {
+	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
+	return anySeq(indexTxn.Iterator())
+}
+
+func (t AnyTable) UnmarshalYAML(data []byte) (any, error) {
+	return t.Meta.unmarshalYAML(data)
+}
+
+func (t AnyTable) Insert(txn WriteTxn, obj any) (old any, hadOld bool, err error) {
+	var iobj object
+	iobj, hadOld, err = txn.getTxn().insert(t.Meta, Revision(0), obj)
+	if hadOld {
+		old = iobj.data
+	}
+	return
+}
+
+func (t AnyTable) Delete(txn WriteTxn, obj any) (old any, hadOld bool, err error) {
+	var iobj object
+	iobj, hadOld, err = txn.getTxn().delete(t.Meta, Revision(0), obj)
+	if hadOld {
+		old = iobj.data
+	}
+	return
+}
+
+func (t AnyTable) Prefix(txn ReadTxn, key string) iter.Seq2[any, Revision] {
+	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
+	iter, _ := indexTxn.Prefix([]byte(key))
+	return anySeq(iter)
+}
+
+func (t AnyTable) LowerBound(txn ReadTxn, key string) iter.Seq2[any, Revision] {
+	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
+	iter := indexTxn.LowerBound([]byte(key))
+	return anySeq(iter)
+}
+
+func (t AnyTable) TableHeader() []string {
+	zero := t.Meta.proto()
+	if tw, ok := zero.(TableWritable); ok {
+		return tw.TableHeader()
+	}
+	return nil
+}
+
+func (t AnyTable) Proto() any {
+	return t.Meta.proto()
+}

db.go

@@ -247,6 +247,25 @@ func (db *DB) WriteTxn(table TableMeta, tables ...TableMeta) WriteTxn {
 	return txn
 }
 
+func (db *DB) GetTables(txn ReadTxn) (tbls []TableMeta) {
+	root := txn.getTxn().root
+	tbls = make([]TableMeta, 0, len(root))
+	for _, table := range root {
+		tbls = append(tbls, table.meta)
+	}
+	return
+}
+
+func (db *DB) GetTable(txn ReadTxn, name string) TableMeta {
+	root := txn.getTxn().root
+	for _, table := range root {
+		if table.meta.Name() == name {
+			return table.meta
+		}
+	}
+	return nil
+}
+
 // Start the background workers for the database.
 //
 // This starts the graveyard worker that deals with garbage collecting

go.mod

@@ -5,11 +5,13 @@ go 1.23
 require (
 	github.com/cilium/hive v0.0.0-20240209163124-bd6ebb4ec11d
 	github.com/cilium/stream v0.0.0-20240209152734-a0792b51812d
+	github.com/rogpeppe/go-internal v1.11.0
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.8.4
 	go.uber.org/goleak v1.3.0
 	golang.org/x/time v0.5.0
+	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
@@ -34,7 +36,7 @@ require (
 	golang.org/x/sys v0.17.0 // indirect
 	golang.org/x/term v0.16.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
+	golang.org/x/tools v0.17.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -77,6 +77,8 @@ golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
 golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
+golang.org/x/tools v0.17.0 h1:FvmRgNOcs3kOa+T20R1uhfP9F6HgG2mfxDv1vrx1Htc=
+golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

iterator.go

@@ -53,7 +53,7 @@ func ToSeq[A, B any](seq iter.Seq2[A, B]) iter.Seq[A] {
 	}
 }
 
-// partSeq returns a sequence of objects from a part Iterator.
+// partSeq returns a casted sequence of objects from a part Iterator.
 func partSeq[Obj any](iter *part.Iterator[object]) iter.Seq2[Obj, Revision] {
 	return func(yield func(Obj, Revision) bool) {
 		// Iterate over a clone of the original iterator to allow the sequence to be iterated
@@ -71,6 +71,24 @@ func partSeq[Obj any](iter *part.Iterator[object]) iter.Seq2[Obj, Revision] {
 	}
 }
 
+// anySeq returns a sequence of objects from a part Iterator.
+func anySeq(iter *part.Iterator[object]) iter.Seq2[any, Revision] {
+	return func(yield func(any, Revision) bool) {
+		// Iterate over a clone of the original iterator to allow the sequence to be iterated
+		// from scratch multiple times.
+		it := iter.Clone()
+		for {
+			_, iobj, ok := it.Next()
+			if !ok {
+				break
+			}
+			if !yield(iobj.data, iobj.revision) {
+				break
+			}
+		}
+	}
+}
+
 // nonUniqueSeq returns a sequence of objects for a non-unique index.
 // Non-unique indexes work by concatenating the secondary key with the
 // primary key and then prefix searching for the items:
@@ -128,6 +146,13 @@ func (it *iterator[Obj]) Next() (obj Obj, revision uint64, ok bool) {
 	return
 }
 
+// Iterator for iterating a sequence objects.
+type Iterator[Obj any] interface {
+	// Next returns the next object and its revision if ok is true, otherwise
+	// zero values to mean that the iteration has finished.
+	Next() (obj Obj, rev Revision, ok bool)
+}
+
 func NewDualIterator[Obj any](left, right Iterator[Obj]) *DualIterator[Obj] {
 	return &DualIterator[Obj]{
 		left:  iterState[Obj]{iter: left},

table.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/cilium/statedb/internal"
 	"github.com/cilium/statedb/part"
+	"gopkg.in/yaml.v3"
 
 	"github.com/cilium/statedb/index"
 )
@@ -184,6 +185,15 @@ func (t *genTable[Obj]) Name() string {
 	return t.table
 }
 
+func (t *genTable[Obj]) Indexes() []string {
+	idxs := make([]string, 0, 1+len(t.secondaryAnyIndexers))
+	idxs = append(idxs, t.primaryAnyIndexer.name)
+	for k := range t.secondaryAnyIndexers {
+		idxs = append(idxs, k)
+	}
+	return idxs
+}
+
 func (t *genTable[Obj]) ToTable() Table[Obj] {
 	return t
 }
@@ -468,5 +478,18 @@ func (t *genTable[Obj]) sortableMutex() internal.SortableMutex {
 	return t.smu
 }
 
+func (t *genTable[Obj]) proto() any {
+	var zero Obj
+	return zero
+}
+
+func (t *genTable[Obj]) unmarshalYAML(data []byte) (any, error) {
+	var obj Obj
+	if err := yaml.Unmarshal(data, &obj); err != nil {
+		return nil, err
+	}
+	return obj, nil
+}
+
 var _ Table[bool] = &genTable[bool]{}
 var _ RWTable[bool] = &genTable[bool]{}