fix consistency level on codecs

Author: shari0311

proxy.go

@@ -24,6 +24,7 @@ import (
 
 func main() {
 	ctx, cancel := signalContext(context.Background(), os.Interrupt, os.Kill)
+
 	defer cancel()
 
 	os.Exit(proxy.Run(ctx, os.Args[1:]))

proxy/codecs.go

@@ -42,7 +42,7 @@ func (c *partialQueryCodec) Decode(source io.Reader, _ primitive.ProtocolVersion
 	if query, err := primitive.ReadLongString(source); err != nil {
 		return nil, err
 	} else {
-		return &partialQuery{query, primitive.ConsistencyLevelAny}, nil
+		return &partialQuery{query, primitive.ConsistencyLevelLocalQuorum}, nil
 	}
 }
 
@@ -105,7 +105,7 @@ func (c *partialExecuteCodec) EncodedLength(_ message.Message, _ primitive.Proto
 }
 
 func (c *partialExecuteCodec) Decode(source io.Reader, _ primitive.ProtocolVersion) (msg message.Message, err error) {
-	execute := &partialExecute{}
+	execute := &partialExecute{Consistency: primitive.ConsistencyLevelLocalQuorum}
 	if execute.queryId, err = primitive.ReadShortBytes(source); err != nil {
 		return nil, fmt.Errorf("cannot read EXECUTE query id: %w", err)
 	} else if len(execute.queryId) == 0 {
@@ -184,7 +184,7 @@ func (p partialBatchCodec) Decode(source io.Reader, version primitive.ProtocolVe
 		}
 		queryOrIds[i] = queryOrId
 	}
-	return &partialBatch{queryOrIds, primitive.ConsistencyLevelAny}, nil
+	return &partialBatch{queryOrIds, primitive.ConsistencyLevelLocalQuorum}, nil
 }
 
 func (p partialBatchCodec) GetOpCode() primitive.OpCode {

proxy/proxy_test.go

@@ -407,47 +407,6 @@ func setupProxyTestWithConfig(ctx context.Context, numNodes int, cfg *proxyTestC
 	return tester, proxyAddr, nil
 }
 
-func TestProxy_Consistency(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
-	tester, proxyContactPoint, err := setupProxyTest(ctx, 3, proxycore.MockRequestHandlers{
-		primitive.OpCodeQuery: func(cl *proxycore.MockClient, frm *frame.Frame) message.Message {
-			qry := frm.Body.Message.(*message.Query)
-			if qry.Query == "SELECT * FROM test.test" {
-				// A read timeout is not retried because of 0 replicas received
-				return &message.ReadTimeout{
-					ErrorMessage: "this is a mock read timeout error",
-					Consistency:  getConsistencyLevel(frm.Body.Message),
-					Received:     0,
-					BlockFor:     1,
-					DataPresent:  false,
-				}
-			}
-			return cl.InterceptQuery(frm.Header, frm.Body.Message.(*message.Query))
-		},
-	})
-	defer func() {
-		cancel()
-		tester.shutdown()
-	}()
-	require.NoError(t, err)
-
-	cl, err := proxycore.ConnectClient(ctx, proxycore.NewEndpoint(proxyContactPoint), proxycore.ClientConnConfig{})
-	require.NoError(t, err)
-
-	version, err := cl.Handshake(ctx, primitive.ProtocolVersion5, nil)
-	require.NoError(t, err)
-	assert.Equal(t, primitive.ProtocolVersion4, version) // Expected to be negotiated to v4
-
-	_, err = cl.Query(ctx, primitive.ProtocolVersion4, &message.Query{
-		Query: "SELECT * FROM test.test",
-		Options: &message.QueryOptions{
-			Consistency: primitive.ConsistencyLevelOne,
-		},
-	})
-
-	print(err)
-}
-
 func connectTestClient(t *testing.T, ctx context.Context, proxyContactPoint string) *proxycore.ClientConn {
 	cl, err := proxycore.ConnectClient(ctx, proxycore.NewEndpoint(proxyContactPoint), proxycore.ClientConnConfig{})
 	require.NoError(t, err)
@@ -459,19 +418,6 @@ func connectTestClient(t *testing.T, ctx context.Context, proxyContactPoint stri
 	return cl
 }
 
-func getConsistencyLevel(message message.Message) primitive.ConsistencyLevel {
-	switch m := message.(type) {
-	case *partialQuery:
-		return m.Consistency
-	case *partialBatch:
-		return m.Consistency
-	case *partialExecute:
-		return m.Consistency
-	default:
-		return primitive.ConsistencyLevelOne
-	}
-}
-
 func waitUntil(d time.Duration, check func() bool) bool {
 	iterations := int(d / (100 * time.Millisecond))
 	for i := 0; i < iterations; i++ {

proxy/request.go

@@ -68,6 +68,8 @@ func (r *request) executeInternal(next bool) {
 			r.done = true
 			r.send(&message.ServerError{ErrorMessage: "Proxy exhausted query plan and there are no more hosts available to try"})
 		} else {
+			r.client.proxy.logger.Info("sending request to host", zap.Stringer("host", r.host),
+				zap.String("consistency", r.getConsistencyLevel(r.msg).String()))
 			err := r.session.Send(r.host, r)
 			if err == nil {
 				break
@@ -78,6 +80,19 @@ func (r *request) executeInternal(next bool) {
 	}
 }
 
+func (r *request) getConsistencyLevel(message message.Message) primitive.ConsistencyLevel {
+	switch m := message.(type) {
+	case *partialQuery:
+		return m.Consistency
+	case *partialBatch:
+		return m.Consistency
+	case *partialExecute:
+		return m.Consistency
+	default:
+		return primitive.ConsistencyLevelOne
+	}
+}
+
 func (r *request) send(msg message.Message) {
 	_ = r.client.conn.Write(proxycore.SenderFunc(func(writer io.Writer) error {
 		return codec.EncodeFrame(frame.NewFrame(r.raw.Header.Version, r.stream, msg), writer)