working tests for varbit

Author: zachmu

server/functions/varbit.go

@@ -16,9 +16,8 @@ package functions
 
 import (
 	"fmt"
-	"strings"
 
-	"github.com/cockroachdb/errors"
+	"github.com/dolthub/doltgresql/postgres/parser/sem/tree"
 	"github.com/dolthub/go-mysql-server/sql"
 
 	"github.com/dolthub/doltgresql/server/functions/framework"
@@ -42,37 +41,23 @@ var varbitin = framework.Function3{
 	Return:     pgtypes.VarBit,
 	Parameters: [3]*pgtypes.DoltgresType{pgtypes.Cstring, pgtypes.Oid, pgtypes.Int32},
 	Strict:     true,
-	Callable: func(ctx *sql.Context, _ [4]*pgtypes.DoltgresType, val1, val2, val3 any) (any, error) {
+	Callable: func(ctx *sql.Context, _ [4]*pgtypes.DoltgresType, val1, _, val3 any) (any, error) {
 		input := val1.(string)
 		typmod := val3.(int32)
-		
-		// Parse bit string - remove leading 'B' or 'b' prefix if present
-		bitStr := strings.TrimSpace(input)
-		if len(bitStr) > 0 && (bitStr[0] == 'B' || bitStr[0] == 'b') {
-			bitStr = bitStr[1:]
-			// Remove quotes if present
-			if len(bitStr) > 0 && (bitStr[0] == '\'' || bitStr[0] == '"') {
-				if len(bitStr) > 1 && bitStr[len(bitStr)-1] == bitStr[0] {
-					bitStr = bitStr[1 : len(bitStr)-1]
-				}
-			}
-		}
-		
-		// Validate that all characters are '0' or '1'
-		for _, r := range bitStr {
-			if r != '0' && r != '1' {
-				return nil, pgtypes.ErrInvalidSyntaxForType.New("varbit", input)
-			}
+
+		bitStr, err := tree.ParseDBitArray(input)
+		if err != nil {
+			return nil, err
 		}
-		
+
 		// Check length against typmod (varbit allows up to typmod length)
 		if typmod != -1 {
 			maxLength := pgtypes.GetCharLengthFromTypmod(typmod)
-			if int32(len(bitStr)) > maxLength {
-				return nil, errors.Wrap(pgtypes.ErrCastOutOfRange, fmt.Sprintf("bit string length %d exceeds maximum length %d for type varbit(%d)", len(bitStr), maxLength, maxLength))
+			if int32(bitStr.BitLen()) > maxLength {
+				return nil, pgtypes.ErrVarBitLengthExceeded.New(maxLength)
 			}
 		}
-		
+
 		return bitStr, nil
 	},
 }
@@ -84,17 +69,8 @@ var varbitout = framework.Function1{
 	Parameters: [1]*pgtypes.DoltgresType{pgtypes.VarBit},
 	Strict:     true,
 	Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
-		bitStr := val.(string)
-		typ := t[0]
-		tm := typ.GetAttTypMod()
-		if tm != -1 {
-			maxLength := pgtypes.GetCharLengthFromTypmod(tm)
-			// Truncate if needed (shouldn't happen normally, but handle it)
-			if int32(len(bitStr)) > maxLength {
-				bitStr = bitStr[:maxLength]
-			}
-		}
-		return bitStr, nil
+		bitStr := val.(*tree.DBitArray)
+		return tree.AsStringWithFlags(bitStr, tree.FmtPgwireText), nil
 	},
 }
 
@@ -110,7 +86,7 @@ var varbitrecv = framework.Function3{
 			return nil, nil
 		}
 		reader := utils.NewReader(data)
-		return reader.String(), nil
+		return tree.ParseDBitArray(reader.String())
 	},
 }
 
@@ -121,9 +97,9 @@ var varbitsend = framework.Function1{
 	Parameters: [1]*pgtypes.DoltgresType{pgtypes.VarBit},
 	Strict:     true,
 	Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) {
-		bitStr := val.(string)
-		writer := utils.NewWriter(uint64(len(bitStr) + 4))
-		writer.String(bitStr)
+		bitStr := val.(*tree.DBitArray)
+		writer := utils.NewWriter(uint64(bitStr.BitLen() + 4))
+		writer.String(tree.AsStringWithFlags(bitStr, tree.FmtPgwireText))
 		return writer.Data(), nil
 	},
 }
