chore(state): Cosmosdb / Postgres - Add query test

Signed-off-by: Luigi Rende <[email protected]>
dapr · Oct 22, 2024 · 4a1380b · 4a1380b
1 parent 241c1d1
commit 4a1380b
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 13 deletions.
diff --git a/common/component/postgresql/v1/postgresql_query_test.go b/common/component/postgresql/v1/postgresql_query_test.go
@@ -18,6 +18,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/dapr/components-contrib/state/utils"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -26,8 +28,9 @@ import (
 
 func TestPostgresqlQueryBuildQuery(t *testing.T) {
 	tests := []struct {
-		input string
-		query string
+		input              string
+		query              string
+		selectedAttributes string
 	}{
 		{
 			input: "../../../../tests/state/query/q1.json",
@@ -61,18 +64,30 @@ func TestPostgresqlQueryBuildQuery(t *testing.T) {
 			input: "../../../../tests/state/query/q8.json",
 			query: "SELECT key, value, xmin as etag FROM state WHERE (value->'person'->>'org'>=$1 OR (value->'person'->>'org'<$2 AND (value->>'state'=$3 OR value->>'state'=$4))) ORDER BY value->>'state' DESC, value->'person'->>'name' LIMIT 2",
 		},
+		{
+			input:              "../../../../tests/state/query/q8.json",
+			query:              "SELECT key, value->>'person' as Person, value->'person'->>'org' as Organization, xmin as etag FROM state WHERE (value->'person'->>'org'>=$1 OR (value->'person'->>'org'<$2 AND (value->>'state'=$3 OR value->>'state'=$4))) ORDER BY value->>'state' DESC, value->'person'->>'name' LIMIT 2",
+			selectedAttributes: `[{"name":"Person", "path":"person", "type":"Object"},{"name":"Organization", "path":"person.org", "type":"Text"}]`,
+		},
 	}
 	for _, test := range tests {
 		data, err := os.ReadFile(test.input)
 		require.NoError(t, err)
 		var qq query.Query
+
 		err = json.Unmarshal(data, &qq)
 		require.NoError(t, err)
 
 		q := &Query{
 			tableName:  defaultTableName,
 			etagColumn: "xmin",
 		}
+
+		if len(test.selectedAttributes) > 0 {
+			q.querySelectedAttributes, err = utils.ParseQuerySelectedAttributes(test.selectedAttributes)
+			require.NoError(t, err)
+		}
+
 		qbuilder := query.NewQueryBuilder(q)
 		err = qbuilder.BuildQuery(&qq)
 		require.NoError(t, err)

diff --git a/state/azure/cosmosdb/cosmosdb_query.go b/state/azure/cosmosdb/cosmosdb_query.go
@@ -231,9 +231,9 @@ func (q *Query) Finalize(filters string, qq *query.Query) error {
 
 	if q.querySelectedAttributes != nil {
 		var columns string
-		columns = "c['id'], c['_etag'] "
+		columns = "c['id'], c['_etag']"
 		for _, item := range q.querySelectedAttributes {
-			columns += ", " + replaceKeywords("c.value."+item.Path) + " as '" + item.Name + "' "
+			columns += ", " + replaceKeywords("c.value."+item.Path) + " as '" + item.Name + "'"
 		}
 		q.query.query = "SELECT " + columns + " FROM c" + filter + orderBy
 	} else {

diff --git a/state/azure/cosmosdb/cosmosdb_query_test.go b/state/azure/cosmosdb/cosmosdb_query_test.go
@@ -18,6 +18,8 @@ import (
 	"os"
 	"testing"
 
+	"github.com/dapr/components-contrib/state/utils"
+
 	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -60,8 +62,9 @@ func TestCosmosDbKeyReplace(t *testing.T) {
 
 func TestCosmosDbQuery(t *testing.T) {
 	tests := []struct {
-		input string
-		query InternalQuery
+		input              string
+		query              InternalQuery
+		selectedAttributes string
 	}{
 		{
 			input: "../../../tests/state/query/q1.json",
@@ -174,6 +177,31 @@ func TestCosmosDbQuery(t *testing.T) {
 				},
 			},
 		},
+		{
+			input: "../../../tests/state/query/q8.json",
+			query: InternalQuery{
+				query: "SELECT c['id'], c['_etag'], c['value']['person'] as 'Person', c['value']['person']['org'] as 'Organization' FROM c WHERE c['value']['person']['org'] >= @__param__0__ OR (c['value']['person']['org'] < @__param__1__ AND c['value']['state'] IN (@__param__2__, @__param__3__)) ORDER BY c['value']['state'] DESC, c['value']['person']['name'] ASC",
+				parameters: []azcosmos.QueryParameter{
+					{
+						Name:  "@__param__0__",
+						Value: 123.0,
+					},
+					{
+						Name:  "@__param__1__",
+						Value: 10.0,
+					},
+					{
+						Name:  "@__param__2__",
+						Value: "CA",
+					},
+					{
+						Name:  "@__param__3__",
+						Value: "WA",
+					},
+				},
+			},
+			selectedAttributes: `[{"name":"Person", "path":"person", "type":"Object"},{"name":"Organization", "path":"person.org", "type":"Text"}]`,
+		},
 	}
 	for _, test := range tests {
 		data, err := os.ReadFile(test.input)
@@ -183,6 +211,10 @@ func TestCosmosDbQuery(t *testing.T) {
 		require.NoError(t, err)
 
 		q := &Query{}
+		if len(test.selectedAttributes) > 0 {
+			q.querySelectedAttributes, err = utils.ParseQuerySelectedAttributes(test.selectedAttributes)
+			require.NoError(t, err)
+		}
 		qbuilder := query.NewQueryBuilder(q)
 		err = qbuilder.BuildQuery(&qq)
 		require.NoError(t, err)

diff --git a/state/utils/query_selected_attributes.go b/state/utils/query_selected_attributes.go
@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"strconv"
+	"strings"
 )
 
 type Type int
@@ -17,11 +19,11 @@ const (
 )
 
 var stringToType = map[string]Type{
-	"Text":    Text,
-	"Numeric": Numeric,
-	"Bool":    Bool,
-	"Object":  Object,
-	"Array":   Array,
+	`"Text"`:    Text,
+	`"Numeric"`: Numeric,
+	`"Bool"`:    Bool,
+	`"Object"`:  Object,
+	`"Array"`:   Array,
 }
 
 func ParseType(typeString string) (Type, error) {
@@ -31,6 +33,16 @@ func ParseType(typeString string) (Type, error) {
 	return Text, fmt.Errorf("invalid type, default text: %s", typeString)
 }
 
+func (t *Type) UnmarshalJSON(p []byte) error {
+	elem := string(p)
+	if elem == `null` || elem == `""` {
+		return nil
+	}
+	var err error
+	*t, err = ParseType(strings.Trim(elem, strconv.Itoa(int('"'))))
+	return err
+}
+
 type Attribute struct {
 	Name string `json:"name,omitempty"`
 	Path string `json:"path,omitempty"`

diff --git a/state/utils/query_selected_attributes_test.go b/state/utils/query_selected_attributes_test.go
@@ -9,9 +9,9 @@ import (
 
 func TestParseQuerySelectedAttributes(t *testing.T) {
 	t.Run("Selected attributes no empty ", func(t *testing.T) {
-		selectedAttributes := `[{"name":"test", "path":"data.test"}]`
+		selectedAttributes := `[{"name":"test", "path":"data.test", "type": "Text"}]`
 
-		attributeArray := []Attribute{{Name: "test", Path: "data.test"}}
+		attributeArray := []Attribute{{Name: "test", Path: "data.test", Type: Text}}
 
 		querySelectedAttributes, err := ParseQuerySelectedAttributes(selectedAttributes)
 		require.NoError(t, err)