Merge pull request #1083 from dolthub/zachmu/copy-tests

Go tests for the COPY command

core/dataloader/csvreader.go

@@ -295,7 +295,10 @@ func (csvr *csvReader) parseQuotedField(rs *recordState) (kontinue bool, err err
 	const quoteLen = len(`"`)
 	dl := len(csvr.delim)
 	recordStartLine := csvr.numLine
-	fullField := rs.line
+	// full copy needed here because we append rs.line to fullField, and this can result in buffer corruption in
+	// some cases (namely when windows line endings are present)
+	fullField := make([]byte, len(rs.line))
+	copy(fullField, rs.line)
 
 	// Quoted string field
 	rs.line = rs.line[quoteLen:]

core/dataloader/dataloader.go

@@ -29,7 +29,7 @@ import (
 // with the incomplete record.
 type DataLoader interface {
 	// LoadChunk reads the records from |data| and inserts them into the previously configured table. Data records
-	// are not guaranteed to stard and end cleanly on chunk boundaries, so implementations must recognize incomplete
+	// are not guaranteed to start and end cleanly on chunk boundaries, so implementations must recognize incomplete
 	// records and save them to prepend on the next processed chunk.
 	LoadChunk(ctx *sql.Context, data *bufio.Reader) error
 

core/dataloader/tabdataloader.go

@@ -111,9 +111,9 @@ func (tdl *TabularDataLoader) LoadChunk(ctx *sql.Context, data *bufio.Reader) er
 			tdl.partialLine = ""
 		}
 
-		// If we see the end of data marker, then jump out of the loop
+		// If we see the end of data marker, return early
 		if line == "\\." {
-			break
+			return nil
 		}
 
 		// Skip over empty lines

testing/go/copy_test.go

@@ -0,0 +1,151 @@
+// Copyright 2024 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package _go
+
+import (
+	"testing"
+
+	"github.com/dolthub/go-mysql-server/sql"
+)
+
+func TestCopy(t *testing.T) {
+	RunScripts(t, []ScriptTest{
+		{
+			Name: "tab delimited with header",
+			SetUpScript: []string{
+				"CREATE TABLE test (pk int primary key);",
+				"INSERT INTO test VALUES (0), (1);",
+				"CREATE TABLE test_info (id int, info varchar(255), test_pk int, primary key(id), foreign key (test_pk) references test(pk));",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:             "COPY test_info FROM STDIN WITH (HEADER);",
+					CopyFromStdInFile: "tab-load-with-header.sql",
+				},
+				{
+					Query: "SELECT * FROM test_info order by 1;",
+					Expected: []sql.Row{
+						{4, "string for 4", 1},
+						{5, "string for 5", 0},
+						{6, "string for 6", 0},
+					},
+				},
+			},
+		},
+		{
+			Name: "tab delimited with quoted column names",
+			SetUpScript: []string{
+				`CREATE TABLE Regions (
+   "Id" SERIAL UNIQUE NOT NULL,
+   "Code" VARCHAR(4) UNIQUE NOT NULL,
+   "Capital" VARCHAR(10) NOT NULL,
+   "Name" VARCHAR(255) UNIQUE NOT NULL
+);`,
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:             "COPY regions (\"Id\", \"Code\", \"Capital\", \"Name\") FROM stdin;\n",
+					CopyFromStdInFile: "tab-load-with-quoted-column-names.sql",
+				},
+			},
+		},
+		{
+			Name: "basic csv",
+			SetUpScript: []string{
+				"CREATE TABLE tbl1 (pk int primary key, c1 varchar(100), c2 varchar(250));",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:             "COPY tbl1 FROM STDIN (FORMAT CSV)",
+					CopyFromStdInFile: "csv-load-basic-cases.sql",
+				},
+				{
+					Query: "select * from tbl1 where pk = 6 order by pk;",
+					Expected: []sql.Row{
+						{6, `foo
+\\.
+bar`, "baz"},
+					},
+				},
+				{
+					Query: "select * from tbl1 where pk = 9;",
+					Expected: []sql.Row{
+						{9, nil, "''"},
+					},
+				},
+			},
+		},
+		{
+			Name: "csv with header",
+			SetUpScript: []string{
+				"CREATE TABLE tbl1 (pk int primary key, c1 varchar(100), c2 varchar(250));",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:             " COPY tbl1 FROM STDIN (FORMAT CSV, HEADER TRUE);",
+					CopyFromStdInFile: "csv-load-with-header.sql",
+				},
+				{
+					Query: "select * from tbl1 where pk = 6 order by pk;",
+					Expected: []sql.Row{
+						{6, `foo
+\\.
+bar`, "baz"},
+					},
+				},
+			},
+		},
+		{
+			Name: "load multiple chunks",
+			SetUpScript: []string{
+				"CREATE TABLE tbl1 (pk int primary key, c1 varchar(100), c2 varchar(250));",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:             "COPY tbl1 FROM STDIN (FORMAT CSV);",
+					CopyFromStdInFile: "csv-load-multi-chunk.sql",
+				},
+				{
+					Query: "select * from tbl1 where pk = 99 order by pk;",
+					Expected: []sql.Row{
+						{99, "foo", "barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash"},
+					},
+				},
+			},
+		},
+		{
+			Name: "load psv with headers",
+			SetUpScript: []string{
+				"CREATE TABLE test (pk int primary key);",
+				"INSERT INTO test VALUES (0), (1);",
+				"CREATE TABLE test_info (id int, info varchar(255), test_pk int, primary key(id), foreign key (test_pk) references test(pk));",
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query:             "COPY test_info FROM STDIN (FORMAT CSV, HEADER TRUE, DELIMITER '|');",
+					CopyFromStdInFile: "psv-load.sql",
+				},
+				{
+					Query: "SELECT * FROM test_info order by 1;",
+					Expected: []sql.Row{
+						{4, "string for 4", 1},
+						{5, "string for 5", 0},
+						{6, "string for 6", 0},
+					},
+				},
+			},
+		},
+	})
+}

