From f87f8bcd979c18ef411685353336f479960e3fbe Mon Sep 17 00:00:00 2001
From: inoth <inoth.work@gmail.com>
Date: Wed, 25 Dec 2024 18:03:47 +0800
Subject: [PATCH] feat: Added `time.Time` support for datetime format (#957)

* feat: Added support for datetime format

* fix: handler "Error return value of `binary.Write` is not checked"

* fix: add toBinaryDateTime test

* fix: update TestToBinaryDateTime

* fix: allow zero time

* fix: code format

---------

Co-authored-by: lance6716 <lance6716@gmail.com>
---
 mysql/resultset_helper.go | 49 ++++++++++++++++++++++++++++++-
 mysql/util_test.go        | 62 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/mysql/resultset_helper.go b/mysql/resultset_helper.go
index 21a120866..21487afdd 100644
--- a/mysql/resultset_helper.go
+++ b/mysql/resultset_helper.go
@@ -1,8 +1,11 @@
 package mysql
 
 import (
+	"bytes"
+	"encoding/binary"
 	"math"
 	"strconv"
+	"time"
 
 	"github.com/pingcap/errors"
 
@@ -39,6 +42,8 @@ func FormatTextValue(value interface{}) ([]byte, error) {
 		return v, nil
 	case string:
 		return utils.StringToByteSlice(v), nil
+	case time.Time:
+		return utils.StringToByteSlice(v.Format(time.DateTime)), nil
 	case nil:
 		return nil, nil
 	default:
@@ -46,6 +51,44 @@ func FormatTextValue(value interface{}) ([]byte, error) {
 	}
 }
 
+func toBinaryDateTime(t time.Time) ([]byte, error) {
+	var buf bytes.Buffer
+
+	if t.IsZero() {
+		return nil, nil
+	}
+
+	year, month, day := t.Year(), t.Month(), t.Day()
+	hour, min, sec := t.Hour(), t.Minute(), t.Second()
+	nanosec := t.Nanosecond()
+
+	if nanosec > 0 {
+		buf.WriteByte(byte(11))
+		_ = binary.Write(&buf, binary.LittleEndian, uint16(year))
+		buf.WriteByte(byte(month))
+		buf.WriteByte(byte(day))
+		buf.WriteByte(byte(hour))
+		buf.WriteByte(byte(min))
+		buf.WriteByte(byte(sec))
+		_ = binary.Write(&buf, binary.LittleEndian, uint32(nanosec/1000))
+	} else if hour > 0 || min > 0 || sec > 0 {
+		buf.WriteByte(byte(7))
+		_ = binary.Write(&buf, binary.LittleEndian, uint16(year))
+		buf.WriteByte(byte(month))
+		buf.WriteByte(byte(day))
+		buf.WriteByte(byte(hour))
+		buf.WriteByte(byte(min))
+		buf.WriteByte(byte(sec))
+	} else {
+		buf.WriteByte(byte(4))
+		_ = binary.Write(&buf, binary.LittleEndian, uint16(year))
+		buf.WriteByte(byte(month))
+		buf.WriteByte(byte(day))
+	}
+
+	return buf.Bytes(), nil
+}
+
 func formatBinaryValue(value interface{}) ([]byte, error) {
 	switch v := value.(type) {
 	case int8:
@@ -76,6 +119,8 @@ func formatBinaryValue(value interface{}) ([]byte, error) {
 		return v, nil
 	case string:
 		return utils.StringToByteSlice(v), nil
+	case time.Time:
+		return toBinaryDateTime(v)
 	default:
 		return nil, errors.Errorf("invalid type %T", value)
 	}
@@ -91,6 +136,8 @@ func fieldType(value interface{}) (typ uint8, err error) {
 		typ = MYSQL_TYPE_DOUBLE
 	case string, []byte:
 		typ = MYSQL_TYPE_VAR_STRING
+	case time.Time:
+		typ = MYSQL_TYPE_DATETIME
 	case nil:
 		typ = MYSQL_TYPE_NULL
 	default:
@@ -110,7 +157,7 @@ func formatField(field *Field, value interface{}) error {
 	case float32, float64:
 		field.Charset = 63
 		field.Flag = BINARY_FLAG | NOT_NULL_FLAG
-	case string, []byte:
+	case string, []byte, time.Time:
 		field.Charset = 33
 	case nil:
 		field.Charset = 33
diff --git a/mysql/util_test.go b/mysql/util_test.go
index 22fbe39e4..175e907e4 100644
--- a/mysql/util_test.go
+++ b/mysql/util_test.go
@@ -2,6 +2,7 @@ package mysql
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -51,3 +52,64 @@ func TestFormatBinaryTime(t *testing.T) {
 		require.Equal(t, test.Expect, string(got), "test case %v", test.Data)
 	}
 }
+
+func TestToBinaryDateTime(t *testing.T) {
+	var (
+		DateTimeNano         = "2006-01-02 15:04:05.000000"
+		formatBinaryDateTime = func(n int, data []byte) string {
+			date, err := FormatBinaryDateTime(n, data)
+			if err != nil {
+				return ""
+			}
+			return string(date)
+		}
+	)
+
+	tests := []struct {
+		Name   string
+		Data   time.Time
+		Expect func(n int, data []byte) string
+		Error  bool
+	}{
+		{
+			Name:   "Zero time",
+			Data:   time.Time{},
+			Expect: nil,
+		},
+		{
+			Name:   "Date with nanoseconds",
+			Data:   time.Date(2023, 10, 10, 10, 10, 10, 123456000, time.UTC),
+			Expect: formatBinaryDateTime,
+		},
+		{
+			Name:   "Date with time",
+			Data:   time.Date(2023, 10, 10, 10, 10, 10, 0, time.UTC),
+			Expect: formatBinaryDateTime,
+		},
+		{
+			Name:   "Date only",
+			Data:   time.Date(2023, 10, 10, 0, 0, 0, 0, time.UTC),
+			Expect: formatBinaryDateTime,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.Name, func(t *testing.T) {
+			got, err := toBinaryDateTime(test.Data)
+			if test.Error {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+			if len(got) == 0 {
+				return
+			}
+			tmp := test.Expect(int(got[0]), got[1:])
+			if int(got[0]) < 11 {
+				require.Equal(t, tmp, test.Data.Format(time.DateTime), "test case %v", test.Data.String())
+			} else {
+				require.Equal(t, tmp, test.Data.Format(DateTimeNano), "test case %v", test.Data.String())
+			}
+		})
+	}
+}