From 176d58bad44a2c92707fee9bbf0c7f1e96c8316d Mon Sep 17 00:00:00 2001
From: WithLin <withlin@yeah.net>
Date: Sat, 27 Oct 2018 09:37:27 +0800
Subject: [PATCH] Fix issue #48

---
 cmd/soar/soar.go |  5 +++++
 common/config.go |  4 ++++
 env/env.go       | 36 ++++++++++++++++++++++++++++++++++--
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/cmd/soar/soar.go b/cmd/soar/soar.go
index 82a17d1..313a6a3 100644
--- a/cmd/soar/soar.go
+++ b/cmd/soar/soar.go
@@ -80,6 +80,11 @@ func main() {
 	// 环境初始化，连接检查线上环境+构建测试环境
 	vEnv, rEnv := env.BuildEnv()
 
+	//如果使用CleanTestDataBase命令，则清除前1小时的数据库
+	if common.Config.CleanTestDataBase {
+		vEnv.CleanTestDataBase()
+	}
+
 	// 如果使用到测试环境，在这里环境清理
 	if common.Config.DropTestTemporary {
 		defer vEnv.CleanUp()
diff --git a/common/config.go b/common/config.go
index 5abc613..541138e 100644
--- a/common/config.go
+++ b/common/config.go
@@ -119,6 +119,7 @@ type Configration struct {
 	Verbose            bool   `yaml:"verbose"`               // verbose模式，会多输出一些信息
 	DryRun             bool   `yaml:"dry-run"`               // 是否在预演环境执行
 	MaxPrettySQLLength int    `yaml:"max-pretty-sql-length"` // 超出该长度的SQL会转换成指纹输出
+	CleanTestDataBase  bool   `yaml:"clean-test-database"`   // 清理数据库  issue #48
 }
 
 // getDefaultLogOutput get default log-output by runtime.GOOS
@@ -219,6 +220,7 @@ var Config = &Configration{
 	ListTestSqls:       false,
 	ListReportTypes:    false,
 	MaxPrettySQLLength: 1024,
+	CleanTestDataBase:  false,
 }
 
 type dsn struct {
@@ -545,6 +547,7 @@ func readCmdFlags() error {
 	verbose := flag.Bool("verbose", Config.Verbose, "Verbose")
 	dryrun := flag.Bool("dry-run", Config.DryRun, "是否在预演环境执行")
 	maxPrettySQLLength := flag.Int("max-pretty-sql-length", Config.MaxPrettySQLLength, "MaxPrettySQLLength, 超出该长度的SQL会转换成指纹输出")
+	cleanTestDataBase := flag.Bool("clean-test-database", Config.CleanTestDataBase, "由于程序异常退出，导致临时数据库的存在，可用此命令删除。")
 	// 一个不存在log-level，用于更新usage。
 	// 因为vitess里面也用了flag，这些vitess的参数我们不需要关注
 	if !Config.Verbose && runtime.GOOS != "windows" {
@@ -648,6 +651,7 @@ func readCmdFlags() error {
 	Config.DryRun = *dryrun
 	Config.MaxPrettySQLLength = *maxPrettySQLLength
 	Config.MaxVarcharLength = *maxVarcharLength
+	Config.CleanTestDataBase = *cleanTestDataBase
 
 	if *ver {
 		version()
diff --git a/env/env.go b/env/env.go
index e4de713..f9466b1 100644
--- a/env/env.go
+++ b/env/env.go
@@ -18,11 +18,12 @@ package env
 
 import (
 	"fmt"
+	"strings"
+	"time"
+
 	"github.com/XiaoMi/soar/ast"
 	"github.com/XiaoMi/soar/common"
 	"github.com/XiaoMi/soar/database"
-	"strings"
-	"time"
 
 	"github.com/dchest/uniuri"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -149,6 +150,37 @@ func (ve VirtualEnv) CleanUp() bool {
 	return true
 }
 
+// CleanTestDataBase 清除一小时前的环境
+func (ve *VirtualEnv) CleanTestDataBase() {
+	common.Log.Debug("CleanTestDataBase ...")
+	dbs, err := ve.Query("show databases like 'optimizer%'")
+	if err == nil {
+		for index, row := range dbs.Rows {
+			s := strings.Split(row.Str(index), "_")
+			pastTime, err := time.Parse("060102150405", s[1])
+			if err != nil {
+				common.Log.Error("CleanTestDataBase compute  pastTime Error: %s", err)
+			}
+			nowTime, err := time.Parse("060102150405", time.Now().Format("060102150405"))
+			if err != nil {
+				common.Log.Error("CleanTestDataBase compute  nowTime Error: %s", err)
+			}
+			subHour := nowTime.Sub(pastTime).Hours()
+			if subHour > 1 {
+				_, err := ve.Query("drop database %s", row.Str(index))
+				if err != nil {
+					common.Log.Error("CleanTestDataBase failed Error: %s", err)
+				}
+			}
+			common.Log.Debug("CleanTestDataBase, done")
+
+		}
+	} else {
+		common.Log.Error("CleanTestDataBase  failed Error:%s", err)
+	}
+
+}
+
 // BuildVirtualEnv rEnv为SQL源环境，DB使用的信息从接口获取
 // 注意：如果是USE,DDL等语句，执行完第一条就会返回，后面的SQL不会执行
 func (ve *VirtualEnv) BuildVirtualEnv(rEnv *database.Connector, SQLs ...string) bool {
-- 
GitLab