Hein
50 lines
1.4 KiB
Go
50 lines
1.4 KiB
Go
package pgx
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
)
|
|
|
|
func RegisterTypes(ctx context.Context, conn *pgx.Conn) error {
|
|
var vectorOid *uint32
|
|
var vectorArrayOid *uint32
|
|
var halfvecOid *uint32
|
|
var halfvecArrayOid *uint32
|
|
var sparsevecOid *uint32
|
|
var sparsevecArrayOid *uint32
|
|
err := conn.QueryRow(ctx, "SELECT to_regtype('vector')::oid, to_regtype('_vector')::oid, to_regtype('halfvec')::oid, to_regtype('_halfvec')::oid, to_regtype('sparsevec')::oid, to_regtype('_sparsevec')::oid").Scan(&vectorOid, &vectorArrayOid, &halfvecOid, &halfvecArrayOid, &sparsevecOid, &sparsevecArrayOid)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if vectorOid == nil {
|
|
return fmt.Errorf("vector type not found in the database")
|
|
}
|
|
|
|
tm := conn.TypeMap()
|
|
registerType(tm, "vector", vectorOid, vectorArrayOid, &VectorCodec{})
|
|
|
|
if halfvecOid != nil {
|
|
registerType(tm, "halfvec", halfvecOid, halfvecArrayOid, &HalfVectorCodec{})
|
|
}
|
|
|
|
if sparsevecOid != nil {
|
|
registerType(tm, "sparsevec", sparsevecOid, sparsevecArrayOid, &SparseVectorCodec{})
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func registerType(tm *pgtype.Map, name string, oid *uint32, arrayOid *uint32, codec pgtype.Codec) {
|
|
t := pgtype.Type{Name: name, OID: *oid, Codec: codec}
|
|
tm.RegisterType(&t)
|
|
|
|
// should never be nil
|
|
if arrayOid != nil {
|
|
tm.RegisterType(&pgtype.Type{Name: "_" + name, OID: *arrayOid, Codec: &pgtype.ArrayCodec{ElementType: &t}})
|
|
}
|
|
}
|