add goccy/json

go.mod

@@ -4,6 +4,7 @@ go 1.20
 
 require (
 	github.com/georgysavva/scany/v2 v2.0.0
+	github.com/goccy/go-json v0.10.2
 	github.com/google/uuid v1.3.0
 	github.com/jackc/pgx/v5 v5.3.1
 )

go.sum

@@ -3,6 +3,8 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/georgysavva/scany/v2 v2.0.0 h1:RGXqxDv4row7/FYoK8MRXAZXqoWF/NM+NP0q50k3DKU=
 github.com/georgysavva/scany/v2 v2.0.0/go.mod h1:sigOdh+0qb/+aOs3TVhehVT10p8qJL7K/Zhyz8vWo38=
+github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
+github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
 github.com/gofrs/flock v0.8.1 h1:+gYjHKf32LDeiEEFhQaotPbLuUXjY5ZqxKgXy7n59aw=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=

pool.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"git.barsukov.pro/barsukov/zdb/ztype"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"log"
 	"reflect"
@@ -107,12 +110,29 @@ func (d *Pool) NewConns(mode connMode, pgConnString ...string) error {
 
 func (d *Pool) newConn(mode connMode, pgConnString string) (q *conn, err error) {
 	var pgxPool *pgxpool.Pool
+	var pgxConfig *pgxpool.Config
 
 	if !strings.Contains(pgConnString, "default_query_exec_mode=") {
 		pgConnString += " default_query_exec_mode=simple_protocol"
 	}
 
-	if pgxPool, err = pgxpool.New(d.ctx, pgConnString); err != nil {
+	if pgxConfig, err = pgxpool.ParseConfig(pgConnString); err != nil {
+		return nil, err
+	}
+
+	pgxConfig.AfterConnect = func(ctx context.Context, p *pgx.Conn) error {
+		j := &pgtype.Type{Name: "jsonb", OID: pgtype.JSONBOID, Codec: ztype.JSONBCodec{}}
+		jb := &pgtype.Type{Name: "json", OID: pgtype.JSONOID, Codec: ztype.JSONCodec{}}
+
+		p.TypeMap().RegisterType(j)
+		p.TypeMap().RegisterType(jb)
+		p.TypeMap().RegisterType(&pgtype.Type{Name: "_json", OID: pgtype.JSONArrayOID, Codec: &pgtype.ArrayCodec{ElementType: j}})
+		p.TypeMap().RegisterType(&pgtype.Type{Name: "_jsonb", OID: pgtype.JSONBArrayOID, Codec: &pgtype.ArrayCodec{ElementType: jb}})
+
+		return nil
+	}
+
+	if pgxPool, err = pgxpool.NewWithConfig(d.ctx, pgxConfig); err != nil {
 		return &conn{Pool: pgxPool, Alive: false, Mode: mode}, err
 	}
 

ztype/json.go

+package ztype
+
+import (
+	"database/sql/driver"
+	"fmt"
+	"github.com/goccy/go-json"
+	pgxtype "github.com/jackc/pgx/v5/pgtype"
+	"reflect"
+)
+
+type JSONCodec struct{}
+
+func (JSONCodec) FormatSupported(format int16) bool {
+	return format == pgxtype.TextFormatCode || format == pgxtype.BinaryFormatCode
+}
+
+func (JSONCodec) PreferredFormat() int16 {
+	return pgxtype.TextFormatCode
+}
+
+func (c JSONCodec) PlanEncode(m *pgxtype.Map, oid uint32, format int16, value any) pgxtype.EncodePlan {
+	switch value.(type) {
+	case string:
+		return encodePlanJSONCodecEitherFormatString{}
+	case []byte:
+		return encodePlanJSONCodecEitherFormatByteSlice{}
+	}
+
+	// Because anything can be marshalled the normal wrapping in Map.PlanScan doesn't get a chance to run. So try the
+	// appropriate wrappers here.
+	for _, f := range []pgxtype.TryWrapEncodePlanFunc{
+		pgxtype.TryWrapDerefPointerEncodePlan,
+		pgxtype.TryWrapFindUnderlyingTypeEncodePlan,
+	} {
+		if wrapperPlan, nextValue, ok := f(value); ok {
+			if nextPlan := c.PlanEncode(m, oid, format, nextValue); nextPlan != nil {
+				wrapperPlan.SetNext(nextPlan)
+				return wrapperPlan
+			}
+		}
+	}
+
+	return encodePlanJSONCodecEitherFormatMarshal{}
+}
+
+type encodePlanJSONCodecEitherFormatString struct{}
+
+func (encodePlanJSONCodecEitherFormatString) Encode(value any, buf []byte) (newBuf []byte, err error) {
+	jsonString := value.(string)
+	buf = append(buf, jsonString...)
+	return buf, nil
+}
+
+type encodePlanJSONCodecEitherFormatByteSlice struct{}
+
+func (encodePlanJSONCodecEitherFormatByteSlice) Encode(value any, buf []byte) (newBuf []byte, err error) {
+	jsonBytes := value.([]byte)
+	if jsonBytes == nil {
+		return nil, nil
+	}
+
+	buf = append(buf, jsonBytes...)
+	return buf, nil
+}
+
+type encodePlanJSONCodecEitherFormatMarshal struct{}
+
+func (encodePlanJSONCodecEitherFormatMarshal) Encode(value any, buf []byte) (newBuf []byte, err error) {
+	jsonBytes, err := json.Marshal(value)
+	if err != nil {
+		return nil, err
+	}
+
+	buf = append(buf, jsonBytes...)
+	return buf, nil
+}
+
+func (JSONCodec) PlanScan(m *pgxtype.Map, oid uint32, format int16, target any) pgxtype.ScanPlan {
+	switch target.(type) {
+	case *string:
+		return scanPlanAnyToString{}
+	case *[]byte:
+		return scanPlanJSONToByteSlice{}
+	case pgxtype.BytesScanner:
+		return scanPlanBinaryBytesToBytesScanner{}
+	}
+
+	return scanPlanJSONToJSONUnmarshal{}
+}
+
+type scanPlanAnyToString struct{}
+
+func (scanPlanAnyToString) Scan(src []byte, dst any) error {
+	p := dst.(*string)
+	*p = string(src)
+	return nil
+}
+
+type scanPlanJSONToByteSlice struct{}
+
+func (scanPlanJSONToByteSlice) Scan(src []byte, dst any) error {
+	dstBuf := dst.(*[]byte)
+	if src == nil {
+		*dstBuf = nil
+		return nil
+	}
+
+	*dstBuf = make([]byte, len(src))
+	copy(*dstBuf, src)
+	return nil
+}
+
+type scanPlanJSONToBytesScanner struct{}
+
+func (scanPlanJSONToBytesScanner) Scan(src []byte, dst any) error {
+	scanner := (dst).(pgxtype.BytesScanner)
+	return scanner.ScanBytes(src)
+}
+
+type scanPlanJSONToJSONUnmarshal struct{}
+
+func (scanPlanJSONToJSONUnmarshal) Scan(src []byte, dst any) error {
+	if src == nil {
+		dstValue := reflect.ValueOf(dst)
+		if dstValue.Kind() == reflect.Ptr {
+			el := dstValue.Elem()
+			switch el.Kind() {
+			case reflect.Ptr, reflect.Slice, reflect.Map:
+				el.Set(reflect.Zero(el.Type()))
+				return nil
+			}
+		}
+
+		return fmt.Errorf("cannot scan NULL into %T", dst)
+	}
+
+	elem := reflect.ValueOf(dst).Elem()
+	elem.Set(reflect.Zero(elem.Type()))
+
+	return json.Unmarshal(src, dst)
+}
+
+func (c JSONCodec) DecodeDatabaseSQLValue(m *pgxtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
+	if src == nil {
+		return nil, nil
+	}
+
+	dstBuf := make([]byte, len(src))
+	copy(dstBuf, src)
+	return dstBuf, nil
+}
+
+func (c JSONCodec) DecodeValue(m *pgxtype.Map, oid uint32, format int16, src []byte) (any, error) {
+	if src == nil {
+		return nil, nil
+	}
+
+	var dst any
+	err := json.Unmarshal(src, &dst)
+	return dst, err
+}
+
+type scanPlanBinaryBytesToBytesScanner struct{}
+
+func (scanPlanBinaryBytesToBytesScanner) Scan(src []byte, dst any) error {
+	scanner := (dst).(pgxtype.BytesScanner)
+	return scanner.ScanBytes(src)
+}

ztype/jsonb.go

+package ztype
+
+import (
+	"database/sql/driver"
+	"fmt"
+	"github.com/goccy/go-json"
+	pgxtype "github.com/jackc/pgx/v5/pgtype"
+)
+
+type JSONBCodec struct{}
+
+func (JSONBCodec) FormatSupported(format int16) bool {
+	return format == pgxtype.TextFormatCode || format == pgxtype.BinaryFormatCode
+}
+
+func (JSONBCodec) PreferredFormat() int16 {
+	return pgxtype.TextFormatCode
+}
+
+func (JSONBCodec) PlanEncode(m *pgxtype.Map, oid uint32, format int16, value any) pgxtype.EncodePlan {
+	switch format {
+	case pgxtype.BinaryFormatCode:
+		plan := JSONCodec{}.PlanEncode(m, oid, pgxtype.TextFormatCode, value)
+		if plan != nil {
+			return &encodePlanJSONBCodecBinaryWrapper{textPlan: plan}
+		}
+	case pgxtype.TextFormatCode:
+		return JSONCodec{}.PlanEncode(m, oid, format, value)
+	}
+
+	return nil
+}
+
+type encodePlanJSONBCodecBinaryWrapper struct {
+	textPlan pgxtype.EncodePlan
+}
+
+func (plan *encodePlanJSONBCodecBinaryWrapper) Encode(value any, buf []byte) (newBuf []byte, err error) {
+	buf = append(buf, 1)
+	return plan.textPlan.Encode(value, buf)
+}
+
+func (JSONBCodec) PlanScan(m *pgxtype.Map, oid uint32, format int16, target any) pgxtype.ScanPlan {
+	switch format {
+	case pgxtype.BinaryFormatCode:
+		plan := JSONCodec{}.PlanScan(m, oid, pgxtype.TextFormatCode, target)
+		if plan != nil {
+			return &scanPlanJSONBCodecBinaryUnwrapper{textPlan: plan}
+		}
+	case pgxtype.TextFormatCode:
+		return JSONCodec{}.PlanScan(m, oid, format, target)
+	}
+
+	return nil
+}
+
+type scanPlanJSONBCodecBinaryUnwrapper struct {
+	textPlan pgxtype.ScanPlan
+}
+
+func (plan *scanPlanJSONBCodecBinaryUnwrapper) Scan(src []byte, dst any) error {
+	if src == nil {
+		return plan.textPlan.Scan(src, dst)
+	}
+
+	if len(src) == 0 {
+		return fmt.Errorf("jsonb too short")
+	}
+
+	if src[0] != 1 {
+		return fmt.Errorf("unknown jsonb version number %d", src[0])
+	}
+
+	return plan.textPlan.Scan(src[1:], dst)
+}
+
+func (c JSONBCodec) DecodeDatabaseSQLValue(m *pgxtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
+	if src == nil {
+		return nil, nil
+	}
+
+	switch format {
+	case pgxtype.BinaryFormatCode:
+		if len(src) == 0 {
+			return nil, fmt.Errorf("jsonb too short")
+		}
+
+		if src[0] != 1 {
+			return nil, fmt.Errorf("unknown jsonb version number %d", src[0])
+		}
+
+		dstBuf := make([]byte, len(src)-1)
+		copy(dstBuf, src[1:])
+		return dstBuf, nil
+	case pgxtype.TextFormatCode:
+		dstBuf := make([]byte, len(src))
+		copy(dstBuf, src)
+		return dstBuf, nil
+	default:
+		return nil, fmt.Errorf("unknown format code: %v", format)
+	}
+}
+
+func (c JSONBCodec) DecodeValue(m *pgxtype.Map, oid uint32, format int16, src []byte) (any, error) {
+	if src == nil {
+		return nil, nil
+	}
+
+	switch format {
+	case pgxtype.BinaryFormatCode:
+		if len(src) == 0 {
+			return nil, fmt.Errorf("jsonb too short")
+		}
+
+		if src[0] != 1 {
+			return nil, fmt.Errorf("unknown jsonb version number %d", src[0])
+		}
+
+		src = src[1:]
+	case pgxtype.TextFormatCode:
+	default:
+		return nil, fmt.Errorf("unknown format code: %v", format)
+	}
+
+	var dst any
+	err := json.Unmarshal(src, &dst)
+	return dst, err
+}