Add support for offset timezones for time and timestamp with timezone

trinodb · Dec 19, 2022 · 565ca64 · 565ca64
1 parent a0416f7
commit 565ca64
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 1 deletion.
diff --git a/trino/trino.go b/trino/trino.go
@@ -1869,6 +1869,12 @@ var timeLayouts = []string{
 	"2006-01-02 15:04:05.000",
 }
 
+// Layout for time and timestamp WITH time zone.
+var timeLayoutsTZ = []string{
+	"15:04:05.000 -07:00",
+	"2006-01-02 15:04:05.000 -07:00",
+}
+
 func scanNullTime(v interface{}) (NullTime, error) {
 	if v == nil {
 		return NullTime{}, nil
@@ -1881,6 +1887,19 @@ func scanNullTime(v interface{}) (NullTime, error) {
 	if len(vparts) > 1 && !unicode.IsDigit(rune(vparts[len(vparts)-1][0])) {
 		return parseNullTimeWithLocation(vv)
 	}
+	// Time literals may not have spaces before the timezone.
+	if strings.ContainsRune(vv, '+') {
+		return parseNullTimeWithLocation(strings.Replace(vv, "+", " +", 1))
+	}
+	hyphenCount := strings.Count(vv, "-")
+	// We need to ensure we don't treat the hyphens in dates as the minus offset sign.
+	// So if there's only one hyphen or more than 2, we have a negative offset.
+	if hyphenCount == 1 || hyphenCount > 2 {
+		// We add a space before the last hyphen to parse properly.
+		i := strings.LastIndex(vv, "-")
+		timestamp := vv[:i] + strings.Replace(vv[i:], "-", " -", 1)
+		return parseNullTimeWithLocation(timestamp)
+	}
 	return parseNullTime(vv)
 }
 
@@ -1902,11 +1921,24 @@ func parseNullTimeWithLocation(v string) (NullTime, error) {
 		return NullTime{}, fmt.Errorf("cannot convert %v (%T) to time+zone", v, v)
 	}
 	stamp, location := v[:idx], v[idx+1:]
+	var t time.Time
+	var err error
+	// Try offset timezones.
+	if strings.HasPrefix(location, "+") || strings.HasPrefix(location, "-") {
+		for _, layout := range timeLayoutsTZ {
+			t, err = time.Parse(layout, v)
+			if err == nil {
+				return NullTime{Valid: true, Time: t}, nil
+			}
+		}
+		return NullTime{}, err
+	}
 	loc, err := time.LoadLocation(location)
+	// Not a named location.
 	if err != nil {
 		return NullTime{}, fmt.Errorf("cannot load timezone %q: %v", location, err)
 	}
-	var t time.Time
+
 	for _, layout := range timeLayouts {
 		t, err = time.ParseInLocation(layout, stamp, loc)
 		if err == nil {

diff --git a/trino/trino_test.go b/trino/trino_test.go
@@ -1010,6 +1010,30 @@ func TestTypeConversion(t *testing.T) {
 			ResponseUnmarshalledSample: "01:02:03.000 UTC",
 			ExpectedGoValue:            time.Date(0, 1, 1, 1, 2, 3, 0, utc),
 		},
+		{
+			DataType:                   "time with time zone",
+			RawType:                    "time with time zone",
+			ResponseUnmarshalledSample: "01:02:03.000 +03:00",
+			ExpectedGoValue:            time.Date(0, 1, 1, 1, 2, 3, 0, time.FixedZone("", 3*3600)),
+		},
+		{
+			DataType:                   "time with time zone",
+			RawType:                    "time with time zone",
+			ResponseUnmarshalledSample: "01:02:03.000+03:00",
+			ExpectedGoValue:            time.Date(0, 1, 1, 1, 2, 3, 0, time.FixedZone("", 3*3600)),
+		},
+		{
+			DataType:                   "time with time zone",
+			RawType:                    "time with time zone",
+			ResponseUnmarshalledSample: "01:02:03.000 -05:00",
+			ExpectedGoValue:            time.Date(0, 1, 1, 1, 2, 3, 0, time.FixedZone("", -5*3600)),
+		},
+		{
+			DataType:                   "time with time zone",
+			RawType:                    "time with time zone",
+			ResponseUnmarshalledSample: "01:02:03.000-05:00",
+			ExpectedGoValue:            time.Date(0, 1, 1, 1, 2, 3, 0, time.FixedZone("", -5*3600)),
+		},
 		{
 			DataType:                   "timestamp",
 			RawType:                    "timestamp",
@@ -1022,6 +1046,30 @@ func TestTypeConversion(t *testing.T) {
 			ResponseUnmarshalledSample: "2017-07-10 01:02:03.000 UTC",
 			ExpectedGoValue:            time.Date(2017, 7, 10, 1, 2, 3, 0, utc),
 		},
+		{
+			DataType:                   "timestamp with time zone",
+			RawType:                    "timestamp with time zone",
+			ResponseUnmarshalledSample: "2017-07-10 01:02:03.000 +03:00",
+			ExpectedGoValue:            time.Date(2017, 7, 10, 1, 2, 3, 0, time.FixedZone("", 3*3600)),
+		},
+		{
+			DataType:                   "timestamp with time zone",
+			RawType:                    "timestamp with time zone",
+			ResponseUnmarshalledSample: "2017-07-10 01:02:03.000+03:00",
+			ExpectedGoValue:            time.Date(2017, 7, 10, 1, 2, 3, 0, time.FixedZone("", 3*3600)),
+		},
+		{
+			DataType:                   "timestamp with time zone",
+			RawType:                    "timestamp with time zone",
+			ResponseUnmarshalledSample: "2017-07-10 01:02:03.000 -04:00",
+			ExpectedGoValue:            time.Date(2017, 7, 10, 1, 2, 3, 0, time.FixedZone("", -4*3600)),
+		},
+		{
+			DataType:                   "timestamp with time zone",
+			RawType:                    "timestamp with time zone",
+			ResponseUnmarshalledSample: "2017-07-10 01:02:03.000-04:00",
+			ExpectedGoValue:            time.Date(2017, 7, 10, 1, 2, 3, 0, time.FixedZone("", -4*3600)),
+		},
 		{
 			DataType: "map(varchar,varchar)",
 			RawType:  "map",