added includeResultSetMetadata to protocol_feature_util

Author: shivam2680

internal/rows/rows.go

@@ -148,6 +148,7 @@ func NewRows(
 		connId,
 		correlationId,
 		r.logger(),
+		config.ThriftProtocolVersion,
 	)
 
 	return r, nil

internal/rows/rowscanner/resultPageIterator.go

@@ -8,6 +8,7 @@ import (
 	"github.com/databricks/databricks-sql-go/driverctx"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
+	"github.com/databricks/databricks-sql-go/internal/thrift_protocol"
 	dbsqllog "github.com/databricks/databricks-sql-go/logger"
 )
 
@@ -53,20 +54,22 @@ func NewResultPageIterator(
 	connectionId string,
 	correlationId string,
 	logger *dbsqllog.DBSQLLogger,
+	thriftProtocolVersion cli_service.TProtocolVersion,
 ) ResultPageIterator {
 
 	// delimiter and hasMoreRows are used to set up the point in the paginated
 	// result set that this iterator starts from.
 	return &resultPageIterator{
-		Delimiter:      delimiter,
-		isFinished:     closedOnServer,
-		maxPageSize:    maxPageSize,
-		opHandle:       opHandle,
-		closedOnServer: closedOnServer,
-		client:         client,
-		connectionId:   connectionId,
-		correlationId:  correlationId,
-		logger:         logger,
+		Delimiter:             delimiter,
+		isFinished:            closedOnServer,
+		maxPageSize:           maxPageSize,
+		opHandle:              opHandle,
+		closedOnServer:        closedOnServer,
+		client:                client,
+		connectionId:          connectionId,
+		correlationId:         correlationId,
+		logger:                logger,
+		thriftProtocolVersion: thriftProtocolVersion,
 	}
 }
 
@@ -106,6 +109,8 @@ type resultPageIterator struct {
 
 	// Hold on to errors so they can be returned by Next()
 	err error
+
+	thriftProtocolVersion cli_service.TProtocolVersion
 }
 
 var _ ResultPageIterator = (*resultPageIterator)(nil)
@@ -184,10 +189,13 @@ func (rpf *resultPageIterator) getNextPage() (*cli_service.TFetchResultsResp, er
 
 		var includeResultSetMetadata = true
 		req := cli_service.TFetchResultsReq{
-			OperationHandle:          rpf.opHandle,
-			MaxRows:                  rpf.maxPageSize,
-			Orientation:              directionToSparkDirection(direction),
-			IncludeResultSetMetadata: &includeResultSetMetadata,
+			OperationHandle: rpf.opHandle,
+			MaxRows:         rpf.maxPageSize,
+			Orientation:     directionToSparkDirection(direction),
+			// IncludeResultSetMetadata: &includeResultSetMetadata,
+		}
+		if thrift_protocol.SupportsIncludeResultSetMetadata(rpf.thriftProtocolVersion) {
+			req.IncludeResultSetMetadata = &includeResultSetMetadata
 		}
 
 		fetchResult, err = rpf.client.FetchResults(ctx, &req)

internal/thrift_protocol/protocol_feature_util.go

@@ -44,3 +44,9 @@ func SupportsParameterizedQueries(version cli_service.TProtocolVersion) bool {
 func SupportsMultipleCatalogs(version cli_service.TProtocolVersion) bool {
 	return version >= cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V4
 }
+
+// SupportsIncludeResultSetMetadata checks if the server protocol version supports IncludeResultSetMetadata
+// Supported in SPARK_CLI_SERVICE_PROTOCOL_V5 and above
+func SupportsIncludeResultSetMetadata(version cli_service.TProtocolVersion) bool {
+	return version >= cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V5
+}