feat: query rewriting for WITH KEY in SHOW TAG KEYS

go.mod

@@ -16,7 +16,7 @@ require (
 	github.com/golang/snappy v0.0.1
 	github.com/google/go-cmp v0.5.0
 	github.com/influxdata/flux v0.65.0
-	github.com/influxdata/influxql v1.1.1-0.20200828144457-65d3ef77d385
+	github.com/influxdata/influxql v1.1.1-0.20210223160523-b6ab99450c93
 	github.com/influxdata/roaring v0.4.13-0.20180809181101-fc520f41fab6
 	github.com/influxdata/usage-client v0.0.0-20160829180054-6d3895376368
 	github.com/jsternberg/zap-logfmt v1.2.0

go.sum

@@ -393,6 +393,8 @@ github.com/influxdata/influxdb1-client v0.0.0-20191209144304-8bf82d3c094d/go.mod
 github.com/influxdata/influxql v1.1.0/go.mod h1:KpVI7okXjK6PRi3Z5B+mtKZli+R1DnZgb3N+tzevNgo=
 github.com/influxdata/influxql v1.1.1-0.20200828144457-65d3ef77d385 h1:ED4e5Cc3z5vSN2Tz2GkOHN7vs4Sxe2yds6CXvDnvZFE=
 github.com/influxdata/influxql v1.1.1-0.20200828144457-65d3ef77d385/go.mod h1:gHp9y86a/pxhjJ+zMjNXiQAA197Xk9wLxaz+fGG+kWk=
+github.com/influxdata/influxql v1.1.1-0.20210223160523-b6ab99450c93 h1:4t/8PcmLnI2vrcaHcEKeeLsGxC0WMRaOQdPX9b7DF8Y=
+github.com/influxdata/influxql v1.1.1-0.20210223160523-b6ab99450c93/go.mod h1:gHp9y86a/pxhjJ+zMjNXiQAA197Xk9wLxaz+fGG+kWk=
 github.com/influxdata/line-protocol v0.0.0-20180522152040-32c6aa80de5e h1:/o3vQtpWJhvnIbXley4/jwzzqNeigJK9z+LZcJZ9zfM=
 github.com/influxdata/line-protocol v0.0.0-20180522152040-32c6aa80de5e/go.mod h1:4kt73NQhadE3daL3WhR5EJ/J2ocX0PZzwxQ0gXJ7oFE=
 github.com/influxdata/promql/v2 v2.12.0/go.mod h1:fxOPu+DY0bqCTCECchSRtWfc+0X19ybifQhZoQNF5D8=

query/statement_rewriter.go

@@ -248,9 +248,12 @@ func rewriteShowSeriesCardinalityStatement(stmt *influxql.ShowSeriesCardinalityS
 	}, nil
 }
 
