fix: set active database from query (#6103)

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## [unreleased]
 
+### Bug Fixes
+
+1. [#6103](https://github.com/influxdata/chronograf/pull/6103): Set active database for InfluxQL meta queries.
+
 ### Other
 
 1. [#6102](https://github.com/influxdata/chronograf/pull/6102): Upgrade golang to 1.21.12.

server/influx.go

@@ -173,25 +173,58 @@ func (s *Service) Write(w http.ResponseWriter, r *http.Request) {
 
 // setupQueryFromCommand set query parameters from its command
 func setupQueryFromCommand(req *chronograf.Query) {
-	// allow to set active database with USE command, examples:
+	// sets active database (and retention policy) from the query
+	useDb := func(dbSpec string) error {
+		dbSpecReader := csv.NewReader(bytes.NewReader(([]byte)(dbSpec)))
+		dbSpecReader.Comma = '.'
+		if dbrp, err := dbSpecReader.Read(); err == nil {
+			if len(dbrp) > 0 {
+				req.DB = dbrp[0]
+			}
+			if len(dbrp) > 1 {
+				req.RP = dbrp[1]
+			}
+			return nil
+		} else {
+			return err
+		}
+	}
+
+	// allow to set active database with USE command or via ON clause, examples:
 	//  use mydb
 	//  use "mydb"
 	//  USE "mydb"."myrp"
 	//  use "mydb.myrp"
 	//  use mydb.myrp
-	if strings.HasPrefix(req.Command, "use ") || strings.HasPrefix(req.Command, "USE ") {
+	//  show tag keys on "mydb"
+	//  SHOW TAG KEYS ON "mydb"
+	command := strings.ToLower(req.Command)
+	if strings.HasPrefix(command, "use ") {
 		if nextCommand := strings.IndexRune(req.Command, ';'); nextCommand > 4 {
 			dbSpec := strings.TrimSpace(req.Command[4:nextCommand])
-			dbSpecReader := csv.NewReader(bytes.NewReader(([]byte)(dbSpec)))
-			dbSpecReader.Comma = '.'
-			if dbrp, err := dbSpecReader.Read(); err == nil {
-				if len(dbrp) > 0 {
-					req.DB = dbrp[0]
+			if useDb(dbSpec) == nil {
+				req.Command = strings.TrimSpace(req.Command[nextCommand+1:])
+			}
+		}
+	} else if strings.Contains(command, " on ") {
+		r := csv.NewReader(strings.NewReader(req.Command))
+		r.Comma = ' '
+		if tokens, err := r.Read(); err == nil {
+			// filter empty tokens (i.e. redundant whitespaces, using https://go.dev/wiki/SliceTricks#filtering-without-allocating)
+			fields := tokens[:0]
+			for _, field := range tokens {
+				if field != "" {
+					fields = append(fields, field)
 				}
-				if len(dbrp) > 1 {
-					req.RP = dbrp[1]
+			}
+			// try to find ON clause and use its value to set the database
+			for i, field := range fields {
+				if strings.ToLower(field) == "on" {
+					if i < len(fields)-1 {
+						_ = useDb(fields[i+1])
+					}
+					break
 				}
-				req.Command = strings.TrimSpace(req.Command[nextCommand+1:])
 			}
 		}
 	}

server/influx_test.go

@@ -217,6 +217,128 @@ func TestService_Influx_UseCommand(t *testing.T) {
 	}
 }
 
+// TestService_Influx_CommandWithOnClause tests preprocessing of command with ON clause
+func TestService_Influx_CommandWithOnClause(t *testing.T) {
+	tests := []struct {
+		name string
+		db   string
+		rp   string
+	}{
+		{
+			name: "/* no command */",
+		},
+		{
+			name: "SHOW MEASUREMENTS",
+		},
+		{
+			name: "SHOW TAG KEYS ON mydb",
+			db:   "mydb",
+		},
+		{
+			name: "SHOW TAG KEYS ON mydb FROM table",
+			db:   "mydb",
+		},
+		{
+			name: "USE anotherdb; SHOW TAG KEYS ON mydb",
+			db:   "anotherdb",
+		},
+		{
+			name: `show tag keys on "mydb"`,
+			db:   "mydb",
+		},
+		{
+			name: `show tag keys oN "mydb"`,
+			db:   "mydb",
+		},
+		{
+			name: `show tag keys on "mydb" from "table"`,
+			db:   "mydb",
+		},
+		{
+			name: `show tag keys on "my_db" from "table"`,
+			db:   "my_db",
+		},
+		{
+			name: `show tag keys on "my-db" from "table"`,
+			db:   "my-db",
+		},
+		{
+			name: `show tag keys on "my/db" from "table"`,
+			db:   "my/db",
+		},
+		{
+			name: `show tag keys on "my db" from "table"`,
+			db:   "my db",
+		},
+		{
+			name: `show tag values    on   "my   db" from "table" with key = "my key"`,
+			db:   "my   db",
+		},
+	}
+
+	h := &Service{
+		Store: &mocks.Store{
+			SourcesStore: &mocks.SourcesStore{
+				GetF: func(ctx context.Context, ID int) (chronograf.Source, error) {
+					return chronograf.Source{
+						ID:  1337,
+						URL: "http://any.url",
+					}, nil
+				},
+			},
+		},
+		TimeSeriesClient: &mocks.TimeSeries{
+			ConnectF: func(ctx context.Context, src *chronograf.Source) error {
+				return nil
+			},
+			QueryF: func(ctx context.Context, query chronograf.Query) (chronograf.Response, error) {
+				return mocks.NewResponse(
+						fmt.Sprintf(`{"db":"%s","rp":"%s"}`, query.DB, query.RP),
+						nil,
+					),
+					nil
+			},
+		},
+		Logger: log.New(log.ErrorLevel),
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			prefixCommand := strings.ReplaceAll(tt.name, "\"", "\\\"")
+			w := httptest.NewRecorder()
+			r := httptest.NewRequest(
+				"POST",
+				"http://any.url",
+				ioutil.NopCloser(
+					bytes.NewReader([]byte(
+						`{"uuid": "tst", "query":"`+prefixCommand+` ; DROP MEASUREMENT test"}`,
+					)),
+				),
+			)
+			r = r.WithContext(httprouter.WithParams(
+				context.Background(),
+				httprouter.Params{
+					{
+						Key:   "id",
+						Value: "1",
+					},
+				},
+			))
+
+			h.Influx(w, r)
+
+			resp := w.Result()
+			body, _ := ioutil.ReadAll(resp.Body)
+
+			want := fmt.Sprintf(`{"results":{"db":"%s","rp":"%s"},"uuid":"tst"}`, tt.db, tt.rp)
+			got := strings.TrimSpace(string(body))
+			if got != want {
+				t.Errorf("%q. Influx() =\ngot  ***%v***\nwant ***%v***\n", tt.name, got, want)
+			}
+
+		})
+	}
+}
+
 func TestService_Influx_Write(t *testing.T) {
 	calledPath := ""
 	ts := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {