internal/database: add DB.WithPGXConn
This function lets code outside the package use pgx-specific features
like CopyFrom.
Trust: Jonathan Amsterdam
Run-TryBot: Jonathan Amsterdam
TryBot-Result: kokoro
@@ -36,15 +36,7 @@ func (db *DB) CopyUpsert(ctx context.Context, table string, columns []string, sr
return errors.New("not in a transaction")
}
- return db.conn.Raw(func(c interface{}) error {
- if w, ok := c.(*wrapConn); ok {
- c = w.underlying
- }
- stdConn, ok := c.(*stdlib.Conn)
- if !ok {
- return fmt.Errorf("DB driver is not pgx or wrapper; conn type is %T", c)
- }
- conn := stdConn.Conn()
+ return db.WithPGXConn(func(conn *pgx.Conn) error {
tempTable := fmt.Sprintf("__%s_copy", table)
stmt := fmt.Sprintf(`
DROP TABLE IF EXISTS %s;
@@ -81,6 +73,22 @@ func (db *DB) CopyUpsert(ctx context.Context, table string, columns []string, sr
})
}
+func (db *DB) WithPGXConn(f func(conn *pgx.Conn) error) error {
+ if !db.InTransaction() {
+ return errors.New("not in a transaction")
+ }
+ return db.conn.Raw(func(c interface{}) error {
+ if w, ok := c.(*wrapConn); ok {
+ c = w.underlying
+ }
+ stdConn, ok := c.(*stdlib.Conn)
+ if !ok {
+ return fmt.Errorf("DB driver is not pgx or wrapper; conn type is %T", c)
+ }
+ return f(stdConn.Conn())
+ })
+}
+
// A RowItem is a row of values or an error.
type RowItem struct {
Values []interface{}