-func rewriteShowTagValuesStatement(stmt *influxql.ShowTagValuesStatement) (influxql.Statement, error) {
+func withKeyExpr(tagKeyExpr influxql.Expr, op influxql.Token) influxql.Expr {
 	var expr influxql.Expr
-	if list, ok := stmt.TagKeyExpr.(*influxql.ListLiteral); ok {
+	if tagKeyExpr == nil {
+		return nil
+	}
+	if list, ok := tagKeyExpr.(*influxql.ListLiteral); ok {
 		for _, tagKey := range list.Vals {
 			tagExpr := &influxql.BinaryExpr{
 				Op:  influxql.EQ,
@@ -270,11 +273,17 @@ func rewriteShowTagValuesStatement(stmt *influxql.ShowTagValuesStatement) (influ
 		}
 	} else {
 		expr = &influxql.BinaryExpr{
-			Op:  stmt.Op,
+			Op:  op,
 			LHS: &influxql.VarRef{Val: "_tagKey"},
-			RHS: stmt.TagKeyExpr,
+			RHS: tagKeyExpr,
 		}
 	}
+	return expr
+}
+
+func rewriteShowTagValuesStatement(stmt *influxql.ShowTagValuesStatement) (influxql.Statement, error) {
+	// parser enforces that TagKeyExpr is non-nil
+	expr := withKeyExpr(stmt.TagKeyExpr, stmt.Op)
 
 	// Set condition or "AND" together.
 	condition := stmt.Condition
@@ -372,9 +381,21 @@ func rewriteShowTagValuesCardinalityStatement(stmt *influxql.ShowTagValuesCardin
 }
 
 func rewriteShowTagKeysStatement(stmt *influxql.ShowTagKeysStatement) (influxql.Statement, error) {
+	condition := rewriteSourcesCondition(stmt.Sources, stmt.Condition)
+	tagExpr := withKeyExpr(stmt.TagKeyExpr, stmt.TagKeyOp)
+	if condition != nil && tagExpr != nil {
+		condition = &influxql.BinaryExpr{
+			LHS: condition,
+			RHS: tagExpr,
+			Op: influxql.AND,
+		}
+	} else if tagExpr != nil {
+		condition = tagExpr
+		// other cases: tagExpr is nil, so condition is already set correctly
+	}
 	return &influxql.ShowTagKeysStatement{
 		Database:   stmt.Database,
-		Condition:  rewriteSourcesCondition(stmt.Sources, stmt.Condition),
+		Condition:  condition,
 		SortFields: stmt.SortFields,
 		Limit:      stmt.Limit,
 		Offset:     stmt.Offset,

tests/server_test.go

@@ -7983,21 +7983,20 @@ func TestServer_Query_ShowTagKeys(t *testing.T) {
 			exp:     `{"results":[{"statement_id":0,"series":[{"name":"cpu","columns":["tagKey"],"values":[["host"],["region"]]},{"name":"gpu","columns":["tagKey"],"values":[["host"],["region"]]}]}]}`,
 			params:  url.Values{"db": []string{"db0"}},
 		},
-		// TODO: WITH KEY + rewriting instead of _tagKey =
 		&Query{
 			name:    `show tag keys on db0 with key`,
-			command: "SHOW TAG KEYS ON db0 where _tagKey =~ /ho/",
+			command: "SHOW TAG KEYS ON db0 WITH KEY =~ /ho/",
 			exp:     `{"results":[{"statement_id":0,"series":[{"name":"cpu","columns":["tagKey"],"values":[["host"]]},{"name":"disk","columns":["tagKey"],"values":[["host"]]},{"name":"gpu","columns":["tagKey"],"values":[["host"]]}]}]}`,
 		},
 		&Query{
 			name:    "show tag keys from with key",
-			command: "SHOW TAG KEYS FROM cpu where _tagKey = 'host'",
+			command: "SHOW TAG KEYS FROM cpu WITH KEY = host",
 			exp:     `{"results":[{"statement_id":0,"series":[{"name":"cpu","columns":["tagKey"],"values":[["host"]]}]}]}`,
 			params:  url.Values{"db": []string{"db0"}},
 		},
 		&Query{
-			name:    "show tag keys from regex with key",
-			command: "SHOW TAG KEYS FROM /[cg]pu/ where _tagKey =~ /[rh]/",
+			name:    "show tag keys from regex with key in",
+			command: "SHOW TAG KEYS FROM /[cg]pu/ WITH KEY IN (host, region) ",
 			exp:     `{"results":[{"statement_id":0,"series":[{"name":"cpu","columns":["tagKey"],"values":[["host"],["region"]]},{"name":"gpu","columns":["tagKey"],"values":[["host"],["region"]]}]}]}`,
 			params:  url.Values{"db": []string{"db0"}},
 		},
@@ -8036,6 +8035,12 @@ func TestServer_Query_ShowTagKeys(t *testing.T) {
 			exp:     `{"results":[{"statement_id":0,"series":[{"name":"disk","columns":["tagKey"],"values":[["host"],["region"]]},{"name":"gpu","columns":["tagKey"],"values":[["host"],["region"]]}]}]}`,
 			params:  url.Values{"db": []string{"db0"}},
 		},
+		&Query{
+			name:    "show tag keys with key with time where",
+			command: "SHOW TAG KEYS WITH KEY = host WHERE host = 'server03' AND time > 0",
+			exp:     `{"results":[{"statement_id":0,"series":[{"name":"disk","columns":["tagKey"],"values":[["host"]]},{"name":"gpu","columns":["tagKey"],"values":[["host"]]}]}]}`,
+			params:  url.Values{"db": []string{"db0"}},
+		},
 		&Query{
 			name:    "show tag keys with time measurement not found",
 			command: "SHOW TAG KEYS FROM doesntexist WHERE time > 0",