feat(orm): ModuleDB JSON import/export/default/validate methods (#11101)

* feat(orm): ModuleDB JSON methods

* WIP

* WIP on JSON

* WIP

* WIP

* tests and docs

* revert

* tests and docs

* docs

* address review comments

orm/encoding/docs.go

@@ -0,0 +1,3 @@
+// Package encoding defines the core types and algorithms for encoding and decoding
+// protobuf objects and values to/from ORM key-value pairs.
+package encoding

orm/internal/codegen/file.go

@@ -138,7 +138,7 @@ func (f fileGen) genStoreInterfaceGuard() {
 }
 
 func (f fileGen) genStoreConstructor(stores []*protogen.Message) {
-	f.P("func New", f.storeInterfaceName(), "(db ", ormdbPkg.Ident("ModuleDB"), ") (", f.storeInterfaceName(), ", error) {")
+	f.P("func New", f.storeInterfaceName(), "(db ", ormTablePkg.Ident("Schema"), ") (", f.storeInterfaceName(), ", error) {")
 	for _, store := range stores {
 		f.P(f.messageStoreReceiverName(store), ", err := ", f.messageConstructorName(store), "(db)")
 		f.P("if err != nil {")

orm/internal/codegen/singleton.go

@@ -75,7 +75,7 @@ func (s singletonGen) genMethods() {
 
 func (s singletonGen) genConstructor() {
 	iface := s.messageStoreInterfaceName(s.msg)
-	s.P("func New", iface, "(db ", ormdbPkg.Ident("ModuleDB"), ") (", iface, ", error) {")
+	s.P("func New", iface, "(db ", ormTablePkg.Ident("Schema"), ") (", iface, ", error) {")
 	s.P("table := db.GetTable(&", s.msg.GoIdent.GoName, "{})")
 	s.P("if table == nil {")
 	s.P("return nil, ", ormErrPkg.Ident("TableNotFound.Wrap"), "(string((&", s.msg.GoIdent.GoName, "{}).ProtoReflect().Descriptor().FullName()))")

orm/internal/codegen/table.go

@@ -270,7 +270,7 @@ func (t tableGen) genStoreImplGuard() {
 
 func (t tableGen) genConstructor() {
 	iface := t.messageStoreInterfaceName(t.msg)
-	t.P("func New", iface, "(db ", ormdbPkg.Ident("ModuleDB"), ") (", iface, ", error) {")
+	t.P("func New", iface, "(db ", ormTablePkg.Ident("Schema"), ") (", iface, ", error) {")
 	t.P("table := db.GetTable(&", t.msg.GoIdent.GoName, "{})")
 	t.P("if table == nil {")
 	t.P("return nil,", ormErrPkg.Ident("TableNotFound.Wrap"), "(string((&", t.msg.GoIdent.GoName, "{}).ProtoReflect().Descriptor().FullName()))")

orm/internal/testkv/compare.go

@@ -0,0 +1,38 @@
+package testkv
+
+import (
+	"bytes"
+
+	"gotest.tools/v3/assert"
+
+	"github.com/cosmos/cosmos-sdk/orm/model/ormtable"
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
+)
+
+func AssertBackendsEqual(t assert.TestingT, b1, b2 ormtable.Backend) {
+	it1, err := b1.CommitmentStoreReader().Iterator(nil, nil)
+	assert.NilError(t, err)
+
+	it2, err := b2.CommitmentStoreReader().Iterator(nil, nil)
+	assert.NilError(t, err)
+
+	AssertIteratorsEqual(t, it1, it2)
+
+	it1, err = b1.IndexStoreReader().Iterator(nil, nil)
+	assert.NilError(t, err)
+
+	it2, err = b2.IndexStoreReader().Iterator(nil, nil)
+	assert.NilError(t, err)
+
+	AssertIteratorsEqual(t, it1, it2)
+}
+
+func AssertIteratorsEqual(t assert.TestingT, it1, it2 kv.Iterator) {
+	for it1.Valid() {
+		assert.Assert(t, it2.Valid())
+		assert.Assert(t, bytes.Equal(it1.Key(), it2.Key()))
+		assert.Assert(t, bytes.Equal(it1.Value(), it2.Value()))
+		it1.Next()
+		it2.Next()
+	}
+}

orm/internal/testkv/debug.go

@@ -3,13 +3,12 @@ package testkv
 import (
 	"fmt"
 
-	"github.com/cosmos/cosmos-sdk/orm/internal/stablejson"
-
 	"google.golang.org/protobuf/proto"
 
 	"github.com/cosmos/cosmos-sdk/orm/encoding/ormkv"
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
+	"github.com/cosmos/cosmos-sdk/orm/internal/stablejson"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormtable"
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
 )
 
 // Debugger is an interface that handles debug info from the debug store wrapper.

orm/internal/testpb/bank.cosmos_orm.go

@@ -5,7 +5,6 @@ package testpb
 import (
 	context "context"
 
-	ormdb "github.com/cosmos/cosmos-sdk/orm/model/ormdb"
 	ormlist "github.com/cosmos/cosmos-sdk/orm/model/ormlist"
 	ormtable "github.com/cosmos/cosmos-sdk/orm/model/ormtable"
 	ormerrors "github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
@@ -131,7 +130,7 @@ func (this balanceStore) doNotImplement() {}
 
 var _ BalanceStore = balanceStore{}
 
-func NewBalanceStore(db ormdb.ModuleDB) (BalanceStore, error) {
+func NewBalanceStore(db ormtable.Schema) (BalanceStore, error) {
 	table := db.GetTable(&Balance{})
 	if table == nil {
 		return nil, ormerrors.TableNotFound.Wrap(string((&Balance{}).ProtoReflect().Descriptor().FullName()))
@@ -241,7 +240,7 @@ func (this supplyStore) doNotImplement() {}
 
 var _ SupplyStore = supplyStore{}
 
-func NewSupplyStore(db ormdb.ModuleDB) (SupplyStore, error) {
+func NewSupplyStore(db ormtable.Schema) (SupplyStore, error) {
 	table := db.GetTable(&Supply{})
 	if table == nil {
 		return nil, ormerrors.TableNotFound.Wrap(string((&Supply{}).ProtoReflect().Descriptor().FullName()))
@@ -273,7 +272,7 @@ func (bankStore) doNotImplement() {}
 
 var _ BankStore = bankStore{}
 
-func NewBankStore(db ormdb.ModuleDB) (BankStore, error) {
+func NewBankStore(db ormtable.Schema) (BankStore, error) {
 	balanceStore, err := NewBalanceStore(db)
 	if err != nil {
 		return nil, err

orm/internal/testpb/test_schema.cosmos_orm.go

@@ -5,7 +5,6 @@ package testpb
 import (
 	context "context"
 
-	ormdb "github.com/cosmos/cosmos-sdk/orm/model/ormdb"
 	ormlist "github.com/cosmos/cosmos-sdk/orm/model/ormlist"
 	ormtable "github.com/cosmos/cosmos-sdk/orm/model/ormtable"
 	ormerrors "github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
@@ -198,7 +197,7 @@ func (this exampleTableStore) doNotImplement() {}
 
 var _ ExampleTableStore = exampleTableStore{}
 
-func NewExampleTableStore(db ormdb.ModuleDB) (ExampleTableStore, error) {
+func NewExampleTableStore(db ormtable.Schema) (ExampleTableStore, error) {
 	table := db.GetTable(&ExampleTable{})
 	if table == nil {
 		return nil, ormerrors.TableNotFound.Wrap(string((&ExampleTable{}).ProtoReflect().Descriptor().FullName()))
@@ -345,7 +344,7 @@ func (this exampleAutoIncrementTableStore) doNotImplement() {}
 
 var _ ExampleAutoIncrementTableStore = exampleAutoIncrementTableStore{}
 
-func NewExampleAutoIncrementTableStore(db ormdb.ModuleDB) (ExampleAutoIncrementTableStore, error) {
+func NewExampleAutoIncrementTableStore(db ormtable.Schema) (ExampleAutoIncrementTableStore, error) {
 	table := db.GetTable(&ExampleAutoIncrementTable{})
 	if table == nil {
 		return nil, ormerrors.TableNotFound.Wrap(string((&ExampleAutoIncrementTable{}).ProtoReflect().Descriptor().FullName()))
@@ -375,7 +374,7 @@ func (x exampleSingletonStore) Save(ctx context.Context, exampleSingleton *Examp
 	return x.table.Save(ctx, exampleSingleton)
 }
 
-func NewExampleSingletonStore(db ormdb.ModuleDB) (ExampleSingletonStore, error) {
+func NewExampleSingletonStore(db ormtable.Schema) (ExampleSingletonStore, error) {
 	table := db.GetTable(&ExampleSingleton{})
 	if table == nil {
 		return nil, ormerrors.TableNotFound.Wrap(string((&ExampleSingleton{}).ProtoReflect().Descriptor().FullName()))
@@ -413,7 +412,7 @@ func (testSchemaStore) doNotImplement() {}
 
 var _ TestSchemaStore = testSchemaStore{}
 
-func NewTestSchemaStore(db ormdb.ModuleDB) (TestSchemaStore, error) {
+func NewTestSchemaStore(db ormtable.Schema) (TestSchemaStore, error) {
 	exampleTableStore, err := NewExampleTableStore(db)
 	if err != nil {
 		return nil, err

orm/model/ormdb/json.go

@@ -0,0 +1,119 @@
+package ormdb
+
+import (
+	"context"
+	"fmt"
+	"sort"
+
+	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
+
+	"github.com/cosmos/cosmos-sdk/orm/types/ormjson"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+
+	"github.com/cosmos/cosmos-sdk/errors"
+)
+
+func (m moduleDB) DefaultJSON(target ormjson.WriteTarget) error {
+	for name, table := range m.tablesByName {
+		w, err := target.OpenWriter(name)
+		if err != nil {
+			return err
+		}
+
+		_, err = w.Write(table.DefaultJSON())
+		if err != nil {
+			return err
+		}
+
+		err = w.Close()
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (m moduleDB) ValidateJSON(source ormjson.ReadSource) error {
+	errMap := map[protoreflect.FullName]error{}
+	for name, table := range m.tablesByName {
+		r, err := source.OpenReader(name)
+		if err != nil {
+			return err
+		}
+
+		err = table.ValidateJSON(r)
+		if err != nil {
+			errMap[name] = err
+		}
+
+		err = r.Close()
+		if err != nil {
+			return err
+		}
+	}
+
+	if len(errMap) != 0 {
+		var allErrors string
+		for name, err := range errMap {
+			allErrors += fmt.Sprintf("Error in JSON for table %s: %v\n", name, err)
+		}
+		return ormerrors.JSONValidationError.Wrap(allErrors)
+	}
+
+	return nil
+}
+
+func (m moduleDB) ImportJSON(ctx context.Context, source ormjson.ReadSource) error {
+	var names []string
+	for name := range m.tablesByName {
+		names = append(names, string(name))
+	}
+	sort.Strings(names)
+
+	for _, name := range names {
+		fullName := protoreflect.FullName(name)
+		table := m.tablesByName[fullName]
+
+		r, err := source.OpenReader(fullName)
+		if err != nil {
+			return errors.Wrapf(err, "table %s", fullName)
+		}
+
+		if r == nil {
+			continue
+		}
+
+		err = table.ImportJSON(ctx, r)
+		if err != nil {
+			return errors.Wrapf(err, "table %s", fullName)
+		}
+
+		err = r.Close()
+		if err != nil {
+			return errors.Wrapf(err, "table %s", fullName)
+		}
+	}
+
+	return nil
+}
+
+func (m moduleDB) ExportJSON(ctx context.Context, sink ormjson.WriteTarget) error {
+	for name, table := range m.tablesByName {
+		w, err := sink.OpenWriter(name)
+		if err != nil {
+			return err
+		}
+
+		err = table.ExportJSON(ctx, w)
+		if err != nil {
+			return err
+		}
+
+		err = w.Close()
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

orm/model/ormdb/module.go

@@ -6,6 +6,8 @@ import (
 	"encoding/binary"
 	"math"
 
+	"github.com/cosmos/cosmos-sdk/orm/types/ormjson"
+
 	"google.golang.org/protobuf/reflect/protodesc"
 
 	"github.com/cosmos/cosmos-sdk/orm/encoding/encodeutil"
@@ -31,7 +33,22 @@ type ModuleSchema struct {
 }
 
 // ModuleDB defines the ORM database type to be used by modules.
-type ModuleDB = ormtable.Schema
+type ModuleDB interface {
+	ormtable.Schema
+
+	// DefaultJSON writes default JSON for each table in the module to the target.
+	DefaultJSON(ormjson.WriteTarget) error
+
+	// ValidateJSON validates JSON for each table in the module.
+	ValidateJSON(ormjson.ReadSource) error
+
+	// ImportJSON imports JSON for each table in the module which has JSON
+	// defined in the read source.
+	ImportJSON(context.Context, ormjson.ReadSource) error
+
+	// ExportJSON exports JSON for each table in the module.
+	ExportJSON(context.Context, ormjson.WriteTarget) error
+}
 
 type moduleDB struct {
 	prefix       []byte

orm/model/ormdb/module_test.go

@@ -3,10 +3,17 @@ package ormdb_test
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"strings"
 	"testing"
 
+	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
+
+	"github.com/cosmos/cosmos-sdk/orm/testing/ormtest"
+
+	"github.com/cosmos/cosmos-sdk/orm/types/ormjson"
+
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"gotest.tools/v3/assert"
 	"gotest.tools/v3/golden"
@@ -157,14 +164,14 @@ func TestModuleDB(t *testing.T) {
 	db, err := ormdb.NewModuleDB(TestBankSchema, ormdb.ModuleDBOptions{})
 	assert.NilError(t, err)
 	debugBuf := &strings.Builder{}
-	store := testkv.NewDebugBackend(
-		testkv.NewSharedMemBackend(),
+	backend := ormtest.NewMemoryBackend()
+	ctx := ormtable.WrapContextDefault(testkv.NewDebugBackend(
+		backend,
 		&testkv.EntryCodecDebugger{
 			EntryCodec: db,
 			Print:      func(s string) { debugBuf.WriteString(s + "\n") },
 		},
-	)
-	ctx := ormtable.WrapContextDefault(store)
+	))
 
 	// create keeper
 	k, err := newKeeper(db)
@@ -205,7 +212,7 @@ func TestModuleDB(t *testing.T) {
 	golden.Assert(t, debugBuf.String(), "bank_scenario.golden")
 
 	// check decode & encode
-	it, err := store.CommitmentStore().Iterator(nil, nil)
+	it, err := backend.CommitmentStore().Iterator(nil, nil)
 	assert.NilError(t, err)
 	for it.Valid() {
 		entry, err := db.DecodeEntry(it.Key(), it.Value())
@@ -216,4 +223,33 @@ func TestModuleDB(t *testing.T) {
 		assert.Assert(t, bytes.Equal(v, it.Value()))
 		it.Next()
 	}
+
+	// check JSON
+	target := ormjson.NewRawMessageTarget()
+	assert.NilError(t, db.DefaultJSON(target))
+	rawJson, err := target.JSON()
+	assert.NilError(t, err)
+	golden.Assert(t, string(rawJson), "default_json.golden")
+
+	target = ormjson.NewRawMessageTarget()
+	assert.NilError(t, db.ExportJSON(ctx, target))
+	rawJson, err = target.JSON()
+	assert.NilError(t, err)
+
+	badJSON := `{
+  "testpb.Balance": 5,
+  "testpb.Supply": {}
+}
+`
+	source, err := ormjson.NewRawMessageSource(json.RawMessage(badJSON))
+	assert.NilError(t, err)
+	assert.ErrorIs(t, db.ValidateJSON(source), ormerrors.JSONValidationError)
+
+	backend2 := ormtest.NewMemoryBackend()
+	ctx2 := ormtable.WrapContextDefault(backend2)
+	source, err = ormjson.NewRawMessageSource(rawJson)
+	assert.NilError(t, err)
+	assert.NilError(t, db.ValidateJSON(source))
+	assert.NilError(t, db.ImportJSON(ctx2, source))
+	testkv.AssertBackendsEqual(t, backend, backend2)
 }

orm/model/ormdb/testdata/default_json.golden

@@ -0,0 +1,4 @@
+{
+  "testpb.Balance": [],
+  "testpb.Supply": []
+}

orm/model/ormtable/auto_increment.go

@@ -10,7 +10,7 @@ import (
 	"google.golang.org/protobuf/reflect/protoreflect"
 
 	"github.com/cosmos/cosmos-sdk/orm/encoding/ormkv"
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
 	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"
 )
 
@@ -125,7 +125,7 @@ func (t autoIncrementTable) ValidateJSON(reader io.Reader) error {
 		messageRef := message.ProtoReflect()
 		id := messageRef.Get(t.autoIncField).Uint()
 		if id > maxID {
-			return fmt.Errorf("invalid ID %d, expected a value <= %d", id, maxID)
+			return fmt.Errorf("invalid ID %d, expected a value <= %d, the highest sequence number", id, maxID)
 		}
 
 		if t.customJSONValidator != nil {
@@ -152,7 +152,7 @@ func (t autoIncrementTable) ImportJSON(ctx context.Context, reader io.Reader) er
 			return err
 		} else {
 			if id > maxID {
-				return fmt.Errorf("invalid ID %d, expected a value <= %d", id, maxID)
+				return fmt.Errorf("invalid ID %d, expected a value <= %d, the highest sequence number", id, maxID)
 			}
 			// we do have an ID and calling Save will fail because it expects
 			// either no ID or SAVE_MODE_UPDATE. So instead we drop one level
@@ -213,18 +213,20 @@ func (t autoIncrementTable) ExportJSON(ctx context.Context, writer io.Writer) er
 		return err
 	}
 
-	bz, err := json.Marshal(seq)
-	if err != nil {
-		return err
-	}
-	_, err = writer.Write(bz)
-	if err != nil {
-		return err
-	}
+	if seq != 0 {
+		bz, err := json.Marshal(seq)
+		if err != nil {
+			return err
+		}
+		_, err = writer.Write(bz)
+		if err != nil {
+			return err
+		}
 
-	_, err = writer.Write([]byte(",\n"))
-	if err != nil {
-		return err
+		_, err = writer.Write([]byte(",\n"))
+		if err != nil {
+			return err
+		}
 	}
 
 	return t.doExportJSON(ctx, writer)

orm/model/ormtable/backend.go

@@ -3,7 +3,7 @@ package ormtable
 import (
 	"context"
 
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
 )
 
 // ReadBackend defines the type used for read-only ORM operations.

orm/model/ormtable/batch.go

@@ -1,6 +1,8 @@
 package ormtable
 
-import "github.com/cosmos/cosmos-sdk/orm/model/kv"
+import (
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
+)
 
 type batchIndexCommitmentWriter struct {
 	Backend

orm/model/ormtable/index.go

@@ -3,11 +3,12 @@ package ormtable
 import (
 	"context"
 
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
+
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 
 	"github.com/cosmos/cosmos-sdk/orm/encoding/ormkv"
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormlist"
 )
 

orm/model/ormtable/index_impl.go

@@ -3,9 +3,10 @@ package ormtable
 import (
 	"context"
 
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
+
 	"github.com/cosmos/cosmos-sdk/orm/internal/fieldnames"
 
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormlist"
 
 	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"

orm/model/ormtable/iterator.go

@@ -4,14 +4,12 @@ import (
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 
-	"github.com/cosmos/cosmos-sdk/orm/encoding/encodeutil"
-
 	queryv1beta1 "github.com/cosmos/cosmos-sdk/api/cosmos/base/query/v1beta1"
-
+	"github.com/cosmos/cosmos-sdk/orm/encoding/encodeutil"
 	"github.com/cosmos/cosmos-sdk/orm/encoding/ormkv"
 	"github.com/cosmos/cosmos-sdk/orm/internal/listinternal"
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormlist"
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
 )
 
 // Iterator defines the interface for iterating over indexes.

orm/model/ormtable/table_test.go

@@ -8,6 +8,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
+
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/testing/protocmp"
@@ -21,7 +23,6 @@ import (
 	"github.com/cosmos/cosmos-sdk/orm/internal/testkv"
 	"github.com/cosmos/cosmos-sdk/orm/internal/testpb"
 	"github.com/cosmos/cosmos-sdk/orm/internal/testutil"
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormlist"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormtable"
 	"github.com/cosmos/cosmos-sdk/orm/types/ormerrors"

orm/model/ormtable/unique.go

@@ -3,9 +3,10 @@ package ormtable
 import (
 	"context"
 
+	"github.com/cosmos/cosmos-sdk/orm/types/kv"
+
 	"github.com/cosmos/cosmos-sdk/orm/internal/fieldnames"
 
-	"github.com/cosmos/cosmos-sdk/orm/model/kv"
 	"github.com/cosmos/cosmos-sdk/orm/model/ormlist"
 
 	"github.com/cosmos/cosmos-sdk/orm/encoding/encodeutil"

orm/testing/ormtest/membackend.go

@@ -0,0 +1,18 @@
+// Package ormtest contains utilities for testing modules built with the ORM.
+package ormtest
+
+import (
+	"github.com/cosmos/cosmos-sdk/orm/internal/testkv"
+	"github.com/cosmos/cosmos-sdk/orm/model/ormtable"
+)
+
+// NewMemoryBackend returns a new ORM memory backend which can be used for
+// testing purposes independent of any storage layer.
+//
+// Example:
+//  backend := ormtest.NewMemoryBackend()
+//  ctx := ormtable.WrapContextDefault()
+//  ...
+func NewMemoryBackend() ormtable.Backend {
+	return testkv.NewSplitMemBackend()
+}

orm/types/docs.go

@@ -0,0 +1,3 @@
+// Package types contains types used in the ORM implementation that don't
+// have another home.
+package types

orm/types/ormerrors/errors.go

@@ -32,4 +32,5 @@ var (
 	InvalidTableDefinition        = errors.New(codespace, 25, "invalid table definition")
 	InvalidFileDescriptorID       = errors.New(codespace, 26, "invalid file descriptor ID")
 	TableNotFound                 = errors.New(codespace, 27, "table not found")
+	JSONValidationError           = errors.New(codespace, 23, "invalid JSON")
 )

orm/types/ormjson/json.go

@@ -0,0 +1,27 @@
+package ormjson
+
+import (
+	"io"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+// ReadSource is a source for reading tables in JSON format. It
+// may abstract over a single JSON object or JSON in separate files that
+// can be streamed over.
+type ReadSource interface {
+	// OpenReader returns an io.ReadCloser for the named table. If there
+	// is no JSON for this table, this method will return nil. It is
+	// important the caller closes the reader when done with it.
+	OpenReader(tableName protoreflect.FullName) (io.ReadCloser, error)
+}
+
+// WriteTarget is a target for writing tables in JSON format. It
+// may abstract over a single JSON object or JSON in separate files that
+// can be written incrementally.
+type WriteTarget interface {
+	// OpenWriter returns an io.WriteCloser for the named table. It is
+	// important the caller closers the writer AND checks the error
+	// when done with it.
+	OpenWriter(tableName protoreflect.FullName) (io.WriteCloser, error)
+}

orm/types/ormjson/raw.go

@@ -0,0 +1,79 @@
+package ormjson
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+)
+
+type rawMessageSource struct {
+	m map[string]json.RawMessage
+}
+
+// NewRawMessageSource returns a new ReadSource for the provided
+// json.RawMessage where it is assumed that the raw message is a JSON
+// map where each table's JSON referenced by the map key corresponding
+// to the tables full protobuf name.
+func NewRawMessageSource(message json.RawMessage) (ReadSource, error) {
+	var m map[string]json.RawMessage
+	err := json.Unmarshal(message, &m)
+	if err != nil {
+		return nil, err
+	}
+	return &rawMessageSource{m}, err
+}
+
+func (r rawMessageSource) OpenReader(tableName protoreflect.FullName) (io.ReadCloser, error) {
+	j, ok := r.m[string(tableName)]
+	if !ok {
+		return nil, nil
+	}
+	return readCloserWrapper{bytes.NewReader(j)}, nil
+}
+
+type readCloserWrapper struct {
+	io.Reader
+}
+
+func (r readCloserWrapper) Close() error { return nil }
+
+var _ ReadSource = rawMessageSource{}
+
+// RawMessageTarget is a WriteTarget wrapping a raw JSON map.
+type RawMessageTarget struct {
+	m map[string]json.RawMessage
+}
+
+// NewRawMessageTarget returns a new WriteTarget where each table's JSON
+// is written to a map key corresponding to the table's full protobuf name.
+func NewRawMessageTarget() *RawMessageTarget {
+	return &RawMessageTarget{}
+}
+
+func (r *RawMessageTarget) OpenWriter(tableName protoreflect.FullName) (io.WriteCloser, error) {
+	if r.m == nil {
+		r.m = map[string]json.RawMessage{}
+	}
+
+	return &rawWriter{Buffer: &bytes.Buffer{}, sink: r, table: tableName}, nil
+}
+
+// JSON returns the JSON map that was written as a json.RawMessage.
+func (r *RawMessageTarget) JSON() (json.RawMessage, error) {
+	return json.MarshalIndent(r.m, "", "  ")
+}
+
+type rawWriter struct {
+	*bytes.Buffer
+	table protoreflect.FullName
+	sink  *RawMessageTarget
+}
+
+func (r rawWriter) Close() error {
+	r.sink.m[string(r.table)] = r.Buffer.Bytes()
+	return nil
+}
+
+var _ WriteTarget = &RawMessageTarget{}