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

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