rows.go

package godatabend

import (
	"context"
	"database/sql/driver"
	"io"
	"reflect"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
)

type resultSchema struct {
	columns []string
	types   []ColumnType
}

type nextRows struct {
	resultSchema
	isClosed int32
	dc       *DatabendConn
	ctx      context.Context
	respData *QueryResponse
}

func waitForData(ctx context.Context, dc *DatabendConn, response *QueryResponse) (*QueryResponse, error) {
	if response.Error != nil {
		return nil, response.Error
	}
	for !response.ReadFinished() && response.bufferedRowCount() == 0 && response.Error == nil {
		nextResponse, err := dc.rest.PollQuery(ctx, response.NextURI)
		if err != nil {
			if errors.Is(err, context.Canceled) {
				// context might be canceled due to timeout or canceled. if it's canceled, we need call
				// the kill url to tell the backend it's killed.
				dc.log("query canceled", response.ID)
				_ = dc.rest.KillQuery(context.Background(), response)
			} else {
				_ = dc.rest.CloseQuery(ctx, response)
			}
			return nil, err
		}
		response = nextResponse
		if response.Error != nil {
			_ = dc.rest.CloseQuery(ctx, response)
			return nil, errors.Errorf("query error: %+v", response.Error)
		}
	}
	return response, nil
}

func parse_schema(fields *[]DataField, opts *ColumnTypeOptions) (*resultSchema, error) {
	if fields == nil {
		return &resultSchema{}, nil
	}

	schema := &resultSchema{
		columns: make([]string, 0, len(*fields)),
		types:   make([]ColumnType, 0, len(*fields)),
	}

	for _, field := range *fields {
		schema.columns = append(schema.columns, field.Name)
		parser, err := NewColumnType(field.Type, opts)
		if err != nil {
			return nil, errors.Wrapf(err, "newTextRows: failed to create a data parser for the type '%s'", field.Type)
		}
		schema.types = append(schema.types, parser)
	}

	return schema, nil
}

func (dc *DatabendConn) newNextRows(ctx context.Context, resp *QueryResponse) (rows *nextRows, err error) {
	if len(resp.Data) != 0 && (resp.Schema == nil || len(*resp.Schema) != len(resp.Data[0])) {
		return nil, errors.New("newNextRows: internal error, data and schema not match")
	}
	if resp.typedRows == nil {
		if err := materializeJSONQueryRows(resp); err != nil {
			return nil, err
		}
	}
	if len(resp.typedRows) != 0 && (resp.Schema == nil || len(*resp.Schema) != len(resp.typedRows[0])) {
		return nil, errors.New("newNextRows: internal error, typed data and schema not match")
	}
	var location *time.Location
	if resp.Settings != nil && resp.Settings.TimeZone != "" {
		location, err = time.LoadLocation(resp.Settings.TimeZone)
		if err != nil {
			return nil, err
		}
	}

	schema, err := parse_schema(resp.Schema, dc.columnTypeOptions(resp.Settings, location))
	if err != nil {
		return nil, err
	}

	rows = &nextRows{
		dc:           dc,
		ctx:          ctx,
		respData:     resp,
		resultSchema: *schema,
	}
	return rows, nil
}

func (r *nextRows) Columns() []string {
	return r.columns
}

// Close will only be called by sql.Rows once.
// But we can doClose internally as soon as EOF.
//
// Not return error for now.
//
// Note it will also be Called by framework when:
//  1. Canceling query/txn Context.
//  2. Next return error other than io.EOF.
func (r *nextRows) Close() error {
	return r.doClose()
}

func (r *nextRows) doClose() error {
	if atomic.CompareAndSwapInt32(&r.isClosed, 0, 1) {
		if r.respData != nil && len(r.respData.FinalURI) != 0 {
			err := r.dc.rest.CloseQuery(r.dc.ctx, r.respData)
			if err != nil {
				return err
			}
			r.respData = nil
		}
		r.dc.cancel = nil
		return nil
	} else {
		// Rows should be safe to close multi times
		return nil
	}
}

func (r *nextRows) Next(dest []driver.Value) error {
	if len(dest) != len(r.columns) {
		return errors.New("query error: Next dest must has same size as the Columns() are wide")
	}

	if atomic.LoadInt32(&r.isClosed) == 1 || r.respData == nil {
		// If user already called Rows.Close(), Rows.Next() will not get here.
		// Get here only because we doClose() internally,
		// only when call Rows.Next() again after it return false.
		return io.EOF
	}
	if r.respData.bufferedRowCount() == 0 {
		var err error
		r.respData, err = waitForData(r.ctx, r.dc, r.respData)
		if err != nil {
			return err
		}
	}

	if r.respData.bufferedRowCount() == 0 {
		_ = r.doClose()
		return io.EOF
	}

	var lineData []*string
	if len(r.respData.Data) > 0 {
		lineData = r.respData.Data[0]
		r.respData.Data = r.respData.Data[1:]
	}
	var typedRow []driver.Value
	if len(r.respData.typedRows) > 0 {
		typedRow = r.respData.typedRows[0]
		r.respData.typedRows = r.respData.typedRows[1:]
	}
	if len(typedRow) != 0 && len(lineData) != 0 {
		return errors.New("query error: internal error, both typed and raw row data are set")
	}
	if len(typedRow) != 0 {
		if len(typedRow) != len(r.columns) {
			return errors.New("query error: internal error, typed data and schema not match")
		}
		copy(dest, typedRow)
		return nil
	}

	if len(lineData) != len(r.columns) {
		return errors.New("query error: internal error, data and schema not match")
	}

	for i, val := range lineData {
		if val == nil {
			dest[i] = nil
			continue
		}
		v, err := r.types[i].Parse(*val)
		if err != nil {
			r.dc.log("fail to parse field", i, ", error: ", err)
			return err
		}
		dest[i] = v
	}
	return nil
}

var _ driver.RowsColumnTypeScanType = (*nextRows)(nil)

func (r *nextRows) ColumnTypeScanType(index int) reflect.Type {
	return r.types[index].ScanType()
}

var _ driver.RowsColumnTypeDatabaseTypeName = (*nextRows)(nil)

func (r *nextRows) ColumnTypeDatabaseTypeName(index int) string {
	return r.types[index].DatabaseTypeName()
}

var _ driver.RowsColumnTypeNullable = (*nextRows)(nil)

func (r *nextRows) ColumnTypeNullable(index int) (bool, bool) {
	return r.types[index].Nullable()
}

var _ driver.RowsColumnTypeLength = (*nextRows)(nil)

func (r *nextRows) ColumnTypeLength(index int) (length int64, ok bool) {
	return r.types[index].Length()
}

var _ driver.RowsColumnTypePrecisionScale = (*nextRows)(nil)

func (r *nextRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) {
	return r.types[index].PrecisionScale()
}