From 2a7ab74005dea5fe7c46f750c990a7f5f117df4b Mon Sep 17 00:00:00 2001
From: Yury Frolov <57130330+EinKrebs@users.noreply.github.com>
Date: Thu, 18 Apr 2024 12:02:24 +0500
Subject: [PATCH] Added processing for statements with subselect compared to
 const (#624)

---
 router/qrouter/proxy_routing.go      | 10 +++++++---
 router/qrouter/proxy_routing_test.go | 30 +++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/router/qrouter/proxy_routing.go b/router/qrouter/proxy_routing.go
index 6f49592f2..84b144761 100644
--- a/router/qrouter/proxy_routing.go
+++ b/router/qrouter/proxy_routing.go
@@ -326,6 +326,10 @@ func (qr *ProxyQrouter) routeByClause(ctx context.Context, expr lyx.Node, meta *
 				default:
 					queue = append(queue, texpr.Left, texpr.Right)
 				}
+			case *lyx.Select:
+				if err := qr.DeparseSelectStmt(ctx, lft, meta); err != nil {
+					return err
+				}
 			default:
 				/* Consider there cases */
 				// if !(texpr.AExpr.Kind == pgquery.A_Expr_Kind_AEXPR_OP || texpr.Kind == pgquery.A_Expr_Kind_AEXPR_BETWEEN) {
@@ -701,7 +705,7 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s
 
 	/*
 	 * Step 1: traverse query tree and deparse mapping from
-	 * columns to their values (either contant or expression).
+	 * columns to their values (either constant or expression).
 	 * Note that exact (routing) value of (sharding) column may not be
 	 * known after this phase, as it can be Parse Step of Extended proto.
 	 */
@@ -713,7 +717,7 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s
 	case *lyx.VariableSetStmt:
 		/* TBD: maybe skip all set stmts? */
 		/*
-		 * SET x = y etc, do not dispatch any statement to shards, just process this in router
+		 * SET x = y etc., do not dispatch any statement to shards, just process this in router
 		 */
 		return routingstate.RandomMatchState{}, nil
 
@@ -746,7 +750,7 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s
 		/*
 		 * Disallow to index on table which does not contain any sharding column
 		 */
-		// XXX: doit
+		// XXX: do it
 		return routingstate.MultiMatchState{}, nil
 
 	case *lyx.Alter, *lyx.Drop, *lyx.Truncate:
diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go
index 9bbd0ab9d..e356f2232 100644
--- a/router/qrouter/proxy_routing_test.go
+++ b/router/qrouter/proxy_routing_test.go
@@ -273,7 +273,7 @@ func TestCTE(t *testing.T) {
 			query: `
 			WITH xxxx AS (
 				SELECT * from t where i = 1
-			) 
+			)
 			SELECT * from xxxx;
 			`,
 			err: nil,
@@ -292,6 +292,34 @@ func TestCTE(t *testing.T) {
 				TargetSessionAttrs: "any",
 			},
 		},
+		{
+			query: `
+			WITH xxxx AS (
+				SELECT * from t where i = 1
+			),
+			zzzz AS (
+				UPDATE t 
+				SET a = 0
+				WHERE i = 1 AND (SELECT COUNT(*) FROM xxxx WHERE b = 0) = 1
+			)	
+			SELECT * FROM xxxx;
+			`,
+			err: nil,
+			exp: routingstate.ShardMatchState{
+				Route: &routingstate.DataShardRoute{
+					Shkey: kr.ShardKey{
+						Name: "sh1",
+					},
+					Matchedkr: &kr.KeyRange{
+						ShardID:      "sh1",
+						ID:           "id1",
+						Distribution: distribution,
+						LowerBound:   []byte("1"),
+					},
+				},
+				TargetSessionAttrs: "any",
+			},
+		},
 	} {
 		parserRes, err := lyx.Parse(tt.query)