From 067f8942027fee3dfbae0d11afcfe18616340bb9 Mon Sep 17 00:00:00 2001
From: Sam Arnold <sarnold@influxdata.com>
Date: Mon, 22 Feb 2021 10:18:56 -0500
Subject: [PATCH] feat: WITH KEY implementation

---
 tests/server_test.go | 18 ++++++++++++++++++
 tsdb/store.go        | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/tests/server_test.go b/tests/server_test.go
index b5d56a66b43..66bfb059b57 100644
--- a/tests/server_test.go
+++ b/tests/server_test.go
@@ -7983,6 +7983,24 @@ 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/",
+			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'",
+			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]/",
+			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"}},
+		},
 		&Query{
 			name:    "show tag keys measurement not found",
 			command: "SHOW TAG KEYS FROM doesntexist",
diff --git a/tsdb/store.go b/tsdb/store.go
index a42ae0f12a0..84fce72671d 100644
--- a/tsdb/store.go
+++ b/tsdb/store.go
@@ -1611,6 +1611,23 @@ func (s *Store) TagKeys(auth query.Authorizer, shardIDs []uint64, cond influxql.
 	if err != nil {
 		return nil, err
 	}
+	// take out the _name = 'mymeasurement' clause from 'FROM' clause
+	tagKeyExpr, filterExpr, err := PartitionExpr(influxql.CloneExpr(cond), func(e influxql.Expr) (bool, error) {
+		switch e := e.(type) {
+		case *influxql.BinaryExpr:
+			switch e.Op {
+			case influxql.EQ, influxql.NEQ, influxql.EQREGEX, influxql.NEQREGEX:
+				tag, ok := e.LHS.(*influxql.VarRef)
+				if ok && tag.Val == "_tagKey" {
+					return true, nil
+				}
+			}
+		}
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
+	}
 
 	// Get all the shards we're interested in.
 	is := IndexSet{Indexes: make([]Index, 0, len(shardIDs))}
@@ -1711,6 +1728,28 @@ func (s *Store) TagKeys(auth query.Authorizer, shardIDs []uint64, cond influxql.
 			Keys:        finalKeys,
 		})
 	}
+
+	if tagKeyExpr != nil {
+		// In-place filtering of results by key filter
+		measurementWrite := 0
+		for _, keys := range results {
+			keyWrite := 0
+			for _, key := range keys.Keys {
+				if influxql.EvalBool(tagKeyExpr, map[string]interface{}{
+					"_tagKey": key,
+				}) {
+					keys.Keys[keyWrite] = key
+					keyWrite++
+				}
+			}
+			if keyWrite != 0 {
+				keys.Keys = keys.Keys[:keyWrite]
+				results[measurementWrite] = keys
+				measurementWrite++
+			}
+		}
+		results = results[:measurementWrite]
+	}
 	return results, nil
 }