Merge pull request #70 from XiaoMi/cleanup-test-database

Add `-cleanup-test-database` arg for database truncate

cmd/soar/soar.go

@@ -80,11 +80,21 @@ func main() {
 	// 环境初始化，连接检查线上环境+构建测试环境
 	vEnv, rEnv := env.BuildEnv()
 
+	// 使用 -cleanup-test-database 命令手动清理残余的 optimizer_xxx 数据库
+	if common.Config.CleanupTestDatabase {
+		vEnv.CleanupTestDatabase()
+	}
+
 	// 如果使用到测试环境，在这里环境清理
 	if common.Config.DropTestTemporary {
 		defer vEnv.CleanUp()
 	}
 
+	// 当程序卡死的时候，或者由于某些原因程序没有退出，可以通过捕获信号量的形式让程序优雅退出并且清理测试环境
+	common.HandleSignal(func() {
+		shutdown(vEnv, rEnv)
+	})
+
 	// 对指定的库表进行索引重复检查
 	if common.Config.ReportType == "duplicate-key-checker" {
 		dupKeySuggest := advisor.DuplicateKeyChecker(rEnv)
@@ -484,3 +494,10 @@ func main() {
 		return
 	}
 }
+
+func shutdown(vEnv *env.VirtualEnv, rEnv *database.Connector) {
+	if common.Config.DropTestTemporary {
+		vEnv.CleanUp()
+	}
+	os.Exit(0)
+}

common/config.go

@@ -43,6 +43,7 @@ type Configration struct {
 	TestDSN                 *dsn   `yaml:"test-dsn"`                  // 测试环境数据库配置
 	AllowOnlineAsTest       bool   `yaml:"allow-online-as-test"`      // 允许Online环境也可以当作Test环境
 	DropTestTemporary       bool   `yaml:"drop-test-temporary"`       // 是否清理Test环境产生的临时库表
+	CleanupTestDatabase     bool   `yaml:"cleanup-test-database"`     // 清理残余的测试数据库（程序异常退出或未开启drop-test-temporary）  issue #48
 	OnlySyntaxCheck         bool   `yaml:"only-syntax-check"`         // 只做语法检查不输出优化建议
 	SamplingStatisticTarget int    `yaml:"sampling-statistic-target"` // 数据采样因子，对应postgres的default_statistics_target
 	Sampling                bool   `yaml:"sampling"`                  // 数据采样开关
@@ -145,6 +146,7 @@ var Config = &Configration{
 	},
 	AllowOnlineAsTest:       false,
 	DropTestTemporary:       true,
+	CleanupTestDatabase:     false,
 	DryRun:                  true,
 	OnlySyntaxCheck:         false,
 	SamplingStatisticTarget: 100,
@@ -478,6 +480,7 @@ func readCmdFlags() error {
 	testDSN := flag.String("test-dsn", FormatDSN(Config.TestDSN), "TestDSN, 测试环境数据库配置, username:password@ip:port/schema")
 	allowOnlineAsTest := flag.Bool("allow-online-as-test", Config.AllowOnlineAsTest, "AllowOnlineAsTest, 允许线上环境也可以当作测试环境")
 	dropTestTemporary := flag.Bool("drop-test-temporary", Config.DropTestTemporary, "DropTestTemporary, 是否清理测试环境产生的临时库表")
+	cleanupTestDatabase := flag.Bool("cleanup-test-database", Config.CleanupTestDatabase, "单次运行清理历史1小时前残余的测试库。")
 	onlySyntaxCheck := flag.Bool("only-syntax-check", Config.OnlySyntaxCheck, "OnlySyntaxCheck, 只做语法检查不输出优化建议")
 	profiling := flag.Bool("profiling", Config.Profiling, "Profiling, 开启数据采样的情况下在测试环境执行Profile")
 	trace := flag.Bool("trace", Config.Trace, "Trace, 开启数据采样的情况下在测试环境执行Trace")
@@ -563,6 +566,7 @@ func readCmdFlags() error {
 	Config.TestDSN = parseDSN(*testDSN, Config.TestDSN)
 	Config.AllowOnlineAsTest = *allowOnlineAsTest
 	Config.DropTestTemporary = *dropTestTemporary
+	Config.CleanupTestDatabase = *cleanupTestDatabase
 	Config.OnlySyntaxCheck = *onlySyntaxCheck
 	Config.Profiling = *profiling
 	Config.Trace = *trace

common/signal.go

+/*
+ * Copyright 2018 Xiaomi, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package common
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+// HandleSignal 当程序卡死的时候，或者由于某些原因程序没有退出，可以通过捕获信号量的形式让程序优雅退出并且清理测试环境
+func HandleSignal(f func()) {
+	sc := make(chan os.Signal, 1)
+	signal.Notify(sc,
+		syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGQUIT)
+
+	go func() {
+		select {
+		case n := <-sc:
+			Log.Info("receive signal %v, closing", n)
+			f()
+		}
+	}()
+}

doc/cheatsheet.md

@@ -8,6 +8,16 @@
 echo "select title from sakila.film" | ./soar -log-output=soar.log
 ```
 
+## 指定输入源
+
+```bash
+# 从文件读取SQL
+./soar -query file.sql
+
+# 从管道读取SQL
+cat file.sql | ./soar
+```
+
 ## 指定配置文件
 
 ```bash
@@ -172,3 +182,13 @@ EOF
 $ cat test.md | soar -report-type md2html > test.html
 ```
 
+## 清理测试环境残余的临时库表
+
+如配置了`-drop-test-temporary=false`或`soar`异常中止，`-test-dsn`中会残余以`optimizer_`为前缀的临时库表。手工清理这些库表可以使用如下命令。
+
+注意：为了不影响正在进行的其他SQL评审，`-cleanup-test-database`中会删除1小时前生成的临时库表。
+
+```bash
+./soar -cleanup-test-database
+```
+

env/env.go

@@ -19,6 +19,7 @@ package env
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/XiaoMi/soar/ast"
 	"github.com/XiaoMi/soar/common"
@@ -149,6 +150,42 @@ func (ve VirtualEnv) CleanUp() bool {
 	return true
 }
 
+// CleanupTestDatabase 清除一小时前的环境
+func (ve *VirtualEnv) CleanupTestDatabase() {
+	common.Log.Debug("CleanupTestDatabase ...")
+	dbs, err := ve.Query("show databases like 'optimizer%'")
+	if err == nil {
+		for _, row := range dbs.Rows {
+			testDatabase := row.Str(0)
+			// test temporary database format `optimizer_YYMMDDHHmmss_randomString(16)`
+			if len(testDatabase) != 39 {
+				common.Log.Debug("CleanupTestDatabase by pass %s", testDatabase)
+				continue
+			}
+			s := strings.Split(testDatabase, "_")
+			pastTime, err := time.Parse("060102150405", s[1])
+			if err != nil {
+				common.Log.Error("CleanupTestDatabase compute  pastTime Error: %s", err.Error())
+				continue
+			}
+			// TODO: 1 hour should be config-able
+			subHour := time.Now().Sub(pastTime).Hours()
+			if subHour > 1 {
+				_, err := ve.Query("drop database %s", testDatabase)
+				if err != nil {
+					common.Log.Error("CleanupTestDatabase failed Error: %s", err.Error())
+					continue
+				}
+				common.Log.Debug("CleanupTestDatabase drop database %s success", testDatabase)
+			} else {
+				common.Log.Debug("CleanupTestDatabase by pass database %s, less than %d hours", testDatabase, subHour)
+			}
+		}
+	} else {
+		common.Log.Error("CleanupTestDatabase failed Error:%s", err.Error())
+	}
+}
+
 // BuildVirtualEnv rEnv为SQL源环境，DB使用的信息从接口获取
 // 注意：如果是USE,DDL等语句，执行完第一条就会返回，后面的SQL不会执行
 func (ve *VirtualEnv) BuildVirtualEnv(rEnv *database.Connector, SQLs ...string) bool {
@@ -294,7 +331,8 @@ func (ve VirtualEnv) createDatabase(rEnv database.Connector, dbName string) erro
 		return nil
 	}
 
-	dbHash := "optimizer_" + uniuri.New()
+	// optimizer_YYMMDDHHmmss_xxxx
+	dbHash := fmt.Sprintf("optimizer_%s_%s", time.Now().Format("060102150405"), uniuri.New())
 	common.Log.Debug("createDatabase, mapping `%s` :`%s`-->`%s`", dbName, dbName, dbHash)
 	ddl, err := rEnv.ShowCreateDatabase(dbName)
 	if err != nil {

env/env_test.go

@@ -113,6 +113,31 @@ func TestNewVirtualEnv(t *testing.T) {
 	}, t.Name(), update)
 }
 
+func TestCleanupTestDatabase(t *testing.T) {
+	vEnv, _ := BuildEnv()
+	vEnv.Query("drop database if exists optimizer_060102150405_xxxxxxxxxxxxxxxx")
+	_, err := vEnv.Query("create database optimizer_060102150405_xxxxxxxxxxxxxxxx")
+	if err != nil {
+		t.Error(err)
+	}
+	vEnv.CleanupTestDatabase()
+	_, err = vEnv.Query("show create database optimizer_060102150405_xxxxxxxxxxxxxxxx")
+	if err == nil {
+		t.Error("optimizer_060102150405_xxxxxxxxxxxxxxxx exist, should be droped")
+	}
+
+	vEnv.Query("drop database if exists optimizer_060102150405")
+	_, err = vEnv.Query("create database optimizer_060102150405")
+	if err != nil {
+		t.Error(err)
+	}
+	vEnv.CleanupTestDatabase()
+	_, err = vEnv.Query("drop database optimizer_060102150405")
+	if err != nil {
+		t.Error("optimizer_060102150405 not exist, should not be droped")
+	}
+}
+
 func TestGenTableColumns(t *testing.T) {
 	vEnv, rEnv := BuildEnv()
 	defer vEnv.CleanUp()