Fix ds.object with_relations (#356)

* Fix `ds.object` with_relations

Passing the `with_relations: true` option to `ds.object` has no
effect because the builtin was assigning the returned object as
the result.

To preserve compatibility with current signature of `ds.object`
this commit only sets the result to the full `GetObjectResponse`
if `with_relations` is true.

* Simplify `ds.object` result

* upd output

* restructure output

* Lint

---------

Co-authored-by: Gert Drapers <gert.drapers@live.com>

builtins/edge/ds/object.go

@@ -3,9 +3,11 @@ package ds
 import (
 	"bytes"
 
+	dsc3 "github.com/aserto-dev/go-directory/aserto/directory/common/v3"
 	dsr3 "github.com/aserto-dev/go-directory/aserto/directory/reader/v3"
 	"github.com/aserto-dev/go-directory/pkg/convert"
 	"github.com/aserto-dev/topaz/resolvers"
+	"google.golang.org/protobuf/encoding/protojson"
 
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
@@ -16,6 +18,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 // RegisterObject - ds.object
@@ -99,25 +102,45 @@ func RegisterObject(logger *zerolog.Logger, fnName string, dr resolvers.Director
 				return nil, err
 			}
 
-			buf := new(bytes.Buffer)
-			var result proto.Message
-
-			if resp.Result != nil {
-				result = resp.Result
-				if outputV2 {
-					result = convert.ObjectToV2(resp.Result)
+			if outputV2 {
+				v, err := v2Value(resp.Result)
+				if err != nil {
+					return nil, err
 				}
+
+				return ast.NewTerm(v), nil
+			}
+
+			buf := new(bytes.Buffer)
+			if err := ProtoToBuf(buf, resp); err != nil {
+				return nil, err
 			}
 
-			if err := ProtoToBuf(buf, result); err != nil {
+			pbs := structpb.Struct{}
+			if err := protojson.Unmarshal(buf.Bytes(), &pbs); err != nil {
 				return nil, err
 			}
 
-			v, err := ast.ValueFromReader(buf)
+			result := pbs.Fields["result"].AsInterface().(map[string]interface{})
+			relations := pbs.Fields["relations"].AsInterface()
+			result["relations"] = relations
+
+			v, err := ast.InterfaceToValue(result)
 			if err != nil {
 				return nil, err
 			}
 
 			return ast.NewTerm(v), nil
 		}
 }
+
+func v2Value(obj *dsc3.Object) (ast.Value, error) {
+	buf := new(bytes.Buffer)
+
+	v2 := convert.ObjectToV2(obj)
+	if err := ProtoToBuf(buf, v2); err != nil {
+		return nil, err
+	}
+
+	return ast.ValueFromReader(buf)
+}