From a25f188aa9a7add50acccedd6edc9bed9267dff5 Mon Sep 17 00:00:00 2001 From: Yury Frolov Date: Wed, 17 Apr 2024 18:56:19 +0500 Subject: [PATCH 1/6] Add another test for CTE --- router/qrouter/proxy_routing.go | 2 +- router/qrouter/proxy_routing_test.go | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/router/qrouter/proxy_routing.go b/router/qrouter/proxy_routing.go index 04f4af212..ca3d2ead9 100644 --- a/router/qrouter/proxy_routing.go +++ b/router/qrouter/proxy_routing.go @@ -105,7 +105,7 @@ func (meta *RoutingMetadataContext) RecordConstExpr(resolvedRelation RelationFQN meta.rels[resolvedRelation] = struct{}{} if _, ok := meta.exprs[resolvedRelation]; !ok { meta.exprs[resolvedRelation] = map[string]string{} - } + } // TODO: else branch delete(meta.unparsed_columns, colname) meta.exprs[resolvedRelation][colname] = expr } diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go index 8c62d3c9e..6897a680d 100644 --- a/router/qrouter/proxy_routing_test.go +++ b/router/qrouter/proxy_routing_test.go @@ -320,6 +320,21 @@ 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 = 12 + ) + SELECT * FROM xxxx; + `, + err: nil, + exp: routingstate.MultiMatchState{}, + }, } { parserRes, err := lyx.Parse(tt.query) From c7cd029dd162fdc5a309c08d460e11630c184da1 Mon Sep 17 00:00:00 2001 From: Yury Frolov Date: Thu, 18 Apr 2024 11:42:48 +0500 Subject: [PATCH 2/6] Forbid setting several values for distribution key --- router/qrouter/proxy_routing.go | 43 ++++++++++++++-------------- router/qrouter/proxy_routing_test.go | 11 +++++-- 2 files changed, 30 insertions(+), 24 deletions(-) diff --git a/router/qrouter/proxy_routing.go b/router/qrouter/proxy_routing.go index ca3d2ead9..fccda4649 100644 --- a/router/qrouter/proxy_routing.go +++ b/router/qrouter/proxy_routing.go @@ -97,17 +97,21 @@ func (meta *RoutingMetadataContext) RFQNIsCTE(resolvedRelation RelationFQN) bool } // TODO : unit tests -func (meta *RoutingMetadataContext) RecordConstExpr(resolvedRelation RelationFQN, colname string, expr string) { +func (meta *RoutingMetadataContext) RecordConstExpr(resolvedRelation RelationFQN, colname string, expr string) error { if meta.RFQNIsCTE(resolvedRelation) { // CTE, skip - return + return nil } meta.rels[resolvedRelation] = struct{}{} if _, ok := meta.exprs[resolvedRelation]; !ok { meta.exprs[resolvedRelation] = map[string]string{} - } // TODO: else branch + } delete(meta.unparsed_columns, colname) + if curExpr, ok := meta.exprs[resolvedRelation][colname]; ok && curExpr != expr { + return spqrerror.Newf(spqrerror.SPQR_COMPLEX_QUERY, "several different values for distribution key") + } meta.exprs[resolvedRelation][colname] = expr + return nil } // TODO : unit tests @@ -185,11 +189,11 @@ func (qr *ProxyQrouter) DeparseKeyWithRangesInternal(_ context.Context, key stri return nil, FailedToFindKeyRange } -func (qr *ProxyQrouter) RecordDistributionKeyColumnValueOnRFQN(meta *RoutingMetadataContext, resolvedRelation RelationFQN, colname, value string) { +func (qr *ProxyQrouter) RecordDistributionKeyColumnValueOnRFQN(meta *RoutingMetadataContext, resolvedRelation RelationFQN, colname, value string) error { /* do not process non-distributed relations or columns not from relation distribution key */ if ds, err := qr.Mgr().GetRelationDistribution(context.TODO(), resolvedRelation.RelationName); err != nil { - return + return nil } else { // TODO: optimize ok := false @@ -201,12 +205,12 @@ func (qr *ProxyQrouter) RecordDistributionKeyColumnValueOnRFQN(meta *RoutingMeta } if !ok { // some junk column - return + return nil } } // will not work not ints - meta.RecordConstExpr(resolvedRelation, colname, value) + return meta.RecordConstExpr(resolvedRelation, colname, value) } // TODO : unit tests @@ -232,30 +236,27 @@ func (qr *ProxyQrouter) RecordDistributionKeyExprOnRFQN(meta *RoutingMetadataCon // ??? protoc violation } - qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, string(routeParam)) - return nil + return qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, string(routeParam)) case *lyx.AExprSConst: - qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, string(e.Value)) - return nil + return qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, string(e.Value)) case *lyx.AExprIConst: val := fmt.Sprintf("%d", e.Value) - qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, string(val)) - return nil + return qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, string(val)) default: return ComplexQuery } } -func (qr *ProxyQrouter) RecordDistributionKeyColumnValue(meta *RoutingMetadataContext, alias, colname, value string) { +func (qr *ProxyQrouter) RecordDistributionKeyColumnValue(meta *RoutingMetadataContext, alias, colname, value string) error { resolvedRelation, err := meta.ResolveRelationByAlias(alias) if err != nil { // failed to resolve relation, skip column meta.unparsed_columns[colname] = struct{}{} - return + return nil } - qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, value) + return qr.RecordDistributionKeyColumnValueOnRFQN(meta, resolvedRelation, colname, value) } // routeByClause de-parses sharding column-value pair from Where clause of the query @@ -281,27 +282,27 @@ func (qr *ProxyQrouter) routeByClause(ctx context.Context, expr lyx.Node, meta * switch rght := texpr.Right.(type) { case *lyx.ParamRef: if rght.Number <= len(meta.params) { - qr.RecordDistributionKeyColumnValue(meta, alias, colname, string(meta.params[rght.Number-1])) + return qr.RecordDistributionKeyColumnValue(meta, alias, colname, string(meta.params[rght.Number-1])) } // else error out? case *lyx.AExprSConst: // TBD: postpone routing from here to root of parsing tree - qr.RecordDistributionKeyColumnValue(meta, alias, colname, rght.Value) + return qr.RecordDistributionKeyColumnValue(meta, alias, colname, rght.Value) case *lyx.AExprIConst: // TBD: postpone routing from here to root of parsing tree // maybe expimely inefficient. Will be fixed in SPQR-2.0 - qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", rght.Value)) + return qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", rght.Value)) case *lyx.AExprList: if len(rght.List) != 0 { expr := rght.List[0] switch bexpr := expr.(type) { case *lyx.AExprSConst: // TBD: postpone routing from here to root of parsing tree - qr.RecordDistributionKeyColumnValue(meta, alias, colname, bexpr.Value) + return qr.RecordDistributionKeyColumnValue(meta, alias, colname, bexpr.Value) case *lyx.AExprIConst: // TBD: postpone routing from here to root of parsing tree // maybe expimely inefficient. Will be fixed in SPQR-2.0 - qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", bexpr.Value)) + return qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", bexpr.Value)) } } case *lyx.FuncApplication: diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go index 6897a680d..7e6e5e68b 100644 --- a/router/qrouter/proxy_routing_test.go +++ b/router/qrouter/proxy_routing_test.go @@ -2,6 +2,7 @@ package qrouter_test import ( "context" + "github.com/pg-sharding/spqr/pkg/models/spqrerror" "testing" "github.com/pg-sharding/spqr/pkg/config" @@ -332,8 +333,8 @@ func TestCTE(t *testing.T) { ) SELECT * FROM xxxx; `, - err: nil, - exp: routingstate.MultiMatchState{}, + err: spqrerror.Newf(spqrerror.SPQR_COMPLEX_QUERY, "several different values for distribution key."), + exp: nil, }, } { parserRes, err := lyx.Parse(tt.query) @@ -342,7 +343,11 @@ func TestCTE(t *testing.T) { tmp, err := pr.Route(context.TODO(), parserRes, session.NewDummyHandler(distribution)) - assert.NoError(err, "query %s", tt.query) + if tt.err == nil { + assert.NoError(err, "query %s", tt.query) + } else { + assert.Error(err, "query %s", tt.query) + } assert.Equal(tt.exp, tmp, tt.query) } From 670fd06e69cf1ea903fa819fa8016ffa9746b240 Mon Sep 17 00:00:00 2001 From: Yury Frolov Date: Thu, 18 Apr 2024 11:59:37 +0500 Subject: [PATCH 3/6] Process several values for each distribution key --- router/qrouter/proxy_routing.go | 64 ++++++++++++++-------------- router/qrouter/proxy_routing_test.go | 33 ++++++++++++-- 2 files changed, 63 insertions(+), 34 deletions(-) diff --git a/router/qrouter/proxy_routing.go b/router/qrouter/proxy_routing.go index fccda4649..cf46a6ea8 100644 --- a/router/qrouter/proxy_routing.go +++ b/router/qrouter/proxy_routing.go @@ -43,7 +43,7 @@ type RoutingMetadataContext struct { // SELECT * FROM a join b WHERE a.c1 = and a.c2 = // can be routed with different rules rels map[RelationFQN]struct{} - exprs map[RelationFQN]map[string]string + exprs map[RelationFQN]map[string][]string // cached CTE names cteNames map[string]struct{} @@ -65,7 +65,7 @@ func NewRoutingMetadataContext(params [][]byte, paramsFormatCodes []int16) *Rout rels: map[RelationFQN]struct{}{}, cteNames: map[string]struct{}{}, tableAliases: map[string]RelationFQN{}, - exprs: map[RelationFQN]map[string]string{}, + exprs: map[RelationFQN]map[string][]string{}, unparsed_columns: map[string]struct{}{}, params: params, } @@ -104,13 +104,13 @@ func (meta *RoutingMetadataContext) RecordConstExpr(resolvedRelation RelationFQN } meta.rels[resolvedRelation] = struct{}{} if _, ok := meta.exprs[resolvedRelation]; !ok { - meta.exprs[resolvedRelation] = map[string]string{} + meta.exprs[resolvedRelation] = map[string][]string{} } delete(meta.unparsed_columns, colname) - if curExpr, ok := meta.exprs[resolvedRelation][colname]; ok && curExpr != expr { - return spqrerror.Newf(spqrerror.SPQR_COMPLEX_QUERY, "several different values for distribution key") + if _, ok := meta.exprs[resolvedRelation][colname]; !ok { + meta.exprs[resolvedRelation][colname] = make([]string, 0) } - meta.exprs[resolvedRelation][colname] = expr + meta.exprs[resolvedRelation][colname] = append(meta.exprs[resolvedRelation][colname], expr) return nil } @@ -901,7 +901,7 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s ok := true - var hashedKey []byte + var hashedKeys [][]byte // TODO: multi-column routing. This works only for one-dim routing for i := 0; i < len(distrKey); i++ { @@ -914,20 +914,22 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s col := distrKey[i].Column - val, valOk := meta.exprs[rfqn][col] + vals, valOk := meta.exprs[rfqn][col] if !valOk { ok = false break } - hashedKey, err = hashfunction.ApplyHashFunction([]byte(val), hf) + hashedKeys = make([][]byte, len(vals)) + for i, val := range vals { + hashedKeys[i], err = hashfunction.ApplyHashFunction([]byte(val), hf) + spqrlog.Zero.Debug().Str("key", meta.exprs[rfqn][col][i]).Str("hashed key", string(hashedKeys[i])).Msg("applying hash function on key") - spqrlog.Zero.Debug().Str("key", meta.exprs[rfqn][col]).Str("hashed key", string(hashedKey)).Msg("applying hash function on key") - - if err != nil { - spqrlog.Zero.Debug().Err(err).Msg("failed to apply hash function") - ok = false - break + if err != nil { + spqrlog.Zero.Debug().Err(err).Msg("failed to apply hash function") + ok = false + break + } } } @@ -935,24 +937,24 @@ func (qr *ProxyQrouter) routeWithRules(ctx context.Context, stmt lyx.Node, sph s // skip this relation continue } + for _, hashedKey := range hashedKeys { + currroute, err := qr.DeparseKeyWithRangesInternal(ctx, string(hashedKey), krs) + if err != nil { + route_err = err + spqrlog.Zero.Debug().Err(route_err).Msg("temporarily skip the route error") + continue + } - currroute, err := qr.DeparseKeyWithRangesInternal(ctx, string(hashedKey), krs) - if err != nil { - route_err = err - spqrlog.Zero.Debug().Err(route_err).Msg("temporarily skip the route error") - continue - } - - spqrlog.Zero.Debug(). - Interface("currroute", currroute). - Str("table", rfqn.RelationName). - Msg("calculated route for table/cols") - - route = routingstate.Combine(route, routingstate.ShardMatchState{ - Route: currroute, - TargetSessionAttrs: tsa, - }) + spqrlog.Zero.Debug(). + Interface("currroute", currroute). + Str("table", rfqn.RelationName). + Msg("calculated route for table/cols") + route = routingstate.Combine(route, routingstate.ShardMatchState{ + Route: currroute, + TargetSessionAttrs: tsa, + }) + } } if route == nil && route_err != nil { return nil, route_err diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go index 7e6e5e68b..3c87bbaa6 100644 --- a/router/qrouter/proxy_routing_test.go +++ b/router/qrouter/proxy_routing_test.go @@ -2,7 +2,6 @@ package qrouter_test import ( "context" - "github.com/pg-sharding/spqr/pkg/models/spqrerror" "testing" "github.com/pg-sharding/spqr/pkg/config" @@ -333,8 +332,36 @@ func TestCTE(t *testing.T) { ) SELECT * FROM xxxx; `, - err: spqrerror.Newf(spqrerror.SPQR_COMPLEX_QUERY, "several different values for distribution key."), - exp: nil, + err: nil, + exp: routingstate.SkipRoutingState{}, + }, + { + query: ` + WITH xxxx AS ( + SELECT * from t where i = 1 + ), + zzzz AS ( + UPDATE t + SET a = 0 + WHERE i = 2 + ) + 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) From ce2645c86471b91d41fa3567ee77ed5ce663cef5 Mon Sep 17 00:00:00 2001 From: Yury Frolov Date: Thu, 18 Apr 2024 14:58:33 +0500 Subject: [PATCH 4/6] Reformatting single_shard_joins regress test --- test/regress/tests/router/expected/single_shard_joins.out | 2 +- test/regress/tests/router/sql/single_shard_joins.sql | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/regress/tests/router/expected/single_shard_joins.out b/test/regress/tests/router/expected/single_shard_joins.out index ed322c854..15d75710c 100644 --- a/test/regress/tests/router/expected/single_shard_joins.out +++ b/test/regress/tests/router/expected/single_shard_joins.out @@ -47,7 +47,7 @@ NOTICE: send query to shard(s) : sh2 12 | 13 (2 rows) -SELECT * FROM sshjt1 WHERE i = 12 AND j =1; +SELECT * FROM sshjt1 WHERE i = 12 AND j = 1; NOTICE: send query to shard(s) : sh2 i | j ---+--- diff --git a/test/regress/tests/router/sql/single_shard_joins.sql b/test/regress/tests/router/sql/single_shard_joins.sql index f78c4bea6..aae5b28f8 100644 --- a/test/regress/tests/router/sql/single_shard_joins.sql +++ b/test/regress/tests/router/sql/single_shard_joins.sql @@ -14,7 +14,7 @@ INSERT INTO sshjt1 (i, j) VALUES(12, 12); INSERT INTO sshjt1 (i, j) VALUES(12, 13); SELECT * FROM sshjt1 WHERE i = 12; -SELECT * FROM sshjt1 WHERE i = 12 AND j =1; +SELECT * FROM sshjt1 WHERE i = 12 AND j = 1; SELECT * FROM sshjt1 a join sshjt1 b WHERE a.i = 12 ON TRUE; SELECT * FROM sshjt1 a join sshjt1 b ON TRUE WHERE a.i = 12; From 16ec98e3e7bb686aa32ee9343ab8beb11c4726e4 Mon Sep 17 00:00:00 2001 From: Yury Frolov Date: Thu, 18 Apr 2024 15:07:45 +0500 Subject: [PATCH 5/6] Fixes --- router/qrouter/proxy_routing.go | 20 +++++++++++++++----- router/qrouter/proxy_routing_test.go | 22 ++++++++++++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/router/qrouter/proxy_routing.go b/router/qrouter/proxy_routing.go index cf46a6ea8..b8802826a 100644 --- a/router/qrouter/proxy_routing.go +++ b/router/qrouter/proxy_routing.go @@ -282,27 +282,37 @@ func (qr *ProxyQrouter) routeByClause(ctx context.Context, expr lyx.Node, meta * switch rght := texpr.Right.(type) { case *lyx.ParamRef: if rght.Number <= len(meta.params) { - return qr.RecordDistributionKeyColumnValue(meta, alias, colname, string(meta.params[rght.Number-1])) + if err := qr.RecordDistributionKeyColumnValue(meta, alias, colname, string(meta.params[rght.Number-1])); err != nil { + return err + } } // else error out? case *lyx.AExprSConst: // TBD: postpone routing from here to root of parsing tree - return qr.RecordDistributionKeyColumnValue(meta, alias, colname, rght.Value) + if err := qr.RecordDistributionKeyColumnValue(meta, alias, colname, rght.Value); err != nil { + return err + } case *lyx.AExprIConst: // TBD: postpone routing from here to root of parsing tree // maybe expimely inefficient. Will be fixed in SPQR-2.0 - return qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", rght.Value)) + if err := qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", rght.Value)); err != nil { + return err + } case *lyx.AExprList: if len(rght.List) != 0 { expr := rght.List[0] switch bexpr := expr.(type) { case *lyx.AExprSConst: // TBD: postpone routing from here to root of parsing tree - return qr.RecordDistributionKeyColumnValue(meta, alias, colname, bexpr.Value) + if err := qr.RecordDistributionKeyColumnValue(meta, alias, colname, bexpr.Value); err != nil { + return err + } case *lyx.AExprIConst: // TBD: postpone routing from here to root of parsing tree // maybe expimely inefficient. Will be fixed in SPQR-2.0 - return qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", bexpr.Value)) + if err := qr.RecordDistributionKeyColumnValue(meta, alias, colname, fmt.Sprintf("%d", bexpr.Value)); err != nil { + return err + } } } case *lyx.FuncApplication: diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go index 3c87bbaa6..75b9ec613 100644 --- a/router/qrouter/proxy_routing_test.go +++ b/router/qrouter/proxy_routing_test.go @@ -659,6 +659,24 @@ func TestSingleShard(t *testing.T) { err: nil, }, + { + query: "SELECT * FROM t WHERE i = 12 AND j = 1;", + 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, + }, { query: "SELECT * FROM t WHERE i = 12 UNION ALL SELECT * FROM xxmixed WHERE i = 22;", exp: routingstate.ShardMatchState{ @@ -1395,3 +1413,7 @@ func TestMiscRouting(t *testing.T) { } } } + +func TestSimple(t *testing.T) { + +} From 01eaf6e0f94e9346835503f20257b61a28608f42 Mon Sep 17 00:00:00 2001 From: Yury Frolov Date: Thu, 25 Apr 2024 15:41:48 +0500 Subject: [PATCH 6/6] Removed redundant test in proxy_routing_test.go --- router/qrouter/proxy_routing_test.go | 4 ---- 1 file changed, 4 deletions(-) diff --git a/router/qrouter/proxy_routing_test.go b/router/qrouter/proxy_routing_test.go index 75b9ec613..a72506229 100644 --- a/router/qrouter/proxy_routing_test.go +++ b/router/qrouter/proxy_routing_test.go @@ -1413,7 +1413,3 @@ func TestMiscRouting(t *testing.T) { } } } - -func TestSimple(t *testing.T) { - -}