Merge pull request #1576 from ydb-platform/pb-value

asmyasnikov · web-flow · commit 2a54081806b7 · 2024-12-05T19:24:24.000+03:00
* Supported raw protobuf typed value using `ydb.ParamsBuilder()`
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,5 @@
+* Supported raw protobuf typed value using `ydb.ParamsBuilder()`
+
 ## v3.93.2
 * Removed experimental helper `ydb.MustParamsFromMap`
 * Changed result of experimental helper `ydb.ParamsFromMap` from tuple <`params.Parameters`, `error`> to `params.Parameters` only 
@@ -19,7 +21,7 @@
 * Fixed connections pool leak on closing
 
 ## v3.92.3
-* Fixed error with incompleted data returen from transaction.ReadQueryResult method
+* Fixed error with incompleted data return from transaction.ReadQueryResult method
 * Added option `query/WithResponsePartLimitSizeBytes(...)` for queries with query service
 
 ## v3.92.2
diff --git a/internal/params/builder.go b/internal/params/builder.go
@@ -6,7 +6,11 @@ type (
 	}
 )
 
-func (b Builder) Build() *Params {
+func (b Builder) Build() Parameters {
+	return &b.params
+}
+
+func (b Builder) build() *Params {
 	return &b.params
 }
 
diff --git a/internal/params/builder_test.go b/internal/params/builder_test.go
@@ -417,7 +417,7 @@ func TestBuilder(t *testing.T) {
 			result, ok := xtest.CallMethod(item, tc.method, tc.args...)[0].(Builder)
 			require.True(t, ok)
 
-			params := result.Build().toYDB(a)
+			params := result.build().toYDB(a)
 
 			require.Equal(t,
 				xtest.ToJSON(
diff --git a/internal/params/dict_test.go b/internal/params/dict_test.go
@@ -440,7 +440,7 @@ func TestDict(t *testing.T) {
 				d, ok := xtest.CallMethod(addedKey, val.method, val.args...)[0].(*dict)
 				require.True(t, ok)
 
-				params := d.EndDict().Build().toYDB(a)
+				params := d.EndDict().build().toYDB(a)
 				require.Equal(t, xtest.ToJSON(
 					map[string]*Ydb.TypedValue{
 						"$x": {
@@ -482,7 +482,7 @@ func TestDict_AddPairs(t *testing.T) {
 		},
 	}
 
-	params := Builder{}.Param("$x").BeginDict().AddPairs(pairs...).EndDict().Build().toYDB(a)
+	params := Builder{}.Param("$x").BeginDict().AddPairs(pairs...).EndDict().build().toYDB(a)
 
 	require.Equal(t, xtest.ToJSON(
 		map[string]*Ydb.TypedValue{
diff --git a/internal/params/list_test.go b/internal/params/list_test.go
@@ -435,7 +435,7 @@ func TestList(t *testing.T) {
 			result, ok := xtest.CallMethod(item, tc.method, tc.args...)[0].(*list)
 			require.True(t, ok)
 
-			params := result.EndList().Build().toYDB(a)
+			params := result.EndList().build().toYDB(a)
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
 					"$x": {
@@ -462,7 +462,7 @@ func TestList_AddItems(t *testing.T) {
 	defer a.Free()
 	params := Builder{}.Param("$x").BeginList().
 		AddItems(value.Uint64Value(123), value.Uint64Value(321)).
-		EndList().Build().toYDB(a)
+		EndList().build().toYDB(a)
 	require.Equal(t, xtest.ToJSON(
 		map[string]*Ydb.TypedValue{
 			"$x": {
diff --git a/internal/params/optional_test.go b/internal/params/optional_test.go
@@ -435,7 +435,7 @@ func TestOptional(t *testing.T) {
 			result, ok := xtest.CallMethod(item, tc.method, tc.args...)[0].(*optionalBuilder)
 			require.True(t, ok)
 
-			params := result.EndOptional().Build().toYDB(a)
+			params := result.EndOptional().build().toYDB(a)
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
 					"$x": {
diff --git a/internal/params/parameters.go b/internal/params/parameters.go
@@ -365,6 +365,13 @@ func (p *Parameter) TzDatetime(v time.Time) Builder {
 	return p.parent
 }
 
+func (p *Parameter) Raw(pb *Ydb.TypedValue) Builder {
+	p.value = value.FromProtobuf(pb)
+	p.parent.params = append(p.parent.params, p)
+
+	return p.parent
+}
+
 func Declare(p *Parameter) string {
 	return fmt.Sprintf(
 		"DECLARE %s AS %s",
diff --git a/internal/params/parameters_test.go b/internal/params/parameters_test.go
@@ -51,7 +51,7 @@ func TestNil(t *testing.T) {
 		},
 		{
 			name: xtest.CurrentFileLine(),
-			p:    Builder{}.Build(),
+			p:    Builder{}.build(),
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
diff --git a/internal/params/pg_test.go b/internal/params/pg_test.go
@@ -86,7 +86,7 @@ func TestPg(t *testing.T) {
 			result, ok := xtest.CallMethod(item, tc.method, tc.args...)[0].(Builder)
 			require.True(t, ok)
 
-			params := result.Build().toYDB(a)
+			params := result.build().toYDB(a)
 
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
diff --git a/internal/params/set_test.go b/internal/params/set_test.go
@@ -435,7 +435,7 @@ func TestSet(t *testing.T) {
 			result, ok := xtest.CallMethod(item, tc.method, tc.args...)[0].(*set)
 			require.True(t, ok)
 
-			params := result.EndSet().Build().toYDB(a)
+			params := result.EndSet().build().toYDB(a)
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
 					"$x": {
@@ -470,7 +470,7 @@ func TestSet_AddItems(t *testing.T) {
 	defer a.Free()
 	params := Builder{}.Param("$x").BeginSet().
 		AddItems(value.Uint64Value(123), value.Uint64Value(321)).
-		EndSet().Build().toYDB(a)
+		EndSet().build().toYDB(a)
 	require.Equal(t, xtest.ToJSON(
 		map[string]*Ydb.TypedValue{
 			"$x": {
diff --git a/internal/params/struct_test.go b/internal/params/struct_test.go
@@ -951,7 +951,7 @@ func TestStruct(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			a := allocator.New()
 			defer a.Free()
-			params := tt.builder.Build().toYDB(a)
+			params := tt.builder.build().toYDB(a)
 			require.Equal(t, xtest.ToJSON(tt.params), xtest.ToJSON(params))
 		})
 	}
diff --git a/internal/params/tuple_test.go b/internal/params/tuple_test.go
@@ -435,7 +435,7 @@ func TestTuple(t *testing.T) {
 			result, ok := xtest.CallMethod(item, tc.method, tc.args...)[0].(*tuple)
 			require.True(t, ok)
 
-			params := result.EndTuple().Build().toYDB(a)
+			params := result.EndTuple().build().toYDB(a)
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
 					"$x": {
@@ -464,7 +464,7 @@ func TestTuple_AddItems(t *testing.T) {
 	defer a.Free()
 	params := Builder{}.Param("$x").BeginTuple().
 		AddItems(value.Uint64Value(123), value.Uint64Value(321)).
-		EndTuple().Build().toYDB(a)
+		EndTuple().build().toYDB(a)
 	require.Equal(t, xtest.ToJSON(
 		map[string]*Ydb.TypedValue{
 			"$x": {
diff --git a/internal/params/variant_struct_test.go b/internal/params/variant_struct_test.go
@@ -467,7 +467,7 @@ func TestVariantStruct(t *testing.T) {
 			builder, ok := xtest.CallMethod(vs.Name("key"), tc.method, tc.itemArgs...)[0].(*variantStructBuilder)
 			require.True(t, ok)
 
-			params := builder.EndStruct().EndVariant().Build().toYDB(a)
+			params := builder.EndStruct().EndVariant().build().toYDB(a)
 
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
@@ -519,7 +519,7 @@ func TestVariantStruct_AddFields(t *testing.T) {
 				T:    types.Text,
 			},
 		}...).Name("key3").Text("Hello, World!").EndStruct().
-		EndVariant().Build().toYDB(a)
+		EndVariant().build().toYDB(a)
 
 	require.Equal(t, xtest.ToJSON(
 		map[string]*Ydb.TypedValue{
diff --git a/internal/params/variant_tuple_test.go b/internal/params/variant_tuple_test.go
@@ -466,7 +466,7 @@ func TestVariantTuple(t *testing.T) {
 			builder, ok := xtest.CallMethod(types.Index(0), tc.method, tc.itemArgs...)[0].(*variantTupleBuilder)
 			require.True(t, ok)
 
-			params := builder.EndTuple().EndVariant().Build().toYDB(a)
+			params := builder.EndTuple().EndVariant().build().toYDB(a)
 
 			require.Equal(t, xtest.ToJSON(
 				map[string]*Ydb.TypedValue{
@@ -503,7 +503,7 @@ func TestVariantTuple_AddTypes(t *testing.T) {
 		Types().AddTypes(types.Int64, types.Bool).
 		Index(1).
 		Bool(true).
-		EndTuple().EndVariant().Build().toYDB(a)
+		EndTuple().EndVariant().build().toYDB(a)
 
 	require.Equal(t, xtest.ToJSON(
 		map[string]*Ydb.TypedValue{
diff --git a/internal/types/types.go b/internal/types/types.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xstring"
@@ -957,3 +958,36 @@ func (Null) ToYDB(*allocator.Allocator) *Ydb.Type {
 func NewNull() Null {
 	return Null{}
 }
+
+var _ Type = (*protobufType)(nil)
+
+type protobufType struct {
+	pb *Ydb.Type
+}
+
+func (v protobufType) Yql() string {
+	return FromYDB([]*Ydb.Type{
+		v.pb,
+	})[0].Yql()
+}
+
+func (v protobufType) String() string {
+	return v.Yql()
+}
+
+func (v protobufType) ToYDB(a *allocator.Allocator) *Ydb.Type {
+	return v.pb
+}
+
+func (v protobufType) equalsTo(rhs Type) bool {
+	switch t := rhs.(type) {
+	case *protobufType:
+		return proto.Equal(v.pb, t.pb)
+	default:
+		return false
+	}
+}
+
+func FromProtobuf(pb *Ydb.Type) *protobufType {
+	return &protobufType{pb: pb}
+}
diff --git a/internal/types/types_test.go b/internal/types/types_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pg"
 )
@@ -273,11 +274,40 @@ func TestTypeToString(t *testing.T) {
 			t: PgType{OID: pg.OIDUnknown},
 			s: "PgType(705)",
 		},
+		{
+			t: FromProtobuf(&Ydb.Type{
+				Type: &Ydb.Type_VariantType{
+					VariantType: &Ydb.VariantType{
+						Type: &Ydb.VariantType_StructItems{
+							StructItems: &Ydb.StructType{
+								Members: []*Ydb.StructMember{
+									{
+										Name: "a",
+										Type: &Ydb.Type{
+											Type: &Ydb.Type_TypeId{
+												TypeId: Ydb.Type_BOOL,
+											},
+										},
+									},
+									{
+										Name: "a",
+										Type: &Ydb.Type{
+											Type: &Ydb.Type_TypeId{
+												TypeId: Ydb.Type_FLOAT,
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			}),
+			s: "Variant<'a':Bool,'a':Float>",
+		},
 	} {
 		t.Run(tt.s, func(t *testing.T) {
-			if got := tt.t.Yql(); got != tt.s {
-				t.Errorf("s representations not equals:\n\n -  got: %s\n\n - want: %s", got, tt.s)
-			}
+			require.Equal(t, tt.s, tt.t.Yql())
 		})
 	}
 }
diff --git a/internal/value/value.go b/internal/value/value.go
@@ -2764,3 +2764,37 @@ func (v bytesValue) toYDB(a *allocator.Allocator) *Ydb.Value {
 func BytesValue(v []byte) bytesValue {
 	return v
 }
+
+var _ Value = (*protobufValue)(nil)
+
+type protobufValue struct {
+	pb *Ydb.TypedValue
+}
+
+func (v protobufValue) Type() types.Type {
+	return types.FromProtobuf(v.pb.GetType())
+}
+
+func (v protobufValue) Yql() string {
+	return FromYDB(v.pb.GetType(), v.pb.GetValue()).Yql()
+}
+
+func (v protobufValue) castTo(dst any) error {
+	switch x := dst.(type) {
+	case *Ydb.TypedValue:
+		x.Type = v.pb.GetType()
+		x.Value = v.pb.GetValue()
+
+		return nil
+	default:
+		return xerrors.WithStackTrace(ErrCannotCast)
+	}
+}
+
+func (v protobufValue) toYDB(a *allocator.Allocator) *Ydb.Value {
+	return v.pb.GetValue()
+}
+
+func FromProtobuf(pb *Ydb.TypedValue) protobufValue {
+	return protobufValue{pb: pb}
+}
diff --git a/internal/value/value_test.go b/internal/value/value_test.go
diff --git a/params_builder_test.go b/params_builder_test.go

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,11 @@ type (`
`6`	`6`	`}`
`7`	`7`	`)`
`8`	`8`
`9`		`-func (b Builder) Build() *Params {`
	`9`	`+func (b Builder) Build() Parameters {`
	`10`	`+ return &b.params`
	`11`	`+}`
	`12`	`+`
	`13`	`+func (b Builder) build() *Params {`
`10`	`14`	`return &b.params`
`11`	`15`	`}`
`12`	`16`
Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ func TestNil(t *testing.T) {`
`51`	`51`	`},`
`52`	`52`	`{`
`53`	`53`	`name: xtest.CurrentFileLine(),`
`54`		`- p: Builder{}.Build(),`
	`54`	`+ p: Builder{}.build(),`
`55`	`55`	`},`
`56`	`56`	`} {`
`57`	`57`	`t.Run(tt.name, func(t *testing.T) {`