Merge pull request #1268 from ydb-platform/range-result

* Added experimental range functions to the `query.Result` and `query.ResultSet` types, available starting with Go version 1.22. These features can be enabled by setting the environment variable `GOEXPERIMENT=rangefunc`.

Author: asmyasnikov

.github/workflows/tests.yml

@@ -78,10 +78,60 @@ jobs:
           go-version: ${{ matrix.go-version }}
           cache: true
       - name: Integration test
-        run: go test -race -tags integration -coverpkg=./... -coverprofile integration-secure.txt -covermode atomic ./tests/integration
+        run: go test -race -tags integration -coverpkg=./... -coverprofile integration.txt -covermode atomic ./tests/integration
       - name: Upload Test secure connection coverage report to Codecov
         uses: codecov/codecov-action@v4
         with:
-          file: ./integration-secure.txt
+          file: ./integration.txt
           flags: integration,${{ matrix.os }},go-${{ matrix.go-version }},ydb-${{ matrix.ydb-version }}
-          name: integration-secure
+          name: integration
+  experiment:
+    concurrency:
+      group: experiment-${{ github.ref }}-${{ matrix.os }}-${{ matrix.go-version }}-${{ matrix.ydb-version }}
+      cancel-in-progress: true
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        go-version: [1.22.x]
+        ydb-version: [24.1]
+    services:
+      ydb:
+        image: ydbplatform/local-ydb:${{ matrix.ydb-version }}
+        ports:
+          - 2135:2135
+          - 2136:2136
+          - 8765:8765
+        volumes:
+          - /tmp/ydb_certs:/ydb_certs
+        env:
+          YDB_LOCAL_SURVIVE_RESTART: true
+          YDB_USE_IN_MEMORY_PDISKS: true
+          YDB_TABLE_ENABLE_PREPARED_DDL: true
+        options: '-h localhost'
+    env:
+      OS: ubuntu-latest
+      GO: ${{ matrix.go-version }}
+      YDB_VERSION: ${{ matrix.ydb-version }}
+      YDB_CONNECTION_STRING: grpc://localhost:2136/local
+      YDB_CONNECTION_STRING_SECURE: grpcs://localhost:2135/local
+      YDB_SSL_ROOT_CERTIFICATES_FILE: /tmp/ydb_certs/ca.pem
+      YDB_SESSIONS_SHUTDOWN_URLS: http://localhost:8765/actors/kqp_proxy?force_shutdown=all
+      HIDE_APPLICATION_OUTPUT: 1
+      GOEXPERIMENT: rangefunc
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ matrix.go-version }}
+          cache: true
+      - name: Integration test
+        run: go test -race -tags integration -coverpkg=./... -coverprofile experiment.txt -covermode atomic ./...
+      - name: Upload Test secure connection coverage report to Codecov
+        uses: codecov/codecov-action@v4
+        with:
+          file: ./integration-secure.txt
+          flags: experiment,${{ matrix.os }},go-${{ matrix.go-version }},ydb-${{ matrix.ydb-version }}
+          name: experiment

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added experimental range functions to the `query.Result` and `query.ResultSet` types, available as for-range loops starting with Go version 1.22. These features can be enabled by setting the environment variable `GOEXPERIMENT=rangefunc`.
 * Added public types for `tx.Option`, `options.DoOption` and `options.DoTxOption`
 
 ## v3.73.1

internal/query/range_experiment.go

@@ -0,0 +1,44 @@
+package query
+
+import (
+	"context"
+	"io"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter"
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+)
+
+func rangeResultSets(ctx context.Context, r query.Result) xiter.Seq2[query.ResultSet, error] {
+	return func(yield func(query.ResultSet, error) bool) {
+		for {
+			rs, err := r.NextResultSet(ctx)
+			if err != nil {
+				if xerrors.Is(err, io.EOF) {
+					return
+				}
+			}
+			cont := yield(rs, err)
+			if !cont || err != nil {
+				return
+			}
+		}
+	}
+}
+
+func rangeRows(ctx context.Context, rs query.ResultSet) xiter.Seq2[query.Row, error] {
+	return func(yield func(query.Row, error) bool) {
+		for {
+			rs, err := rs.NextRow(ctx)
+			if err != nil {
+				if xerrors.Is(err, io.EOF) {
+					return
+				}
+			}
+			cont := yield(rs, err)
+			if !cont || err != nil {
+				return
+			}
+		}
+	}
+}

internal/query/result.go

@@ -11,6 +11,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xiter"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
@@ -37,6 +38,14 @@ type (
 	}
 )
 
+func (r *materializedResult) Range(ctx context.Context) xiter.Seq2[query.ResultSet, error] {
+	return rangeResultSets(ctx, r)
+}
+
+func (r *result) Range(ctx context.Context) xiter.Seq2[query.ResultSet, error] {
+	return rangeResultSets(ctx, r)
+}
+
 func (r *materializedResult) Close(ctx context.Context) error {
 	return nil
 }