Merge pull request #1256 from ydb-platform/result-set-columns

* Added `query/ResultSet.{Columns,ColumnTypes}` methods for get column names and types from query result set

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Added `query/ResultSet.{Columns,ColumnTypes}` methods for get column names and types from query result set
 * Added experimental `retry.RetryWithResult` helper for retry lambda and return value from lambda
 
 ## v3.70.0

internal/query/result.go

@@ -315,7 +315,7 @@ func exactlyOneResultSetFromResult(ctx context.Context, r query.Result) (rs quer
 		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return NewMaterializedResultSet(rows), nil
+	return NewMaterializedResultSet(rs.Columns(), rs.ColumnTypes(), rows), nil
 }
 
 func resultToMaterializedResult(ctx context.Context, r query.Result) (query.Result, error) {
@@ -345,7 +345,7 @@ func resultToMaterializedResult(ctx context.Context, r query.Result) (query.Resu
 			rows = append(rows, row)
 		}
 
-		resultSets = append(resultSets, NewMaterializedResultSet(rows))
+		resultSets = append(resultSets, NewMaterializedResultSet(rs.Columns(), rs.ColumnTypes(), rows))
 	}
 
 	return newMaterializedResult(resultSets), nil

internal/query/result_set.go

@@ -9,6 +9,7 @@ import (
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/types"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 	"github.com/ydb-platform/ydb-go-sdk/v3/query"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
@@ -21,8 +22,10 @@ var (
 
 type (
 	materializedResultSet struct {
-		rows []query.Row
-		idx  int
+		columnNames []string
+		columnTypes []types.Type
+		rows        []query.Row
+		idx         int
 	}
 	resultSet struct {
 		index       int64
@@ -35,6 +38,32 @@ type (
 	}
 )
 
+func (rs *materializedResultSet) Columns() (columnNames []string) {
+	return rs.columnNames
+}
+
+func (rs *materializedResultSet) ColumnTypes() []types.Type {
+	return rs.columnTypes
+}
+
+func (rs *resultSet) ColumnTypes() (columnTypes []types.Type) {
+	columnTypes = make([]types.Type, len(rs.columns))
+	for i := range rs.columns {
+		columnTypes[i] = types.TypeFromYDB(rs.columns[i].GetType())
+	}
+
+	return columnTypes
+}
+
+func (rs *resultSet) Columns() (columnNames []string) {
+	columnNames = make([]string, len(rs.columns))
+	for i := range rs.columns {
+		columnNames[i] = rs.columns[i].GetName()
+	}
+
+	return columnNames
+}
+
 func (rs *materializedResultSet) NextRow(ctx context.Context) (query.Row, error) {
 	if rs.idx == len(rs.rows) {
 		return nil, xerrors.WithStackTrace(io.EOF)
@@ -47,9 +76,15 @@ func (rs *materializedResultSet) NextRow(ctx context.Context) (query.Row, error)
 	return rs.rows[rs.idx], nil
 }
 
-func NewMaterializedResultSet(rows []query.Row) *materializedResultSet {
+func NewMaterializedResultSet(
+	columnNames []string,
+	columnTypes []types.Type,
+	rows []query.Row,
+) *materializedResultSet {
 	return &materializedResultSet{
-		rows: rows,
+		columnNames: columnNames,
+		columnTypes: columnTypes,
+		rows:        rows,
 	}
 }