From c542df4fb4cfaa8126143de5dd8c69f84a9c28cb Mon Sep 17 00:00:00 2001
From: Klaus <klaus_hybrid33@yahoo.com>
Date: Mon, 12 Jun 2023 01:36:15 +0330
Subject: [PATCH] added MarshalJSON and UnmarshalJSON to timestamp and added
 their tests (based on timestamptz implementation)

---
 pgtype/timestamp.go      | 50 ++++++++++++++++++++++++++++++++++++++++
 pgtype/timestamp_test.go | 47 +++++++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/pgtype/timestamp.go b/pgtype/timestamp.go
index 9f3de2c5..35d73956 100644
--- a/pgtype/timestamp.go
+++ b/pgtype/timestamp.go
@@ -3,6 +3,7 @@ package pgtype
 import (
 	"database/sql/driver"
 	"encoding/binary"
+	"encoding/json"
 	"fmt"
 	"strings"
 	"time"
@@ -66,6 +67,55 @@ func (ts Timestamp) Value() (driver.Value, error) {
 	return ts.Time, nil
 }
 
+func (ts Timestamp) MarshalJSON() ([]byte, error) {
+	if !ts.Valid {
+		return []byte("null"), nil
+	}
+
+	var s string
+
+	switch ts.InfinityModifier {
+	case Finite:
+		s = ts.Time.Format(time.RFC3339Nano)
+	case Infinity:
+		s = "infinity"
+	case NegativeInfinity:
+		s = "-infinity"
+	}
+
+	return json.Marshal(s)
+}
+
+func (ts *Timestamp) UnmarshalJSON(b []byte) error {
+	var s *string
+	err := json.Unmarshal(b, &s)
+	if err != nil {
+		return err
+	}
+
+	if s == nil {
+		*ts = Timestamp{}
+		return nil
+	}
+
+	switch *s {
+	case "infinity":
+		*ts = Timestamp{Valid: true, InfinityModifier: Infinity}
+	case "-infinity":
+		*ts = Timestamp{Valid: true, InfinityModifier: -Infinity}
+	default:
+		// PostgreSQL uses ISO 8601 for to_json function and casting from a string to timestamptz
+		tim, err := time.Parse(time.RFC3339Nano, *s)
+		if err != nil {
+			return err
+		}
+
+		*ts = Timestamp{Time: tim, Valid: true}
+	}
+
+	return nil
+}
+
 type TimestampCodec struct{}
 
 func (TimestampCodec) FormatSupported(format int16) bool {
diff --git a/pgtype/timestamp_test.go b/pgtype/timestamp_test.go
index 849f55f6..24f229d5 100644
--- a/pgtype/timestamp_test.go
+++ b/pgtype/timestamp_test.go
@@ -62,3 +62,50 @@ func TestTimestampCodecDecodeTextInvalid(t *testing.T) {
 	err := plan.Scan([]byte(`eeeee`), &ts)
 	require.Error(t, err)
 }
+
+func TestTimestampMarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source pgtype.Timestamp
+		result string
+	}{
+		{source: pgtype.Timestamp{}, result: "null"},
+		{source: pgtype.Timestamp{Time: time.Date(2012, 3, 29, 10, 5, 45, 0, time.UTC), Valid: true}, result: "\"2012-03-29T10:05:45Z\""},
+		{source: pgtype.Timestamp{Time: time.Date(2012, 3, 29, 10, 5, 45, 555*1000*1000, time.UTC), Valid: true}, result: "\"2012-03-29T10:05:45.555Z\""},
+		{source: pgtype.Timestamp{InfinityModifier: pgtype.Infinity, Valid: true}, result: "\"infinity\""},
+		{source: pgtype.Timestamp{InfinityModifier: pgtype.NegativeInfinity, Valid: true}, result: "\"-infinity\""},
+	}
+	for i, tt := range successfulTests {
+		r, err := tt.source.MarshalJSON()
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if string(r) != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, string(r))
+		}
+	}
+}
+
+func TestTimestampUnmarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source string
+		result pgtype.Timestamp
+	}{
+		{source: "null", result: pgtype.Timestamp{}},
+		{source: "\"2012-03-29T10:05:45Z\"", result: pgtype.Timestamp{Time: time.Date(2012, 3, 29, 10, 5, 45, 0, time.UTC), Valid: true}},
+		{source: "\"2012-03-29T10:05:45.555Z\"", result: pgtype.Timestamp{Time: time.Date(2012, 3, 29, 10, 5, 45, 555*1000*1000, time.UTC), Valid: true}},
+		{source: "\"infinity\"", result: pgtype.Timestamp{InfinityModifier: pgtype.Infinity, Valid: true}},
+		{source: "\"-infinity\"", result: pgtype.Timestamp{InfinityModifier: pgtype.NegativeInfinity, Valid: true}},
+	}
+	for i, tt := range successfulTests {
+		var r pgtype.Timestamp
+		err := r.UnmarshalJSON([]byte(tt.source))
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if !r.Time.Equal(tt.result.Time) || r.Valid != tt.result.Valid || r.InfinityModifier != tt.result.InfinityModifier {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r)
+		}
+	}
+}