Repository: mailru/go-clickhouse
Pull request #149 (2021-12-07 09:30:22)

implement DateTime64
Added implementation of DateTime64 that fixes issue #133 


Patch:
diff --git a/dataparser.go b/dataparser.go
index cb7ba62..ea7b205 100644
--- a/dataparser.go
+++ b/dataparser.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"reflect"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -136,10 +137,10 @@ func (p *nullableParser) Parse(s io.RuneScanner) (driver.Value, error) {
 		isNotString := false
 		for {
 			r, size, err := s.ReadRune()
-			if err == io.EOF && size == 0 { 
+			if err == io.EOF && size == 0 {
 				break
 			}
-			
+
 			if err != nil {
 				return nil, fmt.Errorf("error: %v", err)
 			}
@@ -243,18 +244,31 @@ func (p *stringParser) Type() reflect.Type {
 }
 
 type dateTimeParser struct {
-	unquote  bool
-	format   string
-	location *time.Location
+	unquote   bool
+	format    string
+	location  *time.Location
+	precision int
 }
 
 func (p *dateTimeParser) Parse(s io.RuneScanner) (driver.Value, error) {
-	str, err := readString(s, len(p.format), p.unquote)
+	l := len(p.format)
+	if p.precision > 0 {
+		if i := strings.Index(p.format, "."); i >= 0 {
+			l = i + p.precision + 1
+		}
+	}
+
+	str, err := readString(s, l, p.unquote)
 	if err != nil {
 		return nil, fmt.Errorf("failed to read the string representation of date or datetime: %v", err)
 	}
 
-	if str == "0000-00-00" || str == "0000-00-00 00:00:00" {
+	test := str
+	if i := strings.Index(str, " "); i >= 0 {
+		test = str[:i]
+	}
+
+	if test == "0000-00-00" {
 		return time.Time{}, nil
 	}
 
@@ -383,11 +397,12 @@ func (p *simpleAggregateFunctionParser) Parse(s io.RuneScanner) (driver.Value, e
 	return p.arg.Parse(s)
 }
 
-func newDateTimeParser(format string, loc *time.Location, unquote bool) (DataParser, error) {
+func newDateTimeParser(format string, loc *time.Location, precision int, unquote bool) (DataParser, error) {
 	return &dateTimeParser{
-		unquote:  unquote,
-		format:   format,
-		location: loc,
+		unquote:   unquote,
+		format:    format,
+		location:  loc,
+		precision: precision,
 	}, nil
 }
 
@@ -538,7 +553,7 @@ func newDataParser(t *TypeDesc, unquote bool, opt *DataParserOptions) (DataParse
 		if opt != nil && opt.Location != nil {
 			loc = opt.Location
 		}
-		return newDateTimeParser(dateFormat, loc, unquote)
+		return newDateTimeParser(dateFormat, loc, 0, unquote)
 	case "DateTime":
 		loc := time.UTC
 		if (opt == nil || opt.Location == nil || opt.UseDBLocation) && len(t.Args) > 0 {
@@ -550,7 +565,33 @@ func newDataParser(t *TypeDesc, unquote bool, opt *DataParserOptions) (DataParse
 		} else if opt != nil && opt.Location != nil {
 			loc = opt.Location
 		}
-		return newDateTimeParser(timeFormat, loc, unquote)
+		return newDateTimeParser(timeFormat, loc, 0, unquote)
+	case "DateTime64":
+		if len(t.Args) < 1 {
+			return nil, fmt.Errorf("tick size not specified for DateTime64")
+		}
+
+		loc := time.UTC
+		if (opt == nil || opt.Location == nil || opt.UseDBLocation) && len(t.Args) > 1 {
+			var err error
+			loc, err = time.LoadLocation(t.Args[1].Name)
+			if err != nil {
+				return nil, err
+			}
+		} else if opt != nil && opt.Location != nil {
+			loc = opt.Location
+		}
+
+		precision, err := strconv.Atoi(t.Args[0].Name)
+		if err != nil {
+			return nil, err
+		}
+
+		if precision < 0 {
+			return nil, fmt.Errorf("malformed tick size specified for DateTime64")
+		}
+
+		return newDateTimeParser(dateTime64Format, loc, precision, unquote)
 	case "UInt8":
 		return &intParser{false, 8}, nil
 	case "UInt16":
diff --git a/helpers.go b/helpers.go
index fdfee89..6422bc5 100644
--- a/helpers.go
+++ b/helpers.go
@@ -8,9 +8,10 @@ import (
 )
 
 var (
-	escaper    = strings.NewReplacer(`\`, `\\`, `'`, `\'`)
-	dateFormat = "2006-01-02"
-	timeFormat = "2006-01-02 15:04:05"
+	escaper          = strings.NewReplacer(`\`, `\\`, `'`, `\'`)
+	dateFormat       = "2006-01-02"
+	timeFormat       = "2006-01-02 15:04:05"
+	dateTime64Format = "2006-01-02 15:04:05.999999999"
 )
 
 func escape(s string) string {


Test patch:
diff --git a/dataparser_test.go b/dataparser_test.go
index 886a91b..9fbb58a 100644
--- a/dataparser_test.go
+++ b/dataparser_test.go
@@ -29,6 +29,7 @@ func TestParseData(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to load time zone Europe/Moscow: %v", err)
 	}
+	utc := time.UTC
 
 	testCases := []*testCase{
 		{
@@ -199,6 +200,68 @@ func TestParseData(t *testing.T) {
 			output:        time.Time{},
 			failParseData: true,
 		},
+		{
+			name:      "DateTime64: correct",
+			inputtype: "DateTime64(3, 'UTC')",
+			inputdata: "2021-12-07 00:00:00.000",
+			output:    time.Date(2021, 12, 7, 0, 0, 0, 0, utc),
+		},
+		{
+			name:      "DateTime64: in Moscow with precision 6",
+			inputtype: "DateTime64(6, 'Europe/Moscow')",
+			inputdata: "2018-01-02 12:34:56.789012",
+			output:    time.Date(2018, 1, 2, 12, 34, 56, 789012000, moscow),
+		},
+		{
+			name:      "DateTime64: in Moscow with precision 3",
+			inputtype: "DateTime64(3, 'Europe/Moscow')",
+			inputdata: "2018-01-02 12:34:56.789012",
+			output:    time.Date(2018, 1, 2, 12, 34, 56, 789000000, moscow),
+		},
+		{
+			name:          "DateTime64: incorrect",
+			inputtype:     "DateTime64(4, 'Europe/Moscow')",
+			inputdata:     "2018-01-02 12:34:56.789",
+			failParseData: true,
+		},
+		{
+			name:      "DateTime64: zero",
+			inputtype: "DateTime64(3, 'Europe/Moscow')",
+			inputdata: "0000-00-00 00:00:00.000",
+			output:    time.Time{},
+		},
+		{
+			name:          "DateTime64: no tick size",
+			inputtype:     "DateTime64",
+			inputdata:     "2018-01-02 12:34:56.789",
+			failNewParser: true,
+		},
+		{
+			name:          "DateTime64: no tick size",
+			inputtype:     "DateTime64(3, 'Nowhere')",
+			inputdata:     "2018-01-02 12:34:56.789",
+			failNewParser: true,
+		},
+		{
+			name:      "DateTime64: with argument and location, prefer argument",
+			inputtype: "DateTime64(3, 'America/Los_Angeles')",
+			inputdata: "2018-01-02 12:34:56.789",
+			inputopt: &DataParserOptions{
+				Location:      moscow,
+				UseDBLocation: true,
+			},
+			output: time.Date(2018, 1, 2, 12, 34, 56, 789000000, losAngeles),
+		},
+		{
+			name:      "DateTime64: with argument and location, prefer argument",
+			inputtype: "DateTime64(3)",
+			inputdata: "2018-01-02 12:34:56.789",
+			inputopt: &DataParserOptions{
+				Location:      moscow,
+				UseDBLocation: true,
+			},
+			output: time.Date(2018, 1, 2, 12, 34, 56, 789000000, moscow),
+		},
 		{
 			name:      "tuple",
 			inputtype: "Tuple(String, Float64, Int16, UInt16, Int64)",