statement.go

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package mysql

import (
	"encoding/binary"
	"fmt"
	"math"
)

import (
	perrors "github.com/pkg/errors"
)

import (
	"github.com/arana-db/arana/pkg/constants/mysql"
	err2 "github.com/arana-db/arana/pkg/mysql/errors"
	"github.com/arana-db/arana/pkg/proto"
)

type BackendStatement struct {
	conn       *BackendConnection
	id         uint32
	paramCount int
	sql        string
}

// prepare Result Packets
// http://dev.mysql.com/doc/internals/en/com-stmt-prepare-response.html
func (stmt *BackendStatement) readPrepareResultPacket() (uint16, error) {
	data, err := stmt.conn.c.readEphemeralPacket()
	if err != nil {
		return 0, err2.NewSQLError(mysql.CRServerLost, mysql.SSUnknownSQLState, "%v", err)
	}
	defer stmt.conn.c.recycleReadPacket()

	// packet indicator [1 byte]
	if data[0] != mysql.OKPacket {
		return 0, ParseErrorPacket(data)
	}

	// statement id [4 bytes]
	stmt.id = binary.LittleEndian.Uint32(data[1:5])

	// Column count [16 bit uint]
	columnCount := binary.LittleEndian.Uint16(data[5:7])

	// Param count [16 bit uint]
	stmt.paramCount = int(binary.LittleEndian.Uint16(data[7:9]))

	// Reserved [8 bit]

	// Warning count [16 bit uint]

	return columnCount, nil
}

// http://dev.mysql.com/doc/internals/en/com-stmt-send-long-data.html
func (stmt *BackendStatement) writeCommandLongData(paramID int, arg []byte) error {
	maxLen := stmt.conn.conf.MaxAllowedPacket - 1
	pktLen := maxLen

	// After the header (bytes 0-3) follows before the data:
	// 1 byte command
	// 4 bytes stmtID
	// 2 bytes paramID
	const dataOffset = 1 + 4 + 2

	// Cannot use the write buffer since
	// a) the buffer is too small
	// b) it is in use
	data := make([]byte, 4+1+4+2+len(arg))

	copy(data[4+dataOffset:], arg)

	for argLen := len(arg); argLen > 0; argLen -= pktLen - dataOffset {
		if dataOffset+argLen < maxLen {
			pktLen = dataOffset + argLen
		}

		stmt.conn.c.sequence = 0
		// Add command byte [1 byte]
		data[4] = mysql.ComStmtSendLongData

		// Add stmtID [32 bit]
		data[5] = byte(stmt.id)
		data[6] = byte(stmt.id >> 8)
		data[7] = byte(stmt.id >> 16)
		data[8] = byte(stmt.id >> 24)

		// Add paramID [16 bit]
		data[9] = byte(paramID)
		data[10] = byte(paramID >> 8)

		// Send CMD packet
		err := stmt.conn.c.writePacket(data[:4+pktLen])
		if err == nil {
			data = data[pktLen-dataOffset:]
			continue
		}
		return err

	}

	// Reset Packet Sequence
	stmt.conn.c.sequence = 0
	return nil
}

