Bump version to 5.11.0

vendor/github.com/lib/pq/copy.go

@@ -1,6 +1,8 @@
 package pq
 
 import (
+	"bytes"
+	"context"
 	"database/sql/driver"
 	"encoding/binary"
 	"errors"
@@ -19,29 +21,35 @@ var (
 // CopyIn creates a COPY FROM statement which can be prepared with
 // Tx.Prepare().  The target table should be visible in search_path.
 func CopyIn(table string, columns ...string) string {
-	stmt := "COPY " + QuoteIdentifier(table) + " ("
+	buffer := bytes.NewBufferString("COPY ")
+	BufferQuoteIdentifier(table, buffer)
+	buffer.WriteString(" (")
+	makeStmt(buffer, columns...)
+	return buffer.String()
+}
+
+// MakeStmt makes the stmt string for CopyIn and CopyInSchema.
+func makeStmt(buffer *bytes.Buffer, columns ...string) {
+	//s := bytes.NewBufferString()
 	for i, col := range columns {
 		if i != 0 {
-			stmt += ", "
+			buffer.WriteString(", ")
 		}
-		stmt += QuoteIdentifier(col)
+		BufferQuoteIdentifier(col, buffer)
 	}
-	stmt += ") FROM STDIN"
-	return stmt
+	buffer.WriteString(") FROM STDIN")
 }
 
 // CopyInSchema creates a COPY FROM statement which can be prepared with
 // Tx.Prepare().
 func CopyInSchema(schema, table string, columns ...string) string {
-	stmt := "COPY " + QuoteIdentifier(schema) + "." + QuoteIdentifier(table) + " ("
-	for i, col := range columns {
-		if i != 0 {
-			stmt += ", "
-		}
-		stmt += QuoteIdentifier(col)
-	}
-	stmt += ") FROM STDIN"
-	return stmt
+	buffer := bytes.NewBufferString("COPY ")
+	BufferQuoteIdentifier(schema, buffer)
+	buffer.WriteRune('.')
+	BufferQuoteIdentifier(table, buffer)
+	buffer.WriteString(" (")
+	makeStmt(buffer, columns...)
+	return buffer.String()
 }
 
 type copyin struct {
@@ -49,12 +57,14 @@ type copyin struct {
 	buffer  []byte
 	rowData chan []byte
 	done    chan bool
-	driver.Result
 
 	closed bool
 
-	sync.Mutex // guards err
-	err        error
+	mu struct {
+		sync.Mutex
+		err error
+		driver.Result
+	}
 }
 
 const ciBufferSize = 64 * 1024
@@ -98,13 +108,13 @@ awaitCopyInResponse:
 			err = parseError(r)
 		case 'Z':
 			if err == nil {
-				ci.setBad()
+				ci.setBad(driver.ErrBadConn)
 				errorf("unexpected ReadyForQuery in response to COPY")
 			}
 			cn.processReadyForQuery(r)
 			return nil, err
 		default:
-			ci.setBad()
+			ci.setBad(driver.ErrBadConn)
 			errorf("unknown response for copy query: %q", t)
 		}
 	}
@@ -123,7 +133,7 @@ awaitCopyInResponse:
 			cn.processReadyForQuery(r)
 			return nil, err
 		default:
-			ci.setBad()
+			ci.setBad(driver.ErrBadConn)
 			errorf("unknown response for CopyFail: %q", t)
 		}
 	}
