From 750fe1830e895efd0211706c763cb4fb5a45b005 Mon Sep 17 00:00:00 2001
From: Yue Zhu <yzhu@digitalocean.com>
Date: Wed, 11 Dec 2019 23:29:16 -0500
Subject: [PATCH] rados: add support for sending PG commands

---
 rados/conn.go       | 57 +++++++++++++++++++++++++++++++++++++++++++++
 rados/rados_test.go | 21 +++++++++++++++--
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/rados/conn.go b/rados/conn.go
index d2b0a25..ce76619 100644
--- a/rados/conn.go
+++ b/rados/conn.go
@@ -361,3 +361,60 @@ func (c *Conn) monCommand(args, inputBuffer []byte) (buffer []byte, info string,
 
 	return
 }
+
+// PGCommand sends a command to one of the PGs
+func (c *Conn) PGCommand(pgid []byte, args [][]byte) (buffer []byte, info string, err error) {
+	return c.pgCommand(pgid, args, nil)
+}
+
+// PGCommand sends a command to one of the PGs, with an input buffer
+func (c *Conn) PGCommandWithInputBuffer(pgid []byte, args [][]byte, inputBuffer []byte) (buffer []byte, info string, err error) {
+	return c.pgCommand(pgid, args, inputBuffer)
+}
+
+func (c *Conn) pgCommand(pgid []byte, args [][]byte, inputBuffer []byte) (buffer []byte, info string, err error) {
+	name := C.CString(string(pgid))
+	defer C.free(unsafe.Pointer(name))
+
+	argc := len(args)
+	argv := make([]*C.char, argc)
+
+	for i, arg := range args {
+		argv[i] = C.CString(string(arg))
+		defer C.free(unsafe.Pointer(argv[i]))
+	}
+
+	var (
+		outs, outbuf       *C.char
+		outslen, outbuflen C.size_t
+	)
+	inbuf := C.CString(string(inputBuffer))
+	inbufLen := len(inputBuffer)
+	defer C.free(unsafe.Pointer(inbuf))
+
+	ret := C.rados_pg_command(c.cluster,
+		name,
+		&argv[0],
+		C.size_t(argc),
+		inbuf,              // bulk input
+		C.size_t(inbufLen), // length inbuf
+		&outbuf,            // buffer
+		&outbuflen,         // buffer length
+		&outs,              // status string
+		&outslen)
+
+	if outslen > 0 {
+		info = C.GoStringN(outs, C.int(outslen))
+		C.free(unsafe.Pointer(outs))
+	}
+	if outbuflen > 0 {
+		buffer = C.GoBytes(unsafe.Pointer(outbuf), C.int(outbuflen))
+		C.free(unsafe.Pointer(outbuf))
+	}
+	if ret != 0 {
+		err = RadosError(int(ret))
+		return nil, info, err
+	}
+
+	return
+}
diff --git a/rados/rados_test.go b/rados/rados_test.go
index 86b8605..fcf386b 100644
--- a/rados/rados_test.go
+++ b/rados/rados_test.go
@@ -5,11 +5,10 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
-	"strconv"
-	//"net"
 	"math/rand"
 	"os"
 	"sort"
+	"strconv"
 	"testing"
 	"time"
 
@@ -595,6 +594,24 @@ func (suite *RadosTestSuite) TestMonCommandWithInputBuffer() {
 		string(buf[:]))
 }
 
+func (suite *RadosTestSuite) TestPGCommand() {
+	suite.SetupConnection()
+
+	pgid := "1.2"
+
+	command, err := json.Marshal(
+		map[string]string{"prefix": "query", "pgid": pgid, "format": "json"})
+	assert.NoError(suite.T(), err)
+
+	buf, info, err := suite.conn.PGCommand([]byte(pgid), [][]byte{[]byte(command)})
+	assert.NoError(suite.T(), err)
+	assert.Equal(suite.T(), info, "")
+
+	var message map[string]interface{}
+	err = json.Unmarshal(buf, &message)
+	assert.NoError(suite.T(), err)
+}
+
 func (suite *RadosTestSuite) TestObjectListObjects() {
 	suite.SetupConnection()