testing/go/framework.go

@@ -15,13 +15,15 @@
 package _go
 
 import (
+	"bufio"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"math"
 	"net"
 	"os"
+	"path/filepath"
 	"strconv"
 	"testing"
 	"time"
@@ -104,6 +106,9 @@ type ScriptTestAssertion struct {
 
 	// Cols is used to check the column names returned from the server.
 	Cols []string
+
+	// CopyFromSTDIN is used to test the COPY FROM STDIN command.
+	CopyFromStdInFile string
 }
 
 // Connection contains the default and current connections.
@@ -182,7 +187,9 @@ func runScript(t *testing.T, ctx context.Context, script ScriptTest, conn *Conne
 				return
 			}
 			// If we're skipping the results check, then we call Execute, as it uses a simplified message model.
-			if assertion.SkipResultsCheck || assertion.ExpectedErr != "" {
+			if assertion.CopyFromStdInFile != "" {
+				copyFromStdin(t, conn.Current, assertion.Query, assertion.CopyFromStdInFile)
+			} else if assertion.SkipResultsCheck || assertion.ExpectedErr != "" {
 				_, err := conn.Exec(ctx, assertion.Query, assertion.BindVars...)
 				if assertion.ExpectedErr != "" {
 					require.Error(t, err)
@@ -220,6 +227,20 @@ func runScript(t *testing.T, ctx context.Context, script ScriptTest, conn *Conne
 	}
 }
 
+func copyFromStdin(t *testing.T, conn *pgx.Conn, query string, filename string) {
+	filePath := filepath.Join("testdata", filename)
+
+	file, err := os.Open(filePath)
+	if err != nil {
+		t.Fatalf("Failed to open file: %v", err)
+	}
+	defer file.Close()
+
+	reader := bufio.NewReader(file)
+	_, err = conn.PgConn().CopyFrom(context.Background(), reader, query)
+	require.NoError(t, err)
+}
+
 // RunScripts runs the given collection of scripts. This normalizes all rows before comparing them.
 func RunScripts(t *testing.T, scripts []ScriptTest) {
 	runScripts(t, scripts, true)

testing/go/testdata/csv-load-basic-cases.sql

@@ -0,0 +1,15 @@
+1,green,
+2,"blue","a
+q
+u
+a"
+3,"brown",
+4,"NULL",NULL
+5,"?",""
+6,"foo
+\\.
+bar","baz"
+7,  ,' '
+8," ",""
+9,,''
+\.

testing/go/testdata/csv-load-multi-chunk.sql

@@ -0,0 +1,102 @@
+0,foo,barbazbashbarbazbashbarbazbashbarbazbash
+1,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+2,foo,barbazbashbarbazbashbarbazbash
+3,foo,barbazbashbarbazbashbarbazbashbarbazbash
+4,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+5,foo,barbazbashbarbazbashbarbazbash
+6,foo,barbazbashbarbazbashbarbazbashbarbazbash
+7,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+8,foo,barbazbashbarbazbashbarbazbash
+9,foo,barbazbashbarbazbashbarbazbashbarbazbash
+10,foo,barbazbashbarbazbashbarbazbashbarbazbash
+11,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+12,foo,barbazbashbarbazbashbarbazbash
+13,foo,barbazbashbarbazbashbarbazbashbarbazbash
+14,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+15,foo,barbazbashbarbazbashbarbazbash
+16,foo,barbazbashbarbazbashbarbazbashbarbazbash
+17,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+18,foo,barbazbashbarbazbashbarbazbash
+19,foo,barbazbashbarbazbashbarbazbashbarbazbash
+20,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+21,foo,barbazbashbarbazbashbarbazbash
+22,foo,barbazbashbarbazbashbarbazbashbarbazbash
+23,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+24,foo,barbazbashbarbazbashbarbazbash
+25,foo,barbazbashbarbazbashbarbazbashbarbazbash
+26,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+27,foo,barbazbashbarbazbashbarbazbash
+28,foo,barbazbashbarbazbashbarbazbashbarbazbash
+29,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+30,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+31,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+32,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+33,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+34,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+35,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+36,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+37,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+38,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+39,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+40,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+41,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+42,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+43,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+44,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+45,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+46,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+47,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+48,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+49,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+50,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+51,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+52,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+53,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+54,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+55,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+56,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+57,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+58,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+59,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+60,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+61,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+62,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+63,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+64,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+65,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+66,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+67,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+68,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+69,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+70,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+71,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+72,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+73,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+74,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+75,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+76,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+77,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+78,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+79,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+80,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+81,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+82,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+83,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+84,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+85,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+86,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+87,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+88,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+89,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+90,foo,"012345678901234567890123456789012345678901234567890123456789012345678901234567890123456789012
+345678901234567"
+91,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+92,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+93,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+94,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+95,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+96,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+97,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+98,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+99,foo,barbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbashbarbazbash
+\.

testing/go/testdata/csv-load-with-header.sql

@@ -0,0 +1,16 @@
+pk,c1,c2
+1,green,
+2,"blue","a
+q
+u
+a"
+3,"brown",
+4,"NULL",NULL
+5,"?",""
+6,"foo
+\\.
+bar","baz"
+7,  ,' '
+8," ",""
+9,,''
+\.

testing/go/testdata/psv-load.sql

@@ -0,0 +1,5 @@
+id|info|test_pk
+4|string for 4|1
+5|string for 5|0
+6|string for 6|0
+\.

testing/go/testdata/tab-load-with-header.sql

@@ -0,0 +1,5 @@
+id	info	test_pk
+4	string for 4	1
+5	string for 5	0
+6	string for 6	0
+\.

testing/go/testdata/tab-load-with-quoted-column-names.sql

@@ -0,0 +1,4 @@
+1	01	97105	Guadeloupe
+2	02	97209	Martinique
+3	03	97302	Guyane
+\.