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

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"crypto/tls"
     9  	"fmt"
    10  	"os"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/ydb-platform/ydb-go-sdk/v3"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/config"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/table"
    21  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    22  )
    23  
    24  func TestScripting(t *testing.T) {
    25  	ctx := xtest.Context(t)
    26  
    27  	db, err := ydb.Open(ctx,
    28  		os.Getenv("YDB_CONNECTION_STRING"),
    29  		ydb.WithAccessTokenCredentials(
    30  			os.Getenv("YDB_ACCESS_TOKEN_CREDENTIALS"),
    31  		),
    32  		ydb.With(
    33  			config.WithOperationTimeout(time.Second*2),
    34  			config.WithOperationCancelAfter(time.Second*2),
    35  		),
    36  		ydb.WithBalancer(balancers.SingleConn()),
    37  		ydb.WithConnectionTTL(time.Millisecond*10000),
    38  		ydb.WithMinTLSVersion(tls.VersionTLS10),
    39  		ydb.WithLogger(
    40  			newLogger(t),
    41  			trace.MatchDetails(`ydb\.(driver|discovery|retry|scheme).*`),
    42  		),
    43  		ydb.WithUserAgent("scripting"),
    44  	)
    45  	if err != nil {
    46  		t.Fatal(err)
    47  	}
    48  	defer func() {
    49  		// cleanup connection
    50  		if e := db.Close(ctx); e != nil {
    51  			t.Fatalf("db close failed: %+v", e)
    52  		}
    53  	}()
    54  	// Execute
    55  	if err = retry.Retry(ctx, func(ctx context.Context) (err error) {
    56  		res, err := db.Scripting().Execute(
    57  			ctx,
    58  			"SELECT 1+1",
    59  			table.NewQueryParameters(),
    60  		)
    61  		if err != nil {
    62  			return err
    63  		}
    64  		defer func() {
    65  			_ = res.Close()
    66  		}()
    67  		if !res.NextResultSet(ctx) {
    68  			return retry.RetryableError(
    69  				fmt.Errorf("no result sets"),
    70  				retry.WithBackoff(retry.TypeNoBackoff),
    71  			)
    72  		}
    73  		if !res.NextRow() {
    74  			return retry.RetryableError(
    75  				fmt.Errorf("no rows"),
    76  				retry.WithBackoff(retry.TypeSlowBackoff),
    77  			)
    78  		}
    79  		var sum int32
    80  		if err = res.Scan(&sum); err != nil {
    81  			return fmt.Errorf("scan failed: %w", err)
    82  		}
    83  		if sum != 2 {
    84  			return fmt.Errorf("unexpected sum: %v", sum)
    85  		}
    86  		return res.Err()
    87  	}, retry.WithIdempotent(true)); err != nil {
    88  		t.Fatalf("Execute failed: %v", err)
    89  	}
    90  	// StreamExecute
    91  	if err = retry.Retry(ctx, func(ctx context.Context) (err error) {
    92  		res, err := db.Scripting().StreamExecute(
    93  			ctx,
    94  			"SELECT 1+1",
    95  			table.NewQueryParameters(),
    96  		)
    97  		if err != nil {
    98  			return err
    99  		}
   100  		defer func() {
   101  			_ = res.Close()
   102  		}()
   103  		if !res.NextResultSet(ctx) {
   104  			return retry.RetryableError(
   105  				fmt.Errorf("no result sets"),
   106  				retry.WithBackoff(retry.TypeNoBackoff),
   107  				retry.WithDeleteSession(),
   108  			)
   109  		}
   110  		if !res.NextRow() {
   111  			return retry.RetryableError(
   112  				fmt.Errorf("no rows"),
   113  				retry.WithBackoff(retry.TypeFastBackoff),
   114  			)
   115  		}
   116  		var sum int32
   117  		if err = res.Scan(&sum); err != nil {
   118  			return fmt.Errorf("scan failed: %w", err)
   119  		}
   120  		if sum != 2 {
   121  			return fmt.Errorf("unexpected sum: %v", sum)
   122  		}
   123  		return res.Err()
   124  	}, retry.WithIdempotent(true)); err != nil {
   125  		t.Fatalf("StreamExecute failed: %v", err)
   126  	}
   127  	// ExplainPlan
   128  	if err = retry.Retry(ctx, func(ctx context.Context) (err error) {
   129  		res, err := db.Scripting().Explain(
   130  			ctx,
   131  			"SELECT 1+1",
   132  			scripting.ExplainModePlan,
   133  		)
   134  		if err != nil {
   135  			return err
   136  		}
   137  		if res.Plan == "" {
   138  			return fmt.Errorf("empty plan")
   139  		}
   140  		return nil
   141  	}, retry.WithIdempotent(true)); err != nil {
   142  		t.Fatalf("Explain failed: %v", err)
   143  	}
   144  	// ExplainValidate
   145  	if err = retry.Retry(ctx, func(ctx context.Context) (err error) {
   146  		res, err := db.Scripting().Explain(
   147  			ctx,
   148  			"SELECT 1+1",
   149  			scripting.ExplainModeValidate,
   150  		)
   151  		if err != nil {
   152  			return err
   153  		}
   154  		if len(res.ParameterTypes) > 0 {
   155  			return fmt.Errorf("unexpected parameter types")
   156  		}
   157  		return nil
   158  	}, retry.WithIdempotent(true)); err != nil {
   159  		t.Fatalf("Explain failed: %v", err)
   160  	}
   161  }