From d95dda2bdca4fb9cbc17feb1bdf4fcbffa196a08 Mon Sep 17 00:00:00 2001
From: Miron <Gezort@users.noreply.github.com>
Date: Thu, 14 Dec 2023 09:39:49 -0300
Subject: [PATCH] Fix bindings for MultiGetWithCF #137 (#138)

---
 transaction.go        |  7 ++++++-
 transactiondb.go      |  7 ++++++-
 transactiondb_test.go | 48 ++++++++++++++++++++++++++++---------------
 3 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/transaction.go b/transaction.go
index f08e568..dc260a1 100644
--- a/transaction.go
+++ b/transaction.go
@@ -255,10 +255,15 @@ func (transaction *Transaction) MultiGetWithCF(opts *ReadOptions, cf *ColumnFami
 	valSizes := make(sizeTSlice, len(keys))
 	rocksErrs := make(charsSlice, len(keys))
 
+	cfs := make(ColumnFamilyHandles, len(keys))
+	for i := range keys {
+		cfs[i] = cf
+	}
+
 	C.rocksdb_transaction_multi_get_cf(
 		transaction.c,
 		opts.c,
-		&cf.c,
+		cfs.toCSlice().c(),
 		C.size_t(len(keys)),
 		cKeys.c(),
 		cKeySizes.c(),
diff --git a/transactiondb.go b/transactiondb.go
index cdfe73e..8b3341b 100644
--- a/transactiondb.go
+++ b/transactiondb.go
@@ -284,10 +284,15 @@ func (db *TransactionDB) MultiGetWithCF(opts *ReadOptions, cf *ColumnFamilyHandl
 	valSizes := make(sizeTSlice, len(keys))
 	rocksErrs := make(charsSlice, len(keys))
 
+	cfs := make(ColumnFamilyHandles, len(keys))
+	for i := range keys {
+		cfs[i] = cf
+	}
+
 	C.rocksdb_transactiondb_multi_get_cf(
 		db.c,
 		opts.c,
-		&cf.c,
+		cfs.toCSlice().c(),
 		C.size_t(len(keys)),
 		cKeys.c(),
 		cKeySizes.c(),
diff --git a/transactiondb_test.go b/transactiondb_test.go
index 702b110..f09913a 100644
--- a/transactiondb_test.go
+++ b/transactiondb_test.go
@@ -221,31 +221,47 @@ func TestTransactionDBColumnFamilyBatchPutGet(t *testing.T) {
 	givenVal0 := []byte("world0")
 	givenKey1 := []byte("hello1")
 	givenVal1 := []byte("world1")
+	givenKey2 := []byte("hello2")
+	givenVal2 := []byte("world2")
+
+	writeReadBatch := func(cf *ColumnFamilyHandle, keys [][]byte, values [][]byte) {
+		b := NewWriteBatch()
+		defer b.Destroy()
+		for i := range keys {
+			b.PutCF(cf, keys[i], values[i])
+		}
+		require.Nil(t, db.Write(wo, b))
+
+		for i := range keys {
+			actualVal, err := db.GetCF(ro, cf, keys[i])
+			require.Nil(t, err)
+			require.EqualValues(t, actualVal.Data(), values[i])
+			actualVal.Free()
+		}
+	}
 
-	b0 := NewWriteBatch()
-	defer b0.Destroy()
-	b0.PutCF(cfh[0], givenKey0, givenVal0)
-	require.Nil(t, db.Write(wo, b0))
-	actualVal0, err := db.GetCF(ro, cfh[0], givenKey0)
-	defer actualVal0.Free()
-	require.Nil(t, err)
-	require.EqualValues(t, actualVal0.Data(), givenVal0)
+	writeReadBatch(cfh[0], [][]byte{givenKey0}, [][]byte{givenVal0})
 
-	b1 := NewWriteBatch()
-	defer b1.Destroy()
-	b1.PutCF(cfh[1], givenKey1, givenVal1)
-	require.Nil(t, db.Write(wo, b1))
-	actualVal1, err := db.GetCF(ro, cfh[1], givenKey1)
-	defer actualVal1.Free()
-	require.Nil(t, err)
-	require.EqualValues(t, actualVal1.Data(), givenVal1)
+	writeReadBatch(cfh[1], [][]byte{givenKey1, givenKey2}, [][]byte{givenVal1, givenVal2})
 
+	// check read from wrong CF returns nil
 	actualVal, err := db.GetCF(ro, cfh[0], givenKey1)
 	require.Nil(t, err)
 	require.EqualValues(t, actualVal.Size(), 0)
+	actualVal.Free()
+
 	actualVal, err = db.GetCF(ro, cfh[1], givenKey0)
 	require.Nil(t, err)
 	require.EqualValues(t, actualVal.Size(), 0)
+	actualVal.Free()
+
+	// check batch read is correct
+	actualVals, err := db.MultiGetWithCF(ro, cfh[1], givenKey1, givenKey2)
+	require.Nil(t, err)
+	require.EqualValues(t, len(actualVals), 2)
+	require.EqualValues(t, actualVals[0].Data(), givenVal1)
+	require.EqualValues(t, actualVals[1].Data(), givenVal2)
+	actualVals.Destroy()
 
 	// trigger flush
 	require.Nil(t, db.FlushCF(cfh[0], NewDefaultFlushOptions()))