github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/internal/xsql/rows.go (about)

     1  package xsql
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"database/sql/driver"
     7  	"fmt"
     8  	"io"
     9  	"strings"
    10  	"sync"
    11  
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/scanner"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/badconn"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/table/result"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/table/result/indexed"
    18  )
    19  
    20  var (
    21  	_ driver.Rows                           = &rows{}
    22  	_ driver.RowsNextResultSet              = &rows{}
    23  	_ driver.RowsColumnTypeDatabaseTypeName = &rows{}
    24  	_ driver.RowsColumnTypeNullable         = &rows{}
    25  	_ driver.Rows                           = &single{}
    26  
    27  	_ scanner.Scanner = &valuer{}
    28  
    29  	ignoreColumnPrefixName = "__discard_column_"
    30  )
    31  
    32  type rows struct {
    33  	conn   *conn
    34  	result result.BaseResult
    35  
    36  	// nextSet once need for get first result set as default.
    37  	// Iterate over many result sets must be with rows.NextResultSet()
    38  	nextSet sync.Once
    39  }
    40  
    41  func (r *rows) LastInsertId() (int64, error) { return 0, ErrUnsupported }
    42  func (r *rows) RowsAffected() (int64, error) { return 0, ErrUnsupported }
    43  
    44  func (r *rows) Columns() []string {
    45  	r.nextSet.Do(func() {
    46  		r.result.NextResultSet(context.Background())
    47  	})
    48  	cs := make([]string, 0, r.result.CurrentResultSet().ColumnCount())
    49  	r.result.CurrentResultSet().Columns(func(m options.Column) {
    50  		if !strings.HasPrefix(m.Name, ignoreColumnPrefixName) {
    51  			cs = append(cs, m.Name)
    52  		}
    53  	})
    54  
    55  	return cs
    56  }
    57  
    58  // TODO: Need to store column types to internal rows cache.
    59  //
    60  //nolint:godox
    61  func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
    62  	r.nextSet.Do(func() {
    63  		r.result.NextResultSet(context.Background())
    64  	})
    65  
    66  	var i int
    67  	yqlTypes := make([]string, r.result.CurrentResultSet().ColumnCount())
    68  	r.result.CurrentResultSet().Columns(func(m options.Column) {
    69  		yqlTypes[i] = m.Type.Yql()
    70  		i++
    71  	})
    72  
    73  	return yqlTypes[index]
    74  }
    75  
    76  // TODO: Need to store column nullables to internal rows cache.
    77  //
    78  //nolint:godox
    79  func (r *rows) ColumnTypeNullable(index int) (nullable, ok bool) {
    80  	r.nextSet.Do(func() {
    81  		r.result.NextResultSet(context.Background())
    82  	})
    83  
    84  	var i int
    85  	nullables := make([]bool, r.result.CurrentResultSet().ColumnCount())
    86  	r.result.CurrentResultSet().Columns(func(m options.Column) {
    87  		_, nullables[i] = m.Type.(interface {
    88  			IsOptional()
    89  		})
    90  		i++
    91  	})
    92  
    93  	return nullables[index], true
    94  }
    95  
    96  func (r *rows) NextResultSet() (finalErr error) {
    97  	r.nextSet.Do(func() {})
    98  	err := r.result.NextResultSetErr(context.Background())
    99  	if err != nil {
   100  		return badconn.Map(xerrors.WithStackTrace(err))
   101  	}
   102  
   103  	return nil
   104  }
   105  
   106  func (r *rows) HasNextResultSet() bool {
   107  	return r.result.HasNextResultSet()
   108  }
   109  
   110  func (r *rows) Next(dst []driver.Value) error {
   111  	var err error
   112  	r.nextSet.Do(func() {
   113  		err = r.result.NextResultSetErr(context.Background())
   114  	})
   115  	if err != nil {
   116  		return badconn.Map(xerrors.WithStackTrace(err))
   117  	}
   118  	if err = r.result.Err(); err != nil {
   119  		return badconn.Map(xerrors.WithStackTrace(err))
   120  	}
   121  	if !r.result.NextRow() {
   122  		return io.EOF
   123  	}
   124  	values := make([]indexed.RequiredOrOptional, len(dst))
   125  	for i := range dst {
   126  		values[i] = &valuer{}
   127  	}
   128  	if err = r.result.Scan(values...); err != nil {
   129  		return badconn.Map(xerrors.WithStackTrace(err))
   130  	}
   131  	for i := range values {
   132  		val, ok := values[i].(*valuer)
   133  		if !ok {
   134  			panic(fmt.Sprintf("unsupported type conversion from %T to *valuer", val))
   135  		}
   136  
   137  		dst[i] = val.Value()
   138  	}
   139  	if err = r.result.Err(); err != nil {
   140  		return badconn.Map(xerrors.WithStackTrace(err))
   141  	}
   142  
   143  	return nil
   144  }
   145  
   146  func (r *rows) Close() error {
   147  	return r.result.Close()
   148  }
   149  
   150  type single struct {
   151  	values  []sql.NamedArg
   152  	readAll bool
   153  }
   154  
   155  func (r *single) Columns() (columns []string) {
   156  	for i := range r.values {
   157  		columns = append(columns, r.values[i].Name)
   158  	}
   159  
   160  	return columns
   161  }
   162  
   163  func (r *single) Close() error {
   164  	return nil
   165  }
   166  
   167  func (r *single) Next(dst []driver.Value) error {
   168  	if r.values == nil || r.readAll {
   169  		return io.EOF
   170  	}
   171  	for i := range r.values {
   172  		dst[i] = r.values[i].Value
   173  	}
   174  	r.readAll = true
   175  
   176  	return nil
   177  }