Skip to content

Commit 0baec99

Browse files
julien040julien040
committed
✨ Add the query command
It's a shell that enables you to run SQL queries on the Anyquery integrations
1 parent a3de328 commit 0baec99

13 files changed

+1316
-70
lines changed

‎cmd/hashdir.go

-22
This file was deleted.

‎cmd/query.go

+17-6
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,34 @@
11
package cmd
22

33
import (
4-
"fmt"
5-
4+
"github.com/julien040/anyquery/controller"
65
"github.com/spf13/cobra"
76
)
87

98
var queryCmd = &cobra.Command{
10-
Use: "query [sql query]",
9+
Use: "query [database file] [sql query]",
1110
Short: "Run a SQL query",
1211
Long: `Run a SQL query on the data sources installed on the system.
1312
The query can be specified as an argument or read from stdin.
1413
If no query is provided, the command will launch an interactive input.`,
15-
Run: func(cmd *cobra.Command, args []string) {
16-
fmt.Println("query called")
17-
},
14+
RunE: controller.Query,
1815
}
1916

2017
func init() {
2118
rootCmd.AddCommand(queryCmd)
19+
addFlag_commandModifiesConfiguration(queryCmd)
2220
addFlag_commandPrintsData(queryCmd)
21+
queryCmd.Flags().StringP("database", "d", "anyquery.db", "Database to connect to (a path or :memory:)")
22+
queryCmd.Flags().Bool("in-memory", false, "Use an in-memory database")
23+
queryCmd.Flags().Bool("readonly", false, "Start the server in read-only mode")
24+
queryCmd.Flags().Bool("read-only", false, "Start the server in read-only mode")
25+
queryCmd.Flags().StringArray("init", []string{}, "Run SQL commands in a file before the query. You can specify multiple files.")
26+
27+
// Query flags
28+
queryCmd.Flags().StringP("query", "q", "", "Query to run")
29+
30+
// Log flags
31+
queryCmd.Flags().String("log-file", "", "Log file")
32+
queryCmd.Flags().String("log-level", "info", "Log level (trace, debug, info, warn, error, off)")
33+
queryCmd.Flags().String("log-format", "text", "Log format (text, json)")
2334
}

‎cmd/root.go

+2
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package cmd
33
import (
44
"os"
55

6+
"github.com/julien040/anyquery/controller"
67
"github.com/spf13/cobra"
78
)
89

