From f19a1ac6c3f19087d5f12beed1dbf344149400fb Mon Sep 17 00:00:00 2001
From: Yury Frolov <57130330+EinKrebs@users.noreply.github.com>
Date: Thu, 25 Apr 2024 15:09:06 +0500
Subject: [PATCH] Add set operations processing (#639)

* Bump lyx

* Add test for UNION routing

* Bump lyx again

* Add processing for set operations
---
 go.mod                               |  2 +-
 go.sum                               |  4 ++++
 router/qrouter/proxy_routing.go      | 22 +++++++++++++++++++++-
 router/qrouter/proxy_routing_test.go | 19 +++++++++++++++++++
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index 11de61f12..0ce1c0da5 100644
--- a/go.mod
+++ b/go.mod
@@ -17,7 +17,7 @@ require (
 	github.com/lib/pq v1.10.9
 	github.com/libp2p/go-reuseport v0.4.0
 	github.com/opentracing/opentracing-go v1.2.0
-	github.com/pg-sharding/lyx v0.0.0-20240417111237-40a0e3320599
+	github.com/pg-sharding/lyx v0.0.0-20240425090312-06d7412dfba8
 	github.com/pkg/errors v0.9.1
 	github.com/rs/zerolog v1.32.0
 	github.com/sevlyar/go-daemon v0.1.6
diff --git a/go.sum b/go.sum
index d4d0e7e03..f1800c9be 100644
--- a/go.sum
+++ b/go.sum
@@ -143,6 +143,10 @@ github.com/opentracing/opentracing-go v1.2.0 h1:uEJPy/1a5RIPAJ0Ov+OIO8OxWu77jEv+
 github.com/opentracing/opentracing-go v1.2.0/go.mod h1:GxEUsuufX4nBwe+T+Wl9TAgYrxe9dPLANfrWvHYVTgc=
 github.com/pg-sharding/lyx v0.0.0-20240417111237-40a0e3320599 h1:CPMrNVoDZkPNjfqaPpPzCr1UN8LZmWw88NQ5aiZ+RrA=
 github.com/pg-sharding/lyx v0.0.0-20240417111237-40a0e3320599/go.mod h1:2dPBQAhqv/30mhzj2yBXQkXhsGJQ8GhM+oWOfbGua58=
+github.com/pg-sharding/lyx v0.0.0-20240419111011-78b6979ac476 h1:PCRfnJMD8QawQoGqtpyNiN6MLRj9WN628dAWBL+zhTw=
+github.com/pg-sharding/lyx v0.0.0-20240419111011-78b6979ac476/go.mod h1:2dPBQAhqv/30mhzj2yBXQkXhsGJQ8GhM+oWOfbGua58=
+github.com/pg-sharding/lyx v0.0.0-20240425090312-06d7412dfba8 h1:1+eQ40tGL8n/Xli/dBx9oW6ISU1NE2WGSnBDpNnRSiA=
+github.com/pg-sharding/lyx v0.0.0-20240425090312-06d7412dfba8/go.mod h1:2dPBQAhqv/30mhzj2yBXQkXhsGJQ8GhM+oWOfbGua58=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
diff --git a/router/qrouter/proxy_routing.go b/router/qrouter/proxy_routing.go
index 84b144761..04f4af212 100644
--- a/router/qrouter/proxy_routing.go
+++ b/router/qrouter/proxy_routing.go
@@ -372,6 +372,17 @@ func (qr *ProxyQrouter) DeparseSelectStmt(ctx context.Context, selectStmt lyx.No
 			}
 		}
 
+		if s.LArg != nil {
+			if err := qr.DeparseSelectStmt(ctx, s.LArg, meta); err != nil {
+				return err
+			}
+		}
+		if s.RArg != nil {
+			if err := qr.DeparseSelectStmt(ctx, s.RArg, meta); err != nil {
+				return err
+			}
+		}
+
 	/* SELECT * FROM VALUES() ... */
 	case *lyx.ValueClause:
 		/* random route */
@@ -780,7 +791,7 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s
 		/* We cannot route SQL stmt with no FROM clause provided, but there is still
 		* a few cases to consider
 		 */
-		if len(node.FromClause) == 0 {
+		if len(node.FromClause) == 0 && (node.LArg == nil || node.RArg == nil) {
 			/* Step 1.4.8: select a_expr is routable to any shard in case when a_expr is some type of
 			data-independent expr */
 
@@ -815,6 +826,15 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s
 			if any_routable {
 				return routingstate.RandomMatchState{}, nil
 			}
+		} else if node.LArg != nil && node.RArg != nil {
+			err := qr.deparseShardingMapping(ctx, node.LArg, meta)
+			if err != nil {
+				return nil, err
+			}
+			err = qr.deparseShardingMapping(ctx, node.RArg, meta)
+			if err != nil {
+				return nil, err
+			}
 		} else {
 			// SELECT stmts, which
 			// would be routed with their WHERE clause
diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go
index e356f2232..8c62d3c9e 100644
--- a/router/qrouter/proxy_routing_test.go
+++ b/router/qrouter/proxy_routing_test.go
@@ -611,6 +611,25 @@ func TestSingleShard(t *testing.T) {
 			},
 			err: nil,
 		},
+
+		{
+			query: "SELECT * FROM t WHERE i = 12 UNION ALL SELECT * FROM xxmixed WHERE i = 22;",
+			exp: routingstate.ShardMatchState{
+				Route: &routingstate.DataShardRoute{
+					Shkey: kr.ShardKey{
+						Name: "sh2",
+					},
+					Matchedkr: &kr.KeyRange{
+						ShardID:      "sh2",
+						ID:           "id2",
+						Distribution: distribution,
+						LowerBound:   []byte("11"),
+					},
+				},
+				TargetSessionAttrs: "any",
+			},
+			err: nil,
+		},
 	} {
 		parserRes, err := lyx.Parse(tt.query)