@@ -135,7 +111,7 @@ var varbittypmodin = framework.Function1{
 	Parameters: [1]*pgtypes.DoltgresType{pgtypes.CstringArray},
 	Strict:     true,
 	Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) {
-		return getTypModFromStringArr("bit", val.([]any))
+		return getTypModFromStringArr("bit varying", val.([]any))
 	},
 }
 

server/types/varbit.go

@@ -16,8 +16,12 @@ package types
 
 import (
 	"github.com/dolthub/doltgresql/core/id"
+	"gopkg.in/src-d/go-errors.v1"
 )
 
+// ErrVarBitLengthExceeded is returned when a varbit value exceeds the defined length.
+var ErrVarBitLengthExceeded = errors.NewKind(`bit string too long for type bit varying(%d)`)
+
 // VarBit is a varying-length bit string.
 var VarBit = &DoltgresType{
 	ID:            toInternal("varbit"),
@@ -33,7 +37,7 @@ var VarBit = &DoltgresType{
 	Elem:          id.NullType,
 	Array:         toInternal("_varbit"),
 	InputFunc:     toFuncID("varbit_in", toInternal("cstring"), toInternal("oid"), toInternal("int4")),
-	OutputFunc:    toFuncID("varbit_out", toInternal("bit")),
+	OutputFunc:    toFuncID("varbit_out", toInternal("varbit")),
 	ReceiveFunc:   toFuncID("varbit_recv", toInternal("internal"), toInternal("oid"), toInternal("int4")),
 	SendFunc:      toFuncID("varbit_send", toInternal("varbit")),
 	ModInFunc:     toFuncID("varbittypmodin", toInternal("varbit")),
@@ -54,8 +58,8 @@ var VarBit = &DoltgresType{
 	CompareFunc:   toFuncID("bttextcmp", toInternal("text"), toInternal("text")),
 }
 
-// NewBitType returns a Bit type with type modifier set
-// representing the number of bits in the string.
+// NewVarBitType returns a VarBit type with type modifier set
+// representing the max number of bits in the string.
 func NewVarBitType(width int32) (*DoltgresType, error) {
 	typmod, err := GetTypModFromCharLength("bit", width)
 	if err != nil {

testing/go/types_test.go

@@ -79,8 +79,7 @@ var typesTests = []ScriptTest{
 		},
 	},
 	{
-		Name:  "Bit type",
-		Focus: true,
+		Name: "Bit type",
 		SetUpScript: []string{
 			"CREATE TABLE t_bit (id INTEGER primary key, v1 BIT(8), v2 BIT(3));",
 			"INSERT INTO t_bit VALUES (1, B'11011010', '101'), (2, B'00101011', '000');",
@@ -303,8 +302,8 @@ var typesTests = []ScriptTest{
 		},
 	},
 	{
-		Name: "Bit varying type",
-		Skip: true,
+		Name:  "Bit varying type",
+		Focus: true,
 		SetUpScript: []string{
 			"CREATE TABLE t_bit_varying (id INTEGER primary key, v1 BIT VARYING(16));",
 			"INSERT INTO t_bit_varying VALUES (1, B'1101101010101010'), (2, B'0010101101010101');",
@@ -313,8 +312,8 @@ var typesTests = []ScriptTest{
 			{
 				Query: "SELECT * FROM t_bit_varying ORDER BY id;",
 				Expected: []sql.Row{
-					{1, []byte{0xDA, 0xAA}},
-					{2, []byte{0x2B, 0xA5}},
+					{1, pgtype.Bits{Bytes: []uint8{0xda, 0xaa}, Len: 16, Valid: true}},
+					{2, pgtype.Bits{Bytes: []uint8{0x2b, 0x55}, Len: 16, Valid: true}},
 				},
 			},
 		},