@@ -15,6 +16,7 @@ by writing SQL queries. It can be extended with plugins`,
1516
// Avoid writing help when an error occurs
1617
// Thanks https://github.com/spf13/cobra/issues/340#issuecomment-243790200
1718
SilenceUsage: true,
19+
RunE: controller.Query,
1820
}
1921

2022
func Execute() {

‎controller/middleware.go

+311
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,311 @@
1+
// This file defines middlewares to be used in the pipeline
2+
// defined in pipeline.go
3+
package controller
4+
5+
import (
6+
"fmt"
7+
"os"
8+
"os/exec"
9+
"strconv"
10+
"strings"
11+
12+
"github.com/julien040/anyquery/namespace"
13+
"vitess.io/vitess/go/vt/sqlparser"
14+
)
15+
16+
func middlewareDotCommand(queryData *QueryData) bool {
17+
// Check if dot command are enabled
18+
if !queryData.Config.GetBool("dot-command", false) {
19+
return true
20+
}
21+
22+
// Check if the query is a dot command
23+
query := queryData.SQLQuery
24+
25+
// Ensure the trimmed query starts with a dot
26+
// If not, we skip to the next middleware
27+
if !strings.HasPrefix(strings.Trim(query, " "), ".") {
28+
return true
29+
}
30+
31+
command, args := parseDotFunc(query)
32+
33+
switch strings.ToLower(command) {
34+
case "cd":
35+
// Get the path
36+
if len(args) == 0 {
37+
queryData.Message = "No path provided"
38+
queryData.StatusCode = 2
39+
} else {
40+
err := os.Chdir(args[0])
41+
if err != nil {
42+
queryData.Message = err.Error()
43+
queryData.StatusCode = 2
44+
} else {
45+
queryData.Message = "Changed directory to " + args[0]
46+
queryData.StatusCode = 0
47+
}
48+
}
49+
50+
/* case "columns":
51+
queryData.Config.SetString("outputMode", "columns") */
52+
53+
case "databases", "database":
54+
queryData.SQLQuery = "SELECT name FROM pragma_database_list;"
55+
return true
56+
57+
case "help":
58+
queryData.Message = "Documentation available at https://docs.anyquery.dev"
59+
queryData.StatusCode = 0
60+
61+
case "indexes", "index":
62+
queryData.SQLQuery = "SELECT * FROM pragma_index_list;"
63+
64+
case "log":
65+
if len(args) == 0 {
66+
queryData.Config.SetBool("logEnabled", false)
67+
} else {
68+
queryData.Config.SetString("logFile", args[0])
69+
queryData.Config.SetBool("logEnabled", true)
70+
}
71+
72+
case "maxrows":
73+
if len(args) == 0 {
74+
queryData.Message = "No maxrows provided"
75+
queryData.StatusCode = 2
76+
} else {
77+
// Parse the maxrows
78+
val, err := strconv.Atoi(args[0])
79+
if err != nil {
80+
queryData.Message = err.Error()
81+
queryData.StatusCode = 2
82+
} else {
83+
queryData.Config.SetInt("maxRows", val)
84+
}
85+
}
86+
case "mode":
87+
if len(args) == 0 {
88+
queryData.Message = "No mode provided"
89+
queryData.StatusCode = 2
90+
} else {
91+
// Check if the mode is valid
92+
if _, ok := formatName[args[0]]; !ok {
93+
queryData.Message = fmt.Sprintf("Invalid mode %s", args[0])
94+
queryData.StatusCode = 2
95+
} else {
96+
queryData.Config.SetString("outputMode", args[0])
97+
queryData.Message = fmt.Sprintf("Output mode set to %s", args[0])
98+
queryData.StatusCode = 0
99+
}
100+
}
101+
102+
case "once":
103+
if len(args) == 0 {
104+
queryData.Message = "No output file provided"
105+
queryData.StatusCode = 2
106+
} else {
107+
queryData.Config.SetString("onceOutputFile", args[0])
108+
// We don't set the message because it would be outputted
109+
// to the file
110+
}
111+
112+
case "output":
113+
if len(args) == 0 {
114+
if queryData.Config.GetString("outputFile", "") != "" {
115+
queryData.Message = "Output now falls back to stdout"
116+
queryData.Config.SetString("outputFile", "")
117+
} else {
118+
queryData.Message = "No output file provided"
119+
queryData.StatusCode = 2
120+
}
121+
} else {
122+
queryData.Config.SetString("outputFile", args[0])
123+
// We don't set the message because it would be outputted
124+
// to the file
125+
}
126+
127+
case "print":
128+
queryData.Message = strings.Join(args, " ")
129+
130+
case "schema":
131+
if len(args) == 0 {
132+
queryData.SQLQuery = "SELECT sql FROM sqlite_schema"
133+
} else {
134+
queryData.SQLQuery = "SELECT sql FROM sqlite_schema WHERE tbl_name = ?"
135+
queryData.Args = append(queryData.Args, args[0])
136+
}
137+
138+
return true
139+
140+
case "separator":
141+
if len(args) == 0 {
142+
queryData.Message = "No separator provided"
143+
queryData.StatusCode = 2
144+
} else if len(args) == 1 {
145+
queryData.Config.SetString("separatorColumn", args[0])
146+
} else {
147+
queryData.Config.SetString("separatorColumn", args[0])
148+
queryData.Config.SetString("separatorRow", args[1])
149+
}
150+
151+
case "shell", "system":
152+
// We run the command
153+
if len(args) == 0 {
154+
queryData.Message = "No command provided"
155+
queryData.StatusCode = 2
156+
} else {
157+
command := args[0]
158+
args = args[1:]
159+
cmd := exec.Command(command, args...)
160+
// Buffer the output
161+
output, err := cmd.CombinedOutput()
162+
if err != nil {
163+
queryData.Message = err.Error()
164+
queryData.StatusCode = 2
165+
} else {
166+
queryData.Message = string(output)
167+
queryData.StatusCode = 0
168+
}
169+
}
170+
case "tables", "table":
171+
queryData.SQLQuery = "SELECT name FROM pragma_table_list UNION SELECT name FROM pragma_module_list" +
172+
" WHERE name NOT LIKE 'fts%' AND name NOT LIKE 'rtree%'"
173+
return true
174+
}
175+
176+
return false
177+
178+
}
179+
180+
// Return the command and the arguments of a dot command
181+
func parseDotFunc(query string) (string, []string) {
182+
// We parse it by splitting it by spaces
183+
// and removing the first element
184+
// which is the dot command itself
185+
command := ""
186+
args := []string{}
187+
188+
tempStr := strings.Builder{}
189+
for i := 1; i < len(query); i++ {
190+
if query[i] == ' ' {
191+
if command == "" {
192+
command = tempStr.String()
193+
} else {
194+
args = append(args, tempStr.String())
195+
}
196+
tempStr.Reset()
197+
} else {
198+
tempStr.WriteByte(query[i])
199+
}
200+
}
201+
202+
// We add the last argument
203+
if command == "" {
204+
command = tempStr.String()
205+
} else {
206+
args = append(args, tempStr.String())
207+
}
208+
209+
return command, args
210+
}
211+
212+
// Rewrite the query
213+
func middlewareMySQL(queryData *QueryData) bool {
214+
// Check if the query is a MySQL query
215+
if !queryData.Config.GetBool("mysql", false) {
216+
return true
217+
}
218+
219+
queryType, stmt, err := namespace.GetQueryType(queryData.SQLQuery)
220+
if err != nil {
221+
// If we wan't parse the query, we just pass it
222+
return true
223+
}
224+
if queryType == sqlparser.StmtShow {
225+
queryData.SQLQuery, queryData.Args = namespace.RewriteShowStatement(stmt.(*sqlparser.Show))
226+
}
227+
228+
return true
229+
}
230+
231+
func middlewareQuery(queryData *QueryData) bool {
232+
// Run the query on the database
233+
if queryData.DB == nil {
234+
return true
235+
}
236+
237+
// We run the query
238+
rows, err := queryData.DB.Query(queryData.SQLQuery, queryData.Args...)
239+
if err != nil {
240+
queryData.Message = err.Error()
241+
queryData.StatusCode = 2
242+
return false
243+
}
244+
245+
queryData.Result = rows
246+
return true
247+
}
248+
249+
func middlewareSlashCommand(queryData *QueryData) bool {
250+
// Check if slash command are enabled
251+
if !queryData.Config.GetBool("slash-command", false) {
252+
return true
253+
}
254+
255+
// Check if the query is a slash command
256+
query := queryData.SQLQuery
257+
258+
// Ensure the trimmed query starts with a slash
259+
// If not, we skip to the next middleware
260+
trimmedQuery := strings.TrimSpace(query)
261+
if !strings.HasPrefix(trimmedQuery, "\\") {
262+
return true
263+
}
264+
265+
splittedCommand := strings.Split(trimmedQuery, " ")
266+
commands := strings.ToLower(splittedCommand[0][1:])
267+
args := splittedCommand[1:]
268+
269+
switch commands {
270+
case "l":
271+
queryData.SQLQuery = "SELECT name FROM pragma_database_list;"
272+
return true
273+
case "d":
274+
if len(args) == 0 {
275+
queryData.SQLQuery = "SELECT name FROM pragma_table_list UNION SELECT name FROM pragma_module_list" +
276+
" WHERE name NOT LIKE 'fts%' AND name NOT LIKE 'rtree%'"
277+
} else {
278+
queryData.SQLQuery = "SELECT name as Column, type as Type, '' as Collation, iif(\"notnull\" = 0, '', 'not null') as \"Null\"," +
279+
" dflt_value as \"Default\", pk as PrimaryKey FROM pragma_table_info(?)"
280+
281+
queryData.Args = append(queryData.Args, args[0])
282+
}
283+
return true
284+
case "d+":
285+
if len(args) == 0 {
286+
queryData.SQLQuery = "SELECT * FROM pragma_table_list"
287+
} else {
288+
queryData.SQLQuery = "SELECT name as Column, type as Type, '' as Collation, iif(\"notnull\" = 0, '', 'not null') as \"Null\"," +
289+
" dflt_value as \"Default\", pk as PrimaryKey FROM pragma_table_info(?)"
290+
queryData.Args = append(queryData.Args, args[0])
291+
}
292+
return true
293+
case "di":
294+
queryData.SQLQuery = "SELECT * FROM pragma_index_list;"
295+
return true
296+
case "dt":
297+
queryData.SQLQuery = "SELECT name FROM pragma_table_list UNION SELECT name FROM pragma_module_list" +
298+
" WHERE name NOT LIKE 'fts%' AND name NOT LIKE 'rtree%'"
299+
return true
300+
case "dv":
301+
queryData.SQLQuery = "SELECT * FROM pragma_table_list WHERE type = 'view';"
302+
return true
303+
304+
default:
305+
queryData.Message = fmt.Sprintf("Unknown command \\%s", commands)
306+
queryData.StatusCode = 2
307+
}
308+
309+
return false
310+
311+
}

0 commit comments

Comments
 (0)