remove-naive-query-framentation

Author: general-kroll-4-life

internal/stackql/acid/tsm_physio/best_effort_orchestrator.go

@@ -2,8 +2,8 @@ package tsm_physio //nolint:revive,stylecheck // prefer this nomenclature
 
 import (
 	"fmt"
-	"strings"
 
+	"github.com/stackql/stackql-parser/go/vt/sqlparser"
 	"github.com/stackql/stackql/internal/stackql/acid/binlog"
 	"github.com/stackql/stackql/internal/stackql/acid/tsm"
 	"github.com/stackql/stackql/internal/stackql/acid/txn_context"
@@ -42,7 +42,8 @@ func (orc *bestEffortOrchestrator) processQueryOrQueries(
 ) ([]internaldto.ExecutorOutput, bool) {
 	var retVal []internaldto.ExecutorOutput
 	cmdString := handlerCtx.GetRawQuery()
-	for _, s := range strings.Split(cmdString, ";") {
+	splitQueries, _ := sqlparser.SplitStatementToPieces(cmdString)
+	for _, s := range splitQueries {
 		response, hasResponse := orc.processQuery(handlerCtx, s)
 		if hasResponse {
 			retVal = append(retVal, response...)

internal/stackql/acid/tsm_physio/txn_orchestrator.go

@@ -2,9 +2,9 @@ package tsm_physio //nolint:stylecheck,revive // prefer this nomenclature
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/stackql/any-sdk/pkg/constants"
+	"github.com/stackql/stackql-parser/go/vt/sqlparser"
 	"github.com/stackql/stackql/internal/stackql/acid/tsm"
 	"github.com/stackql/stackql/internal/stackql/acid/txn_context"
 	"github.com/stackql/stackql/internal/stackql/handler"
@@ -68,7 +68,8 @@ func (orc *standardOrchestrator) processQueryOrQueries(
 ) ([]internaldto.ExecutorOutput, bool) {
 	var retVal []internaldto.ExecutorOutput
 	cmdString := handlerCtx.GetRawQuery()
-	for _, s := range strings.Split(cmdString, ";") {
+	splitQueries, _ := sqlparser.SplitStatementToPieces(cmdString)
+	for _, s := range splitQueries {
 		response, hasResponse := orc.processQuery(handlerCtx, s)
 		if hasResponse {
 			retVal = append(retVal, response...)

internal/stackql/cmd/shell.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/stackql/any-sdk/pkg/dto"
 	"github.com/stackql/any-sdk/pkg/logging"
+	"github.com/stackql/stackql-parser/go/vt/sqlparser"
 	"github.com/stackql/stackql/internal/stackql/config"
 	"github.com/stackql/stackql/internal/stackql/driver"
 	"github.com/stackql/stackql/internal/stackql/entryutil"
@@ -225,20 +226,24 @@ var shellCmd = &cobra.Command{
 				if inlineCommentIdx > -1 {
 					line = line[:inlineCommentIdx]
 				}
-				semiColonIdx := strings.Index(line, ";")
-				if semiColonIdx > -1 {
-					line = strings.TrimSpace(line[:semiColonIdx+1])
-					subSemiColonIdx := strings.Index(line, ";")
-					sb.WriteString(" " + line[:subSemiColonIdx+1])
-					rawQuery := sb.String()
-					queryToExecute, qErr := entryutil.PreprocessInline(runtimeCtx, rawQuery)
-					if qErr != nil {
-						io.WriteString(outErrFile, "\r\n"+qErr.Error()+"\r\n") //nolint:errcheck // TODO: investigate
+				hasRHSSemiColon := strings.HasSuffix(strings.TrimSpace(line), ";")
+				splitQueries, _ := sqlparser.SplitStatementToPieces(line)
+				if len(splitQueries) > 0 {
+					for i, s := range splitQueries {
+						if i == len(splitQueries)-1 && !hasRHSSemiColon {
+							sb.Reset()
+							sb.WriteString(s)
+						}
+						line = s
+						sb.WriteString(" " + line)
+						rawQuery := sb.String()
+						queryToExecute, qErr := entryutil.PreprocessInline(runtimeCtx, rawQuery)
+						if qErr != nil {
+							io.WriteString(outErrFile, "\r\n"+qErr.Error()+"\r\n") //nolint:errcheck // TODO: investigate
+						}
+						l.WriteToHistory(rawQuery) //nolint:errcheck // TODO: investigate
+						sessionRunnerInstance.RunCommand(queryToExecute)
 					}
-					l.WriteToHistory(rawQuery) //nolint:errcheck // TODO: investigate
-					sessionRunnerInstance.RunCommand(queryToExecute)
-					sb.Reset()
-					sb.WriteString(line[subSemiColonIdx+1:])
 				} else {
 					sb.WriteString(" " + line)
 				}

internal/stackql/driver/driver.go

@@ -147,9 +147,6 @@ func (dr *basicStackQLDriver) CloneSQLBackend() sqlbackend.ISQLBackend {
 //nolint:revive // TODO: review
 func (dr *basicStackQLDriver) HandleSimpleQuery(ctx context.Context, query string) (sqldata.ISQLResultStream, error) {
 	dr.handlerCtx.SetRawQuery(query)
-	// if strings.Count(query, ";") > 1 {
-	// 	return nil, fmt.Errorf("only support single queries in server mode at this time")
-	// }
 	res, ok := dr.processQueryOrQueries(dr.handlerCtx)
 	if !ok {
 		return nil, fmt.Errorf("no SQLresults available")