addressed feedback

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

Author: rcypher-databricks

connection.go

@@ -102,13 +102,14 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 		// we have an operation id so update the logger
 		log = logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID))
 
-		// since we have an operation handle we can close the operation
-		if opStatusResp == nil || opStatusResp.GetOperationState() != cli_service.TOperationState_CLOSED_STATE {
+		// since we have an operation handle we can close the operation if necessary
+		alreadyClosed := exStmtResp.DirectResults != nil && exStmtResp.DirectResults.CloseOperation != nil
+		if !alreadyClosed && (opStatusResp == nil || opStatusResp.GetOperationState() != cli_service.TOperationState_CLOSED_STATE) {
 			_, err1 := c.client.CloseOperation(ctx, &cli_service.TCloseOperationReq{
 				OperationHandle: exStmtResp.OperationHandle,
 			})
 			if err1 != nil {
-				log.Err(err1).Msg("failed to close operation after executing statement")
+				log.Err(err1).Msg("databricks: failed to close operation after executing statement")
 			}
 		}
 	}

connection_test.go

@@ -167,6 +167,18 @@ func TestConn_executeStatement(t *testing.T) {
 			assert.Equal(t, 1, executeStatementCount)
 			assert.Equal(t, opTest.closeOperationCount, closeOperationCount)
 		}
+
+		// if the execute statement response contains direct results with a non-nil CloseOperation member
+		// we shouldn't call close
+		closeOperationCount = 0
+		executeStatementCount = 0
+		executeStatementResp.DirectResults.CloseOperation = &cli_service.TCloseOperationResp{}
+		finished := cli_service.TOperationState_FINISHED_STATE
+		executeStatementResp.DirectResults.OperationStatus.OperationState = &finished
+		_, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{})
+		assert.NoError(t, err)
+		assert.Equal(t, 1, executeStatementCount)
+		assert.Equal(t, 0, closeOperationCount)
 	})
 
 }