Add parsing of query response columns

kostyantyn · kostyantyn · commit 316d3cad1498 · 2018-02-25T21:42:52.000+01:00
diff --git a/com_query_response.go b/com_query_response.go
@@ -1,9 +1,31 @@
 package mysqlproto
 
+import (
+	"errors"
+	"fmt"
+)
+
 type ResultSet struct {
+	Columns []Column
+
 	conn Conn
 }
 
+// https://dev.mysql.com/doc/internals/en/com-query-response.html#column-definition
+type Column struct {
+	Catalog      string
+	Schema       string
+	Table        string
+	OrgTable     string
+	Name         string
+	OrgName      string
+	CharacterSet uint16
+	ColumnLength uint64
+	ColumnType   Type
+	Flags        uint16
+	Decimals     byte
+}
+
 func (r ResultSet) Row() ([]byte, error) {
 	packet, err := r.conn.NextPacket()
 	if err != nil {
@@ -17,28 +39,79 @@ func (r ResultSet) Row() ([]byte, error) {
 	return packet.Payload, nil
 }
 
+// https://dev.mysql.com/doc/internals/en/com-query-response.html
 func ComQueryResponse(conn Conn) (ResultSet, error) {
-	packet, err := conn.NextPacket()
-	if err != nil {
-		return ResultSet{}, err
+	read := func() ([]byte, error) {
+		packet, err := conn.NextPacket()
+		if err != nil {
+			return nil, err
+		}
+		if len(packet.Payload) == 0 {
+			return nil, errors.New("mysqlproto: empty payload")
+		}
+		if packet.Payload[0] == ERR_PACKET {
+			return nil, parseError(packet.Payload, conn.CapabilityFlags)
+		}
+		return packet.Payload, nil
 	}
 
-	if packet.Payload[0] == ERR_PACKET {
-		return ResultSet{}, parseError(packet.Payload, conn.CapabilityFlags)
+	payload, err := read()
+	if err != nil {
+		return ResultSet{}, err
 	}
 
-	columns, _, _ := lenDecInt(packet.Payload)
-	skip := int(columns) + 1 // skip column definition + first EOF
-	for i := 0; i < skip; i++ {
-		packet, err := conn.NextPacket()
+	colCount, _, _ := lenDecInt(payload)
+	columns := make([]Column, int(colCount))
+	for i := 0; i < int(colCount); i++ {
+		payload, err := read()
 		if err != nil {
 			return ResultSet{}, err
 		}
 
-		if packet.Payload[0] == ERR_PACKET {
-			return ResultSet{}, parseError(packet.Payload, conn.CapabilityFlags)
+		column := Column{}
+		bytes, offset, _ := ReadRowValue(payload, 0)
+		column.Catalog = string(bytes)
+
+		bytes, offset, _ = ReadRowValue(payload, offset)
+		column.Schema = string(bytes)
+
+		bytes, offset, _ = ReadRowValue(payload, offset)
+		column.Table = string(bytes)
+
+		bytes, offset, _ = ReadRowValue(payload, offset)
+		column.OrgTable = string(bytes)
+
+		bytes, offset, _ = ReadRowValue(payload, offset)
+		column.Name = string(bytes)
+
+		bytes, offset, _ = ReadRowValue(payload, offset)
+		column.OrgName = string(bytes)
+
+		bytes, _, _ = ReadRowValue(payload, offset)
+		if len(bytes) < 10 {
+			return ResultSet{}, fmt.Errorf("mysqlproto: invalid column payload: %x", bytes)
 		}
+
+		column.CharacterSet = uint16(bytes[0]) | uint16(bytes[1])<<8
+		column.ColumnLength = uint64(bytes[2]) | uint64(bytes[3])<<8 | uint64(bytes[4])<<16 | uint64(bytes[5])<<32
+		column.ColumnType = Type(bytes[6])
+		column.Flags = uint16(bytes[7]) | uint16(bytes[9])<<8
+		column.Decimals = bytes[9]
+
+		columns[i] = column
 	}
 
-	return ResultSet{conn}, nil
+	payload, err = read()
+	if err != nil {
+		return ResultSet{}, err
+	}
+	if payload[0] != EOF_PACKET {
+		return ResultSet{}, parseError(payload, conn.CapabilityFlags)
+	}
+
+	rs := ResultSet{
+		Columns: columns,
+		conn:    conn,
+	}
+	return rs, nil
 }
diff --git a/com_query_response_test.go b/com_query_response_test.go
@@ -0,0 +1,105 @@
+package mysqlproto
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestComQueryResponseColumnReader(t *testing.T) {
+	buf := newBuffer([]byte{
+		// DB name "test"
+		// table name "people" AS "p"
+
+		// total records
+		0x01, 0x00, 0x00, 0x01, 0x05,
+
+		// id INT
+		0x25, 0x00, 0x00, 0x02, 0x03, 0x64, 0x65, 0x66, 0x04, 0x74, 0x65, 0x73, 0x74, 0x01, 0x70, 0x06, 0x70, 0x65, 0x6f, 0x70, 0x6c, 0x65, 0x02, 0x69, 0x64, 0x02, 0x69, 0x64, 0x0c, 0x3f, 0x00, 0x0b, 0x00, 0x00, 0x00, 0x03, 0x03, 0x42, 0x00, 0x00, 0x00,
+
+		// firstname VARCHAR(255) AS name
+		0x2e, 0x00, 0x00, 0x03, 0x03, 0x64, 0x65, 0x66, 0x04, 0x74, 0x65, 0x73, 0x74, 0x01, 0x70, 0x06, 0x70, 0x65, 0x6f, 0x70, 0x6c, 0x65, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x09, 0x66, 0x69, 0x72, 0x73, 0x74, 0x6e, 0x61, 0x6d, 0x65, 0x0c, 0x21, 0x00, 0xfd, 0x02, 0x00, 0x00, 0xfd, 0x00, 0x00, 0x00, 0x00, 0x00,
+
+		// married TINYINT
+		0x2f, 0x00, 0x00, 0x04, 0x03, 0x64, 0x65, 0x66, 0x04, 0x74, 0x65, 0x73, 0x74, 0x01, 0x70, 0x06, 0x70, 0x65, 0x6f, 0x70, 0x6c, 0x65, 0x07, 0x6d, 0x61, 0x72, 0x72, 0x69, 0x65, 0x64, 0x07, 0x6d, 0x61, 0x72, 0x72, 0x69, 0x65, 0x64, 0x0c, 0x3f, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
+
+		// score DECIMAL(6,2)
+		0x2b, 0x00, 0x00, 0x05, 0x03, 0x64, 0x65, 0x66, 0x04, 0x74, 0x65, 0x73, 0x74, 0x01, 0x70, 0x06, 0x70, 0x65, 0x6f, 0x70, 0x6c, 0x65, 0x05, 0x73, 0x63, 0x6f, 0x72, 0x65, 0x05, 0x73, 0x63, 0x6f, 0x72, 0x65, 0x0c, 0x3f, 0x00, 0x08, 0x00, 0x00, 0x00, 0xf6, 0x00, 0x00, 0x02, 0x00, 0x00,
+
+		// note TEXT
+		0x29, 0x00, 0x00, 0x06, 0x03, 0x64, 0x65, 0x66, 0x04, 0x74, 0x65, 0x73, 0x74, 0x01, 0x70, 0x06, 0x70, 0x65, 0x6f, 0x70, 0x6c, 0x65, 0x04, 0x6e, 0x6f, 0x74, 0x65, 0x04, 0x6e, 0x6f, 0x74, 0x65, 0x0c, 0x21, 0x00, 0xfd, 0xff, 0x02, 0x00, 0xfc, 0x10, 0x00, 0x00, 0x00, 0x00,
+
+		// EOF
+		0x05, 0x00, 0x00, 0x07, 0xfe, 0x00, 0x00, 0x22, 0x00,
+	})
+
+	conn := Conn{Stream: &Stream{stream: buf}}
+	rs, err := ComQueryResponse(conn)
+	assert.NoError(t, err)
+	assert.Len(t, rs.Columns, 5)
+
+	id := rs.Columns[0]
+	assert.Equal(t, id.Catalog, "def")
+	assert.Equal(t, id.Schema, "test")
+	assert.Equal(t, id.Table, "p")
+	assert.Equal(t, id.OrgTable, "people")
+	assert.Equal(t, id.Name, "id")
+	assert.Equal(t, id.OrgName, "id")
+	assert.Equal(t, id.CharacterSet, uint16(63))
+	assert.Equal(t, id.ColumnLength, uint64(11))
+	assert.Equal(t, id.ColumnType.String(), "LONG")
+	assert.Equal(t, id.Flags, uint16(3))
+	assert.Equal(t, id.Decimals, uint8(0))
+
+	name := rs.Columns[1]
+	assert.Equal(t, name.Catalog, "def")
+	assert.Equal(t, name.Schema, "test")
+	assert.Equal(t, name.Table, "p")
+	assert.Equal(t, name.OrgTable, "people")
+	assert.Equal(t, name.Name, "name")
+	assert.Equal(t, name.OrgName, "firstname")
+	assert.Equal(t, name.CharacterSet, uint16(33))
+	assert.Equal(t, name.ColumnLength, uint64(765))
+	assert.Equal(t, name.ColumnType.String(), "VAR_STRING")
+	assert.Equal(t, name.Flags, uint16(0))
+	assert.Equal(t, name.Decimals, uint8(0))
+
+	married := rs.Columns[2]
+	assert.Equal(t, married.Catalog, "def")
+	assert.Equal(t, married.Schema, "test")
+	assert.Equal(t, married.Table, "p")
+	assert.Equal(t, married.OrgTable, "people")
+	assert.Equal(t, married.Name, "married")
+	assert.Equal(t, married.OrgName, "married")
+	assert.Equal(t, married.CharacterSet, uint16(63))
+	assert.Equal(t, married.ColumnLength, uint64(4))
+	assert.Equal(t, married.ColumnType.String(), "TINY")
+	assert.Equal(t, married.Flags, uint16(0))
+	assert.Equal(t, married.Decimals, uint8(0))
+
+	score := rs.Columns[3]
+	assert.Equal(t, score.Catalog, "def")
+	assert.Equal(t, score.Schema, "test")
+	assert.Equal(t, score.Table, "p")
+	assert.Equal(t, score.OrgTable, "people")
+	assert.Equal(t, score.Name, "score")
+	assert.Equal(t, score.OrgName, "score")
+	assert.Equal(t, score.CharacterSet, uint16(63))
+	assert.Equal(t, score.ColumnLength, uint64(8))
+	assert.Equal(t, score.ColumnType.String(), "NEWDECIMAL")
+	assert.Equal(t, score.Flags, uint16(512))
+	assert.Equal(t, score.Decimals, uint8(2))
+
+	note := rs.Columns[4]
+	assert.Equal(t, note.Catalog, "def")
+	assert.Equal(t, note.Schema, "test")
+	assert.Equal(t, note.Table, "p")
+	assert.Equal(t, note.OrgTable, "people")
+	assert.Equal(t, note.Name, "note")
+	assert.Equal(t, note.OrgName, "note")
+	assert.Equal(t, note.CharacterSet, uint16(33))
+	assert.Equal(t, note.ColumnLength, uint64(196605))
+	assert.Equal(t, note.ColumnType.String(), "BLOB")
+	assert.Equal(t, note.Flags, uint16(16))
+	assert.Equal(t, note.Decimals, uint8(0))
+}
diff --git a/packet_test.go b/packet_test.go
@@ -6,7 +6,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestParseOKPacketInvalidPayout(t *testing.T) {
+func TestParseOKPacketInvalidPayload(t *testing.T) {
 	data := []byte{0xff}
 	_, err := ParseOKPacket(data, 0)
 	assert.Equal(t, err.Error(), "mysqlproto: invalid OK_PACKET payload: ff")
diff --git a/types.go b/types.go
@@ -0,0 +1,104 @@
+package mysqlproto
+
+// https://dev.mysql.com/doc/internals/en/com-query-response.html#column-type
+type Type byte
+
+const (
+	TypeDecimal    Type = 0x00
+	TypeTiny       Type = 0x01
+	TypeShort      Type = 0x02
+	TypeLong       Type = 0x03
+	TypeFloat      Type = 0x04
+	TypeDouble     Type = 0x05
+	TypeNULL       Type = 0x06
+	TypeTimestamp  Type = 0x07
+	TypeLongLong   Type = 0x08
+	TypeInt24      Type = 0x09
+	TypeDate       Type = 0x0a
+	TypeTime       Type = 0x0b
+	TypeDateTime   Type = 0x0c
+	TypeYear       Type = 0x0d
+	TypeNewDate    Type = 0x0e
+	TypeVarchar    Type = 0x0f
+	TypeBit        Type = 0x10
+	TypeTimestamp2 Type = 0x11
+	TypeDateTime2  Type = 0x12
+	TypeTime2      Type = 0x13
+	TypeNewDecimal Type = 0xf6
+	TypeEnum       Type = 0xf7
+	TypeSet        Type = 0xf8
+	TypeTinyBLOB   Type = 0xf9
+	TypeMediumBLOB Type = 0xfa
+	TypeLongBLOB   Type = 0xfb
+	TypeBLOB       Type = 0xfc
+	TypeVarString  Type = 0xfd
+	TypeString     Type = 0xfe
+	TypeGEOMETRY   Type = 0xff
+)
+
+func (t Type) String() string {
+	switch t {
+	case TypeDecimal:
+		return "DECIMAL"
+	case TypeTiny:
+		return "TINY"
+	case TypeShort:
+		return "SHORT"
+	case TypeLong:
+		return "LONG"
+	case TypeFloat:
+		return "FLOAT"
+	case TypeDouble:
+		return "DOUBLE"
+	case TypeNULL:
+		return "NULL"
+	case TypeTimestamp:
+		return "TIMESTAMP"
+	case TypeLongLong:
+		return "LONGLONG"
+	case TypeInt24:
+		return "INT24"
+	case TypeDate:
+		return "DATE"
+	case TypeTime:
+		return "TIME"
+	case TypeDateTime:
+		return "DATETIME"
+	case TypeYear:
+		return "YEAR"
+	case TypeNewDate:
+		return "NEWDATE"
+	case TypeVarchar:
+		return "VARCHAR"
+	case TypeBit:
+		return "BIT"
+	case TypeTimestamp2:
+		return "TIMESTAMP2"
+	case TypeDateTime2:
+		return "DATETIME2"
+	case TypeTime2:
+		return "TIME2"
+	case TypeNewDecimal:
+		return "NEWDECIMAL"
+	case TypeEnum:
+		return "ENUM"
+	case TypeSet:
+		return "SET"
+	case TypeTinyBLOB:
+		return "TINY_BLOB"
+	case TypeMediumBLOB:
+		return "MEDIUM_BLOB"
+	case TypeLongBLOB:
+		return "LONG_BLOB"
+	case TypeBLOB:
+		return "BLOB"
+	case TypeVarString:
+		return "VAR_STRING"
+	case TypeString:
+		return "STRING"
+	case TypeGEOMETRY:
+		return "GEOMETRY"
+	default:
+		return "UNKNOWN"
+	}
+}
diff --git a/types_test.go b/types_test.go
@@ -0,0 +1,51 @@
+package mysqlproto
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestTypeString(t *testing.T) {
+	testCases := []struct {
+		typ Type
+		hex byte
+		str string
+	}{
+		{typ: TypeDecimal, hex: 0x00, str: "DECIMAL"},
+		{typ: TypeTiny, hex: 0x01, str: "TINY"},
+		{typ: TypeShort, hex: 0x02, str: "SHORT"},
+		{typ: TypeLong, hex: 0x03, str: "LONG"},
+		{typ: TypeFloat, hex: 0x04, str: "FLOAT"},
+		{typ: TypeDouble, hex: 0x05, str: "DOUBLE"},
+		{typ: TypeNULL, hex: 0x06, str: "NULL"},
+		{typ: TypeTimestamp, hex: 0x07, str: "TIMESTAMP"},
+		{typ: TypeLongLong, hex: 0x08, str: "LONGLONG"},
+		{typ: TypeInt24, hex: 0x09, str: "INT24"},
+		{typ: TypeDate, hex: 0x0a, str: "DATE"},
+		{typ: TypeTime, hex: 0x0b, str: "TIME"},
+		{typ: TypeDateTime, hex: 0x0c, str: "DATETIME"},
+		{typ: TypeYear, hex: 0x0d, str: "YEAR"},
+		{typ: TypeNewDate, hex: 0x0e, str: "NEWDATE"},
+		{typ: TypeVarchar, hex: 0x0f, str: "VARCHAR"},
+		{typ: TypeBit, hex: 0x10, str: "BIT"},
+		{typ: TypeTimestamp2, hex: 0x11, str: "TIMESTAMP2"},
+		{typ: TypeDateTime2, hex: 0x12, str: "DATETIME2"},
+		{typ: TypeTime2, hex: 0x13, str: "TIME2"},
+		{typ: TypeNewDecimal, hex: 0xf6, str: "NEWDECIMAL"},
+		{typ: TypeEnum, hex: 0xf7, str: "ENUM"},
+		{typ: TypeSet, hex: 0xf8, str: "SET"},
+		{typ: TypeTinyBLOB, hex: 0xf9, str: "TINY_BLOB"},
+		{typ: TypeMediumBLOB, hex: 0xfa, str: "MEDIUM_BLOB"},
+		{typ: TypeLongBLOB, hex: 0xfb, str: "LONG_BLOB"},
+		{typ: TypeBLOB, hex: 0xfc, str: "BLOB"},
+		{typ: TypeVarString, hex: 0xfd, str: "VAR_STRING"},
+		{typ: TypeString, hex: 0xfe, str: "STRING"},
+		{typ: TypeGEOMETRY, hex: 0xff, str: "GEOMETRY"},
+	}
+
+	for i, tc := range testCases {
+		assert.Equal(t, byte(tc.typ), tc.hex, "test case %v", i)
+		assert.Equal(t, tc.typ.String(), tc.str, "test case %v", i)
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@ import (`
`6`	`6`	`"github.com/stretchr/testify/assert"`
`7`	`7`	`)`
`8`	`8`
`9`		`-func TestParseOKPacketInvalidPayout(t *testing.T) {`
	`9`	`+func TestParseOKPacketInvalidPayload(t *testing.T) {`
`10`	`10`	`data := []byte{0xff}`
`11`	`11`	`_, err := ParseOKPacket(data, 0)`
`12`	`12`	`assert.Equal(t, err.Error(), "mysqlproto: invalid OK_PACKET payload: ff")`