Merge pull request #985 from philipglazman/copy-return-command-tag

Return rows copied for COPY command

Author: madelynnblue

copy.go

@@ -49,6 +49,7 @@ type copyin struct {
 	buffer  []byte
 	rowData chan []byte
 	done    chan bool
+	driver.Result
 
 	closed bool
 
@@ -151,6 +152,8 @@ func (ci *copyin) resploop() {
 		switch t {
 		case 'C':
 			// complete
+			res, _ := ci.cn.parseComplete(r.string())
+			ci.setResult(res)
 		case 'N':
 			if n := ci.cn.noticeHandler; n != nil {
 				n(parseError(&r))
@@ -201,6 +204,22 @@ func (ci *copyin) setError(err error) {
 	ci.Unlock()
 }
 
+func (ci *copyin) setResult(result driver.Result) {
+	ci.Lock()
+	ci.Result = result
+	ci.Unlock()
+}
+
+func (ci *copyin) getResult() driver.Result {
+	ci.Lock()
+	result := ci.Result
+	if result == nil {
+		return driver.RowsAffected(0)
+	}
+	ci.Unlock()
+	return result
+}
+
 func (ci *copyin) NumInput() int {
 	return -1
 }
@@ -231,7 +250,11 @@ func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
 	}
 
 	if len(v) == 0 {
-		return driver.RowsAffected(0), ci.Close()
+		if err := ci.Close(); err != nil {
+			return driver.RowsAffected(0), err
+		}
+
+		return ci.getResult(), nil
 	}
 
 	numValues := len(v)

copy_test.go

@@ -73,11 +73,20 @@ func TestCopyInMultipleValues(t *testing.T) {
 		}
 	}
 
-	_, err = stmt.Exec()
+	result, err := stmt.Exec()
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if rowsAffected != 500 {
+		t.Fatalf("expected 500 rows affected, not %d", rowsAffected)
+	}
+
 	err = stmt.Close()
 	if err != nil {
 		t.Fatal(err)