Skip to content

Commit

Permalink
copy: enhance copyfrom tests with kvtrace feature and more tests
Browse files Browse the repository at this point in the history 
Epic: CRDB-18892
Informs: cockroachdb#91831
Release note: None
  • Loading branch information
@cucaroach
cucaroach committed Mar 14, 2023
1 parent e54b589 commit 89ae338
Show file tree
Hide file tree
Showing 2 changed files with 263 additions and 66 deletions.
179 changes: 113 additions & 66 deletions pkg/sql/copy/copy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,80 +75,127 @@ const lineitemSchema string = `CREATE TABLE lineitem (

const csvData = `%d|155190|7706|1|17|21168.23|0.04|0.02|N|O|1996-03-13|1996-02-12|1996-03-22|DELIVER IN PERSON|TRUCK|egular courts above the`

func TestCopy(t *testing.T) {
func TestDataDriven(t *testing.T) {
defer leaktest.AfterTest(t)()
defer log.Scope(t).Close(t)
ctx := context.Background()

datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) {
s, _, _ := serverutils.StartServer(t, base.TestServerArgs{
Settings: cluster.MakeTestingClusterSettings(),
})
defer s.Stopper().Stop(ctx)
for _, vectorize := range []string{"on", "off"} {
for _, atomic := range []string{"on", "off"} {
for _, fastPath := range []string{"on", "off"} {
datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) {
s, _, _ := serverutils.StartServer(t, base.TestServerArgs{
Settings: cluster.MakeTestingClusterSettings(),
})
defer s.Stopper().Stop(ctx)

url, cleanup := sqlutils.PGUrl(t, s.ServingSQLAddr(), t.Name(), url.User(username.RootUser))
defer cleanup()
var sqlConnCtx clisqlclient.Context
conn := sqlConnCtx.MakeSQLConn(io.Discard, io.Discard, url.String())

err := conn.Exec(ctx, fmt.Sprintf(`SET VECTORIZE='%s'`, vectorize))
require.NoError(t, err)

url, cleanup := sqlutils.PGUrl(t, s.ServingSQLAddr(), t.Name(), url.User(username.RootUser))
defer cleanup()
var sqlConnCtx clisqlclient.Context
conn := sqlConnCtx.MakeSQLConn(io.Discard, io.Discard, url.String())

datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
switch d.Cmd {
case "exec-ddl":
err := conn.Exec(ctx, d.Input)
if err != nil {
require.NoError(t, err, "%s: %s", d.Pos, d.Cmd)
}
return ""
case "copy-from", "copy-from-error":
lines := strings.Split(d.Input, "\n")
stmt := lines[0]
data := strings.Join(lines[1:], "\n")
rows, err := conn.GetDriverConn().CopyFrom(ctx, strings.NewReader(data), stmt)
if d.Cmd == "copy-from" {
require.NoError(t, err, "%s\n%s\n", d.Cmd, d.Input)
require.Equal(t, int(rows), len(lines)-1, "not all rows were inserted")
} else {
require.Error(t, err)
return err.Error()
}
return fmt.Sprintf("%d", rows)
case "copy-to", "copy-to-error":
var buf bytes.Buffer
err := conn.GetDriverConn().CopyTo(ctx, &buf, d.Input)
if d.Cmd == "copy-to" {
err = conn.Exec(ctx, fmt.Sprintf(`SET COPY_FAST_PATH_ENABLED='%s'`, fastPath))
require.NoError(t, err)
} else {
require.Error(t, err)
return expandErrorString(err)
}
return buf.String()
case "query":
rows, err := conn.Query(ctx, d.Input)
require.NoError(t, err)
vals := make([]driver.Value, len(rows.Columns()))
var results string
for {
if err := rows.Next(vals); err == io.EOF {
break
} else if err != nil {
require.NoError(t, err)
}
for i, v := range vals {
if i > 0 {
results += "|"

err = conn.Exec(ctx, fmt.Sprintf(`SET COPY_FROM_ATOMIC_ENABLED='%s'`, atomic))
require.NoError(t, err)

datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
switch d.Cmd {
case "exec-ddl":
err := conn.Exec(ctx, d.Input)
if err != nil {
require.NoError(t, err, "%s: %s", d.Pos, d.Cmd)
}
return ""
case "copy-from", "copy-from-error", "copy-from-kvtrace":
kvtrace := d.Cmd == "copy-from-kvtrace"
lines := strings.Split(d.Input, "\n")
stmt := lines[0]
data := strings.Join(lines[1:], "\n")
if kvtrace {
err := conn.Exec(ctx, "SET TRACING=on,kv")
require.NoError(t, err)
}
rows, err := conn.GetDriverConn().CopyFrom(ctx, strings.NewReader(data), stmt)
if kvtrace {
err := conn.Exec(ctx, "SET TRACING=off")
require.NoError(t, err)
}
switch d.Cmd {
case "copy-from":
require.NoError(t, err, "%s\n%s\n", d.Cmd, d.Input)
require.Equal(t, int(rows), len(lines)-1, "Not all rows were inserted")
return fmt.Sprintf("%d", rows)
case "copy-from-error":
require.Error(t, err, "copy-from-error didn't return and error!")
return err.Error()
case "copy-from-kvtrace":
rows, err := conn.Query(ctx,
`SELECT
regexp_replace(message, '/Table/[0-9]*/', '/Table/<>/')
FROM [SHOW KV TRACE FOR SESSION]
WHERE message LIKE '%Put % -> %'`)
defer func() {
_ = rows.Close()
}()
require.NoError(t, err)
vals := make([]driver.Value, 1)
var results strings.Builder
for err = nil; err == nil; {
err = rows.Next(vals)
if err == io.EOF {
break
}
require.NoError(t, err)
results.WriteString(fmt.Sprintf("%v\n", vals[0]))
}
return results.String()
}
case "copy-to", "copy-to-error":
var buf bytes.Buffer
err := conn.GetDriverConn().CopyTo(ctx, &buf, d.Input)
if d.Cmd == "copy-to" {
require.NoError(t, err)
} else {
require.Error(t, err)
return expandErrorString(err)
}
return buf.String()
case "query":
rows, err := conn.Query(ctx, d.Input)
require.NoError(t, err)
vals := make([]driver.Value, len(rows.Columns()))
var results string
for {
if err := rows.Next(vals); err == io.EOF {
break
} else if err != nil {
require.NoError(t, err)
}
for i, v := range vals {
if i > 0 {
results += "|"
}
results += fmt.Sprintf("%v", v)
}
results += "\n"
}
err = rows.Close()
require.NoError(t, err)
return results
default:
return fmt.Sprintf("unknown command: %s\n", d.Cmd)
}
results += fmt.Sprintf("%v", v)
}
results += "\n"
}
err = rows.Close()
require.NoError(t, err)
return results
default:
return fmt.Sprintf("unknown command: %s\n", d.Cmd)
return ""
})
})
}
})
})
}
}
}

var issueLinkRE = regexp.MustCompile("https://go.crdb.dev/issue-v/([0-9]+)/.*")
Expand Down
150 changes: 150 additions & 0 deletions pkg/sql/copy/testdata/copy_from
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,11 @@ COPY t3 FROM STDIN
----
ERROR: failed to satisfy CHECK constraint (i > 0:::INT8) (SQLSTATE 23514)

copy-from
COPY t3 FROM STDIN
1
----
1

# Foreign key checks happen
exec-ddl
Expand Down Expand Up @@ -510,3 +515,148 @@ COPY tab_child FROM STDIN
'tower' 'former' 'mainly' 'point' 'class' 'idea'
----
ERROR: insert on table "tab_child" violates foreign key constraint "tab_child_col5_col6_fkey" (SQLSTATE 23503)

exec-ddl
CREATE TABLE tabnn (i INT NOT NULL)
----

copy-from-error
COPY tabnn FROM STDIN
\N
----
ERROR: null value in column "i" violates not-null constraint (SQLSTATE 23502)

exec-ddl
CREATE TABLE tinet (i INET CHECK (i < b), b INET)
----

# Put to rest concerns about Datum types on both sides of check constraint expr.
copy-from-error
COPY tinet FROM STDIN
192.168.100.128/25 0.0.0.0/0
----
ERROR: failed to satisfy CHECK constraint (i < b) (SQLSTATE 23514)

# Put to rest concerns about Datum types on both sides of check constraint expr.
copy-from
COPY tinet FROM STDIN
0.0.0.0/0 192.168.100.128/25
----
1

exec-ddl
CREATE TABLE tpartial (i INT PRIMARY KEY, index(i) WHERE i > 0)
----

copy-from-kvtrace
COPY tpartial FROM STDIN
-1
0
1
----
CPut /Table/<>/1/-1/0 -> /TUPLE/
CPut /Table/<>/1/0/0 -> /TUPLE/
CPut /Table/<>/1/1/0 -> /TUPLE/
InitPut /Table/<>/2/1/0 -> /BYTES/

exec-ddl
CREATE TABLE tpartial2 (i INT PRIMARY KEY, b INT, index(b) WHERE i > 0, FAMILY (i), FAMILY (b))
----

# Unfortunately we have to do the inserts 1 row at a time because the row inserter does it
# that way and the vector inserter doesn't.
copy-from-kvtrace
COPY tpartial2 FROM STDIN
-1 -2
----
CPut /Table/<>/1/-1/0 -> /TUPLE/
CPut /Table/<>/1/-1/1/1 -> /INT/-2

copy-from-kvtrace
COPY tpartial2 FROM STDIN
0 0
----
CPut /Table/<>/1/0/0 -> /TUPLE/
CPut /Table/<>/1/0/1/1 -> /INT/0

copy-from-kvtrace
COPY tpartial2 FROM STDIN
1 2
----
CPut /Table/<>/1/1/0 -> /TUPLE/
CPut /Table/<>/1/1/1/1 -> /INT/2
InitPut /Table/<>/2/2/1/0 -> /BYTES/

exec-ddl
CREATE TYPE testenum AS ENUM('cat','dog','bear');
CREATE TABLE tenum (i INT PRIMARY KEY,c1 testenum)
----

copy-from
COPY tenum FROM STDIN
0 cat
1 dog
2 bear
----
3

copy-from-kvtrace
COPY tenum FROM STDIN
0 cat
1 dog
2 bear
----
CPut /Table/<>/1/0/0 -> /TUPLE/2:2:Bytes/@
CPut /Table/<>/1/1/0 -> /TUPLE/2:2:Bytes/0x80
CPut /Table/<>/1/2/0 -> /TUPLE/2:2:Bytes/0xc0

exec-ddl
CREATE TABLE tenum2 (i INT PRIMARY KEY, c1 testenum, INDEX(c1))
----

copy-from
COPY tenum2 FROM STDIN
0 cat
----
1

copy-from-kvtrace
COPY tenum2 FROM STDIN
0 cat
----
CPut /Table/<>/1/0/0 -> /TUPLE/2:2:Bytes/@
InitPut /Table/<>/2/"@"/0/0 -> /BYTES/

exec-ddl
CREATE TABLE tenum3 (i INT PRIMARY KEY, c1 testenum, UNIQUE INDEX(c1))
----

copy-from
COPY tenum3 FROM STDIN
0 cat
----
1

copy-from-kvtrace
COPY tenum3 FROM STDIN
0 cat
----
CPut /Table/<>/1/0/0 -> /TUPLE/2:2:Bytes/@
InitPut /Table/<>/2/"@"/0 -> /BYTES/0x88

exec-ddl
CREATE TYPE comp AS (a INT, b INT);
CREATE TABLE tcomp (i INT PRIMARY KEY, c comp)
----

copy-from
COPY tcomp FROM STDIN
0 (1, 2)
----
1

copy-from-kvtrace
COPY tcomp FROM STDIN
0 (1, 2)
----
CPut /Table/<>/1/0/0 -> /TUPLE/

0 comments on commit 89ae338

Please sign in to comment.