From f8e8c63b05069a793599bbfaaf28a8bc377c59cc Mon Sep 17 00:00:00 2001
From: yisaer <disxiaofei@163.com>
Date: Mon, 29 Aug 2022 12:41:22 +0800
Subject: [PATCH] add dbvars and tiflash

Signed-off-by: yisaer <disxiaofei@163.com>
---
 collector/plan_replayer.go               | 76 ++++++++++++++++++++++++
 configs/profiles/tidb-plan-replayer.toml |  1 -
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/collector/plan_replayer.go b/collector/plan_replayer.go
index acb7d4b1..3e8ce706 100644
--- a/collector/plan_replayer.go
+++ b/collector/plan_replayer.go
@@ -11,6 +11,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/BurntSushi/toml"
 	"github.com/joomcode/errorx"
 	"github.com/pingcap/diag/pkg/models"
 	perrs "github.com/pingcap/errors"
@@ -135,9 +136,84 @@ func (c *PlanReplayerCollectorOptions) collectPlanReplayer(ctx context.Context,
 	if len(errs) > 0 {
 		logger.Warnf("error happened during explain sql, err:%v", strings.Join(errs, ","))
 	}
+	errs = c.collectDBVars(zw, db)
+	if len(errs) > 0 {
+		logger.Warnf("error happened during collect dbvars, err:%v", strings.Join(errs, ","))
+	}
+	errs = c.collectTiflashReplicas(zw, db)
+	if len(errs) > 0 {
+		logger.Warnf("error happened during collect tiflash replica info, err:%v", strings.Join(errs, ","))
+	}
 	return nil
 }
 
+/*
+ |-table_tiflash_replica.txt
+*/
+func (c *PlanReplayerCollectorOptions) collectTiflashReplicas(zw *zip.Writer, db *sql.DB) (errs []string) {
+	vf, err := zw.Create("table_tiflash_replica.txt")
+	if err != nil {
+		errs = append(errs, fmt.Errorf("create table_tiflash_replica.txt failed, err:%v", err).Error())
+		return errs
+	}
+	for table := range c.tables {
+		err := func() error {
+			sql := fmt.Sprintf("SELECT TABLE_SCHEMA,TABLE_NAME,REPLICA_COUNT FROM INFORMATION_SCHEMA.TIFLASH_REPLICA WHERE TABLE_SCHEMA='%s' AND TABLE_NAME ='%s' AND REPLICA_COUNT >0", table.dbName, table.tableName)
+			rows, err := db.Query(sql)
+			if err != nil {
+				return fmt.Errorf("failed to query tiflash replicas, err:%v", err)
+			}
+			for rows.Next() {
+				var dbName, tableName, count string
+				err = rows.Scan(&dbName, &tableName, &count)
+				r := []string{
+					dbName, tableName, count,
+				}
+				fmt.Fprintf(vf, "%s\n", strings.Join(r, "\t"))
+			}
+			return nil
+		}()
+		if err != nil {
+			errs = append(errs, err.Error())
+		}
+	}
+	return errs
+}
+
+/*
+ |-variables.toml
+*/
+func (c *PlanReplayerCollectorOptions) collectDBVars(zw *zip.Writer, db *sql.DB) (errs []string) {
+	err := func() error {
+		vf, err := zw.Create("variables.toml")
+		if err != nil {
+			return fmt.Errorf("create variables.toml failed, err:%v", err)
+		}
+		rows, err := db.Query("show variables")
+		if err != nil {
+			return fmt.Errorf("show variables failed, err:%v", err)
+		}
+		defer rows.Close()
+		varMap := make(map[string]string)
+		var name, value string
+		for rows.Next() {
+			err := rows.Scan(&name, &value)
+			if err != nil {
+				return fmt.Errorf("show variables failed, err:%v", err)
+			}
+			varMap[name] = value
+		}
+		if err := toml.NewEncoder(vf).Encode(varMap); err != nil {
+			return fmt.Errorf("dump varMap failed, err:%v", err)
+		}
+		return nil
+	}()
+	if err != nil {
+		errs = append(errs, err.Error())
+	}
+	return errs
+}
+
 /*
  |-meta.txt
  |-schema
diff --git a/configs/profiles/tidb-plan-replayer.toml b/configs/profiles/tidb-plan-replayer.toml
index 9d5d8274..6c099f9d 100644
--- a/configs/profiles/tidb-plan-replayer.toml
+++ b/configs/profiles/tidb-plan-replayer.toml
@@ -7,7 +7,6 @@ maintainers = [
 # list of data types to collect
 collectors = [
     "config",
-    "db_vars",
     "sql_bind",
     "plan_replayer"
 ]