Make transactions work in Batch

Signed-off-by: Piotr Jastrzebski <piotr@chiselstrike.com>
tursodatabase · Apr 10, 2024 · 6fedad9 · 6fedad9
1 parent f087768
commit 6fedad9
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/libsql/internal/hrana/stream_request.go b/libsql/internal/hrana/stream_request.go
@@ -49,8 +49,20 @@ func BatchStream(sqls []string, params []shared.Params, wantRows bool) (*StreamR
 		if err := stmt.AddArgs(params[idx]); err != nil {
 			return nil, err
 		}
-		batch.Add(*stmt, nil)
+		var condition *BatchCondition
+		if idx > 0 {
+			prev_idx := int32(idx - 1)
+			condition = &BatchCondition{
+				Type: "ok",
+				Step: &prev_idx,
+			}
+		}
+		batch.Add(*stmt, condition)
 	}
+	rollback := "ROLLBACK"
+	last_idx := int32(len(sqls) - 1)
+	batch.Add(Stmt{Sql: &rollback, WantRows: false},
+		&BatchCondition{Type: "not", Cond: &BatchCondition{Type: "ok", Step: &last_idx}})
 	return &StreamRequest{Type: "batch", Batch: batch}, nil
 }
 

diff --git a/libsql/internal/http/hranaV2/hranaV2.go b/libsql/internal/http/hranaV2/hranaV2.go
@@ -338,7 +338,8 @@ func (h *hranaV2Conn) ExecContext(ctx context.Context, query string, args []driv
 		}
 		lastInsertRowId := int64(0)
 		affectedRowCount := int64(0)
-		for _, r := range res.StepResults {
+		for idx := 0; idx < len(res.StepResults)-1; idx++ {
+			r := res.StepResults[idx]
 			rowId := r.GetLastInsertRowId()
 			if rowId > 0 {
 				lastInsertRowId = rowId
@@ -457,6 +458,8 @@ func (h *hranaV2Conn) QueryContext(ctx context.Context, query string, args []dri
 		if err != nil {
 			return nil, err
 		}
+		res.StepResults = res.StepResults[:len(res.StepResults)-1]
+		res.StepErrors = res.StepErrors[:len(res.StepErrors)-1]
 		return shared.NewRows(&BatchResultRowsProvider{res}), nil
 	default:
 		return nil, fmt.Errorf("failed to execute SQL: %s\n%s", query, "unknown response type")