rows.go

// Copyright 2026 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package embedded

import (
	"database/sql/driver"
	"errors"
	"fmt"
	"io"

	gms "github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/types"
)

// doltMultiRows implements driver.RowsNextResultSet by aggregating a
// set of individual doltRows instances. Each doltRows is a lazy
// producer to get the row iter, and they are run on demand. The
// current result set is in currentRowSet, with the index of its
// producer in currentIdx.
type doltMultiRows struct {
	rowSets       []func() (*doltRows, error)
	currentIdx    int
	currentRowSet *doltRows
}

var _ driver.RowsNextResultSet = (*doltMultiRows)(nil)
var _ driver.RowsColumnTypeDatabaseTypeName = (*doltMultiRows)(nil)

func (d *doltMultiRows) Columns() []string {
	if d.currentRowSet == nil {
		return nil
	}
	return d.currentRowSet.Columns()
}

func (d *doltMultiRows) ColumnTypeDatabaseTypeName(i int) string {
	if d.currentRowSet == nil {
		return ""
	}
	return d.currentRowSet.ColumnTypeDatabaseTypeName(i)
}

// Close implements the driver.Rows interface. When Close is called on
// a doltMultiRows instance, it will close the open-and-unclosed
// doltRows instance that it contains, if there is one. If any errors
// are encountered while closing the row set, that error will be
// returned.
func (d *doltMultiRows) Close() error {
	rowSet := d.currentRowSet
	d.currentRowSet = nil
	d.rowSets = nil
	if rowSet != nil {
		return rowSet.Close()
	}
	return nil
}

// Return the next row from the current result set.
func (d *doltMultiRows) Next(dest []driver.Value) error {
	if d.currentRowSet == nil {
		return io.EOF
	}
	return d.currentRowSet.Next(dest)
}

// Called by the database/sql implementation at the end of a result
// set. Simply returns |true| if there are more queries we need to
// run. It is OK if this returns |true| but |NextResultSet| ends up
// returning |io.EOF| after we run all the remaining queries.
func (d *doltMultiRows) HasNextResultSet() bool {
	return (d.currentIdx + 1) < len(d.rowSets)
}

// Called anytime we are going on to the next result set, including at
// the end of an existing result set or in the middle of one at the
// request of the caller.
//
// This is responsible for closing the existing result set and for
// running the remaining queries until we get to an iterator that has
// a result set.
//
// Returns io.EOF if we have run all the queries and there are no
// result sets to iterate.
func (d *doltMultiRows) NextResultSet() error {
	if d.currentRowSet == nil {
		return io.EOF
	}
	err := d.currentRowSet.Close()
	if err != nil {
		return err
	}
	d.currentRowSet = nil
	for d.currentIdx = d.currentIdx + 1; d.currentIdx < len(d.rowSets); d.currentIdx++ {
		iter, err := d.rowSets[d.currentIdx]()
		if err != nil {
			return err
		} else if iter.isQueryResultSet() {
			d.currentRowSet = iter
			return nil
		} else {
			err = iter.Close()
			if err != nil {
				return err
			}
		}
	}
	return io.EOF
}

type doltRows struct {
	sch     gms.Schema
	rowIter gms.RowIter
	conn    *DoltConn
	gmsCtx  *gms.Context // per-query context; used for Next and Close calls on the iter.

	columns []string

	// True once rowIter has seen io.EOF.
	drained bool
}

var _ driver.Rows = (*doltRows)(nil)
var _ driver.RowsColumnTypeDatabaseTypeName = (*doltRows)(nil)

// ColumnTypeDatabaseTypeName returns the database system type name for column i.
func (rows *doltRows) ColumnTypeDatabaseTypeName(i int) string {
	if i >= 0 && i < len(rows.sch) {
		return rows.sch[i].Type.String()
	}
	return ""
}

// Columns returns the names of the columns. The number of columns of the result is inferred from the length of the
// slice. If a particular column name isn't known, an empty string should be returned for that entry.
func (rows *doltRows) Columns() []string {
	if rows.columns == nil {
		rows.columns = make([]string, len(rows.sch))
		for i, col := range rows.sch {
			rows.columns[i] = col.Name
		}
	}

	return rows.columns
}

// isQueryResultSet indicates if this result set was generated by a statement that doesn't produce a result set. For
// example, an INSERT or DML statement doesn't return a result set. This is used to skip over this doltRows
// instance if it does not have a result set when calling NextResultSet() on a doltMultiRows instance.
func (rows *doltRows) isQueryResultSet() bool {
	return isQueryResultSet(rows.sch)
}

// Close closes the rows iterator.
func (rows *doltRows) Close() error {
	if !rows.isQueryResultSet() {
		// DML and DDL want Next() to be called on their iters in order to run.
		err := rows.drain()
		if err != nil {
			_ = rows.rowIter.Close(rows.gmsCtx)
			return err
		}
	}
	err := translateError(rows.rowIter.Close(rows.gmsCtx))
	return err
}

// drain consumes all remaining rows in the iterator without converting them.
// This is used by NextResultSet to mirror MySQL client/server behavior, where
// the client must consume all pending rows before advancing to the next result set.
func (rows *doltRows) drain() error {
	if !rows.drained {
		rows.drained = true
		for {
			_, err := rows.rowIter.Next(rows.gmsCtx)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return translateError(err)
			}
		}
	}
	return nil
}

// Next is called to populate the next row of data into the provided slice. The provided slice will be the same size as
// the Columns() are wide. Next returns io.EOF when there are no more rows.
func (rows *doltRows) Next(dest []driver.Value) error {
	nextRow, err := rows.rowIter.Next(rows.gmsCtx)
	if err != nil {
		if err == io.EOF {
			rows.drained = true
			return io.EOF
		}
		return translateError(err)
	}

	if len(dest) != len(nextRow) {
		return errors.New("mismatch between expected column count and actual column count")
	}

	for i := range nextRow {
		if v, ok := nextRow[i].(driver.Valuer); ok {
			dest[i], err = v.Value()

			if err != nil {
				return fmt.Errorf("error processing column %d: %w", i, err)
			}
		} else if geomValue, ok := nextRow[i].(types.GeometryValue); ok {
			dest[i] = geomValue.Serialize()
		} else if enumType, ok := rows.sch[i].Type.(gms.EnumType); ok {
			if v, _, err := enumType.Convert(rows.gmsCtx, nextRow[i]); err != nil {
				return fmt.Errorf("could not convert to expected enum type for column %d: %w", i, err)
			} else if enumStr, ok := enumType.At(int(v.(uint16))); !ok {
				return fmt.Errorf("not a valid enum index for column %d: %v", i, v)
			} else {
				dest[i] = enumStr
			}
		} else if setType, ok := rows.sch[i].Type.(gms.SetType); ok {
			if v, _, err := setType.Convert(rows.gmsCtx, nextRow[i]); err != nil {
				return fmt.Errorf("could not convert to expected set type for column %d: %w", i, err)
			} else if setStr, err := setType.BitsToString(v.(uint64)); err != nil {
				return fmt.Errorf("could not convert value to set string for column %d: %w", i, err)
			} else {
				dest[i] = setStr
			}
		} else {
			dest[i] = nextRow[i]
		}
	}

	return nil
}