Merge #82457

82457: sql: various COPY improvements r=rafiss a=otan

refs  #41608

Release note: None

Co-authored-by: Oliver Tan <[email protected]>

docs/generated/sql/bnf/stmt_block.bnf

@@ -1860,6 +1860,7 @@ copy_options ::=
 	| 'CSV'
 	| 'DELIMITER' string_or_placeholder
 	| 'NULL' string_or_placeholder
+	| 'HEADER'
 
 db_object_name_component ::=
 	name

pkg/sql/copy.go

@@ -67,6 +67,8 @@ type copyMachine struct {
 	textDelim   []byte
 	null        string
 	binaryState binaryState
+	// csvExpectHeader is true if we are expecting a header for the CSV input.
+	csvExpectHeader bool
 	// forceNotNull disables converting values matching the null string to
 	// NULL. The spec says this is only supported for CSV, and also must specify
 	// which columns it applies to.
@@ -121,10 +123,11 @@ func newCopyMachine(
 		conn: conn,
 		// TODO(georgiah): Currently, insertRows depends on Table and Columns,
 		//  but that dependency can be removed by refactoring it.
-		table:   &n.Table,
-		columns: n.Columns,
-		format:  n.Options.CopyFormat,
-		txnOpt:  txnOpt,
+		table:           &n.Table,
+		columns:         n.Columns,
+		format:          n.Options.CopyFormat,
+		txnOpt:          txnOpt,
+		csvExpectHeader: n.Options.Header,
 		// The planner will be prepared before use.
 		p:              planner{execCfg: execCfg, alloc: &tree.DatumAlloc{}},
 		execInsertPlan: execInsertPlan,
@@ -147,9 +150,16 @@ func newCopyMachine(
 		c.delimiter = ','
 	}
 
+	if n.Options.Header && c.format != tree.CopyFormatCSV {
+		return nil, pgerror.Newf(pgcode.FeatureNotSupported, "HEADER only supported with CSV format")
+	}
+
 	if n.Options.Delimiter != nil {
 		if c.format == tree.CopyFormatBinary {
-			return nil, errors.Newf("DELIMITER unsupported in BINARY format")
+			return nil, pgerror.Newf(
+				pgcode.Syntax,
+				"DELIMITER unsupported in BINARY format",
+			)
 		}
 		fn, err := c.p.TypeAsString(ctx, n.Options.Delimiter, "COPY")
 		if err != nil {
@@ -160,13 +170,19 @@ func newCopyMachine(
 			return nil, err
 		}
 		if len(delim) != 1 || !utf8.ValidString(delim) {
-			return nil, errors.Newf("delimiter must be a single-byte character")
+			return nil, pgerror.Newf(
+				pgcode.FeatureNotSupported,
+				"delimiter must be a single-byte character",
+			)
 		}
 		c.delimiter = delim[0]
 	}
 	if n.Options.Null != nil {
 		if c.format == tree.CopyFormatBinary {
-			return nil, errors.Newf("NULL unsupported in BINARY format")
+			return nil, pgerror.Newf(
+				pgcode.Syntax,
+				"NULL unsupported in BINARY format",
+			)
 		}
 		fn, err := c.p.TypeAsString(ctx, n.Options.Null, "COPY")
 		if err != nil {
@@ -471,6 +487,13 @@ func (c *copyMachine) readCSVData(ctx context.Context, final bool) (brk bool, er
 		}
 	}
 
+	// If we are using COPY FROM and expecting a header, PostgreSQL ignores
+	// the header row in all circumstances. Do the same.
+	if c.csvExpectHeader {
+		c.csvExpectHeader = false
+		return false, nil
+	}
+
 	c.csvInput.Write(fullLine)
 	record, err := c.csvReader.Read()
 	// Look for end of data before checking for errors, since a field count

pkg/sql/parser/parse_test.go

@@ -402,7 +402,6 @@ func TestUnimplementedSyntax(t *testing.T) {
 
 		{`COPY t FROM STDIN OIDS`, 41608, `oids`, ``},
 		{`COPY t FROM STDIN FREEZE`, 41608, `freeze`, ``},
-		{`COPY t FROM STDIN HEADER`, 41608, `header`, ``},
 		{`COPY t FROM STDIN ENCODING 'utf-8'`, 41608, `encoding`, ``},
 		{`COPY t FROM STDIN QUOTE 'x'`, 41608, `quote`, ``},
 		{`COPY t FROM STDIN FORCE QUOTE *`, 41608, `quote`, ``},

pkg/sql/parser/sql.y

@@ -3560,9 +3560,9 @@ copy_options:
   {
     return unimplementedWithIssueDetail(sqllex, 41608, "freeze")
   }
-| HEADER error
+| HEADER
   {
-    return unimplementedWithIssueDetail(sqllex, 41608, "header")
+    $$.val = &tree.CopyOptions{Header: true}
   }
 | QUOTE SCONST
   {

pkg/sql/parser/testdata/copy

@@ -79,9 +79,9 @@ COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER '_' destination = '_' -- literals
 COPY _ (_, _, _) FROM STDIN WITH CSV DELIMITER ' ' destination = 'filename' -- identifiers removed
 
 parse
-COPY t (a, b, c) FROM STDIN destination = 'filename' CSV DELIMITER ' ' ESCAPE 'x'
+COPY t (a, b, c) FROM STDIN destination = 'filename' CSV DELIMITER ' ' ESCAPE 'x' HEADER
 ----
-COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER ' ' destination = 'filename' ESCAPE 'x' -- normalized!
-COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER (' ') destination = ('filename') ESCAPE ('x') -- fully parenthesized
-COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER '_' destination = '_' ESCAPE '_' -- literals removed
-COPY _ (_, _, _) FROM STDIN WITH CSV DELIMITER ' ' destination = 'filename' ESCAPE 'x' -- identifiers removed
+COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER ' ' destination = 'filename' ESCAPE 'x' HEADER -- normalized!
+COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER (' ') destination = ('filename') ESCAPE ('x') HEADER -- fully parenthesized
+COPY t (a, b, c) FROM STDIN WITH CSV DELIMITER '_' destination = '_' ESCAPE '_' HEADER -- literals removed
+COPY _ (_, _, _) FROM STDIN WITH CSV DELIMITER ' ' destination = 'filename' ESCAPE 'x' HEADER -- identifiers removed

pkg/sql/pgwire/testdata/pgtest/copy

@@ -212,6 +212,18 @@ ReadyForQuery
 {"Type":"ErrorResponse","Code":"0A000"}
 {"Type":"ReadyForQuery","TxStatus":"I"}
 
+# Header without CSV.
+send
+Query {"String": "COPY t FROM STDIN HEADER"}
+----
+
+until
+ErrorResponse
+ReadyForQuery
+----
+{"Type":"ErrorResponse","Code":"0A000"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
 # Wrong number of columns.
 send
 Query {"String": "COPY t FROM STDIN"}
@@ -322,6 +334,45 @@ ReadyForQuery
 {"Type":"CommandComplete","CommandTag":"SELECT 2"}
 {"Type":"ReadyForQuery","TxStatus":"I"}
 
+send
+Query {"String": "DELETE FROM t"}
+Query {"String": "COPY t FROM STDIN HEADER CSV"}
+CopyData {"Data": "\\.\n"}
+CopyDone
+Query {"String": "COPY t FROM STDIN HEADER CSV"}
+CopyData {"Data": "justonelinewithheader\n"}
+CopyData {"Data": "\\.\n"}
+CopyDone
+Query {"String": "COPY t FROM STDIN HEADER CSV"}
+CopyData {"Data": "ignore,this,entire,line\n"}
+CopyData {"Data": "1,blah\n"}
+CopyData {"Data": "\\.\n"}
+CopyDone
+Query {"String": "SELECT * FROM t ORDER BY i"}
+----
+
+until ignore=RowDescription
+ReadyForQuery
+ReadyForQuery
+ReadyForQuery
+ReadyForQuery
+ReadyForQuery
+----
+{"Type":"CommandComplete","CommandTag":"DELETE 2"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+{"Type":"CopyInResponse","ColumnFormatCodes":[0,0]}
+{"Type":"CommandComplete","CommandTag":"COPY 0"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+{"Type":"CopyInResponse","ColumnFormatCodes":[0,0]}
+{"Type":"CommandComplete","CommandTag":"COPY 0"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+{"Type":"CopyInResponse","ColumnFormatCodes":[0,0]}
+{"Type":"CommandComplete","CommandTag":"COPY 1"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+{"Type":"DataRow","Values":[{"text":"1"},{"text":"blah"}]}
+{"Type":"CommandComplete","CommandTag":"SELECT 1"}
+{"Type":"ReadyForQuery","TxStatus":"I"}
+
 send
 Query {"String": "DELETE FROM t"}
 Query {"String": "COPY t FROM STDIN NULL 'NS'"}
@@ -338,7 +389,7 @@ ReadyForQuery
 ReadyForQuery
 ReadyForQuery
 ----
-{"Type":"CommandComplete","CommandTag":"DELETE 2"}
+{"Type":"CommandComplete","CommandTag":"DELETE 1"}
 {"Type":"ReadyForQuery","TxStatus":"I"}
 {"Type":"CopyInResponse","ColumnFormatCodes":[0,0]}
 {"Type":"CommandComplete","CommandTag":"COPY 3"}

pkg/sql/sem/tree/copy.go

@@ -30,6 +30,7 @@ type CopyOptions struct {
 	Delimiter   Expr
 	Null        Expr
 	Escape      *StrVal
+	Header      bool
 }
 
 var _ NodeFormatter = &CopyOptions{}
@@ -97,6 +98,10 @@ func (o *CopyOptions) Format(ctx *FmtCtx) {
 		ctx.WriteString("ESCAPE ")
 		ctx.FormatNode(o.Escape)
 	}
+	if o.Header {
+		maybeAddSep()
+		ctx.WriteString("HEADER")
+	}
 }
 
 // IsDefault returns true if this struct has default value.
@@ -137,6 +142,9 @@ func (o *CopyOptions) CombineWith(other *CopyOptions) error {
 		}
 		o.Escape = other.Escape
 	}
+	if other.Header {
+		o.Header = true
+	}
 	return nil
 }