fix(orm): codegen issues (#11061)

orm/internal/codegen/index.go

@@ -20,9 +20,9 @@ func (t tableGen) genIndexKeys() {
 
 	// start with primary key..
 	t.P("// primary key starting index..")
-	t.genIndex(t.table.PrimaryKey.Fields, t.ormTable.ID())
+	t.genIndex(t.table.PrimaryKey.Fields, 0, true)
 	for _, idx := range t.table.Index {
-		t.genIndex(idx.Fields, idx.Id)
+		t.genIndex(idx.Fields, idx.Id, false)
 	}
 }
 
@@ -94,9 +94,15 @@ func (t tableGen) indexStructName(fields []string) string {
 	return t.msg.GoIdent.GoName + joinedNames + "IndexKey"
 }
 
-func (t tableGen) genIndex(fields string, id uint32) {
+func (t tableGen) genIndex(fields string, id uint32, isPrimaryKey bool) {
 	fieldsSlc := strings.Split(fields, ",")
 	idxKeyName := t.indexStructName(fieldsSlc)
+
+	if isPrimaryKey {
+		t.P("type ", t.msg.GoIdent.GoName, "PrimaryKey = ", idxKeyName)
+		t.P()
+	}
+
 	t.P("type ", idxKeyName, " struct {")
 	t.P("vs []interface{}")
 	t.P("}")

orm/internal/codegen/singleton.go

@@ -60,12 +60,9 @@ func (s singletonGen) genMethods() {
 	varName := s.param(s.msg.GoIdent.GoName)
 	// Get
 	s.P(receiver, "Get(ctx ", contextPkg.Ident("Context"), ") (*", s.msg.GoIdent.GoName, ", error) {")
-	s.P("var ", varName, " ", s.msg.GoIdent.GoName)
-	s.P("found, err := x.table.Get(ctx, &", varName, ")")
-	s.P("if !found {")
-	s.P("return nil, err")
-	s.P("}")
-	s.P("return &", varName, ", err")
+	s.P(varName, " := &", s.msg.GoIdent.GoName, "{}")
+	s.P("_, err := x.table.Get(ctx, ", varName, ")")
+	s.P("return ", varName, ", err")
 	s.P("}")
 	s.P()
 

orm/internal/codegen/table.go

@@ -4,8 +4,6 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/iancoleman/strcase"
-
 	"google.golang.org/protobuf/compiler/protogen"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/dynamicpb"
@@ -189,43 +187,45 @@ func (t tableGen) genStoreImpl() {
 
 		// has
 		t.P("func (", receiverVar, " ", t.messageStoreReceiverName(t.msg), ") ", hasName, "{")
-		t.P("return ", receiverVar, ".table.Has(ctx, &", t.msg.GoIdent.GoName, "{")
+		t.P("return ", receiverVar, ".table.GetIndexByID(", idx.Id, ").(",
+			tablePkg.Ident("UniqueIndex"), ").Has(ctx,")
 		for _, field := range fields {
-			t.P(strcase.ToCamel(field), ": ", field, ",")
+			t.P(field, ",")
 		}
-		t.P("})")
+		t.P(")")
 		t.P("}")
 		t.P()
 
 		// get
 		varName := t.param(t.msg.GoIdent.GoName)
 		varTypeName := t.msg.GoIdent.GoName
 		t.P("func (", receiverVar, " ", t.messageStoreReceiverName(t.msg), ") ", getName, "{")
-		t.P(varName, " := &", varTypeName, "{")
+		t.P("var ", varName, " ", varTypeName)
+		t.P("found, err := ", receiverVar, ".table.GetIndexByID(", idx.Id, ").(",
+			tablePkg.Ident("UniqueIndex"), ").Get(ctx, &", varName, ",")
 		for _, field := range fields {
-			t.P(strcase.ToCamel(field), ": ", field, ",")
+			t.P(field, ",")
 		}
-		t.P("}")
-		t.P("found, err := ", receiverVar, ".table.Get(ctx, ", varName, ")")
+		t.P(")")
 		t.P("if !found {")
 		t.P("return nil, err")
 		t.P("}")
-		t.P("return ", varName, ", nil")
+		t.P("return &", varName, ", nil")
 		t.P("}")
 		t.P()
 	}
 
 	// List
 	t.P(receiver, "List(ctx ", contextPkg.Ident("Context"), ", prefixKey ", t.indexKeyInterfaceName(), ", opts ...", ormListPkg.Ident("Option"), ") (", t.iteratorName(), ", error) {")
-	t.P("opts = append(opts, ", ormListPkg.Ident("Prefix"), "(prefixKey.values()))")
+	t.P("opts = append(opts, ", ormListPkg.Ident("Prefix"), "(prefixKey.values()...))")
 	t.P("it, err := ", receiverVar, ".table.GetIndexByID(prefixKey.id()).Iterator(ctx, opts...)")
 	t.P("return ", t.iteratorName(), "{it}, err")
 	t.P("}")
 	t.P()
 
 	// ListRange
 	t.P(receiver, "ListRange(ctx ", contextPkg.Ident("Context"), ", from, to ", t.indexKeyInterfaceName(), ", opts ...", ormListPkg.Ident("Option"), ") (", t.iteratorName(), ", error) {")
-	t.P("opts = append(opts, ", ormListPkg.Ident("Start"), "(from.values()), ", ormListPkg.Ident("End"), "(to))")
+	t.P("opts = append(opts, ", ormListPkg.Ident("Start"), "(from.values()...), ", ormListPkg.Ident("End"), "(to.values()...))")
 	t.P("it, err := ", receiverVar, ".table.GetIndexByID(from.id()).Iterator(ctx, opts...)")
 	t.P("return ", t.iteratorName(), "{it}, err")
 	t.P("}")

orm/model/ormdb/module.go

@@ -31,12 +31,7 @@ type ModuleSchema struct {
 }
 
 // ModuleDB defines the ORM database type to be used by modules.
-type ModuleDB interface {
-	ormkv.EntryCodec
-
-	// GetTable returns the table for the provided message type or nil.
-	GetTable(message proto.Message) ormtable.Table
-}
+type ModuleDB = ormtable.Schema
 
 type moduleDB struct {
 	prefix       []byte

orm/model/ormtable/auto_increment.go

@@ -216,3 +216,10 @@ func (t autoIncrementTable) ExportJSON(ctx context.Context, writer io.Writer) er
 
 	return t.doExportJSON(ctx, writer)
 }
+
+func (t *autoIncrementTable) GetTable(message proto.Message) Table {
+	if message.ProtoReflect().Descriptor().FullName() == t.MessageType().Descriptor().FullName() {
+		return t
+	}
+	return nil
+}