driver.Valuer existence takes precedence over codec for nil values #2611

package main

import (
	"context"
	"database/sql/driver"
	"errors"
	"fmt"
	"log"
	"os"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/pborman/uuid"
)

type Result struct {
	ID   int64
	Area pgtype.Box
}

func main() {
	conn, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatal(err)
	}
	defer conn.Close(context.Background())

	m := conn.TypeMap()

	m.RegisterType(&pgtype.Type{
		Name:  "uuid",
		OID:   pgtype.UUIDOID,
		Codec: UUIDCodec{},
	})

	_, err = conn.Exec(context.Background(),
		`CREATE OR REPLACE FUNCTION do_something( in_value      UUID[] )
    		RETURNS TABLE (out_value UUID)
		AS
		$$
		BEGIN
    		RETURN QUERY SELECT * FROM unnest(in_value);
		END;
		$$
		LANGUAGE plpgsql SECURITY DEFINER IMMUTABLE;`)
	if err != nil {
		log.Fatal(err)
	}

	input := []uuid.UUID{uuid.NewRandom(), nil, uuid.NewRandom()}
	_, err = conn.Exec(context.Background(), "SELECT * FROM do_something( $1 )", input)
	if err != nil {
		log.Fatal(err)
	}
}

type UUIDCodec struct {
	pgtype.UUIDCodec
}

func (c UUIDCodec) DecodeDatabaseSQLValue(tm *pgtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
	if src == nil {
		return nil, nil
	}

	var target uuid.UUID
	scanPlan := tm.PlanScan(oid, format, &target)
	if scanPlan == nil {
		return nil, errors.New("PlanScan did not find a plan")
	}

	err := scanPlan.Scan(src, &target)
	if err != nil {
		return nil, err
	}

	return target.Value()
}

func (c UUIDCodec) PlanEncode(m *pgtype.Map, oid uint32, format int16, value any) pgtype.EncodePlan {
	switch value.(type) {
	default:
		// If this is not the pborman uuid, then fallback to the pgtype implementation
		return c.UUIDCodec.PlanEncode(m, oid, format, value)

	case uuid.UUID, *uuid.UUID:
		switch format {
		case pgtype.BinaryFormatCode:
			return binaryEncodePlan{}
		case pgtype.TextFormatCode:
			return textEncodePlan{}
		}
	}

	return nil
}

type binaryEncodePlan struct{}

func (plan binaryEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) {
	switch src := value.(type) {
	case *uuid.UUID:
		if src == nil {
			return nil, nil
		}
		return plan.encodeUUID(*src, buf)

	case uuid.UUID:
		return plan.encodeUUID(src, buf)

	default:
		return nil, fmt.Errorf("cannot encode %T as uuid", src)
	}
}

func (binaryEncodePlan) encodeUUID(value uuid.UUID, buf []byte) (newBuf []byte, err error) {
	if value == nil {
		return nil, nil
	}

	return append(buf, value[:]...), nil
}

type textEncodePlan struct{}

func (plan textEncodePlan) Encode(value any, buf []byte) (newBuf []byte, err error) {
	switch src := value.(type) {
	case *uuid.UUID:
		if src == nil {
			return nil, nil
		}
		return plan.encodeUUID(*src, buf)

	case uuid.UUID:
		return plan.encodeUUID(src, buf)

	default:
		return nil, fmt.Errorf("cannot encode %T as uuid", src)
	}
}

func (textEncodePlan) encodeUUID(value uuid.UUID, buf []byte) (newBuf []byte, err error) {
	if value == nil {
		return nil, nil
	}

	return append(buf, value.String()...), nil
}