Execute big batches in chunks

Fixes tursodatabase/turso-cli#839
Refs tursodatabase/turso-cli#735

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>

Author: haaawk

libsql/internal/http/hranaV2/hranaV2.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/tursodatabase/libsql-client-go/sqliteparserutils"
 	"io"
 	"net/http"
 	net_url "net/url"
@@ -303,10 +304,114 @@ func (h *hranaV2Conn) executeMsg(ctx context.Context, msg *hrana.PipelineRequest
 	return result, nil
 }
 
+type chunker struct {
+	chunk    []string
+	iterator *sqliteparserutils.StatementIterator
+	limit    int
+}
+
+func newChunker(iterator *sqliteparserutils.StatementIterator, limit int) *chunker {
+	return &chunker{iterator: iterator, chunk: make([]string, 0, limit), limit: limit}
+}
+
+func isTransactionStatement(stmt string) bool {
+	patterns := [][]byte{[]byte("begin"), []byte("commit"), []byte("end"), []byte("rollback")}
+	for _, p := range patterns {
+		if len(stmt) >= len(p) && bytes.Equal(bytes.ToLower([]byte(stmt[0:len(p)])), p) {
+			return true
+		}
+	}
+	return false
+}
+
+func (c *chunker) Next() (chunk []string, isEOF bool) {
+	c.chunk = c.chunk[:0]
+	var stmt string
+	for !isEOF && len(c.chunk) < c.limit {
+		stmt, _, isEOF = c.iterator.Next()
+		// We need to skip transaction statements. Chunks run in a transaction by default.
+		if stmt != "" && !isTransactionStatement(stmt) {
+			c.chunk = append(c.chunk, stmt)
+		}
+	}
+	return c.chunk, isEOF
+}
+
+func (h *hranaV2Conn) executeSingleStmt(ctx context.Context, stmt string, wantRows bool) (*hrana.PipelineResponse, error) {
+	msg := &hrana.PipelineRequest{}
+	executeStream, err := hrana.ExecuteStream(stmt, nil, wantRows)
+	if err != nil {
+		return nil, fmt.Errorf("failed to execute SQL: %s\n%w", stmt, err)
+	}
+	msg.Add(*executeStream)
+	res, err := h.executeMsg(ctx, msg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to execute SQL: %s\n%w", stmt, err)
+	}
+	return res, nil
+}
+
+func (h *hranaV2Conn) executeInChunks(ctx context.Context, query string, wantRows bool) (*hrana.PipelineResponse, error) {
+	const chunkSize = 4096
+	iterator := sqliteparserutils.CreateStatementIterator(query)
+	chunker := newChunker(iterator, chunkSize)
+
+	chunk, isEOF := chunker.Next()
+	if isEOF && len(chunk) == 1 {
+		return h.executeSingleStmt(ctx, chunk[0], wantRows)
+	}
+
+	_, err := h.executeSingleStmt(ctx, "BEGIN", false)
+	if err != nil {
+		return nil, err
+	}
+
+	batch := &hrana.Batch{Steps: make([]hrana.BatchStep, chunkSize)}
+	msg := &hrana.PipelineRequest{}
+	msg.Add(hrana.StreamRequest{Type: "batch", Batch: batch})
+	for idx := range batch.Steps {
+		batch.Steps[idx].Stmt.WantRows = wantRows
+	}
+
+	result := &hrana.PipelineResponse{}
+	for {
+		for idx := range chunk {
+			batch.Steps[idx].Stmt.Sql = &chunk[idx]
+		}
+		if len(chunk) < chunkSize {
+			// We can trim batch.Steps because this is the last chunk anyway.
+			// isEOF has to be true at this point.
+			batch.Steps = batch.Steps[:len(chunk)]
+		}
+		res, err := h.executeMsg(ctx, msg)
+		if err != nil {
+			h.closeStream()
+			return nil, fmt.Errorf("failed to execute SQL:\n%w", err)
+		}
+		result.Baton = res.Baton
+		result.BaseUrl = res.BaseUrl
+		result.Results = append(result.Results, res.Results...)
+		if isEOF {
+			break
+		}
+		chunk, isEOF = chunker.Next()
+	}
+	_, err = h.executeSingleStmt(ctx, "COMMIT", false)
+	if err != nil {
+		h.closeStream()
+		return nil, err
+	}
+	return result, nil
+}
+
 func (h *hranaV2Conn) executeStmt(ctx context.Context, query string, args []driver.NamedValue, wantRows bool) (*hrana.PipelineResponse, error) {
+	const querySizeLimitForChunking = 20 * 1024 * 1024
+	if len(args) == 0 && len(query) > querySizeLimitForChunking && !h.schemaDb {
+		return h.executeInChunks(ctx, query, wantRows)
+	}
 	stmts, params, err := shared.ParseStatementAndArgs(query, args)
 	if err != nil {
-		return nil, fmt.Errorf("failed to execute SQL: %s\n%w", query, err)
+		return nil, fmt.Errorf("failed to execute SQL:\n%w", err)
 	}
 	msg := &hrana.PipelineRequest{}
 	if len(stmts) == 1 {
@@ -316,20 +421,20 @@ func (h *hranaV2Conn) executeStmt(ctx context.Context, query string, args []driv
 		}
 		executeStream, err := hrana.ExecuteStream(stmts[0], p, wantRows)
 		if err != nil {
-			return nil, fmt.Errorf("failed to execute SQL: %s\n%w", query, err)
+			return nil, fmt.Errorf("failed to execute SQL:\n%w", err)
 		}
 		msg.Add(*executeStream)
 	} else {
 		batchStream, err := hrana.BatchStream(stmts, params, wantRows, !h.schemaDb)
 		if err != nil {
-			return nil, fmt.Errorf("failed to execute SQL: %s\n%w", query, err)
+			return nil, fmt.Errorf("failed to execute SQL:\n%w", err)
 		}
 		msg.Add(*batchStream)
 	}
 
 	resp, err := h.executeMsg(ctx, msg)
 	if err != nil {
-		return nil, fmt.Errorf("failed to execute SQL: %s\n%w", query, err)
+		return nil, fmt.Errorf("failed to execute SQL:\n%w", err)
 	}
 	return resp, nil
 }