Reimplement script commands against hive script

any_table.go

@@ -1,6 +1,10 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Cilium
+
 package statedb
 
 import (
+	"fmt"
 	"iter"
 )
 
@@ -12,8 +16,13 @@ type AnyTable struct {
 }
 
 func (t AnyTable) All(txn ReadTxn) iter.Seq2[any, Revision] {
+	all, _ := t.AllWatch(txn)
+	return all
+}
+
+func (t AnyTable) AllWatch(txn ReadTxn) (iter.Seq2[any, Revision], <-chan struct{}) {
 	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
-	return partSeq[any](indexTxn.Iterator())
+	return partSeq[any](indexTxn.Iterator()), indexTxn.RootWatch()
 }
 
 func (t AnyTable) UnmarshalYAML(data []byte) (any, error) {
@@ -38,22 +47,85 @@ func (t AnyTable) Delete(txn WriteTxn, obj any) (old any, hadOld bool, err error
 	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))
-	if indexTxn.unique {
-		return partSeq[any](iter)
+func (t AnyTable) Get(txn ReadTxn, index string, key string) (any, Revision, bool, error) {
+	itxn, rawKey, err := t.queryIndex(txn, index, key)
+	if err != nil {
+		return nil, 0, false, err
+	}
+	if itxn.unique {
+		obj, _, ok := itxn.Get(rawKey)
+		return obj.data, obj.revision, ok, nil
 	}
-	return nonUniqueSeq[any](iter, true, []byte(key))
+	// For non-unique indexes we need to prefix search and make sure to fully
+	// match the secondary key.
+	iter, _ := itxn.Prefix(rawKey)
+	for {
+		k, obj, ok := iter.Next()
+		if !ok {
+			break
+		}
+		secondary, _ := decodeNonUniqueKey(k)
+		if len(secondary) == len(rawKey) {
+			return obj.data, obj.revision, true, nil
+		}
+	}
+	return nil, 0, false, nil
 }
 
-func (t AnyTable) LowerBound(txn ReadTxn, key string) iter.Seq2[any, Revision] {
-	indexTxn := txn.getTxn().mustIndexReadTxn(t.Meta, PrimaryIndexPos)
-	iter := indexTxn.LowerBound([]byte(key))
-	if indexTxn.unique {
-		return partSeq[any](iter)
+func (t AnyTable) Prefix(txn ReadTxn, index string, key string) (iter.Seq2[any, Revision], error) {
+	itxn, rawKey, err := t.queryIndex(txn, index, key)
+	if err != nil {
+		return nil, err
+	}
+	iter, _ := itxn.Prefix(rawKey)
+	if itxn.unique {
+		return partSeq[any](iter), nil
+	}
+	return nonUniqueSeq[any](iter, true, rawKey), nil
+}
+
+func (t AnyTable) LowerBound(txn ReadTxn, index string, key string) (iter.Seq2[any, Revision], error) {
+	itxn, rawKey, err := t.queryIndex(txn, index, key)
+	if err != nil {
+		return nil, err
+	}
+	iter := itxn.LowerBound(rawKey)
+	if itxn.unique {
+		return partSeq[any](iter), nil
+	}
+	return nonUniqueLowerBoundSeq[any](iter, rawKey), nil
+}
+
+func (t AnyTable) List(txn ReadTxn, index string, key string) (iter.Seq2[any, Revision], error) {
+	itxn, rawKey, err := t.queryIndex(txn, index, key)
+	if err != nil {
+		return nil, err
+	}
+	iter, _ := itxn.Prefix(rawKey)
+	if itxn.unique {
+		// Unique index means that there can be only a single matching object.
+		// Doing a Get() is more efficient than constructing an iterator.
+		value, _, ok := itxn.Get(rawKey)
+		return func(yield func(any, Revision) bool) {
+			if ok {
+				yield(value.data, value.revision)
+			}
+		}, nil
+	}
+	return nonUniqueSeq[any](iter, false, rawKey), nil
+}
+
+func (t AnyTable) queryIndex(txn ReadTxn, index string, key string) (indexReadTxn, []byte, error) {
+	indexer := t.Meta.getIndexer(index)
+	if indexer == nil {
+		return indexReadTxn{}, nil, fmt.Errorf("invalid index %q", index)
+	}
+	rawKey, err := indexer.fromString(key)
+	if err != nil {
+		return indexReadTxn{}, nil, err
 	}
-	return nonUniqueLowerBoundSeq[any](iter, []byte(key))
+	itxn, err := txn.getTxn().indexReadTxn(t.Meta, indexer.pos)
+	return itxn, rawKey, err
 }
 
 func (t AnyTable) TableHeader() []string {

benchmarks_test.go

@@ -6,6 +6,7 @@ package statedb
 import (
 	"context"
 	"iter"
+	"log/slog"
 	"math/rand"
 	"sort"
 	"testing"
@@ -16,6 +17,7 @@ import (
 
 	"github.com/cilium/hive"
 	"github.com/cilium/hive/cell"
+	"github.com/cilium/hive/hivetest"
 	"github.com/cilium/statedb/index"
 	"github.com/cilium/statedb/part"
 )
@@ -471,19 +473,18 @@ func BenchmarkDB_PropagationDelay(b *testing.B) {
 		table2 = MustNewTable("test2", id2Index)
 	)
 
-	h := hive.NewWithOptions(
-		hive.Options{Logger: logger},
-
+	h := hive.New(
 		Cell, // DB
 		cell.Invoke(func(db_ *DB) error {
 			db = db_
 			return db.RegisterTable(table1, table2)
 		}),
 	)
 
-	require.NoError(b, h.Start(context.TODO()))
+	log := hivetest.Logger(b, hivetest.LogLevel(slog.LevelError))
+	require.NoError(b, h.Start(log, context.TODO()))
 	b.Cleanup(func() {
-		assert.NoError(b, h.Stop(context.TODO()))
+		assert.NoError(b, h.Stop(log, context.TODO()))
 	})
 
 	b.ResetTimer()

cell.go

@@ -16,6 +16,7 @@ var Cell = cell.Module(
 
 	cell.Provide(
 		newHiveDB,
+		ScriptCommands,
 	),
 )
 

db.go

@@ -208,9 +208,20 @@ func (db *DB) WriteTxn(table TableMeta, tables ...TableMeta) WriteTxn {
 	lockAt := time.Now()
 	smus.Lock()
 	acquiredAt := time.Now()
-
 	root := *db.root.Load()
 	tableEntries := make([]*tableEntry, len(root))
+
+	txn := &txn{
+		db:         db,
+		root:       root,
+		handle:     db.handleName,
+		acquiredAt: time.Now(),
+		writeTxn: writeTxn{
+			modifiedTables: tableEntries,
+			smus:           smus,
+		},
+	}
+
 	var tableNames []string
 	for _, table := range allTables {
 		tableEntry := root[table.tablePos()]
@@ -223,26 +234,19 @@ func (db *DB) WriteTxn(table TableMeta, tables ...TableMeta) WriteTxn {
 			table.Name(),
 			table.sortableMutex().AcquireDuration(),
 		)
+		table.acquired(txn)
 	}
 
 	// Sort the table names so they always appear ordered in metrics.
 	sort.Strings(tableNames)
+	txn.tableNames = tableNames
 
 	db.metrics.WriteTxnTotalAcquisition(
 		db.handleName,
 		tableNames,
 		acquiredAt.Sub(lockAt),
 	)
 
-	txn := &txn{
-		db:             db,
-		root:           root,
-		modifiedTables: tableEntries,
-		smus:           smus,
-		acquiredAt:     acquiredAt,
-		tableNames:     tableNames,
-		handle:         db.handleName,
-	}
 	runtime.SetFinalizer(txn, txnFinalizer)
 	return txn
 }

db_test.go

@@ -9,9 +9,10 @@ import (
 	"expvar"
 	"fmt"
 	"log/slog"
-	"os"
 	"runtime"
 	"slices"
+	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -21,6 +22,7 @@ import (
 
 	"github.com/cilium/hive"
 	"github.com/cilium/hive/cell"
+	"github.com/cilium/hive/hivetest"
 	"github.com/cilium/statedb/index"
 	"github.com/cilium/statedb/part"
 	"github.com/cilium/stream"
@@ -47,36 +49,54 @@ func (t testObject) String() string {
 	return fmt.Sprintf("testObject{ID: %d, Tags: %v}", t.ID, t.Tags)
 }
 
+func (t testObject) TableHeader() []string {
+	return []string{"ID", "Tags"}
+}
+
+func (t testObject) TableRow() []string {
+	return []string{
+		strconv.FormatUint(uint64(t.ID), 10),
+		strings.Join(slices.Collect(t.Tags.All()), ", "),
+	}
+}
+
 var (
 	idIndex = Index[testObject, uint64]{
 		Name: "id",
 		FromObject: func(t testObject) index.KeySet {
 			return index.NewKeySet(index.Uint64(t.ID))
 		},
-		FromKey: index.Uint64,
-		Unique:  true,
+		FromKey:    index.Uint64,
+		FromString: index.Uint64String,
+		Unique:     true,
 	}
 
 	tagsIndex = Index[testObject, string]{
 		Name: "tags",
 		FromObject: func(t testObject) index.KeySet {
 			return index.Set(t.Tags)
 		},
-		FromKey: index.String,
-		Unique:  false,
+		FromKey:    index.String,
+		FromString: index.FromString,
+		Unique:     false,
 	}
 )
 
+func newTestObjectTable(t testing.TB, name string, secondaryIndexers ...Indexer[testObject]) RWTable[testObject] {
+	table, err := NewTable(
+		name,
+		idIndex,
+		secondaryIndexers...,
+	)
+	require.NoError(t, err, "NewTable[testObject]")
+	return table
+}
+
 const (
 	INDEX_TAGS    = true
 	NO_INDEX_TAGS = false
 )
 
-// Do not log debug&info level logs in tests.
-var logger = slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
-	Level: slog.LevelError,
-}))
-
 func newTestDB(t testing.TB, secondaryIndexers ...Indexer[testObject]) (*DB, RWTable[testObject], *ExpVarMetrics) {
 	metrics := NewExpVarMetrics(false)
 	db, table := newTestDBWithMetrics(t, metrics, secondaryIndexers...)
@@ -87,16 +107,9 @@ func newTestDBWithMetrics(t testing.TB, metrics Metrics, secondaryIndexers ...In
 	var (
 		db *DB
 	)
-	table, err := NewTable(
-		"test",
-		idIndex,
-		secondaryIndexers...,
-	)
-	require.NoError(t, err, "NewTable[testObject]")
-
-	h := hive.NewWithOptions(
-		hive.Options{Logger: logger},
+	table := newTestObjectTable(t, "test", secondaryIndexers...)
 
+	h := hive.New(
 		cell.Provide(func() Metrics { return metrics }),
 		Cell, // DB
 		cell.Invoke(func(db_ *DB) {
@@ -110,9 +123,10 @@ func newTestDBWithMetrics(t testing.TB, metrics Metrics, secondaryIndexers ...In
 		}),
 	)
 
-	require.NoError(t, h.Start(context.TODO()))
+	log := hivetest.Logger(t, hivetest.LogLevel(slog.LevelError))
+	require.NoError(t, h.Start(log, context.TODO()))
 	t.Cleanup(func() {
-		assert.NoError(t, h.Stop(context.TODO()))
+		assert.NoError(t, h.Stop(log, context.TODO()))
 	})
 	return db, table
 }
@@ -243,7 +257,7 @@ func TestDB_Prefix(t *testing.T) {
 	txn := db.ReadTxn()
 
 	iter, watch := table.PrefixWatch(txn, tagsIndex.Query("ab"))
-	require.Equal(t, Collect(Map(iter, testObject.getID)), []uint64{71, 82})
+	require.Equal(t, []uint64{71, 82}, Collect(Map(iter, testObject.getID)))
 
 	select {
 	case <-watch:

derive_test.go

@@ -5,6 +5,7 @@ package statedb
 
 import (
 	"context"
+	"log/slog"
 	"slices"
 	"testing"
 	"time"
@@ -13,6 +14,7 @@ import (
 
 	"github.com/cilium/hive"
 	"github.com/cilium/hive/cell"
+	"github.com/cilium/hive/hivetest"
 	"github.com/cilium/hive/job"
 	"github.com/cilium/statedb/index"
 	"github.com/cilium/statedb/part"
@@ -52,6 +54,8 @@ func (*nopHealth) OK(status string) {
 func (*nopHealth) Stopped(reason string) {
 }
 
+func (*nopHealth) Close() {}
+
 func newNopHealth() (cell.Health, *nopHealth) {
 	h := &nopHealth{}
 	return h, h
@@ -103,7 +107,8 @@ func TestDerive(t *testing.T) {
 			cell.Invoke(Derive("testObject-to-derived", transform)),
 		),
 	)
-	require.NoError(t, h.Start(context.TODO()), "Start")
+	log := hivetest.Logger(t, hivetest.LogLevel(slog.LevelError))
+	require.NoError(t, h.Start(log, context.TODO()), "Start")
 
 	getDerived := func() []derived {
 		txn := db.ReadTxn()
@@ -174,5 +179,5 @@ func TestDerive(t *testing.T) {
 		"expected 1 to be gone, and 2 mark deleted",
 	)
 
-	require.NoError(t, h.Stop(context.TODO()), "Stop")
+	require.NoError(t, h.Stop(log, context.TODO()), "Stop")
 }