feat(db): support key in nested message

Resolve key fields declared in nested messages (for example metadata.id) across key lookup, indexability checks, order planning, and field-reader path resolution. Add recursive protoc-gen validation to reject invalid/ambiguous nested key declarations and extend tests for nested-key behavior.

Signed-off-by: Adphi <philippe.adrien.nousse@gmail.com>

Author: Adphi

cmd/protoc-gen-protodb/main.go

@@ -42,6 +42,8 @@ type mod struct {
 	oneOfs  map[string]struct{}
 }
 
+type keyPath []pgs.Field
+
 func (p *mod) Name() string {
 	return "protodb"
 }
@@ -85,28 +87,87 @@ func (p *mod) Execute(targets map[string]pgs.File, _ map[string]pgs.Package) []p
 			if !ok || !enabled {
 				continue
 			}
-			var keyField string
-			for _, f := range m.Fields() {
-				var k bool
-				f.Extension(protodb.E_Key, &k)
-				if !k {
-					continue
-				}
-				if keyField != "" {
-					p.Failf("%s: %s: key already defined for message: %s", m.Name(), f.Name(), keyField)
-				}
-				keyField = f.Name().String()
-				if f.Type().IsMap() || f.Type().IsRepeated() || f.InOneOf() || f.Descriptor().GetType() == descriptorpb.FieldDescriptorProto_TYPE_MESSAGE {
-					p.Failf("%s: %s: only non repeated and not oneof scalar types are supported as key, got %T", m.Name(), f.Name(), f.Type())
-				}
-			}
+			p.validateMessageKey(m)
 			msgs = append(msgs, m)
 		}
 		p.generate(f, msgs)
 	}
 	return p.Artifacts()
 }
 
