Skip to content

Commit

Permalink
Simplify custom type autoloading with pgxpool
Browse files Browse the repository at this point in the history
Provide a backwards-compatible configuration option for pgxpool
which streamlines the use of the bulk loading and registration of types:
- ReuseTypeMaps: if enabled, pgxpool will cache the typemap information,
  avoiding the need to perform any further queries as new connections
  are created.

ReuseTypeMaps is disabled by default as in some situations, a
connection string might resolve to a pool of servers which do not share
the same type name -> OID mapping.
  • Loading branch information
nicois committed Jun 27, 2024
1 parent b8b9930 commit cc208ca
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 0 deletions.
6 changes: 6 additions & 0 deletions conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,12 @@ type ConnConfig struct {

// automatically call LoadTypes with these values
AutoLoadTypes []string

// TypeRegistrationMap is used to register types which require special operations.
// The type name is the key, the value is a function which will be called for each
// connection, providing the OID of that type name for that connection.
// The function will manipulate conn.TypeMap() in some way.
TypeRegistrationMap map[string]CustomRegistrationFunction
}

// ParseConfigOptions contains options that control how a config is built such as getsslpassword.
Expand Down
30 changes: 30 additions & 0 deletions derived_types.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@ import (
"github.com/jackc/pgx/v5/pgtype"
)

// CustomRegistrationFunction is capable of registering whatever is necessary for
// a custom type. It is provided with the backend's OID for this type.
type CustomRegistrationFunction func(ctx context.Context, m *pgtype.Map, oid uint32) error

/*
buildLoadDerivedTypesSQL generates the correct query for retrieving type information.
Expand Down Expand Up @@ -254,3 +258,29 @@ func serverVersion(c *Conn) (int64, error) {
}
return serverVersion, nil
}

func fetchOidMapForCustomRegistration(ctx context.Context, conn *Conn) (map[string]uint32, error) {
sql := `
SELECT oid, typname
FROM pg_type
WHERE typname = ANY($1)`
result := make(map[string]uint32)
typeNames := make([]string, 0, len(conn.config.TypeRegistrationMap))
for typeName := range conn.config.TypeRegistrationMap {
typeNames = append(typeNames, typeName)
}
rows, err := conn.Query(ctx, sql, typeNames)
if err != nil {
return nil, fmt.Errorf("While collecting OIDs for custom registrations: %w", err)
}
defer rows.Close()
var typeName string
var oid uint32
for rows.Next() {
if err := rows.Scan(&typeName, &oid); err != nil {
return nil, fmt.Errorf("While scanning a row for custom registrations: %w", err)
}
result[typeName] = oid
}
return result, nil
}

0 comments on commit cc208ca

Please sign in to comment.