feat(sqltype/time): AddSupportedLayout

Author: morlay

pkg/sqltype/time/timestamp.go

@@ -1,6 +1,7 @@
 package time
 
 import (
+	"cmp"
 	"database/sql"
 	"database/sql/driver"
 	"encoding"
@@ -17,8 +18,7 @@ var (
 )
 
 var (
-	CST          = time.UTC
-	OutputLayout = time.RFC3339
+	CST = time.UTC
 )
 
 func init() {
@@ -28,8 +28,19 @@ func init() {
 	}
 }
 
-func SetOutput(layout string, location *time.Location) {
-	OutputLayout = layout
+var (
+	outputLayout     = time.RFC3339
+	supportedLayouts = map[string]*time.Location{
+		time.RFC3339: nil,
+	}
+)
+
+func AddSupportedLayout(layout string, location *time.Location) {
+	supportedLayouts[layout] = cmp.Or(location, CST)
+}
+
+func SetOutputLayout(layout string, location *time.Location) {
+	outputLayout = layout
 
 	if location != nil {
 		CST = location
@@ -62,19 +73,17 @@ func (Timestamp) DataType(engine string) string {
 	return "bigint"
 }
 
-func ParseTimestampFromString(s string) (Timestamp, error) {
-	if OutputLayout != time.RFC3339 {
-		t, err := time.ParseInLocation(OutputLayout, s, CST)
-		if err == nil {
+func ParseTimestampFromString(s string) (d Timestamp, err error) {
+	for layout, cst := range supportedLayouts {
+		// fallback
+		t, e := time.ParseInLocation(layout, s, cst)
+		if e == nil {
 			return Timestamp(t), nil
 		}
+		err = e
 	}
-	// fallback
-	t, err := time.Parse(time.RFC3339, s)
-	if err != nil {
-		return TimestampUnixZero, err
-	}
-	return Timestamp(t), nil
+
+	return
 }
 
 func ParseTimestampFromStringWithLayout(input, layout string) (Timestamp, error) {
@@ -124,7 +133,7 @@ func (dt Timestamp) String() string {
 	if dt.IsZero() {
 		return ""
 	}
-	return time.Time(dt).In(CST).Format(OutputLayout)
+	return time.Time(dt).In(CST).Format(outputLayout)
 }
 
 func (dt Timestamp) Format(layout string) string {

pkg/sqltype/time/timestamp_test.go

@@ -40,7 +40,7 @@ func TestTimestamp(t *testing.T) {
 	})
 
 	b.Given("time string", func(b bdd.T) {
-		t0, _ := time.Parse(time.RFC3339, "2017-03-27T23:58:59+08:00")
+		t0, _ := time.ParseInLocation(time.RFC3339, "2017-03-27T23:58:59+08:00", nil)
 		dt := Timestamp(t0)
 
 		b.Then("output RFC3339 string",
@@ -70,11 +70,12 @@ func TestTimestamp(t *testing.T) {
 	})
 
 	b.Given("time string for custom output layout", func(b bdd.T) {
-		t0, _ := time.Parse(time.RFC3339, "2017-03-27T23:58:59+08:00")
+		t0, _ := time.ParseInLocation(time.RFC3339, "2017-03-27T23:58:59+08:00", nil)
 		dt := Timestamp(t0)
 
 		b.When("marshal text", func(b bdd.T) {
-			SetOutput("2006-01-02 15:04:05", nil)
+			AddSupportedLayout("2006-01-02 15:04:05", nil)
+			SetOutputLayout("2006-01-02 15:04:05", nil)
 
 			data := bdd.Must(dt.MarshalText())