From d2cf33ed40645729a0439cc1b33eb05c3ed0c17f Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 1 Jan 2022 11:25:26 -0600
Subject: [PATCH] Add UnmarshalJSON to generated ints

---
 pgtype/int.go          | 17 +++++++++++++++++
 pgtype/int.go.erb      | 16 ++++++++++++++++
 pgtype/int_test.go     | 41 +++++++++++++++++++++++++++++++++++++++++
 pgtype/int_test.go.erb | 41 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 115 insertions(+)

diff --git a/pgtype/int.go b/pgtype/int.go
index 4d7dea83..9da79a89 100644
--- a/pgtype/int.go
+++ b/pgtype/int.go
@@ -4,6 +4,7 @@ package pgtype
 import (
 	"database/sql/driver"
 	"encoding/binary"
+	"encoding/json"
 	"fmt"
 	"math"
 	"strconv"
@@ -92,6 +93,22 @@ func (src Int2) MarshalJSON() ([]byte, error) {
 	return []byte(strconv.FormatInt(int64(src.Int), 10)), nil
 }
 
+func (dst *Int2) UnmarshalJSON(b []byte) error {
+	var n *int16
+	err := json.Unmarshal(b, &n)
+	if err != nil {
+		return err
+	}
+
+	if n == nil {
+		*dst = Int2{}
+	} else {
+		*dst = Int2{Int: *n, Valid: true}
+	}
+
+	return nil
+}
+
 type Int2Codec struct{}
 
 func (Int2Codec) FormatSupported(format int16) bool {
diff --git a/pgtype/int.go.erb b/pgtype/int.go.erb
index 152dabed..5c8e44fa 100644
--- a/pgtype/int.go.erb
+++ b/pgtype/int.go.erb
@@ -94,6 +94,22 @@ func (src Int<%= pg_byte_size %>) MarshalJSON() ([]byte, error) {
 	return []byte(strconv.FormatInt(int64(src.Int), 10)), nil
 }
 
+func (dst *Int<%= pg_byte_size %>) UnmarshalJSON(b []byte) error {
+	var n *int<%= pg_bit_size %>
+	err := json.Unmarshal(b, &n)
+	if err != nil {
+		return err
+	}
+
+	if n == nil {
+		*dst = Int<%= pg_byte_size %>{}
+	} else {
+		*dst = Int<%= pg_byte_size %>{Int: *n, Valid: true}
+	}
+
+	return nil
+}
+
 type Int<%= pg_byte_size %>Codec struct{}
 
 func (Int<%= pg_byte_size %>Codec) FormatSupported(format int16) bool {
diff --git a/pgtype/int_test.go b/pgtype/int_test.go
index 3ba1306b..272ad69b 100644
--- a/pgtype/int_test.go
+++ b/pgtype/int_test.go
@@ -41,3 +41,44 @@ func TestInt2Codec(t *testing.T) {
 		{nil, new(*int16), isExpectedEq((*int16)(nil))},
 	})
 }
+
+func TestInt2MarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source pgtype.Int2
+		result string
+	}{
+		{source: pgtype.Int2{Int: 0}, result: "null"},
+		{source: pgtype.Int2{Int: 1, Valid: true}, result: "1"},
+	}
+	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 TestInt2UnmarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source string
+		result pgtype.Int2
+	}{
+		{source: "null", result: pgtype.Int2{Int: 0}},
+		{source: "1", result: pgtype.Int2{Int: 1, Valid: true}},
+	}
+	for i, tt := range successfulTests {
+		var r pgtype.Int2
+		err := r.UnmarshalJSON([]byte(tt.source))
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if r != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r)
+		}
+	}
+}
diff --git a/pgtype/int_test.go.erb b/pgtype/int_test.go.erb
index be1f5358..28847a11 100644
--- a/pgtype/int_test.go.erb
+++ b/pgtype/int_test.go.erb
@@ -42,4 +42,45 @@ func TestInt<%= pg_byte_size %>Codec(t *testing.T) {
 		{nil, new(*int<%= pg_bit_size %>), isExpectedEq((*int<%= pg_bit_size %>)(nil))},
 	})
 }
+
+func TestInt<%= pg_byte_size %>MarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source pgtype.Int<%= pg_byte_size %>
+		result string
+	}{
+		{source: pgtype.Int<%= pg_byte_size %>{Int: 0}, result: "null"},
+		{source: pgtype.Int<%= pg_byte_size %>{Int: 1, Valid: true}, result: "1"},
+	}
+	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 TestInt<%= pg_byte_size %>UnmarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source string
+		result pgtype.Int<%= pg_byte_size %>
+	}{
+		{source: "null", result: pgtype.Int<%= pg_byte_size %>{Int: 0}},
+		{source: "1", result: pgtype.Int<%= pg_byte_size %>{Int: 1, Valid: true}},
+	}
+	for i, tt := range successfulTests {
+		var r pgtype.Int<%= pg_byte_size %>
+		err := r.UnmarshalJSON([]byte(tt.source))
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if r != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r)
+		}
+	}
+}
 <% end %>