PMM-11808 Fix query results deserialization bug (#1877)

* PMM-11808 Fix * PMM-11808 Format * PMM-11808 Fix clickhouse query * PMM-11808 Fix * PMM-11808 Fix * Update managed/cmd/pmm-managed-starlark/main.go Co-authored-by: Alex Tymchuk <[email protected]> * Update managed/services/checks/checks.go Co-authored-by: Alex Tymchuk <[email protected]> --------- Co-authored-by: Alex Tymchuk <[email protected]>
percona · Mar 17, 2023 · 8269774 · 8269774
1 parent 3d2e6ed
commit 8269774
Show file tree

Hide file tree

Showing 2 changed files with 133 additions and 70 deletions.
diff --git a/managed/cmd/pmm-managed-starlark/main.go b/managed/cmd/pmm-managed-starlark/main.go
@@ -16,6 +16,7 @@
 package main
 
 import (
+	"encoding/base64"
 	"encoding/json"
 	"log"
 	"os"
@@ -30,6 +31,7 @@ import (
 	"golang.org/x/sys/unix"
 	"gopkg.in/alecthomas/kingpin.v2"
 
+	"github.com/percona/pmm/api/agentpb"
 	"github.com/percona/pmm/managed/services/checks"
 	"github.com/percona/pmm/managed/utils/logger"
 	"github.com/percona/pmm/version"
@@ -117,17 +119,55 @@ func runChecks(l *logrus.Entry, data *checks.StarlarkScriptData) ([]check.Result
 		return nil, errors.Wrap(err, "error initializing starlark env")
 	}
 
+	res := make([]any, len(data.QueriesResults))
+	for i, queryResult := range data.QueriesResults {
+		switch qr := queryResult.(type) {
+		case map[string]any: // used for PG multidb results where key is database name and value is rows
+			dbRes := make(map[string]any, len(qr))
+			for dbName, dbQr := range qr {
+				s, ok := dbQr.(string)
+				if !ok {
+					return nil, errors.Errorf("unexpected query result type: %T", dbQr)
+				}
+				if dbRes[dbName], err = unmarshallQueryResult(s); err != nil {
+					return nil, err
+				}
+			}
+			res[i] = dbRes
+		case string: // used for all other databases
+			if res[i], err = unmarshallQueryResult(qr); err != nil {
+				return nil, err
+			}
+		default:
+			return nil, errors.Errorf("unknown query result type %T", qr)
+		}
+	}
+
 	var results []check.Result
 	contextFuncs := checks.GetAdditionalContext()
 	switch data.Version {
 	case 1:
-		results, err = env.Run(data.Name, data.QueriesResults[0], contextFuncs, l.Debugln)
+		results, err = env.Run(data.Name, res[0], contextFuncs, l.Debugln)
 	case 2:
-		results, err = env.Run(data.Name, data.QueriesResults, contextFuncs, l.Debugln)
+		results, err = env.Run(data.Name, res, contextFuncs, l.Debugln)
 	}
 	if err != nil {
 		return nil, errors.Wrap(err, "error running starlark env")
 	}
 
 	return results, nil
 }
+
+func unmarshallQueryResult(qr string) ([]map[string]any, error) {
+	b, err := base64.StdEncoding.DecodeString(qr)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to decode base64 encoded query result")
+	}
+
+	res, err := agentpb.UnmarshalActionQueryResult(b)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to unmarshal query result")
+	}
+
+	return res, nil
+}