github.com/snowflakedb/gosnowflake@v1.9.0/chunk_downloader_test.go (about)

     1  package gosnowflake
     2  
     3  import (
     4  	"context"
     5  	"database/sql/driver"
     6  	"testing"
     7  )
     8  
     9  func TestChunkDownloaderDoesNotStartWhenArrowParsingCausesError(t *testing.T) {
    10  	tcs := []string{
    11  		"invalid base64",
    12  		"aW52YWxpZCBhcnJvdw==", // valid base64, but invalid arrow
    13  	}
    14  	for _, tc := range tcs {
    15  		t.Run(tc, func(t *testing.T) {
    16  			scd := snowflakeChunkDownloader{
    17  				ctx:               context.Background(),
    18  				QueryResultFormat: "arrow",
    19  				RowSet: rowSetType{
    20  					RowSetBase64: tc,
    21  				},
    22  			}
    23  
    24  			err := scd.start()
    25  
    26  			assertNotNilF(t, err)
    27  		})
    28  	}
    29  }
    30  
    31  func TestWithArrowBatchesWhenQueryReturnsNoRowsWhenUsingNativeGoSQLInterface(t *testing.T) {
    32  	runDBTest(t, func(dbt *DBTest) {
    33  		var rows driver.Rows
    34  		var err error
    35  		err = dbt.conn.Raw(func(x interface{}) error {
    36  			rows, err = x.(driver.QueryerContext).QueryContext(WithArrowBatches(context.Background()), "SELECT 1 WHERE 0 = 1", nil)
    37  			return err
    38  		})
    39  		assertNilF(t, err)
    40  		rows.Close()
    41  	})
    42  }
    43  
    44  func TestWithArrowBatchesWhenQueryReturnsRowsAndReadingRows(t *testing.T) {
    45  	runDBTest(t, func(dbt *DBTest) {
    46  		rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
    47  		defer rows.Close()
    48  		assertFalseF(t, rows.Next())
    49  	})
    50  }
    51  
    52  func TestWithArrowBatchesWhenQueryReturnsNoRowsAndReadingRows(t *testing.T) {
    53  	runDBTest(t, func(dbt *DBTest) {
    54  		rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1 WHERE 1 = 0")
    55  		defer rows.Close()
    56  		assertFalseF(t, rows.Next())
    57  	})
    58  }
    59  
    60  func TestWithArrowBatchesWhenQueryReturnsNoRowsAndReadingArrowBatches(t *testing.T) {
    61  	runDBTest(t, func(dbt *DBTest) {
    62  		var rows driver.Rows
    63  		var err error
    64  		err = dbt.conn.Raw(func(x any) error {
    65  			rows, err = x.(driver.QueryerContext).QueryContext(WithArrowBatches(context.Background()), "SELECT 1 WHERE 1 = 0", nil)
    66  			return err
    67  		})
    68  		assertNilF(t, err)
    69  		defer rows.Close()
    70  		batches, err := rows.(SnowflakeRows).GetArrowBatches()
    71  		assertNilF(t, err)
    72  		assertEmptyE(t, batches)
    73  	})
    74  }
    75  
    76  func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormatUsingNativeGoSQLInterface(t *testing.T) {
    77  	for _, tc := range []struct {
    78  		useJSON bool
    79  		desc    string
    80  	}{
    81  		{
    82  			useJSON: true,
    83  			desc:    "json",
    84  		},
    85  		{
    86  			useJSON: false,
    87  			desc:    "arrow",
    88  		},
    89  	} {
    90  		t.Run(tc.desc, func(t *testing.T) {
    91  			runDBTest(t, func(dbt *DBTest) {
    92  				if tc.useJSON {
    93  					dbt.mustExec(forceJSON)
    94  				}
    95  				var rows driver.Rows
    96  				var err error
    97  				err = dbt.conn.Raw(func(x interface{}) error {
    98  					rows, err = x.(driver.QueryerContext).QueryContext(WithArrowBatches(context.Background()), "SELECT 1", nil)
    99  					return err
   100  				})
   101  				assertNilF(t, err)
   102  				defer rows.Close()
   103  				values := make([]driver.Value, 1)
   104  				rows.Next(values)
   105  				assertEqualE(t, values[0], nil)
   106  			})
   107  		})
   108  	}
   109  }
   110  
   111  func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormat(t *testing.T) {
   112  	for _, tc := range []struct {
   113  		useJSON bool
   114  		desc    string
   115  	}{
   116  		{
   117  			useJSON: true,
   118  			desc:    "json",
   119  		},
   120  		{
   121  			useJSON: false,
   122  			desc:    "arrow",
   123  		},
   124  	} {
   125  		t.Run(tc.desc, func(t *testing.T) {
   126  			runDBTest(t, func(dbt *DBTest) {
   127  				if tc.useJSON {
   128  					dbt.mustExec(forceJSON)
   129  				}
   130  				rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
   131  				defer rows.Close()
   132  				assertFalseF(t, rows.Next())
   133  			})
   134  		})
   135  	}
   136  }