@@ -21,69 +21,48 @@ import (
2121 "fmt"
2222 "github.com/linuxsuren/api-testing/pkg/server"
2323 "gorm.io/gorm"
24+ "log"
2425 "reflect"
26+ "sort"
27+ "strings"
2528 "time"
2629)
2730
2831func (s * dbserver ) Query (ctx context.Context , query * server.DataQuery ) (result * server.DataQueryResult , err error ) {
2932 var db * gorm.DB
30- if db , err = s .getClientWithDatabase (ctx , query .Key ); err != nil {
33+ var dbQuery DataQuery
34+ if dbQuery , err = s .getClientWithDatabase (ctx , query .Key ); err != nil {
3135 return
3236 }
3337
38+ db = dbQuery .GetClient ()
39+
3440 result = & server.DataQueryResult {
3541 Data : []* server.Pair {},
3642 Items : make ([]* server.Pairs , 0 ),
3743 Meta : & server.DataMeta {},
3844 }
3945
4046 // query database and tables
41- queryDatabaseSql := "show databases"
42- var databaseResult * server.DataQueryResult
43- if databaseResult , err = sqlQuery (ctx , queryDatabaseSql , db ); err == nil {
44- for _ , table := range databaseResult .Items {
45- for _ , item := range table .GetData () {
46- if item .Key == "Database" {
47- var found bool
48- for _ , name := range result .Meta .Databases {
49- if name == item .Value {
50- found = true
51- }
52- }
53- if ! found {
54- result .Meta .Databases = append (result .Meta .Databases , item .Value )
55- }
56- }
57- }
58- }
47+ if result .Meta .Databases , err = dbQuery .GetDatabases (ctx ); err != nil {
48+ log .Printf ("failed to query databases: %v\n " , err )
5949 }
6050
61- var row * sql.Row
62- if row = db .Raw ("SELECT DATABASE() as name" ).Row (); row != nil {
63- if err = row .Scan (& result .Meta .CurrentDatabase ); err == nil {
64- queryTableSql := "show tables"
65- var tableResult * server.DataQueryResult
66- if tableResult , err = sqlQuery (ctx , queryTableSql , db ); err == nil {
67- for _ , table := range tableResult .Items {
68- for _ , item := range table .GetData () {
69- if item .Key == fmt .Sprintf ("Tables_in_%s" , result .Meta .CurrentDatabase ) {
70- var found bool
71- for _ , name := range result .Meta .Tables {
72- if name == item .Value {
73- found = true
74- }
75- }
76- if ! found {
77- result .Meta .Tables = append (result .Meta .Tables , item .Value )
78- }
79- }
80- }
81- }
82- }
51+ if result .Meta .CurrentDatabase = query .Key ; query .Key == "" {
52+ if result .Meta .CurrentDatabase , err = dbQuery .GetCurrentDatabase (); err != nil {
53+ log .Printf ("failed to query current database: %v\n " , err )
8354 }
8455 }
8556
57+ if result .Meta .Tables , err = dbQuery .GetTables (ctx , result .Meta .CurrentDatabase ); err != nil {
58+ log .Printf ("failed to query tables: %v\n " , err )
59+ }
60+
8661 // query data
62+ if query .Sql == "" {
63+ return
64+ }
65+
8766 var dataResult * server.DataQueryResult
8867 if dataResult , err = sqlQuery (ctx , query .Sql , db ); err == nil {
8968 result .Items = dataResult .Items
@@ -177,3 +156,91 @@ func sqlQuery(ctx context.Context, sql string, db *gorm.DB) (result *server.Data
177156 }
178157 return
179158}
159+
160+ const queryDatabaseSql = "show databases"
161+
162+ type DataQuery interface {
163+ GetDatabases (context.Context ) (databases []string , err error )
164+ GetTables (ctx context.Context , currentDatabase string ) (tables []string , err error )
165+ GetCurrentDatabase () (string , error )
166+ GetClient () * gorm.DB
167+ }
168+
169+ type commonDataQuery struct {
170+ showDatabases , showTables , currentDatabase string
171+ db * gorm.DB
172+ }
173+
174+ var _ DataQuery = & commonDataQuery {}
175+
176+ func NewCommonDataQuery (showDatabases , showTables , currentDatabase string , db * gorm.DB ) DataQuery {
177+ return & commonDataQuery {
178+ showDatabases : showDatabases ,
179+ showTables : showTables ,
180+ currentDatabase : currentDatabase ,
181+ db : db ,
182+ }
183+ }
184+
185+ func (q * commonDataQuery ) GetDatabases (ctx context.Context ) (databases []string , err error ) {
186+ var databaseResult * server.DataQueryResult
187+ if databaseResult , err = sqlQuery (ctx , q .showDatabases , q .db ); err == nil {
188+ for _ , table := range databaseResult .Items {
189+ for _ , item := range table .GetData () {
190+ if item .Key == "Database" || item .Key == "name" {
191+ var found bool
192+ for _ , name := range databases {
193+ if name == item .Value {
194+ found = true
195+ }
196+ }
197+ if ! found {
198+ databases = append (databases , item .Value )
199+ }
200+ }
201+ }
202+ }
203+ sort .Strings (databases )
204+ }
205+ return
206+ }
207+
208+ func (q * commonDataQuery ) GetTables (ctx context.Context , currentDatabase string ) (tables []string , err error ) {
209+ showTables := q .showTables
210+ if strings .Contains (showTables , "%s" ) {
211+ showTables = fmt .Sprintf (showTables , currentDatabase )
212+ }
213+
214+ var tableResult * server.DataQueryResult
215+ if tableResult , err = sqlQuery (ctx , showTables , q .db ); err == nil {
216+ for _ , table := range tableResult .Items {
217+ for _ , item := range table .GetData () {
218+ if item .Key == fmt .Sprintf ("Tables_in_%s" , currentDatabase ) || item .Key == "table_name" ||
219+ item .Key == "Tables" || item .Key == "tablename" {
220+ var found bool
221+ for _ , name := range tables {
222+ if name == item .Value {
223+ found = true
224+ }
225+ }
226+ if ! found {
227+ tables = append (tables , item .Value )
228+ }
229+ }
230+ }
231+ }
232+ sort .Strings (tables )
233+ }
234+ return
235+ }
236+ func (q * commonDataQuery ) GetCurrentDatabase () (current string , err error ) {
237+ var row * sql.Row
238+ if row = q .db .Raw (q .currentDatabase ).Row (); row != nil {
239+ err = row .Scan (& current )
240+ }
241+ return
242+ }
243+
244+ func (q * commonDataQuery ) GetClient () * gorm.DB {
245+ return q .db
246+ }
0 commit comments