// Execute Prepared Statement
// http://dev.mysql.com/doc/internals/en/com-stmt-execute.html
func (stmt *BackendStatement) writeExecutePacket(args []proto.Value) error {
	if len(args) != stmt.paramCount {
		return fmt.Errorf(
			"argument count mismatch (got: %d; has: %d)",
			len(args),
			stmt.paramCount,
		)
	}

	const minPktLen = 4 + 1 + 4 + 1 + 4
	bc := stmt.conn

	// Determine threshold dynamically to avoid packet size shortage.
	longDataSize := bc.conf.MaxAllowedPacket / (stmt.paramCount + 1)
	if longDataSize < 64 {
		longDataSize = 64
	}

	// Reset packet-sequence
	bc.c.sequence = 0

	var data []byte
	var err error

	if len(args) == 0 {
		data = stmt.conn.c.startEphemeralPacket(minPktLen)
	} else {
		data = stmt.conn.c.startEphemeralPacket(mysql.MaxPacketSize)
		// In this case the len(data) == cap(data) which is used to optimize the flow below.
	}
	defer bc.c.recycleWritePacket()

	// command [1 byte]
	data[4] = mysql.ComStmtExecute

	// statement_id [4 bytes]
	data[5] = byte(stmt.id)
	data[6] = byte(stmt.id >> 8)
	data[7] = byte(stmt.id >> 16)
	data[8] = byte(stmt.id >> 24)

	// flags (0: CURSOR_TYPE_NO_CURSOR) [1 byte]
	data[9] = 0x00

	// iteration_count (uint32(1)) [4 bytes]
	data[10] = 0x01
	data[11] = 0x00
	data[12] = 0x00
	data[13] = 0x00

	if len(args) > 0 {
		pos := minPktLen

		var nullMask []byte
		if maskLen, typesLen := (len(args)+7)/8, 1+2*len(args); pos+maskLen+typesLen >= cap(data) {
			// buffer has to be extended but we don't know by how much so
			// we depend on append after all data with known sizes fit.
			// We stop at that because we deal with a lot of columns here
			// which makes the required allocation size hard to guess.
			tmp := make([]byte, pos+maskLen+typesLen)
			copy(tmp[:pos], data[:pos])
			data = tmp
			nullMask = data[pos : pos+maskLen]
			// No need to clean nullMask as make ensures that.
			pos += maskLen
		} else {
			nullMask = data[pos : pos+maskLen]
			for i := range nullMask {
				nullMask[i] = 0
			}
			pos += maskLen
		}

		// newParameterBoundFlag 1 [1 byte]
		data[pos] = 0x01
		pos++

		// type of each parameter [len(args)*2 bytes]
		paramTypes := data[pos:]
		pos += len(args) * 2

		// value of each parameter [n bytes]
		paramValues := data[pos:pos]
		valuesCap := cap(paramValues)

		for i, arg := range args {
			// build NULL-bitmap
			if arg == nil {
				nullMask[i/8] |= 1 << (uint(i) & 7)
				paramTypes[i+i] = byte(mysql.FieldTypeNULL)
				paramTypes[i+i+1] = 0x00
				continue
			}

			// cache types and values
			switch arg.Family() {
			case proto.ValueFamilySign:
				v, _ := arg.Int64()

				paramTypes[i+i] = byte(mysql.FieldTypeLongLong)
				paramTypes[i+i+1] = 0x00

				if cap(paramValues)-len(paramValues)-8 >= 0 {
					paramValues = paramValues[:len(paramValues)+8]
					binary.LittleEndian.PutUint64(
						paramValues[len(paramValues)-8:],
						uint64(v),
					)
				} else {
					paramValues = append(paramValues,
						uint64ToBytes(uint64(v))...,
					)
				}

			case proto.ValueFamilyUnsigned:
				v, _ := arg.Uint64()

				paramTypes[i+i] = byte(mysql.FieldTypeLongLong)
				paramTypes[i+i+1] = 0x80 // type is unsigned

				if cap(paramValues)-len(paramValues)-8 >= 0 {
					paramValues = paramValues[:len(paramValues)+8]
					binary.LittleEndian.PutUint64(
						paramValues[len(paramValues)-8:],
						v,
					)
				} else {
					paramValues = append(paramValues,
						uint64ToBytes(v)...,
					)
				}

			case proto.ValueFamilyFloat:
				v, _ := arg.Float64()

				paramTypes[i+i] = byte(mysql.FieldTypeDouble)
				paramTypes[i+i+1] = 0x00

				if cap(paramValues)-len(paramValues)-8 >= 0 {
					paramValues = paramValues[:len(paramValues)+8]
					binary.LittleEndian.PutUint64(
						paramValues[len(paramValues)-8:],
						math.Float64bits(v),
					)
				} else {
					paramValues = append(paramValues,
						uint64ToBytes(math.Float64bits(v))...,
					)
				}

			case proto.ValueFamilyBool:
				v, _ := arg.Bool()

				paramTypes[i+i] = byte(mysql.FieldTypeTiny)
				paramTypes[i+i+1] = 0x00

				if v {
					paramValues = append(paramValues, 0x01)
				} else {
					paramValues = append(paramValues, 0x00)
				}

			//case []byte:
			//	// Common case (non-nil value) first
			//	if v != nil {
			//		paramTypes[i+i] = byte(mysql.FieldTypeString)
			//		paramTypes[i+i+1] = 0x00
			//
			//		if len(v) < longDataSize {
			//			paramValues = appendLengthEncodedInteger(paramValues,
			//				uint64(len(v)),
			//			)
			//			paramValues = append(paramValues, v...)
			//		} else {
			//			if err := stmt.writeCommandLongData(i, v); err != nil {
			//				return err
			//			}
			//		}
			//		continue
			//	}
			//
			//	// Handle []byte(nil) as a NULL value
			//	nullMask[i/8] |= 1 << (uint(i) & 7)
			//	paramTypes[i+i] = byte(mysql.FieldTypeNULL)
			//	paramTypes[i+i+1] = 0x00

			case proto.ValueFamilyString, proto.ValueFamilyDecimal:
				v := arg.String()

				paramTypes[i+i] = byte(mysql.FieldTypeString)
				paramTypes[i+i+1] = 0x00

				if len(v) < longDataSize {
					paramValues = appendLengthEncodedInteger(paramValues,
						uint64(len(v)),
					)
					paramValues = append(paramValues, v...)
				} else {
					if err := stmt.writeCommandLongData(i, []byte(v)); err != nil {
						return err
					}
				}

			case proto.ValueFamilyTime:
				v, _ := arg.Time()

				paramTypes[i+i] = byte(mysql.FieldTypeString)
				paramTypes[i+i+1] = 0x00

				var a [64]byte
				b := a[:0]

				if v.IsZero() {
					b = append(b, "0000-00-00"...)
				} else {
					b, err = AppendDateTime(b, v.In(bc.conf.Loc))
					if err != nil {
						return err
					}
				}

				paramValues = appendLengthEncodedInteger(paramValues,
					uint64(len(b)),
				)
				paramValues = append(paramValues, b...)

			default:
				return perrors.Errorf("cannot convert type: %T", arg.Family())
			}
		}

		// Check if param values exceeded the available buffer
		// In that case we must build the data packet with the new values buffer
		if valuesCap != cap(paramValues) {
			data = append(data[:pos], paramValues...)
		}

		pos += len(paramValues)
		data = data[:pos]
	}

	return bc.c.writePacket(data[4:])
}

