From d3704524490a981e567add197c01b00711e4f9b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kuba=20Podg=C3=B3rski?= <kuba--@users.noreply.github.com>
Date: Tue, 16 Mar 2021 16:32:35 +0100
Subject: [PATCH] Add helper function FirstStringArg

---
 executor.go | 52 ++++++++++++++++------------------------------------
 pql/ast.go  | 15 +++++++++++++++
 2 files changed, 31 insertions(+), 36 deletions(-)

diff --git a/executor.go b/executor.go
index c6b5326d3..9ee05633e 100644
--- a/executor.go
+++ b/executor.go
@@ -1125,12 +1125,9 @@ func (e *executor) executeSum(ctx context.Context, qcx *Qcx, index string, c *pq
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeSum")
 	defer span.Finish()
 
-	fieldName := c.ArgString("field")
-	if fieldName == "" {
-		fieldName = c.ArgString("_field")
-	}
-	if fieldName == "" {
-		return ValCount{}, errors.New("Sum(): field required")
+	fieldName, err := c.FirstStringArg("field", "_field")
+	if err != nil {
+		return ValCount{}, errors.Wrap(err, "Sum(): field required")
 	}
 
 	if len(c.Children) > 1 {
@@ -1233,12 +1230,8 @@ func (e *executor) executeMin(ctx context.Context, qcx *Qcx, index string, c *pq
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMin")
 	defer span.Finish()
 
-	field := c.Args["field"]
-	if field == "" {
-		field = c.ArgString("_field")
-	}
-	if field == "" {
-		return ValCount{}, errors.New("Min(): field required")
+	if _, err := c.FirstStringArg("field", "_field"); err != nil {
+		return ValCount{}, errors.Wrap(err, "Min(): field required")
 	}
 
 	if len(c.Children) > 1 {
@@ -1273,12 +1266,8 @@ func (e *executor) executeMax(ctx context.Context, qcx *Qcx, index string, c *pq
 	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeMax")
 	defer span.Finish()
 
-	field := c.ArgString("field")
-	if field == "" {
-		field = c.ArgString("_field")
-	}
-	if field == "" {
-		return ValCount{}, errors.New("Max(): field required")
+	if _, err := c.FirstStringArg("field", "_field"); err != nil {
+		return ValCount{}, errors.Wrap(err, "Max(): field required")
 	}
 
 	if len(c.Children) > 1 {
@@ -1901,12 +1890,9 @@ func (e *executor) executeSumCountShard(ctx context.Context, qcx *Qcx, index str
 		filter = row
 	}
 
-	fieldName := c.ArgString("field")
-	if fieldName == "" {
-		fieldName = c.ArgString("_field")
-	}
-	if fieldName == "" {
-		return ValCount{}, errors.New("Sum(): field required")
+	fieldName, err := c.FirstStringArg("field", "_field")
+	if err != nil {
+		return ValCount{}, errors.Wrap(err, "Sum(): field required")
 	}
 
 	field := e.Holder.Field(index, fieldName)
@@ -1958,12 +1944,9 @@ func (e *executor) executeMinShard(ctx context.Context, qcx *Qcx, index string,
 		filter = row
 	}
 
-	fieldName, ok := c.Args["field"].(string)
-	if !ok || fieldName == "" {
-		fieldName = c.ArgString("_field")
-	}
-	if fieldName == "" {
-		return ValCount{}, errors.New("Min(): field required")
+	fieldName, err := c.FirstStringArg("field", "_field")
+	if err != nil {
+		return ValCount{}, errors.Wrap(err, "Min(): field required")
 	}
 
 	field := e.Holder.Field(index, fieldName)
@@ -1993,12 +1976,9 @@ func (e *executor) executeMaxShard(ctx context.Context, qcx *Qcx, index string,
 		filter = row
 	}
 
-	fieldName := c.ArgString("field")
-	if fieldName == "" {
-		fieldName = c.ArgString("_field")
-	}
-	if fieldName == "" {
-		return ValCount{}, errors.New("Max(): field required")
+	fieldName, err := c.FirstStringArg("field", "_field")
+	if err != nil {
+		return ValCount{}, errors.Wrap(err, "Max(): field required")
 	}
 
 	field := e.Holder.Field(index, fieldName)
diff --git a/pql/ast.go b/pql/ast.go
index 729393a86..e6155ac7f 100644
--- a/pql/ast.go
+++ b/pql/ast.go
@@ -723,6 +723,21 @@ func (c *Call) StringArg(key string) (string, bool, error) {
 	}
 }
 
+func (c *Call) FirstStringArg(keys ...string) (string, error) {
+	for _, k := range keys {
+		val, ok, err := c.StringArg(k)
+		if err != nil {
+			return "", err
+		}
+		if !ok {
+			continue
+		}
+		return val, nil
+	}
+
+	return "", fmt.Errorf("keys: %v not found", keys)
+}
+
 // CallArg is for reading the value at key from call.Args as a Call. If the
 // key is not in Call.Args, the value of the returned value will be nil, and
 // the error will be nil. An error is returned if the value is not a Call.