Close optimization for Rows

When a query returns direct results it is possible that all the query result rows are contained in the direct response.  In this case the server can automatically close the operation.  This is indicated by a non-null CloseOperation member in the TSparkDirectResults structure.
In this case calls to Rows.Close don't require any interaction with the server as the operation is already closed.
- added a NewRows function and updated connection to use it when creating Rows
- added logic in NewRows to check for an already closed operation
- updated Rows.Close so we don't make a server round trip when not necessary

Signed-off-by: Raymond Cypher <[email protected]>

Author: rcypher-databricks

connection.go

@@ -140,22 +140,9 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 	// hold on to the operation handle
 	opHandle := exStmtResp.OperationHandle
 
-	rows := rows{
-		connId:        c.id,
-		correlationId: corrId,
-		client:        c.client,
-		opHandle:      opHandle,
-		pageSize:      int64(c.cfg.MaxRows),
-		location:      c.cfg.Location,
-	}
-
-	if exStmtResp.DirectResults != nil {
-		// return results
-		rows.fetchResults = exStmtResp.DirectResults.ResultSet
-		rows.fetchResultsMetadata = exStmtResp.DirectResults.ResultSetMetadata
+	rows := NewRows(c.id, corrId, c.client, opHandle, int64(c.cfg.MaxRows), c.cfg.Location, exStmtResp.DirectResults)
 
-	}
-	return &rows, nil
+	return rows, nil
 
 }
 

rows.go

@@ -28,6 +28,7 @@ type rows struct {
 	fetchResultsMetadata *cli_service.TGetResultSetMetadataResp
 	nextRowIndex         int64
 	nextRowNumber        int64
+	closed               bool
 }
 
 var _ driver.Rows = (*rows)(nil)
@@ -41,6 +42,27 @@ var errRowsNoSchemaAvailable = "no schema in result set metadata response"
 var errRowsNoClient = "instance of Rows missing client"
 var errRowsNilRows = "nil Rows instance"
 
+func NewRows(connID string, corrId string, client cli_service.TCLIService, opHandle *cli_service.TOperationHandle, pageSize int64, location *time.Location, directResults *cli_service.TSparkDirectResults) driver.Rows {
+	r := &rows{
+		connId:        connID,
+		correlationId: corrId,
+		client:        client,
+		opHandle:      opHandle,
+		pageSize:      pageSize,
+		location:      location,
+	}
+
+	if directResults != nil {
+		r.fetchResults = directResults.ResultSet
+		r.fetchResultsMetadata = directResults.ResultSetMetadata
+		if directResults.CloseOperation != nil {
+			r.closed = true
+		}
+	}
+
+	return r
+}
+
 // 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
@@ -72,20 +94,23 @@ func (r *rows) Columns() []string {
 
 // Close closes the rows iterator.
 func (r *rows) Close() error {
-	err := isValidRows(r)
-	if err != nil {
-		return err
-	}
+	if !r.closed {
+		err := isValidRows(r)
+		if err != nil {
+			return err
+		}
 
-	req := cli_service.TCloseOperationReq{
-		OperationHandle: r.opHandle,
-	}
-	ctx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), r.connId), r.correlationId)
+		req := cli_service.TCloseOperationReq{
+			OperationHandle: r.opHandle,
+		}
+		ctx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), r.connId), r.correlationId)
 
-	_, err1 := r.client.CloseOperation(ctx, &req)
-	if err1 != nil {
-		return err1
+		_, err1 := r.client.CloseOperation(ctx, &req)
+		if err1 != nil {
+			return err1
+		}
 	}
+
 	return nil
 }
 

rows_test.go

@@ -799,6 +799,40 @@ func TestColumnTypeDatabaseTypeName(t *testing.T) {
 	assert.Equal(t, expectedScanTypes, scanTypes)
 }
 
+func TestRowsCloseOptimization(t *testing.T) {
+	t.Parallel()
+
+	var closeCount int
+	client := &client.TestClient{
+		FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (_r *cli_service.TCloseOperationResp, _err error) {
+			closeCount++
+			return nil, nil
+		},
+	}
+
+	rowSet := NewRows("", "", client, &cli_service.TOperationHandle{}, 1, nil, nil)
+
+	// rowSet has no direct results calling Close should result in call to client to close operation
+	err := rowSet.Close()
+	assert.Nil(t, err, "rows.Close should not throw an error")
+	assert.Equal(t, 1, closeCount)
+
+	// rowSet has direct results, but operation was not closed so it should call client to close operation
+	closeCount = 0
+	rowSet = NewRows("", "", client, &cli_service.TOperationHandle{}, 1, nil, &cli_service.TSparkDirectResults{})
+	err = rowSet.Close()
+	assert.Nil(t, err, "rows.Close should not throw an error")
+	assert.Equal(t, 1, closeCount)
+
+	// rowSet has direct results which include a close operation response.  rowSet should be marked as closed
+	// and calling Close should not call into the client.
+	closeCount = 0
+	rowSet = NewRows("", "", client, &cli_service.TOperationHandle{}, 1, nil, &cli_service.TSparkDirectResults{CloseOperation: &cli_service.TCloseOperationResp{}})
+	err = rowSet.Close()
+	assert.Nil(t, err, "rows.Close should not throw an error")
+	assert.Equal(t, 0, closeCount)
+}
+
 type rowTestPagingResult struct {
 	getMetadataCount  int
 	fetchResultsCount int