func (stmt *BackendStatement) execArgs(args []proto.Value) (proto.Result, error) {
	err := stmt.writeExecutePacket(args)
	if err != nil {
		return nil, err
	}

	//if colNumber > 0 {
	//	// columns
	//	if err = stmt.conn.DrainResults(); err != nil {
	//		return nil, 0, err
	//	}
	//	// rows
	//	if err = stmt.conn.DrainResults(); err != nil {
	//		return nil, 0, err
	//	}
	//}

	return stmt.conn.ReadQueryRow(), nil
}

func (stmt *BackendStatement) queryArgs(args []proto.Value) (proto.Result, error) {
	err := stmt.writeExecutePacket(args)
	if err != nil {
		return nil, err
	}

	res := stmt.conn.ReadQueryRow()
	res.setWantFields(true)
	res.setBinaryProtocol()

	return res, nil
}

func (stmt *BackendStatement) exec(args []byte) (proto.Result, error) {
	args[1] = byte(stmt.id)
	args[2] = byte(stmt.id >> 8)
	args[3] = byte(stmt.id >> 16)
	args[4] = byte(stmt.id >> 24)

	// Reset packet-sequence
	stmt.conn.c.sequence = 0

	err := stmt.conn.c.writePacket(args)
	if err != nil {
		return nil, err
	}
	stmt.conn.c.recycleWritePacket()

	res := stmt.conn.ReadQueryRow()

	//if colNumber > 0 {
	//	// columns
	//	if err = stmt.conn.DrainResults(); err != nil {
	//		return nil, 0, err
	//	}
	//	// rows
	//	if err = stmt.conn.DrainResults(); err != nil {
	//		return nil, 0, err
	//	}
	//}

	res.setBinaryProtocol()
	res.setWantFields(true)
	return res, nil
}

func (stmt *BackendStatement) query(args []byte) (proto.Result, error) {
	args[1] = byte(stmt.id)
	args[2] = byte(stmt.id >> 8)
	args[3] = byte(stmt.id >> 16)
	args[4] = byte(stmt.id >> 24)

	// Reset packet-sequence
	stmt.conn.c.sequence = 0

	err := stmt.conn.c.writePacket(args)
	if err != nil {
		return nil, err
	}

	return stmt.conn.ReadQueryResult(true), nil
}