github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/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  				if err != nil {
    70  					return err
    71  				}
    72  
    73  				require.True(t, isPk)
    74  			}
    75  			return nil
    76  		}, retry.WithIdempotent(true))
    77  
    78  		require.NoError(t, err)
    79  	})
    80  
    81  	t.Run("is-not-primary-key", func(t *testing.T) {
    82  		err := retry.Do(scope.Ctx, db, func(ctx context.Context, cc *sql.Conn) (err error) {
    83  			for _, pk := range []string{"title", "air_date", "views"} {
    84  				isPk := false
    85  				err = cc.Raw(func(drvConn interface{}) (err error) {
    86  					q, ok := drvConn.(interface {
    87  						IsPrimaryKey(context.Context, string, string) (bool, error)
    88  					})
    89  
    90  					if !ok {
    91  						return errors.New("drvConn does not implement extended API")
    92  					}
    93  
    94  					isPk, err = q.IsPrimaryKey(ctx, "episodes", pk)
    95  					return err
    96  				})
    97  				if err != nil {
    98  					return err
    99  				}
   100  
   101  				require.False(t, isPk)
   102  			}
   103  			return nil
   104  		}, retry.WithIdempotent(true))
   105  
   106  		require.NoError(t, err)
   107  	})
   108  }