@@ -144,7 +154,7 @@ func (ci *copyin) resploop() {
 		var r readBuf
 		t, err := ci.cn.recvMessage(&r)
 		if err != nil {
-			ci.setBad()
+			ci.setBad(driver.ErrBadConn)
 			ci.setError(err)
 			ci.done <- true
 			return
@@ -166,7 +176,7 @@ func (ci *copyin) resploop() {
 			err := parseError(&r)
 			ci.setError(err)
 		default:
-			ci.setBad()
+			ci.setBad(driver.ErrBadConn)
 			ci.setError(fmt.Errorf("unknown response during CopyIn: %q", t))
 			ci.done <- true
 			return
@@ -174,46 +184,41 @@ func (ci *copyin) resploop() {
 	}
 }
 
-func (ci *copyin) setBad() {
-	ci.Lock()
-	ci.cn.setBad()
-	ci.Unlock()
+func (ci *copyin) setBad(err error) {
+	ci.cn.err.set(err)
 }
 
-func (ci *copyin) isBad() bool {
-	ci.Lock()
-	b := ci.cn.getBad()
-	ci.Unlock()
-	return b
+func (ci *copyin) getBad() error {
+	return ci.cn.err.get()
 }
 
-func (ci *copyin) isErrorSet() bool {
-	ci.Lock()
-	isSet := (ci.err != nil)
-	ci.Unlock()
-	return isSet
+func (ci *copyin) err() error {
+	ci.mu.Lock()
+	err := ci.mu.err
+	ci.mu.Unlock()
+	return err
 }
 
 // setError() sets ci.err if one has not been set already.  Caller must not be
 // holding ci.Mutex.
 func (ci *copyin) setError(err error) {
-	ci.Lock()
-	if ci.err == nil {
-		ci.err = err
+	ci.mu.Lock()
+	if ci.mu.err == nil {
+		ci.mu.err = err
 	}
-	ci.Unlock()
+	ci.mu.Unlock()
 }
 
 func (ci *copyin) setResult(result driver.Result) {
-	ci.Lock()
-	ci.Result = result
-	ci.Unlock()
+	ci.mu.Lock()
+	ci.mu.Result = result
+	ci.mu.Unlock()
 }
 
 func (ci *copyin) getResult() driver.Result {
-	ci.Lock()
-	result := ci.Result
-	ci.Unlock()
+	ci.mu.Lock()
+	result := ci.mu.Result
+	ci.mu.Unlock()
 	if result == nil {
 		return driver.RowsAffected(0)
 	}
@@ -240,13 +245,13 @@ func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
 		return nil, errCopyInClosed
 	}
 
-	if ci.isBad() {
-		return nil, driver.ErrBadConn
+	if err := ci.getBad(); err != nil {
+		return nil, err
 	}
 	defer ci.cn.errRecover(&err)
 
-	if ci.isErrorSet() {
-		return nil, ci.err
+	if err := ci.err(); err != nil {
+		return nil, err
 	}
 
 	if len(v) == 0 {
@@ -276,14 +281,51 @@ func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
 	return driver.RowsAffected(0), nil
 }
 
+// CopyData inserts a raw string into the COPY stream. The insert is
+// asynchronous and CopyData can return errors from previous CopyData calls to
+// the same COPY stmt.
+//
+// You need to call Exec(nil) to sync the COPY stream and to get any
+// errors from pending data, since Stmt.Close() doesn't return errors
+// to the user.
+func (ci *copyin) CopyData(ctx context.Context, line string) (r driver.Result, err error) {
+	if ci.closed {
+		return nil, errCopyInClosed
+	}
+
+	if finish := ci.cn.watchCancel(ctx); finish != nil {
+		defer finish()
+	}
+
+	if err := ci.getBad(); err != nil {
+		return nil, err
+	}
+	defer ci.cn.errRecover(&err)
+
+	if err := ci.err(); err != nil {
+		return nil, err
+	}
+
+	ci.buffer = append(ci.buffer, []byte(line)...)
+	ci.buffer = append(ci.buffer, '\n')
+
+	if len(ci.buffer) > ciBufferFlushSize {
+		ci.flush(ci.buffer)
+		// reset buffer, keep bytes for message identifier and length
+		ci.buffer = ci.buffer[:5]
+	}
+
+	return driver.RowsAffected(0), nil
+}
+
 func (ci *copyin) Close() (err error) {
 	if ci.closed { // Don't do anything, we're already closed
 		return nil
 	}
 	ci.closed = true
 
-	if ci.isBad() {
-		return driver.ErrBadConn
+	if err := ci.getBad(); err != nil {
+		return err
 	}
 	defer ci.cn.errRecover(&err)
 
@@ -299,8 +341,7 @@ func (ci *copyin) Close() (err error) {
 	<-ci.done
 	ci.cn.inCopy = false
 
-	if ci.isErrorSet() {
-		err = ci.err
+	if err := ci.err(); err != nil {
 		return err
 	}
 	return nil