Support set statement routing + unit test

router/qrouter/proxy_routing.go

@@ -729,12 +729,11 @@ func (qr *ProxyQrouter) Route(ctx context.Context, stmt lyx.Node, dataspace stri
 		return nil, err
 	}
 
-	spqrlog.Zero.Debug().
-		Interface("route", route).
-		Msg("parsed shard route")
 	switch v := route.(type) {
 	case routingstate.ShardMatchState:
 		return v, nil
+	case routingstate.RandomMatchState:
+		return v, nil
 	case routingstate.MultiMatchState:
 		switch qr.cfg.DefaultRouteBehaviour {
 		case "BLOCK":

router/qrouter/proxy_routing_test.go

@@ -949,3 +949,99 @@ func TestInsertMultiDataspace(t *testing.T) {
 		assert.Equal(tt.exp, tmp, tt.query)
 	}
 }
+
+func TestSetStmt(t *testing.T) {
+	assert := assert.New(t)
+
+	type tcase struct {
+		query     string
+		dataspace string
+		exp       routingstate.RoutingState
+		err       error
+	}
+	db, _ := qdb.NewMemQDB(MemQDBPath)
+	dataspace1 := "ds1"
+	dataspace2 := "ds2"
+
+	_ = db.AddShardingRule(context.TODO(), &qdb.ShardingRule{
+		ID:          "id1",
+		DataspaceId: dataspace1,
+		TableName:   "",
+		Entries: []qdb.ShardingRuleEntry{
+			{
+				Column: "i",
+			},
+		},
+	})
+
+	_ = db.AddShardingRule(context.TODO(), &qdb.ShardingRule{
+		ID:          "id1",
+		DataspaceId: dataspace2,
+		TableName:   "",
+		Entries: []qdb.ShardingRuleEntry{
+			{
+				Column: "i",
+			},
+		},
+	})
+
+	err := db.AddKeyRange(context.TODO(), &qdb.KeyRange{
+		ShardID:     "sh1",
+		DataspaceId: dataspace1,
+		KeyRangeID:  "id1",
+		LowerBound:  []byte("1"),
+		UpperBound:  []byte("11"),
+	})
+
+	assert.NoError(err)
+
+	err = db.AddKeyRange(context.TODO(), &qdb.KeyRange{
+		ShardID:     "sh2",
+		DataspaceId: dataspace2,
+		KeyRangeID:  "id2",
+		LowerBound:  []byte("1"),
+		UpperBound:  []byte("11"),
+	})
+
+	assert.NoError(err)
+
+	lc := local.NewLocalCoordinator(db)
+
+	pr, err := qrouter.NewProxyRouter(map[string]*config.Shard{
+		"sh1": {
+			Hosts: nil,
+		},
+		"sh2": {
+			Hosts: nil,
+		},
+	}, lc, &config.QRouter{
+		DefaultRouteBehaviour: "BLOCK",
+	})
+
+	assert.NoError(err)
+
+	for _, tt := range []tcase{
+		{
+			query:     "SET extra_float_digits = 3",
+			dataspace: dataspace1,
+			exp:       routingstate.RandomMatchState{},
+			err:       nil,
+		},
+		{
+			query:     "SET application_name = 'jiofewjijiojioji';",
+			dataspace: dataspace2,
+			exp:       routingstate.RandomMatchState{},
+			err:       nil,
+		},
+	} {
+		parserRes, err := lyx.Parse(tt.query)
+
+		assert.NoError(err, "query %s", tt.query)
+
+		tmp, err := pr.Route(context.TODO(), parserRes, tt.dataspace, nil, routehint.EmptyRouteHint{})
+
+		assert.NoError(err, "query %s", tt.query)
+
+		assert.Equal(tt.exp, tmp, tt.query)
+	}
+}

router/relay/relay.go

@@ -328,6 +328,13 @@ func (rst *RelayStateImpl) Reroute(params [][]byte, rh routehint.RouteHint) erro
 		Msg("rerouting the client connection, resolving shard")
 
 	routingState, err := rst.Qr.Route(context.TODO(), rst.qp.Stmt(), rst.Cl.DS(), params, rh)
+
+	spqrlog.Zero.Debug().
+		Uint("client", rst.Client().ID()).
+		Interface("routing-state", routingState).
+		Err(err).
+		Msg("parsing routing state done")
+
 	if err != nil {
 		return fmt.Errorf("error processing query '%v': %v", rst.plainQ, err)
 	}