github.com/datastax/go-cassandra-native-protocol@v0.0.0-20220706104457-5e8aad05cf90/client/prepare.go (about)

     1  // Copyright 2020 DataStax
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package client
    16  
    17  import (
    18  	"github.com/rs/zerolog/log"
    19  
    20  	"github.com/datastax/go-cassandra-native-protocol/frame"
    21  	"github.com/datastax/go-cassandra-native-protocol/message"
    22  )
    23  
    24  // A RequestHandler to handle PREPARE and EXECUTE requests for the given query string, effectively emulating the
    25  // behavior of a statement being prepared, then executed.
    26  // When a PREPARE request targets the query string, it is intercepted and the handler replies with a PreparedResult.
    27  // The prepared id is simply the query string bytes, and the metadata is the metadata provided to the function.
    28  // When an EXECUTE request targets the same query string:
    29  // - If the request was previously prepared, returns a Rows RESULT response; the actual data returned is produced by
    30  // invoking the provided rows factory function, which allows the result to be customized according to the bound
    31  // variables provided with the EXECUTE message.
    32  // - If the request was not prepared, returns an Unprepared ERROR response.
    33  func NewPreparedStatementHandler(
    34  	query string,
    35  	variables *message.VariablesMetadata,
    36  	columns *message.RowsMetadata,
    37  	rows func(options *message.QueryOptions) message.RowSet,
    38  ) RequestHandler {
    39  	prepared := false
    40  	return func(request *frame.Frame, conn *CqlServerConnection, ctx RequestHandlerContext) (response *frame.Frame) {
    41  		version := request.Header.Version
    42  		id := request.Header.StreamId
    43  		switch msg := request.Body.Message.(type) {
    44  		case *message.Prepare:
    45  			if msg.Query == query {
    46  				log.Debug().Msgf("%v: [prepare handler]: intercepted PREPARE", conn)
    47  				result := &message.PreparedResult{
    48  					PreparedQueryId:   []byte(query),
    49  					VariablesMetadata: variables,
    50  					ResultMetadata:    columns,
    51  				}
    52  				prepared = true
    53  				response = frame.NewFrame(version, id, result)
    54  				log.Debug().Msgf("%v: [prepare handler]: returning %v", conn, response)
    55  			}
    56  		case *message.Execute:
    57  			if string(msg.QueryId) == query {
    58  				log.Debug().Msgf("%v: [prepare handler]: intercepted EXECUTE", conn)
    59  				if prepared {
    60  					result := &message.RowsResult{
    61  						Metadata: columns,
    62  						Data:     rows(msg.Options),
    63  					}
    64  					response = frame.NewFrame(version, id, result)
    65  				} else {
    66  					result := &message.Unprepared{
    67  						ErrorMessage: "Unprepared query: " + query,
    68  						Id:           []byte(query),
    69  					}
    70  					response = frame.NewFrame(version, id, result)
    71  				}
    72  				log.Debug().Msgf("%v: [prepare handler]: returning %v", conn, response)
    73  			}
    74  		}
    75  		return
    76  	}
    77  }