From b3b833258663afbe55669e5b53e5d62f7e1231bd Mon Sep 17 00:00:00 2001
From: Oliver Tan <otan@cockroachlabs.com>
Date: Wed, 13 Apr 2022 06:05:56 +1000
Subject: [PATCH] expose raw CopyData command (#1077)

---
 copy.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/copy.go b/copy.go
index c072bc3b..a24ea3f5 100644
--- a/copy.go
+++ b/copy.go
@@ -1,6 +1,7 @@
 package pq
 
 import (
+	"context"
 	"database/sql/driver"
 	"encoding/binary"
 	"errors"
@@ -273,6 +274,38 @@ func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
 	return driver.RowsAffected(0), nil
 }
 
+// CopyData executes a raw CopyData command using the PostgreSQL Frontend/Backend
+// protocol. Use Exec(nil) to finish the command.
+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