internal/database: add DB.WithPGXConn

This function lets code outside the package use pgx-specific features
like CopyFrom.

Change-Id: I40d2c8b5d9a92f61f535fc25059fb976baf5d178
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/307870
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>

internal/database/copy.go

@@ -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{}