Fix for result.AffectedRows (#79)

There were a few cases where we were not returning the correct
AffectedRows in the result. Silly mistake.

Author: rcypher-databricks

connection.go

@@ -208,18 +208,18 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 			// good
 			case cli_service.TOperationState_FINISHED_STATE:
 				// return handle to fetch results later
-				return exStmtResp, opStatus, nil
+				return exStmtResp, statusResp, nil
 			// bad
 			case cli_service.TOperationState_CANCELED_STATE,
 				cli_service.TOperationState_CLOSED_STATE,
 				cli_service.TOperationState_ERROR_STATE,
 				cli_service.TOperationState_TIMEDOUT_STATE:
 				logBadQueryState(log, statusResp)
-				return exStmtResp, opStatus, errors.New(statusResp.GetDisplayMessage())
+				return exStmtResp, statusResp, errors.New(statusResp.GetDisplayMessage())
 				// live states
 			default:
 				logBadQueryState(log, statusResp)
-				return exStmtResp, opStatus, errors.New("invalid operation state. This should not have happened")
+				return exStmtResp, statusResp, errors.New("invalid operation state. This should not have happened")
 			}
 		// weird states
 		default:

connection_test.go

@@ -731,11 +731,13 @@ func TestConn_runQuery(t *testing.T) {
 			}
 			return executeStatementResp, nil
 		}
+		var numModRows int64 = 2
 
 		getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) {
 			getOperationStatusCount++
 			getOperationStatusResp := &cli_service.TGetOperationStatusResp{
-				OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE),
+				OperationState:  cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE),
+				NumModifiedRows: &numModRows,
 			}
 			return getOperationStatusResp, nil
 		}
@@ -756,6 +758,7 @@ func TestConn_runQuery(t *testing.T) {
 		assert.Equal(t, 1, getOperationStatusCount)
 		assert.NotNil(t, exStmtResp)
 		assert.NotNil(t, opStatusResp)
+		assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows)
 	})
 
 	t.Run("runQuery should return resp and error when query is canceled", func(t *testing.T) {
@@ -775,11 +778,13 @@ func TestConn_runQuery(t *testing.T) {
 			}
 			return executeStatementResp, nil
 		}
+		var numModRows int64 = 3
 
 		getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) {
 			getOperationStatusCount++
 			getOperationStatusResp := &cli_service.TGetOperationStatusResp{
-				OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE),
+				OperationState:  cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE),
+				NumModifiedRows: &numModRows,
 			}
 			return getOperationStatusResp, nil
 		}
@@ -800,6 +805,7 @@ func TestConn_runQuery(t *testing.T) {
 		assert.Equal(t, 1, getOperationStatusCount)
 		assert.NotNil(t, exStmtResp)
 		assert.NotNil(t, opStatusResp)
+		assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows)
 	})
 
 	t.Run("runQuery should return resp when query is finished with DirectResults", func(t *testing.T) {
@@ -931,11 +937,12 @@ func TestConn_runQuery(t *testing.T) {
 			}
 			return executeStatementResp, nil
 		}
-
+		var numModRows int64 = 3
 		getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) {
 			getOperationStatusCount++
 			getOperationStatusResp := &cli_service.TGetOperationStatusResp{
-				OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE),
+				OperationState:  cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE),
+				NumModifiedRows: &numModRows,
 			}
 			return getOperationStatusResp, nil
 		}
@@ -953,6 +960,7 @@ func TestConn_runQuery(t *testing.T) {
 
 		assert.NoError(t, err)
 		assert.Equal(t, 1, executeStatementCount)
+		assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows)
 		assert.Equal(t, 1, getOperationStatusCount)
 		assert.NotNil(t, exStmtResp)
 		assert.NotNil(t, opStatusResp)