/
profiling.go
135 lines (118 loc) · 3.54 KB
/
profiling.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
/*
* 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 database
import (
"errors"
"fmt"
"io"
"strings"
"time"
"github.com/XiaoMi/soar/common"
"vitess.io/vitess/go/vt/sqlparser"
)
// Profiling show profile输出的结果
type Profiling struct {
Rows []ProfilingRow
}
// ProfilingRow show profile每一行信息
type ProfilingRow struct {
Status string
Duration float64
// TODO: 支持show profile all,不过目前看all的信息过多有点眼花缭乱
}
// Profiling 执行SQL,并对其Profiling
func (db *Connector) Profiling(sql string, params ...interface{}) (*QueryResult, error) {
// 过滤不需要profiling的SQL
switch sqlparser.Preview(sql) {
case sqlparser.StmtSelect, sqlparser.StmtUpdate, sqlparser.StmtDelete:
default:
return nil, errors.New("no need profiling")
}
// 测试环境如果检查是关闭的,则SQL不会被执行
if common.Config.TestDSN.Disable {
return nil, errors.New("TestDsn Disable")
}
// 数据库安全性检查:如果Connector的IP端口与TEST环境不一致,则启用SQL白名单
// 不在白名单中的SQL不允许执行
// 执行环境与test环境不相同
if db.Addr != common.Config.TestDSN.Addr && db.dangerousQuery(sql) {
return nil, fmt.Errorf("query execution deny: Execute SQL with DSN(%s/%s) '%s'",
db.Addr, db.Database, fmt.Sprintf(sql, params...))
}
common.Log.Debug("Execute SQL with DSN(%s/%s) : %s", db.Addr, db.Database, sql)
conn := db.NewConnection()
// 设置SQL连接超时时间
conn.SetTimeout(time.Duration(common.Config.ConnTimeOut) * time.Second)
defer conn.Close()
err := conn.Connect()
if err != nil {
return nil, err
}
// 添加SQL执行超时限制
ch := make(chan QueryResult, 1)
go func() {
// 开启Profiling
_, _, err = conn.Query("set @@profiling=1")
common.LogIfError(err, "")
// 执行SQL,抛弃返回结果
result, err := conn.Start(sql, params...)
if err != nil {
ch <- QueryResult{
Error: err,
}
return
}
row := result.MakeRow()
for {
err = result.ScanRow(row)
if err == io.EOF {
break
}
}
// 返回Profiling结果
res := QueryResult{}
res.Rows, res.Result, res.Error = conn.Query("show profile")
_, _, err = conn.Query("set @@profiling=0")
common.LogIfError(err, "")
ch <- res
}()
select {
case res := <-ch:
return &res, res.Error
case <-time.After(time.Duration(common.Config.QueryTimeOut) * time.Second):
return nil, errors.New("query execution timeout")
}
}
func getProfiling(res *QueryResult) Profiling {
var rows []ProfilingRow
for _, row := range res.Rows {
rows = append(rows, ProfilingRow{
Status: row.Str(0),
Duration: row.Float(1),
})
}
return Profiling{Rows: rows}
}
// FormatProfiling 格式化输出Profiling信息
func FormatProfiling(res *QueryResult) string {
profiling := getProfiling(res)
str := []string{"| Status | Duration |"}
str = append(str, "| --- | --- |")
for _, row := range profiling.Rows {
str = append(str, fmt.Sprintf("| %s | %f |", row.Status, row.Duration))
}
return strings.Join(str, "\n")
}