+func (p *mod) validateMessageKey(m pgs.Message) {
+	keys := p.collectMessageKeys(m)
+	if len(keys) == 0 {
+		return
+	}
+	for _, path := range keys {
+		leaf := path.leaf()
+		if leaf.Type().IsMap() || leaf.Type().IsRepeated() || leaf.InOneOf() || leaf.Descriptor().GetType() == descriptorpb.FieldDescriptorProto_TYPE_MESSAGE {
+			p.Failf("%s: %s: only non repeated and not oneof scalar types are supported as key, got %T", m.Name(), path.String(), leaf.Type())
+		}
+		for i := 0; i < len(path)-1; i++ {
+			f := path[i]
+			if !f.Type().IsEmbed() || f.Type().IsMap() || f.Type().IsRepeated() || f.InOneOf() {
+				p.Failf("%s: %s: nested key path must traverse singular message fields", m.Name(), path.String())
+			}
+		}
+	}
+	if len(keys) == 1 {
+		return
+	}
+	paths := make([]string, 0, len(keys))
+	for _, path := range keys {
+		paths = append(paths, path.String())
+	}
+	p.Failf("%s: multiple keys found for message: %s", m.Name(), strings.Join(paths, ", "))
+}
+
+func (p *mod) collectMessageKeys(m pgs.Message) []keyPath {
+	stack := map[string]bool{m.FullyQualifiedName(): true}
+	return p.collectMessageKeysIn(m, nil, stack)
+}
+
+func (p *mod) collectMessageKeysIn(m pgs.Message, path keyPath, stack map[string]bool) []keyPath {
+	var out []keyPath
+	for _, f := range m.Fields() {
+		next := make(keyPath, 0, len(path)+1)
+		next = append(next, path...)
+		next = append(next, f)
+		var key bool
+		f.Extension(protodb.E_Key, &key)
+		if key {
+			out = append(out, next)
+		}
+		if !f.Type().IsEmbed() || f.Type().IsMap() || f.Type().IsRepeated() || f.InOneOf() {
+			continue
+		}
+		embed := f.Type().Embed()
+		if embed == nil {
+			continue
+		}
+		name := embed.FullyQualifiedName()
+		if stack[name] {
+			continue
+		}
+		stack[name] = true
+		out = append(out, p.collectMessageKeysIn(embed, next, stack)...)
+		delete(stack, name)
+	}
+	return out
+}
+
+func (p keyPath) leaf() pgs.Field {
+	return p[len(p)-1]
+}
+
+func (p keyPath) String() string {
+	parts := make([]string, 0, len(p))
+	for _, f := range p {
+		parts = append(parts, f.Name().String())
+	}
+	return strings.Join(parts, ".")
+}
+
 func (p *mod) generate(f pgs.File, msgs []pgs.Message) {
 	if len(msgs) == 0 {
 		return

internal/db/tx.go

@@ -1174,7 +1174,7 @@ func buildOrderPlan(msg protoreflect.Message, orderBy *v1alpha1.OrderBy) (orderF
 	if !isOrderableFieldPath(fds) {
 		return orderField{}, fmt.Errorf("order_by field %q is not sortable", fieldPath)
 	}
-	isKey := hasKey && len(fds) == 1 && string(fds[0].Name()) == keyField
+	isKey := hasKey && nameFieldPath(fds) == keyField
 	if !isKey {
 		fd := fds[len(fds)-1]
 		if !proto.HasExtension(fd.Options(), protopts.E_Index) {
@@ -1319,3 +1319,11 @@ func hash(f protodb.Filter) (hash string, err error) {
 	h := sha512.Sum512(b)
 	return base64.StdEncoding.EncodeToString(h[:]), nil
 }
+
+func nameFieldPath(fds []protoreflect.FieldDescriptor) string {
+	parts := make([]string, 0, len(fds))
+	for _, fd := range fds {
+		parts = append(parts, string(fd.Name()))
+	}
+	return strings.Join(parts, ".")
+}

internal/db/tx_test.go

@@ -6,13 +6,18 @@ import (
 	"testing"
 
 	"github.com/dgraph-io/badger/v3"
+	"github.com/jhump/protoreflect/v2/protobuilder"
 	"github.com/stretchr/testify/require"
 	"go.linka.cloud/protofilters/filters"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protodesc"
 	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/types/descriptorpb"
 	"google.golang.org/protobuf/types/dynamicpb"
 
 	"go.linka.cloud/protodb/internal/protodb"
 	"go.linka.cloud/protodb/internal/token"
+	protopts "go.linka.cloud/protodb/protodb"
 	v1alpha1 "go.linka.cloud/protodb/protodb/v1alpha1"
 )
 
@@ -55,6 +60,18 @@ func TestBuildOrderPlan(t *testing.T) {
 		_, err := buildOrderPlan(m.ProtoReflect(), &v1alpha1.OrderBy{Field: "status", Direction: v1alpha1.OrderDirection(99)})
 		require.ErrorContains(t, err, "invalid direction")
 	})
+
+	require.NoError(t, d.RegisterProto(ctx, buildNestedKeyOrderFileDescriptor(t)))
+	nmd, err := lookupMessage(d, "tests.order.OrderDoc")
+	require.NoError(t, err)
+	nm := dynamicpb.NewMessage(nmd)
+
+	t.Run("nested key field", func(t *testing.T) {
+		pl, err := buildOrderPlan(nm.ProtoReflect(), &v1alpha1.OrderBy{Field: "metadata.id"})
+		require.NoError(t, err)
+		require.Equal(t, "metadata.id", pl.fieldPath)
+		require.Equal(t, v1alpha1.OrderDirectionAsc, pl.direction)
+	})
 }
 
 func TestGetContinuationValidation(t *testing.T) {
@@ -297,6 +314,30 @@ func lookupNumberField(md protoreflect.MessageDescriptor, n string) (protoreflec
 	return nil, fmt.Errorf("field number %s not found", n)
 }
 
+func buildNestedKeyOrderFileDescriptor(t *testing.T) *descriptorpb.FileDescriptorProto {
+	t.Helper()
+	keyOpts := &descriptorpb.FieldOptions{}
+	proto.SetExtension(keyOpts, protopts.E_Key, true)
+
+	meta := protobuilder.NewMessage("Metadata").
+		AddField(protobuilder.NewField("id", protobuilder.FieldTypeString()).SetNumber(1).SetOptions(keyOpts))
+
+	msg := protobuilder.NewMessage("OrderDoc").
+		AddField(protobuilder.NewField("metadata", protobuilder.FieldTypeMessage(meta)).SetNumber(1)).
+		AddField(protobuilder.NewField("status", protobuilder.FieldTypeString()).SetNumber(2)).
+		AddNestedMessage(meta)
+
+	file := protobuilder.NewFile("tests/order_nested_key.proto").
+		SetPackageName(protoreflect.FullName("tests.order")).
+		SetSyntax(protoreflect.Proto3).
+		AddMessage(msg).
+		AddImportedDependency(protopts.File_protodb_protodb_proto)
+
+	fd, err := file.Build()
+	require.NoError(t, err)
+	return protodesc.ToFileDescriptorProto(fd)
+}
+
 func TestGetWithMalformedTokenReturnsInvalid(t *testing.T) {
 	ctx := context.Background()
 	dbif, err := Open(ctx, WithPath(t.TempDir()), WithApplyDefaults(true))

internal/index/indexer.go

@@ -984,14 +984,14 @@ func isKeyField(md protoreflect.MessageDescriptor, fds []protoreflect.FieldDescr
 	if md == nil {
 		return false
 	}
-	if len(fds) != 1 {
+	if len(fds) == 0 {
 		return false
 	}
 	field, ok := protodb.KeyFieldName(md)
 	if !ok {
 		return false
 	}
-	return field == string(fds[0].Name())
+	return field == fieldPathFromNames(fds)
 }
 
 func isIndexableFieldPathDescriptors(fds []protoreflect.FieldDescriptor) bool {

internal/index/indexer_test.go

@@ -52,6 +52,12 @@ func TestIndexableFilter(t *testing.T) {
 
 	_, err = idx.IndexableFilter(m, filters.Where("meta.nope").StringEquals("x"))
 	require.Error(t, err)
+
+	nestedKey := buildIndexerNestedKeyDescriptor(t)
+	nestedMsg := dynamicpb.NewMessage(nestedKey)
+	ok, err = idx.IndexableFilter(nestedMsg, filters.Where("metadata.id").StringEquals("x"))
+	require.NoError(t, err)
+	require.True(t, ok)
 }
 
 func TestEnforceUnique(t *testing.T) {
@@ -349,6 +355,32 @@ func buildUniqueDocDescriptor(t *testing.T) protoreflect.MessageDescriptor {
 	return md
 }
 
+func buildIndexerNestedKeyDescriptor(t *testing.T) protoreflect.MessageDescriptor {
+	t.Helper()
+	keyOpts := &descriptorpb.FieldOptions{}
+	proto.SetExtension(keyOpts, protopts.E_Key, true)
+
+	meta := protobuilder.NewMessage("Metadata").
+		AddField(protobuilder.NewField("id", protobuilder.FieldTypeString()).SetNumber(1).SetOptions(keyOpts))
+
+	msg := protobuilder.NewMessage("NestedKeyDoc").
+		AddField(protobuilder.NewField("metadata", protobuilder.FieldTypeMessage(meta)).SetNumber(1)).
+		AddField(protobuilder.NewField("status", protobuilder.FieldTypeString()).SetNumber(2).SetOptions(&descriptorpb.FieldOptions{})).
+		AddNestedMessage(meta)
+
+	file := protobuilder.NewFile("tests/indexer_nested_key.proto").
+		SetPackageName(protoreflect.FullName("tests.index")).
+		SetSyntax(protoreflect.Proto3).
+		AddMessage(msg).
+		AddImportedDependency(protopts.File_protodb_protodb_proto)
+
+	fd, err := file.Build()
+	require.NoError(t, err)
+	md := fd.Messages().ByName("NestedKeyDoc")
+	require.NotNil(t, md)
+	return md
+}
+
 func buildUniqueTagsDescriptor(t *testing.T) protoreflect.MessageDescriptor {
 	t.Helper()
 	keyOpts := &descriptorpb.FieldOptions{}

internal/index/store.go

@@ -321,7 +321,7 @@ func (f *fieldReader) Get(_ context.Context, n protoreflect.Name) iter.Seq2[pfin
 
 type keyField struct {
 	name       protoreflect.Name
-	fd         protoreflect.FieldDescriptor
+	fds        []protoreflect.FieldDescriptor
 	dataPrefix []byte
 	txn        badgerd.Tx
 	entries    []keyEntry
@@ -371,7 +371,7 @@ func (k *keyField) iterate(yield func(pfindex.Field, error) bool) {
 	}
 	for _, e := range k.entries {
 		entry := &field{
-			descriptors: []protoreflect.FieldDescriptor{k.fd},
+			descriptors: k.fds,
 			value:       protoreflect.ValueOfString(e.value),
 			bitmap:      e.bitmap,
 		}
@@ -558,13 +558,7 @@ func fieldPathFromNames(fds []protoreflect.FieldDescriptor) string {
 }
 
 func lookupByNumberPath(md protoreflect.MessageDescriptor, fieldPath string) ([]protoreflect.FieldDescriptor, error) {
-	parts := strings.Split(fieldPath, ".")
-	if len(parts) == 0 {
-		return nil, fmt.Errorf("empty field path")
-	}
-	var fds []protoreflect.FieldDescriptor
-	cur := md
-	for _, part := range parts {
+	return lookupByPath(md, fieldPath, func(cur protoreflect.MessageDescriptor, part string) (protoreflect.FieldDescriptor, error) {
 		num, err := strconv.Atoi(part)
 		if err != nil {
 			return nil, fmt.Errorf("invalid field number %q", part)
@@ -573,15 +567,40 @@ func lookupByNumberPath(md protoreflect.MessageDescriptor, fieldPath string) ([]
 		if fd == nil {
 			return nil, fmt.Errorf("%s does not contain field number %d", cur.FullName(), num)
 		}
+		return fd, nil
+	})
+}
+
+func lookupByNamePath(md protoreflect.MessageDescriptor, fieldPath string) ([]protoreflect.FieldDescriptor, error) {
+	return lookupByPath(md, fieldPath, func(cur protoreflect.MessageDescriptor, part string) (protoreflect.FieldDescriptor, error) {
+		fd := cur.Fields().ByName(protoreflect.Name(part))
+		if fd == nil {
+			return nil, fmt.Errorf("%s does not contain field '%s'", cur.FullName(), part)
+		}
+		return fd, nil
+	})
+}
+
+func lookupByPath(md protoreflect.MessageDescriptor, fieldPath string, lookup func(protoreflect.MessageDescriptor, string) (protoreflect.FieldDescriptor, error)) ([]protoreflect.FieldDescriptor, error) {
+	if fieldPath == "" {
+		return nil, fmt.Errorf("empty field path")
+	}
+	parts := strings.Split(fieldPath, ".")
+	var fds []protoreflect.FieldDescriptor
+	cur := md
+	for i, part := range parts {
+		fd, err := lookup(cur, part)
+		if err != nil {
+			return nil, err
+		}
 		fds = append(fds, fd)
-		if fd.Kind() == protoreflect.MessageKind {
-			cur = fd.Message()
-		} else {
-			cur = nil
+		if i == len(parts)-1 {
+			continue
 		}
-		if cur == nil && part != parts[len(parts)-1] {
+		if fd.Kind() != protoreflect.MessageKind {
 			return nil, fmt.Errorf("%s does not contain '%s'", md.FullName(), fieldPath)
 		}
+		cur = fd.Message()
 	}
 	return fds, nil
 }
@@ -598,11 +617,11 @@ func buildFieldReader(txn badgerd.Tx, resolver protodesc.Resolver, name protoref
 	var key *keyField
 	keyName, ok := protodb.KeyFieldName(md)
 	if ok {
-		fd := md.Fields().ByName(protoreflect.Name(keyName))
-		if fd != nil {
+		fds, err := lookupByNamePath(md, keyName)
+		if err == nil {
 			key = &keyField{
-				name:       fd.Name(),
-				fd:         fd,
+				name:       protoreflect.Name(fieldPathFromNames(fds)),
+				fds:        fds,
 				dataPrefix: []byte(protodb.Data + "/" + string(name) + "/"),
 				txn:        txn,
 			}