From b6f88edb73d1e281dec72fe235c17d38d92f1360 Mon Sep 17 00:00:00 2001
From: diphantxm <diphantxm@yandex-team.ru>
Date: Mon, 26 Aug 2024 16:49:00 +0300
Subject: [PATCH] add tests

---
 router/relay/relay.go                         | 10 +++---
 .../tests/router/expected/copy_routing.out    | 36 +++++++++++++------
 .../regress/tests/router/sql/copy_routing.sql | 19 ++++++----
 3 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/router/relay/relay.go b/router/relay/relay.go
index eae4e97b0..05851f2b7 100644
--- a/router/relay/relay.go
+++ b/router/relay/relay.go
@@ -665,8 +665,6 @@ func (rst *RelayStateImpl) ProcCopy(stmt *lyx.Copy, data *pgproto3.CopyData) err
 	rst.Client().RLock()
 	defer rst.Client().RUnlock()
 
-	spqrlog.Zero.Debug().Interface("copy stmt", stmt).Msg("here201")
-
 	// Read delimiter from COPY options
 	delimiter := byte(';')
 	for _, opt := range stmt.Options {
@@ -698,8 +696,6 @@ func (rst *RelayStateImpl) ProcCopy(stmt *lyx.Copy, data *pgproto3.CopyData) err
 			return err
 		}
 
-		spqrlog.Zero.Debug().Interface("value", valueClause.Values).Msg("here")
-
 		smt, ok := r.(routingstate.ShardMatchState)
 		if !ok {
 			return fmt.Errorf("multishard copy is not supported")
@@ -809,6 +805,7 @@ func (rst *RelayStateImpl) ProcQuery(query pgproto3.FrontendMessage, waitForResp
 			q := rst.qp.Stmt().(*lyx.Copy)
 
 			if err := func() error {
+				buf := []byte{}
 				for {
 					cpMsg, err := rst.Client().Receive()
 					if err != nil {
@@ -817,10 +814,11 @@ func (rst *RelayStateImpl) ProcQuery(query pgproto3.FrontendMessage, waitForResp
 
 					switch msg := cpMsg.(type) {
 					case *pgproto3.CopyData:
-						if err := rst.ProcCopy(q, msg); err != nil {
+						buf = append(buf, msg.Data...)
+					case *pgproto3.CopyDone, *pgproto3.CopyFail:
+						if err := rst.ProcCopy(q, &pgproto3.CopyData{Data: buf}); err != nil {
 							return err
 						}
-					case *pgproto3.CopyDone, *pgproto3.CopyFail:
 						if err := rst.ProcCopyComplete(&cpMsg); err != nil {
 							return err
 						}
diff --git a/test/regress/tests/router/expected/copy_routing.out b/test/regress/tests/router/expected/copy_routing.out
index 89e5df1b7..114a33580 100644
--- a/test/regress/tests/router/expected/copy_routing.out
+++ b/test/regress/tests/router/expected/copy_routing.out
@@ -34,8 +34,8 @@ ALTER DISTRIBUTION ds1 ATTACH RELATION copy_test DISTRIBUTION KEY id;
 \c regress
 CREATE TABLE copy_test (id int);
 NOTICE: send query to shard(s) : sh1,sh2
-COPY copy_test FROM STDIN WHERE id <= 10;
-NOTICE: send query to shard(s) : sh1
+COPY copy_test(id) FROM STDIN WHERE id <= 10;
+NOTICE: send query to shard(s) : sh1,sh2
 SELECT * FROM copy_test WHERE id <= 10;
 NOTICE: send query to shard(s) : sh1
  id 
@@ -47,10 +47,24 @@ NOTICE: send query to shard(s) : sh1
   5
 (5 rows)
 
-COPY copy_test FROM STDIN WHERE id <= 30;
-NOTICE: send query to shard(s) : sh2
-SELECT * FROM copy_test WHERE id <= 30 ORDER BY copy_test;
-NOTICE: send query to shard(s) : sh2
+COPY copy_test(id) FROM STDIN;
+NOTICE: send query to shard(s) : sh1,sh2
+ERROR:  client processing error: multishard copy is not supported, tx status IDLE
+SELECT * FROM copy_test;
+NOTICE: send query to shard(s) : sh1,sh2
+ id 
+----
+  1
+  2
+  3
+  4
+  5
+(5 rows)
+
+COPY copy_test(id) FROM STDIN;
+NOTICE: send query to shard(s) : sh1,sh2
+SELECT * FROM copy_test;
+NOTICE: send query to shard(s) : sh1,sh2
  id 
 ----
   1
@@ -58,10 +72,12 @@ NOTICE: send query to shard(s) : sh2
   3
   4
   5
- 12
- 22
- 23
-(8 rows)
+ 41
+ 42
+ 43
+ 44
+ 45
+(10 rows)
 
 DROP TABLE copy_test;
 NOTICE: send query to shard(s) : sh1,sh2
diff --git a/test/regress/tests/router/sql/copy_routing.sql b/test/regress/tests/router/sql/copy_routing.sql
index 8e4010484..8b8ef16dc 100644
--- a/test/regress/tests/router/sql/copy_routing.sql
+++ b/test/regress/tests/router/sql/copy_routing.sql
@@ -7,20 +7,17 @@ ALTER DISTRIBUTION ds1 ATTACH RELATION copy_test DISTRIBUTION KEY id;
 \c regress
 CREATE TABLE copy_test (id int);
 
-COPY copy_test FROM STDIN WHERE id <= 10;
+COPY copy_test(id) FROM STDIN WHERE id <= 10;
 1
 2
 3
 4
 5
-12
-3434
-43
 \.
 
 SELECT * FROM copy_test WHERE id <= 10;
 
-COPY copy_test FROM STDIN WHERE id <= 30;
+COPY copy_test(id) FROM STDIN;
 1
 2
 3
@@ -35,7 +32,17 @@ COPY copy_test FROM STDIN WHERE id <= 30;
 43
 \.
 
-SELECT * FROM copy_test WHERE id <= 30 ORDER BY copy_test;
+SELECT * FROM copy_test;
+
+COPY copy_test(id) FROM STDIN;
+41
+42
+43
+44
+45
+\.
+
+SELECT * FROM copy_test;
 
 DROP TABLE copy_test;