OAS-10935 Get Raw JSON bytes from query cursor (#683)

Author: enikon

v2/arangodb/cursor.go

@@ -23,6 +23,8 @@ package arangodb
 import (
 	"context"
 	"io"
+
+	"github.com/arangodb/go-driver/v2/connection"
 )
 
 // Cursor is returned from a query, used to iterate over a list of documents.
@@ -77,6 +79,14 @@ type CursorBatch interface {
 	// E.g. `var result []MyStruct{}`.
 	RetryReadBatch(ctx context.Context, result interface{}) error
 
+	// ReadNextRawBatch reads the next batch of documents from the cursor.
+	// The result must be a pointer to a byte array *[]bytes.
+	ReadNextRawBatch(ctx context.Context, result *connection.RawObject) error
+
+	// RetryReadRawBatch retries the last batch read made by ReadNextRawBatch.
+	// The result must be a pointer to a byte array *[]bytes.
+	RetryReadRawBatch(ctx context.Context, result *connection.RawObject) error
+
 	// Count returns the total number of result documents available.
 	// A valid return value is only available when the cursor has been created with `Count` and not with `Stream`.
 	Count() int64

v2/arangodb/cursor_impl.go

@@ -115,22 +115,37 @@ func (c *cursor) ReadDocument(ctx context.Context, result interface{}) (Document
 	return c.readDocument(ctx, result)
 }
 
-func (c *cursor) ReadNextBatch(ctx context.Context, result interface{}) error {
-	err := c.getNextBatch(ctx, "")
+func (c *cursor) readBatch(ctx context.Context, result interface{}, retryBatchID string) error {
+	err := c.getNextBatch(ctx, retryBatchID)
 	if err != nil {
 		return err
 	}
-
 	return json.Unmarshal(c.data.Result.in, result)
 }
 
-func (c *cursor) RetryReadBatch(ctx context.Context, result interface{}) error {
-	err := c.getNextBatch(ctx, c.retryData.currentBatchID)
+func (c *cursor) readRawBatch(ctx context.Context, result *connection.RawObject, retryBatchID string) error {
+	err := c.getNextBatch(ctx, retryBatchID)
 	if err != nil {
 		return err
 	}
+	*result = append([]byte(nil), c.data.Result.in...)
+	return nil
+}
 
-	return json.Unmarshal(c.data.Result.in, result)
+func (c *cursor) ReadNextBatch(ctx context.Context, result interface{}) error {
+	return c.readBatch(ctx, result, "")
+}
+
+func (c *cursor) RetryReadBatch(ctx context.Context, result interface{}) error {
+	return c.readBatch(ctx, result, c.retryData.currentBatchID)
+}
+
+func (c *cursor) ReadNextRawBatch(ctx context.Context, result *connection.RawObject) error {
+	return c.readRawBatch(ctx, result, "")
+}
+
+func (c *cursor) RetryReadRawBatch(ctx context.Context, result *connection.RawObject) error {
+	return c.readRawBatch(ctx, result, c.retryData.currentBatchID)
 }
 
 func (c *cursor) readDocument(ctx context.Context, result interface{}) (DocumentMeta, error) {

v2/connection/connection.go

@@ -148,3 +148,7 @@ type Response interface {
 
 	RawResponse() *http.Response
 }
+
+// RawObject is a raw encoded object.
+// Connection implementations must be able to unmarshal *RawObject into Go objects.
+type RawObject []byte

v2/tests/cursor_test.go

@@ -0,0 +1,78 @@
+//
+// DISCLAIMER
+//
+// Copyright 2025 ArangoDB GmbH, Cologne, Germany
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Copyright holder is ArangoDB GmbH, Cologne, Germany
+//
+
+package tests
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/arangodb/go-driver/v2/arangodb"
+	"github.com/arangodb/go-driver/v2/connection"
+	"github.com/stretchr/testify/require"
+)
+
+// Test_ExplainQuery tries to explain several AQL queries.
+func Test_CursorRawResult(t *testing.T) {
+	Wrap(t, func(t *testing.T, client arangodb.Client) {
+		WithDatabase(t, client, nil, func(db arangodb.Database) {
+			WithCollectionV2(t, db, nil, func(col arangodb.Collection) {
+				WithUserDocs(t, col, func(docs []UserDoc) {
+					withContextT(t, defaultTestTimeout, func(ctx context.Context, tb testing.TB) {
+						skipBelowVersion(client, ctx, "3.11", t)
+
+						query := fmt.Sprintf("FOR d IN `%s` SORT d.Title RETURN d", col.Name())
+
+						t.Run("Test retry read when reading raw JSON batch", func(t *testing.T) {
+							opts := arangodb.QueryOptions{
+								Count:     true,
+								BatchSize: 2,
+								Options: arangodb.QuerySubOptions{
+									AllowRetry: true,
+								},
+							}
+
+							cursor, err := db.QueryBatch(ctx, query, &opts, nil)
+							require.NoError(t, err)
+							for {
+								if !cursor.HasMoreBatches() {
+									break
+								}
+								var result connection.RawObject
+								require.NoError(t, cursor.ReadNextRawBatch(ctx, &result))
+
+								var resultRetry connection.RawObject
+								require.NoError(t, cursor.RetryReadRawBatch(ctx, &resultRetry))
+
+								require.Equal(t, len(result), len(resultRetry))
+								require.Equal(t, result, resultRetry)
+
+							}
+
+							err = cursor.Close()
+							require.NoError(t, err)
+						})
+					})
+				})
+			})
+		})
+	})
+}