Close optimization for Rows (#64)

When a query returns direct results it is possible that all the query
result rows are contained therein. 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
- added a unit test for the close optimization logic

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