diff --git a/v2/arangodb/cursor.go b/v2/arangodb/cursor.go index f7ae8f79..1eee54fe 100644 --- a/v2/arangodb/cursor.go +++ b/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 diff --git a/v2/arangodb/cursor_impl.go b/v2/arangodb/cursor_impl.go index c57cf7dd..9d4559a7 100644 --- a/v2/arangodb/cursor_impl.go +++ b/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) { diff --git a/v2/connection/connection.go b/v2/connection/connection.go index 06ea6ff2..52d6c271 100644 --- a/v2/connection/connection.go +++ b/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 diff --git a/v2/tests/cursor_test.go b/v2/tests/cursor_test.go new file mode 100644 index 00000000..7791037a --- /dev/null +++ b/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) + }) + }) + }) + }) + }) + }) +}