Repository: mailru/go-clickhouse
Pull request #78 (2019-07-05 10:13:39)

Add DateTime default location and override options
This PR fixes the FIXME comment on supporting custom locations on DateTime and Date fields 


Patch:
diff --git a/dataparser.go b/dataparser.go
index 07e164a..330ea15 100644
--- a/dataparser.go
+++ b/dataparser.go
@@ -242,11 +242,7 @@ func (p *lowCardinalityParser) Parse(s io.RuneScanner) (driver.Value, error) {
 	return p.arg.Parse(s)
 }
 
-func newDateTimeParser(format, locname string, unquote bool) (DataParser, error) {
-	loc, err := time.LoadLocation(locname)
-	if err != nil {
-		return nil, err
-	}
+func newDateTimeParser(format string, loc *time.Location, unquote bool) (DataParser, error) {
 	return &dateTimeParser{
 		unquote:  unquote,
 		format:   format,
@@ -368,28 +364,45 @@ func (p *nothingParser) Type() reflect.Type {
 	return reflectTypeEmptyStruct
 }
 
+// DataParserOptions describes DataParser options.
+// Ex.: Fields Location and UseDBLocation specify timezone options.
+type DataParserOptions struct {
+	// Location describes default location for DateTime and Date field without Timezone argument.
+	Location *time.Location
+	// UseDBLocation if false: always use Location, ignore DateTime argument.
+	UseDBLocation bool
+}
+
 // NewDataParser creates a new DataParser based on the
 // given TypeDesc.
-func NewDataParser(t *TypeDesc) (DataParser, error) {
-	return newDataParser(t, false)
+func NewDataParser(t *TypeDesc, opt *DataParserOptions) (DataParser, error) {
+	return newDataParser(t, false, opt)
 }
 
-func newDataParser(t *TypeDesc, unquote bool) (DataParser, error) {
+func newDataParser(t *TypeDesc, unquote bool, opt *DataParserOptions) (DataParser, error) {
 	switch t.Name {
 	case "Nothing":
 		return &nothingParser{}, nil
 	case "Nullable":
 		return nil, fmt.Errorf("Nullable types are not supported")
 	case "Date":
-		// FIXME: support custom default/override location
-		return newDateTimeParser(dateFormat, "UTC", unquote)
+		loc := time.UTC
+		if opt != nil && opt.Location != nil {
+			loc = opt.Location
+		}
+		return newDateTimeParser(dateFormat, loc, unquote)
 	case "DateTime":
-		// FIXME: support custom default/override location
-		locname := "UTC"
-		if len(t.Args) > 0 {
-			locname = t.Args[0].Name
+		loc := time.UTC
+		if (opt == nil || opt.Location == nil || opt.UseDBLocation) && len(t.Args) > 0 {
+			var err error
+			loc, err = time.LoadLocation(t.Args[0].Name)
+			if err != nil {
+				return nil, err
+			}
+		} else if opt != nil && opt.Location != nil {
+			loc = opt.Location
 		}
-		return newDateTimeParser(timeFormat, locname, unquote)
+		return newDateTimeParser(timeFormat, loc, unquote)
 	case "UInt8":
 		return &intParser{false, 8}, nil
 	case "UInt16":
@@ -425,7 +438,7 @@ func newDataParser(t *TypeDesc, unquote bool) (DataParser, error) {
 		if len(t.Args) != 1 {
 			return nil, fmt.Errorf("element type not specified for Array")
 		}
-		subParser, err := newDataParser(t.Args[0], true)
+		subParser, err := newDataParser(t.Args[0], true, opt)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create parser for array elements: %v", err)
 		}
@@ -436,7 +449,7 @@ func newDataParser(t *TypeDesc, unquote bool) (DataParser, error) {
 		}
 		subParsers := make([]DataParser, len(t.Args), len(t.Args))
 		for i, arg := range t.Args {
-			subParser, err := newDataParser(arg, true)
+			subParser, err := newDataParser(arg, true, opt)
 			if err != nil {
 				return nil, fmt.Errorf("failed to create parser for tuple element: %v", err)
 			}
@@ -447,7 +460,7 @@ func newDataParser(t *TypeDesc, unquote bool) (DataParser, error) {
 		if len(t.Args) != 1 {
 			return nil, fmt.Errorf("element type not specified for LowCardinality")
 		}
-		subParser, err := newDataParser(t.Args[0], unquote)
+		subParser, err := newDataParser(t.Args[0], unquote, opt)
 		if err != nil {
 			return nil, fmt.Errorf("failed to create parser for LowCardinality elements: %v", err)
 		}
diff --git a/rows.go b/rows.go
index c2947a3..a29c708 100644
--- a/rows.go
+++ b/rows.go
@@ -38,7 +38,10 @@ func newTextRows(c *conn, body io.ReadCloser, location *time.Location, useDBLoca
 			return nil, err
 		}
 
-		parsers[i], err = NewDataParser(desc)
+		parsers[i], err = NewDataParser(desc, &DataParserOptions{
+			Location:      location,
+			UseDBLocation: useDBLocation,
+		})
 		if err != nil {
 			return nil, err
 		}


Test patch:
diff --git a/dataparser_test.go b/dataparser_test.go
index 9694784..bd8cfe0 100644
--- a/dataparser_test.go
+++ b/dataparser_test.go
@@ -13,6 +13,7 @@ func TestParseData(t *testing.T) {
 	type testCase struct {
 		name          string
 		inputtype     string
+		inputopt      *DataParserOptions
 		inputdata     string
 		output        interface{}
 		failParseDesc bool
@@ -22,7 +23,11 @@ func TestParseData(t *testing.T) {
 
 	losAngeles, err := time.LoadLocation("America/Los_Angeles")
 	if err != nil {
-		t.Fatalf("failed to load time zone: %v", err)
+		t.Fatalf("failed to load time zone America/Los_Angeles: %v", err)
+	}
+	moscow, err := time.LoadLocation("Europe/Moscow")
+	if err != nil {
+		t.Fatalf("failed to load time zone Europe/Moscow: %v", err)
 	}
 
 	testCases := []*testCase{
@@ -86,6 +91,15 @@ func TestParseData(t *testing.T) {
 			inputdata: "0000-00-00",
 			output:    time.Time{},
 		},
+		{
+			name:      "with special timezone",
+			inputtype: "Date",
+			inputdata: "2019-06-29",
+			inputopt: &DataParserOptions{
+				Location: losAngeles,
+			},
+			output: time.Date(2019, 6, 29, 0, 0, 0, 0, losAngeles),
+		},
 		{
 			name:      "enum",
 			inputtype: "Enum8('hello' = 1, 'world' = 2)",
@@ -99,17 +113,54 @@ func TestParseData(t *testing.T) {
 			output:    "c79a9747-7cef-4b11-8177-380f7ed462a4",
 		},
 		{
-			name:      "datetime",
+			name:      "datetime, without options and argument",
 			inputtype: "DateTime",
 			inputdata: "2018-01-02 12:34:56",
 			output:    time.Date(2018, 1, 2, 12, 34, 56, 0, time.UTC),
 		},
 		{
-			name:      "datetime in Los Angeles",
+			name:      "datetime, with argument",
 			inputtype: "DateTime('America/Los_Angeles')",
 			inputdata: "2018-01-02 12:34:56",
 			output:    time.Date(2018, 1, 2, 12, 34, 56, 0, losAngeles),
 		},
+		{
+			name:      "datetime with argument, but location nil",
+			inputtype: "DateTime('America/Los_Angeles')",
+			inputdata: "2018-01-02 12:34:56",
+			inputopt: &DataParserOptions{
+				Location: nil,
+			},
+			output: time.Date(2018, 1, 2, 12, 34, 56, 0, losAngeles),
+		},
+		{
+			name:      "datetime without argument, but use location",
+			inputtype: "DateTime",
+			inputdata: "2018-01-02 12:34:56",
+			inputopt: &DataParserOptions{
+				Location: moscow,
+			},
+			output: time.Date(2018, 1, 2, 12, 34, 56, 0, moscow),
+		},
+		{
+			name:      "datetime with argument and location, ingnore argument",
+			inputtype: "DateTime('America/Los_Angeles')",
+			inputdata: "2018-01-02 12:34:56",
+			inputopt: &DataParserOptions{
+				Location: moscow,
+			},
+			output: time.Date(2018, 1, 2, 12, 34, 56, 0, moscow),
+		},
+		{
+			name:      "datetime with argument and location, prefer argument",
+			inputtype: "DateTime('America/Los_Angeles')",
+			inputdata: "2018-01-02 12:34:56",
+			inputopt: &DataParserOptions{
+				Location:      moscow,
+				UseDBLocation: true,
+			},
+			output: time.Date(2018, 1, 2, 12, 34, 56, 0, losAngeles),
+		},
 		{
 			name:          "datetime in nowhere",
 			inputtype:     "DateTime('Nowhere')",
@@ -290,7 +341,7 @@ func TestParseData(t *testing.T) {
 				return
 			}
 
-			parser, err := newDataParser(desc, false)
+			parser, err := newDataParser(desc, false, tc.inputopt)
 			if tc.failNewParser {
 				assert.Error(tt, err)
 				return