Repository: vitessio/vitess
Pull request #8384 (2021-06-24 16:31:01)

query projection test
Signed-off-by: Harshit Gangal <harshit@planetscale.com>

<!--
  Thank you for your contribution to the Vitess project.
  How to contribute: https://vitess.io/docs/contributing/
  Please first make sure there is an open Issue to discuss the feature/fix suggested in this PR.
  If this is a new feature, please mark the Issue as "RFC".
 -->

<!-- if this PR is Work in Progress please create it as a Draft Pull Request -->

## Description
<!-- A few sentences describing the overall goals of the pull request's commits. -->
<!-- If this is a bug fix and you think the fix should be backported, please write so. -->

## Related Issue(s)
<!-- List related issues and pull requests. If this PR fixes an issue, please add it using Fixes #????  -->


## Checklist
- [x] Tests were added or are not required
- [x] Documentation was added or is not required

## Deployment Notes
<!-- Notes regarding deployment of the contained body of work. These should note any db migrations, etc. -->


Patch:
diff --git a/go/vt/vtgate/planbuilder/queryprojection.go b/go/vt/vtgate/planbuilder/queryprojection.go
index 2bfb526481..1096a7d6a7 100644
--- a/go/vt/vtgate/planbuilder/queryprojection.go
+++ b/go/vt/vtgate/planbuilder/queryprojection.go
@@ -69,17 +69,23 @@ func createQPFromSelect(sel *sqlparser.Select) (*queryProjection, error) {
 	allExpr := append(qp.selectExprs, qp.aggrExprs...)
 
 	for _, order := range sel.OrderBy {
-		qp.addOrderBy(order, allExpr)
+		err := qp.addOrderBy(order, allExpr)
+		if err != nil {
+			return nil, err
+		}
 	}
 	return qp, nil
 }
 
-func (qp *queryProjection) addOrderBy(order *sqlparser.Order, allExpr []*sqlparser.AliasedExpr) {
+func (qp *queryProjection) addOrderBy(order *sqlparser.Order, allExpr []*sqlparser.AliasedExpr) error {
 	// Order by is the column offset to be used from the select expressions
 	// Eg - select id from music order by 1
 	literalExpr, isLiteral := order.Expr.(*sqlparser.Literal)
 	if isLiteral && literalExpr.Type == sqlparser.IntVal {
 		num, _ := strconv.Atoi(literalExpr.Val)
+		if num > len(allExpr) {
+			return vterrors.NewErrorf(vtrpcpb.Code_INVALID_ARGUMENT, vterrors.BadFieldError, "Unknown column '%d' in 'order clause'", num)
+		}
 		aliasedExpr := allExpr[num-1]
 		expr := aliasedExpr.Expr
 		if !aliasedExpr.As.IsEmpty() {
@@ -95,7 +101,7 @@ func (qp *queryProjection) addOrderBy(order *sqlparser.Order, allExpr []*sqlpars
 			},
 			weightStrExpr: aliasedExpr.Expr,
 		})
-		return
+		return nil
 	}
 
 	// If the ORDER BY is against a column alias, we need to remember the expression
@@ -110,7 +116,7 @@ func (qp *queryProjection) addOrderBy(order *sqlparser.Order, allExpr []*sqlpars
 					inner:         order,
 					weightStrExpr: expr.Expr,
 				})
-				return
+				return nil
 			}
 		}
 	}
@@ -119,4 +125,5 @@ func (qp *queryProjection) addOrderBy(order *sqlparser.Order, allExpr []*sqlpars
 		inner:         order,
 		weightStrExpr: order.Expr,
 	})
+	return nil
 }


