github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/tests/integration/query_multi_result_sets_test.go (about)

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"errors"
     9  	"fmt"
    10  	"io"
    11  	"os"
    12  	"testing"
    13  
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/version"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/query"
    16  )
    17  
    18  func TestQueryMultiResultSets(t *testing.T) {
    19  	if version.Lt(os.Getenv("YDB_VERSION"), "24.1") {
    20  		t.Skip("query service not allowed in YDB version '" + os.Getenv("YDB_VERSION") + "'")
    21  	}
    22  
    23  	scope := newScope(t)
    24  	var i, j, k int
    25  	db := scope.Driver()
    26  	err := db.Query().Do(scope.Ctx, func(ctx context.Context, s query.Session) (err error) {
    27  		res, err := s.Query(ctx, `SELECT 42; SELECT 43, 44;`,
    28  			query.WithIdempotent(),
    29  		)
    30  		if err != nil {
    31  			return fmt.Errorf("can't get result: %w", err)
    32  		}
    33  		set, err := res.NextResultSet(ctx)
    34  		if err != nil {
    35  			return fmt.Errorf("set 0: get next result set: %w", err)
    36  		}
    37  
    38  		for row, err := set.NextRow(ctx); !errors.Is(err, io.EOF); row, err = set.NextRow(ctx) {
    39  			if err != nil {
    40  				return fmt.Errorf("set 0: get next row: %w", err)
    41  			}
    42  			if err := row.Scan(&i); err != nil {
    43  				return fmt.Errorf("set 0: scan row: %w", err)
    44  			}
    45  			fmt.Println(i)
    46  		}
    47  
    48  		set, err = res.NextResultSet(ctx)
    49  		if err != nil {
    50  			return err
    51  		}
    52  
    53  		for row, err := set.NextRow(ctx); !errors.Is(err, io.EOF); row, err = set.NextRow(ctx) {
    54  			if err != nil {
    55  				return fmt.Errorf("set 1: get next row: %w", err)
    56  			}
    57  
    58  			if err := row.Scan(&j, &k); err != nil {
    59  				return fmt.Errorf("set 1: scan row: %w", err)
    60  			}
    61  			fmt.Println(j, k)
    62  		}
    63  		_, err = res.NextResultSet(ctx)
    64  		if !errors.Is(err, io.EOF) {
    65  			return fmt.Errorf("get next result set: %w", err)
    66  		}
    67  
    68  		return nil
    69  	}, query.WithIdempotent())
    70  
    71  	scope.Require.NoError(err)
    72  
    73  	scope.Require.Equal(42, i)
    74  	scope.Require.Equal(43, j)
    75  	scope.Require.Equal(44, k)
    76  }