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

     1  //go:build integration
     2  // +build integration
     3  
     4  package integration
     5  
     6  import (
     7  	"context"
     8  	"database/sql"
     9  	"errors"
    10  	"testing"
    11  
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/ydb-platform/ydb-go-sdk/v3"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
    16  )
    17  
    18  func TestDatabaseSqlIsPrimaryKey(t *testing.T) {
    19  	var (
    20  		scope = newScope(t)
    21  		db    = scope.SQLDriverWithFolder()
    22  	)
    23  
    24  	defer func() {
    25  		_ = db.Close()
    26  	}()
    27  
    28  	t.Run("create-tables", func(t *testing.T) {
    29  		err := retry.Do(scope.Ctx, db, func(ctx context.Context, cc *sql.Conn) (err error) {
    30  			_, err = cc.ExecContext(
    31  				ydb.WithQueryMode(ctx, ydb.SchemeQueryMode),
    32  				`CREATE TABLE episodes (
    33  				series_id Uint64,
    34  				season_id Uint64,
    35  				episode_id Uint64,
    36  				title UTF8,
    37  				air_date Date,
    38  				views Uint64,
    39  				PRIMARY KEY (
    40  					series_id,
    41  					season_id,
    42  					episode_id
    43  				)
    44  			)`,
    45  			)
    46  
    47  			return err
    48  		}, retry.WithIdempotent(true))
    49  
    50  		require.NoError(t, err)
    51  	})
    52  
    53  	t.Run("is-primary-key", func(t *testing.T) {
    54  		err := retry.Do(scope.Ctx, db, func(ctx context.Context, cc *sql.Conn) (err error) {
    55  			for _, pk := range []string{"series_id", "season_id", "episode_id"} {
    56  				isPk := false
    57  				err = cc.Raw(func(drvConn interface{}) (err error) {
    58  					q, ok := drvConn.(interface {
    59  						IsPrimaryKey(context.Context, string, string) (bool, error)
    60  					})
    61  
    62  					if !ok {
    63  						return errors.New("drvConn does not implement extended API")
    64  					}
    65  
    66  					isPk, err = q.IsPrimaryKey(ctx, "episodes", pk)
    67  					return err
    68  				})
    69  
    70  				if err != nil {
    71  					return err
    72  				}
    73  
    74  				require.True(t, isPk)
    75  			}
    76  			return nil
    77  		}, retry.WithIdempotent(true))
    78  
    79  		require.NoError(t, err)
    80  	})
    81  
    82  	t.Run("is-not-primary-key", func(t *testing.T) {
    83  		err := retry.Do(scope.Ctx, db, func(ctx context.Context, cc *sql.Conn) (err error) {
    84  			for _, pk := range []string{"title", "air_date", "views"} {
    85  				isPk := false
    86  				err = cc.Raw(func(drvConn interface{}) (err error) {
    87  					q, ok := drvConn.(interface {
    88  						IsPrimaryKey(context.Context, string, string) (bool, error)
    89  					})
    90  
    91  					if !ok {
    92  						return errors.New("drvConn does not implement extended API")
    93  					}
    94  
    95  					isPk, err = q.IsPrimaryKey(ctx, "episodes", pk)
    96  					return err
    97  				})
    98  
    99  				if err != nil {
   100  					return err
   101  				}
   102  
   103  				require.False(t, isPk)
   104  			}
   105  			return nil
   106  		}, retry.WithIdempotent(true))
   107  
   108  		require.NoError(t, err)
   109  	})
   110  }