✨ Add support for querying files in MCP #47

Author: julien040

controller/llm.go

@@ -11,7 +11,6 @@ import (
 	"os"
 	"os/signal"
 	"runtime"
-	"strconv"
 	"strings"
 	"time"
 
@@ -20,7 +19,6 @@ import (
 	"github.com/charmbracelet/lipgloss"
 	"github.com/julien040/anyquery/controller/config/model"
 	"github.com/julien040/anyquery/namespace"
-	"github.com/julien040/anyquery/other/sqlparser"
 	ws_tunnel "github.com/julien040/anyquery/other/websocket_tunnel/client"
 	"github.com/julien040/anyquery/rpc"
 	"github.com/mark3labs/mcp-go/mcp"
@@ -183,108 +181,22 @@ func executeQueryLLM(
 	query string,
 	w io.Writer,
 ) error {
-	stmt, _, err := namespace.GetQueryType(query)
-	if err != nil { // If we can't determine the query type, we assume it's a SELECT
-		stmt = sqlparser.StmtSelect
-	}
-
-	if stmt == sqlparser.StmtSelect {
-		// Make a context that'll cancel the query after 40 seconds
-		ctx, cancel := context.WithTimeout(context.Background(), 40*time.Second)
-		defer cancel()
-
-		rows, err := db.QueryContext(ctx, query)
-		if err != nil {
-			return fmt.Errorf("failed to run the query: %w", err)
-		}
-
-		defer rows.Close()
-
-		columns, err := rows.Columns()
-		if err != nil {
-			return fmt.Errorf("failed to get the columns: %w", err)
-		}
-
-		// Write the columns, and table data as a markdown table
-		w.Write([]byte("|"))
-		for _, column := range columns {
-			w.Write([]byte(" "))
-			w.Write([]byte(column))
-			w.Write([]byte(" |"))
-		}
-		w.Write([]byte("\n|"))
-		for _, column := range columns {
-			w.Write([]byte(" "))
-			for i := 0; i < len(column); i++ {
-				w.Write([]byte("-"))
-			}
-			w.Write([]byte(" |"))
-		}
-		w.Write([]byte("\n"))
-
-		for rows.Next() {
-			values := make([]interface{}, len(columns))
-			for i := range values {
-				values[i] = new(interface{})
-			}
-
-			if err := rows.Scan(values...); err != nil {
-				return fmt.Errorf("failed to scan the row: %w", err)
-			}
-
-			w.Write([]byte("|"))
-			for _, value := range values {
-				w.Write([]byte(" "))
-				unknown, ok := value.(*interface{})
-				if ok && unknown != nil && *unknown != nil {
-					switch parsed := (*unknown).(type) {
-					case []byte:
-						w.Write([]byte(fmt.Sprintf("%x", parsed)))
-					case string:
-						w.Write([]byte(fmt.Sprintf("%s", parsed)))
-					case int64:
-						w.Write([]byte(strconv.FormatInt(parsed, 10)))
-					case float64:
-						w.Write([]byte(strconv.FormatFloat(parsed, 'f', -1, 64)))
-					case bool:
-						if parsed {
-							w.Write([]byte("true"))
-						} else {
-							w.Write([]byte("false"))
-						}
-					case time.Time:
-						w.Write([]byte(parsed.Format(time.RFC3339)))
-
-					default:
-						w.Write([]byte(fmt.Sprintf("%v", *unknown)))
-					}
-
-				} else {
-					w.Write([]byte("NULL"))
-				}
-
-				w.Write([]byte(" |"))
-			}
-			w.Write([]byte("\n"))
-		}
-
-		if rows.Err() != nil {
-			return fmt.Errorf("failed to iterate over the rows %w", rows.Err())
-		}
-
-	} else {
-		res, err := db.Exec(query)
-		if err != nil {
-			return fmt.Errorf("failed to execute the query: %w", err)
-		}
-
-		rowsAffected, err := res.RowsAffected()
-		if err != nil {
-			return fmt.Errorf("failed to get the number of rows affected: %w", err)
-		}
-
-		w.Write([]byte(fmt.Sprintf("Query executed, %d rows affected", rowsAffected)))
-	}
+	sh := shell{
+		DB:             db,
+		OutputFileDesc: w,
+		Middlewares: []middleware{
+			middlewareMySQL,
+			middlewareFileQuery,
+			middlewareQuery,
+		},
+		Config: middlewareConfiguration{
+			"mysql":             true,
+			"doNotModifyOutput": true, // Do not modify the output, just keep w
+		},
+	}
+
+	// Run the query
+	sh.Run(query)
 
 	return nil
 }

controller/middleware.go

@@ -59,7 +59,7 @@ func middlewareDotCommand(queryData *QueryData) bool {
 		return true
 
 	case "help":
-		queryData.Message = "Documentation available at https://docs.anyquery.dev"
+		queryData.Message = "Documentation available at https://anyquery.dev/docs/usage/running-queries#commands"
 		queryData.StatusCode = 0
 
 	case "indexes", "index":

controller/query.go

@@ -153,7 +153,6 @@ func Query(cmd *cobra.Command, args []string) error {
 			"mysql":         true,
 			"slash-command": true,
 		},
-		Namespace:      namespace,
 		OutputFile:     "stdout",
 		OutputFileDesc: os.Stdout,
 	}
@@ -190,7 +189,7 @@ func Query(cmd *cobra.Command, args []string) error {
 
 	// Check if the output file is a tty
 	// If not, we set the output mode to plain
-	if !term.IsTerminal(int(shell.OutputFileDesc.Fd())) {
+	if file, ok := shell.OutputFileDesc.(*os.File); !ok || !term.IsTerminal(int(file.Fd())) {
 		shell.Config["outputMode"] = "plain"
 	}
 

controller/run.go

@@ -509,7 +509,6 @@ func Run(cmd *cobra.Command, args []string) error {
 			"mysql":         true,
 			"slash-command": true,
 		},
-		Namespace:      namespace,
 		OutputFile:     "stdout",
 		OutputFileDesc: os.Stdout,
 	}
@@ -527,7 +526,7 @@ func Run(cmd *cobra.Command, args []string) error {
 
 	// Check if the output file is a tty
 	// If not, we set the output mode to plain
-	if !term.IsTerminal(int(shell.OutputFileDesc.Fd())) {
+	if file, ok := shell.OutputFileDesc.(*os.File); !ok || !term.IsTerminal(int(file.Fd())) {
 		shell.Config["outputMode"] = "plain"
 	}
 

controller/shell.go

@@ -18,7 +18,6 @@ import (
 	"github.com/briandowns/spinner"
 	"github.com/charmbracelet/lipgloss"
 	"github.com/elk-language/go-prompt"
-	"github.com/julien040/anyquery/namespace"
 	"golang.org/x/term"
 )
 
@@ -127,10 +126,7 @@ type shell struct {
 
 	// Where to output the result
 	OutputFile     string
-	OutputFileDesc *os.File
-
-	// The namespace to use
-	Namespace *namespace.Namespace
+	OutputFileDesc io.Writer
 
 	// The history of the shell
 	History []string
@@ -208,41 +204,45 @@ func (p *shell) Run(rawQuery string, args ...interface{}) bool {
 		}
 		s.Stop()
 
-		// Where the output will be written for this loop
-		tempOutput := os.Stdout
+		var tempOutput io.Writer = p.OutputFileDesc
 		/* tempOutputMustClose := false */
 
 		// If the output file is specified, we write the result to it
 		// and save it for later execution in p.OutputFileDesc
 		// We also set it to tempOutput to write the result to it
 		confPath := p.Config.GetString("outputFile", "")
-		if confPath == "" || confPath == "stdout" {
-			p.OutputFileDesc = os.Stdout
-			p.OutputFile = "stdout"
-		} else {
-			// We check if the file is already open
-			// If not, open it
-			if p.OutputFile != confPath {
-				if p.OutputFileDesc != nil && p.OutputFileDesc != os.Stdout {
-					p.OutputFileDesc.Close()
-				}
-				file, err := os.OpenFile(confPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
-				if err != nil {
-					queryData = QueryData{
-						Message:    fmt.Sprintf("Error opening file %s: %s", confPath, err.Error()),
-						StatusCode: 2,
+		doNotModifyOutput := p.Config.GetBool("doNotModifyOutput", false)
+		if !doNotModifyOutput {
+			// Where the output will be written for this loop
+			if confPath == "" || confPath == "stdout" {
+				p.OutputFileDesc = os.Stdout
+				p.OutputFile = "stdout"
+			} else {
+				// We check if the file is already open
+				// If not, open it
+				if p.OutputFile != confPath {
+					// Close previous file if it's a file and not stdout
+					if fileDesc, ok := p.OutputFileDesc.(*os.File); ok && fileDesc != os.Stdout {
+						fileDesc.Close()
+					}
+					file, err := os.OpenFile(confPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
+					if err != nil {
+						queryData = QueryData{
+							Message:    fmt.Sprintf("Error opening file %s: %s", confPath, err.Error()),
+							StatusCode: 2,
+						}
+					} else {
+						p.OutputFileDesc = file
+						p.OutputFile = confPath
 					}
-				} else {
-					p.OutputFileDesc = file
-					p.OutputFile = confPath
+				} else if fileDesc, ok := p.OutputFileDesc.(*os.File); ok && fileDesc != os.Stdout {
+					// If the file is already open and it's a file (not stdout), we seek to the end
+					// just to be sure to append the result
+					fileDesc.Seek(0, 2)
 				}
-			} else {
-				// If the file is already open, we seek to the end
-				// just to be sure to append the result
-				p.OutputFileDesc.Seek(0, 2)
 			}
+			tempOutput = p.OutputFileDesc
 		}
-		tempOutput = p.OutputFileDesc
 
 		// We check also if the output must be displayed only once somewhere
 		/* if p.Config.GetString("onceOutputFile", "") != "" {