Test patch:
diff --git a/go/vt/vtgate/planbuilder/queryprojection_test.go b/go/vt/vtgate/planbuilder/queryprojection_test.go
new file mode 100644
index 0000000000..5c2f98d715
--- /dev/null
+++ b/go/vt/vtgate/planbuilder/queryprojection_test.go
@@ -0,0 +1,114 @@
+/*
+Copyright 2021 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package planbuilder
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/vt/sqlparser"
+)
+
+func TestQP(t *testing.T) {
+	tcases := []struct {
+		sql string
+
+		expErr   string
+		expOrder []orderBy
+	}{
+		{
+			sql:    "select * from user",
+			expErr: "gen4 does not yet support: *sqlparser.StarExpr in select list",
+		},
+		{
+			sql:    "select next value from user_seq",
+			expErr: "gen4 does not yet support: *sqlparser.Nextval in select list",
+		},
+		{
+			sql: "select (select 1) from user",
+		},
+		{
+			sql:    "select 1, count(1) from user",
+			expErr: "gen4 does not yet support: aggregation and non-aggregation expressions, together are not supported in cross-shard query",
+		},
+		{
+			sql: "select max(id) from user",
+		},
+		{
+			sql:    "select max(a, b) from user",
+			expErr: "aggregate functions take a single argument 'max(a, b)'",
+		},
+		{
+			sql:    "select func(max(id)) from user",
+			expErr: "unsupported: in scatter query: complex aggregate expression",
+		},
+		{
+			sql:    "select 1, count(1) from user order by 1",
+			expErr: "gen4 does not yet support: aggregation and non-aggregation expressions, together are not supported in cross-shard query",
+		},
+		{
+			sql: "select id from user order by col, id, 1",
+			expOrder: []orderBy{
+				{inner: &sqlparser.Order{Expr: sqlparser.NewColName("col")}, weightStrExpr: sqlparser.NewColName("col")},
+				{inner: &sqlparser.Order{Expr: sqlparser.NewColName("id")}, weightStrExpr: sqlparser.NewColName("id")},
+				{inner: &sqlparser.Order{Expr: sqlparser.NewColName("id")}, weightStrExpr: sqlparser.NewColName("id")},
+			},
+		},
+		{
+			sql:    "select id from user order by 2", // positional order not supported
+			expErr: "Unknown column '2' in 'order clause'",
+		},
+		{
+			sql: "SELECT CONCAT(last_name,', ',first_name) AS full_name FROM mytable, tbl2 ORDER BY full_name", // alias in order not supported
+			expOrder: []orderBy{
+				{
+					inner: &sqlparser.Order{Expr: sqlparser.NewColName("full_name")},
+					weightStrExpr: &sqlparser.FuncExpr{
+						Name: sqlparser.NewColIdent("CONCAT"),
+						Exprs: sqlparser.SelectExprs{
+							&sqlparser.AliasedExpr{Expr: sqlparser.NewColName("last_name")},
+							&sqlparser.AliasedExpr{Expr: sqlparser.NewStrLiteral(", ")},
+							&sqlparser.AliasedExpr{Expr: sqlparser.NewColName("first_name")},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for _, tcase := range tcases {
+		t.Run(tcase.sql, func(t *testing.T) {
+			stmt, err := sqlparser.Parse(tcase.sql)
+			require.NoError(t, err)
+
+			sel := stmt.(*sqlparser.Select)
+			qp, err := createQPFromSelect(sel)
+			if tcase.expErr != "" {
+				require.EqualError(t, err, tcase.expErr)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, len(sel.SelectExprs), len(qp.selectExprs)+len(qp.aggrExprs))
+				for index, expOrder := range tcase.expOrder {
+					assert.True(t, sqlparser.EqualsSQLNode(expOrder.inner, qp.orderExprs[index].inner), "want: %+v, got %+v", sqlparser.String(expOrder.inner), sqlparser.String(qp.orderExprs[index].inner))
+					assert.True(t, sqlparser.EqualsSQLNode(expOrder.weightStrExpr, qp.orderExprs[index].weightStrExpr), "want: %v, got %v", sqlparser.String(expOrder.weightStrExpr), sqlparser.String(qp.orderExprs[index].weightStrExpr))
+				}
+			}
+		})
+	}
+}