clisqlshell: implement COPY ... FROM STDIN for CLI

DEPS.bzl

@@ -5059,10 +5059,10 @@ def go_deps():
         name = "com_github_lib_pq",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/lib/pq",
-        sha256 = "0f50cfc8d4ed4bbb39767aacc04d6b23e1105d2fa50dcb8e4ae204b2c90018f0",
-        strip_prefix = "github.com/lib/[email protected].2",
+        sha256 = "5bca281c55dd8918e49a7e68d562eefb37f2cf17f7d45e1f3bd77e8eae49eb6e",
+        strip_prefix = "github.com/lib/[email protected].6-0.20220412200556-b3b833258663",
         urls = [
-            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/lib/pq/com_github_lib_pq-v1.10.2.zip",
+            "https://storage.googleapis.com/cockroach-godeps/gomod/github.com/lib/pq/com_github_lib_pq-v1.10.6-0.20220412200556-b3b833258663.zip",
         ],
     )
     go_repository(

go.mod

@@ -105,7 +105,7 @@ require (
 	github.com/kr/text v0.2.0
 	github.com/kylelemons/godebug v1.1.0
 	github.com/leanovate/gopter v0.2.5-0.20190402064358-634a59d12406
-	github.com/lib/pq v1.10.2
+	github.com/lib/pq v1.10.6-0.20220412200556-b3b833258663
 	github.com/lib/pq/auth/kerberos v0.0.0-20200720160335-984a6aa1ca46
 	github.com/linkedin/goavro/v2 v2.10.0
 	github.com/lufia/iostat v1.2.1

go.sum

@@ -1482,8 +1482,9 @@ github.com/lib/pq v1.2.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.3.0/go.mod h1:5WUZQaWbwv1U+lTReE5YruASi9Al49XbQIvNi/34Woo=
 github.com/lib/pq v1.8.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lib/pq v1.10.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
-github.com/lib/pq v1.10.2 h1:AqzbZs4ZoCBp+GtejcpCpcxM3zlSMx29dXbUSeVtJb8=
 github.com/lib/pq v1.10.2/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
+github.com/lib/pq v1.10.6-0.20220412200556-b3b833258663 h1:zQ4V0s+y+YrjamtmcRoERUVfsN/4jEb/pVdxrpL6zjU=
+github.com/lib/pq v1.10.6-0.20220412200556-b3b833258663/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lib/pq/auth/kerberos v0.0.0-20200720160335-984a6aa1ca46 h1:q7hY+WNJTcSqJNGwJzXZYL++nWBaoKlKdgZOyY6jxz4=
 github.com/lib/pq/auth/kerberos v0.0.0-20200720160335-984a6aa1ca46/go.mod h1:jydegJvs5JvVcuFD/YAT8JRmRVeOoRhtnGEgRnAoPpE=
 github.com/lightstep/lightstep-tracer-common/golang/gogo v0.0.0-20190605223551-bc2310a04743/go.mod h1:qklhhLq1aX+mtWk9cPHPzaBjWImj5ULL6C7HFJtXQMM=

pkg/cli/clisqlclient/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "api.go",
         "conn.go",
         "context.go",
+        "copy.go",
         "doc.go",
         "init_conn_error.go",
         "make_query.go",

pkg/cli/clisqlclient/copy.go

@@ -0,0 +1,116 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package clisqlclient
+
+import (
+	"context"
+	"database/sql/driver"
+	"io"
+	"reflect"
+	"strings"
+
+	"github.com/cockroachdb/errors"
+)
+
+type copyFromer interface {
+	CopyData(ctx context.Context, line string) (r driver.Result, err error)
+	Exec(v []driver.Value) (r driver.Result, err error)
+	Close() error
+}
+
+// CopyFromState represents an in progress COPY FROM.
+type CopyFromState struct {
+	driver.Tx
+	copyFromer
+}
+
+// BeginCopyFrom starts a COPY FROM query.
+func BeginCopyFrom(ctx context.Context, conn Conn, query string) (*CopyFromState, error) {
+	txn, err := conn.(*sqlConn).conn.(driver.ConnBeginTx).BeginTx(ctx, driver.TxOptions{})
+	if err != nil {
+		return nil, err
+	}
+	stmt, err := txn.(driver.Conn).Prepare(query)
+	if err != nil {
+		return nil, errors.CombineErrors(err, txn.Rollback())
+	}
+	return &CopyFromState{Tx: txn, copyFromer: stmt.(copyFromer)}, nil
+}
+
+// copyFromRows is a mock Rows interface for COPY results.
+type copyFromRows struct {
+	r driver.Result
+}
+
+func (c copyFromRows) Close() error {
+	return nil
+}
+
+func (c copyFromRows) Columns() []string {
+	return nil
+}
+
+func (c copyFromRows) ColumnTypeScanType(index int) reflect.Type {
+	return nil
+}
+
+func (c copyFromRows) ColumnTypeDatabaseTypeName(index int) string {
+	return ""
+}
+
+func (c copyFromRows) ColumnTypeNames() []string {
+	return nil
+}
+
+func (c copyFromRows) Result() driver.Result {
+	return c.r
+}
+
+func (c copyFromRows) Tag() string {
+	return "COPY"
+}
+
+func (c copyFromRows) Next(values []driver.Value) error {
+	return io.EOF
+}
+
+func (c copyFromRows) NextResultSet() (bool, error) {
+	return false, nil
+}
+
+// Cancel cancels a COPY FROM query from completing.
+func (c *CopyFromState) Cancel() error {
+	return errors.CombineErrors(c.copyFromer.Close(), c.Tx.Rollback())
+}
+
+// Commit completes a COPY FROM query by committing lines to the database.
+func (c *CopyFromState) Commit(ctx context.Context, cleanupFunc func(), lines string) QueryFn {
+	return func(ctx context.Context, conn Conn) (Rows, bool, error) {
+		defer cleanupFunc()
+		rows, isMulti, err := func() (Rows, bool, error) {
+			for _, l := range strings.Split(lines, "\n") {
+				_, err := c.copyFromer.CopyData(ctx, l)
+				if err != nil {
+					return nil, false, err
+				}
+			}
+			r, err := c.copyFromer.Exec(nil)
+			if err != nil {
+				return nil, false, err
+			}
+			return copyFromRows{r: r}, false, c.Tx.Commit()
+		}()
+		if err != nil {
+			return rows, isMulti, errors.CombineErrors(err, errors.CombineErrors(c.copyFromer.Close(), c.Tx.Rollback()))
+		}
+		return rows, isMulti, err
+	}
+}

pkg/cli/clisqlclient/make_query.go

@@ -13,10 +13,8 @@ package clisqlclient
 import (
 	"context"
 	"database/sql/driver"
-	"strings"
 
 	"github.com/cockroachdb/cockroach/pkg/sql/scanner"
-	"github.com/cockroachdb/errors"
 )
 
 // QueryFn is the type of functions produced by MakeQuery.
@@ -28,7 +26,6 @@ func MakeQuery(query string, parameters ...interface{}) QueryFn {
 	return func(ctx context.Context, conn Conn) (Rows, bool, error) {
 		isMultiStatementQuery, _ := scanner.HasMultipleStatements(query)
 		rows, err := conn.Query(ctx, query, parameters...)
-		err = handleCopyError(conn.(*sqlConn), err)
 		return rows, isMultiStatementQuery, err
 	}
 }
@@ -51,22 +48,3 @@ func convertArgs(parameters []interface{}) ([]driver.NamedValue, error) {
 	}
 	return dVals, nil
 }
-
-// handleCopyError ensures the user is properly informed when they issue
-// a COPY statement somewhere in their input.
-func handleCopyError(conn *sqlConn, err error) error {
-	if err == nil {
-		return nil
-	}
-
-	if !strings.HasPrefix(err.Error(), "pq: unknown response for simple query: 'G'") {
-		return err
-	}
-
-	// The COPY statement has hosed the connection by putting the
-	// protocol in a state that lib/pq cannot understand any more. Reset
-	// it.
-	_ = conn.Close()
-	conn.reconnecting = true
-	return errors.New("woops! COPY has confused this client! Suggestion: use 'psql' for COPY")
-}

pkg/cli/clisqlexec/run_query.go

@@ -282,6 +282,7 @@ var tagsWithRowsAffected = map[string]struct{}{
 	"DELETE":    {},
 	"MOVE":      {},
 	"DROP USER": {},
+	"COPY":      {},
 	// This one is used with e.g. CREATE TABLE AS (other SELECT
 	// statements have type Rows, not RowsAffected).
 	"SELECT": {},

pkg/cli/clisqlshell/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
         "//pkg/sql/scanner",
         "//pkg/sql/sqlfsm",
         "//pkg/util/envutil",
+        "//pkg/util/errorutil/unimplemented",
         "//pkg/util/syncutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_knz_go_libedit//:go-libedit",