write_batch.go

package gorocksdb

// #include "rocksdb/c.h"
// #include "extended.h"
import "C"
import "io"

// WriteBatch is a batching of Puts, Merges and Deletes.
type WriteBatch struct {
	c *C.rocksdb_writebatch_t
}

// NewWriteBatch create a WriteBatch object.
func NewWriteBatch() *WriteBatch {
	return NewNativeWriteBatch(C.rocksdb_writebatch_create())
}

// NewNativeWriteBatch create a WriteBatch object.
func NewNativeWriteBatch(c *C.rocksdb_writebatch_t) *WriteBatch {
	return &WriteBatch{c}
}

// WriteBatchFrom creates a write batch from a serialized WriteBatch.
func WriteBatchFrom(data []byte) *WriteBatch {
	return NewNativeWriteBatch(C.rocksdb_writebatch_create_from(byteToChar(data), C.size_t(len(data))))
}

func (wb *WriteBatch) BatchedPut(key1, value1, key2, value2, key3, value3, key4, value4,
	key5, value5, key6, value6, key7, value7, key8, value8 []byte) {
	C.rocksdb_writebatch_batched_put(wb.c,
		byteToChar(key1), C.size_t(len(key1)), byteToChar(value1), C.size_t(len(value1)),
		byteToChar(key2), C.size_t(len(key2)), byteToChar(value2), C.size_t(len(value2)),
		byteToChar(key3), C.size_t(len(key3)), byteToChar(value3), C.size_t(len(value3)),
		byteToChar(key4), C.size_t(len(key4)), byteToChar(value4), C.size_t(len(value4)),
		byteToChar(key5), C.size_t(len(key5)), byteToChar(value5), C.size_t(len(value5)),
		byteToChar(key6), C.size_t(len(key6)), byteToChar(value6), C.size_t(len(value6)),
		byteToChar(key7), C.size_t(len(key7)), byteToChar(value7), C.size_t(len(value7)),
		byteToChar(key8), C.size_t(len(key8)), byteToChar(value8), C.size_t(len(value8)))
}

// Put queues a key-value pair.
func (wb *WriteBatch) Put(key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)
	C.rocksdb_writebatch_put(wb.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// PutCF queues a key-value pair in a column family.
func (wb *WriteBatch) PutCF(cf *ColumnFamilyHandle, key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)
	C.rocksdb_writebatch_put_cf(wb.c, cf.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// Merge queues a merge of "value" with the existing value of "key".
func (wb *WriteBatch) Merge(key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)
	C.rocksdb_writebatch_merge(wb.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// MergeCF queues a merge of "value" with the existing value of "key" in a
// column family.
func (wb *WriteBatch) MergeCF(cf *ColumnFamilyHandle, key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)
	C.rocksdb_writebatch_merge_cf(wb.c, cf.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// Delete queues a deletion of the data at key.
func (wb *WriteBatch) Delete(key []byte) {
	cKey := byteToChar(key)
	C.rocksdb_writebatch_delete(wb.c, cKey, C.size_t(len(key)))
}

// DeleteRange queues a deletion of the specified key range.
func (wb *WriteBatch) DeleteRange(startKey []byte, endKey []byte) {
	cStartKey := byteToChar(startKey)
	cEndKey := byteToChar(endKey)
	C.rocksdb_writebatch_delete_range(wb.c, cStartKey, C.size_t(len(startKey)),
		cEndKey, C.size_t(len(endKey)))
}

// DeleteCF queues a deletion of the data at key in a column family.
func (wb *WriteBatch) DeleteCF(cf *ColumnFamilyHandle, key []byte) {
	cKey := byteToChar(key)
	C.rocksdb_writebatch_delete_cf(wb.c, cf.c, cKey, C.size_t(len(key)))
}

// Data returns the serialized version of this batch.
func (wb *WriteBatch) Data() []byte {
	var cSize C.size_t
	cValue := C.rocksdb_writebatch_data(wb.c, &cSize)
	return charToByte(cValue, cSize)
}

// Count returns the number of updates in the batch.
func (wb *WriteBatch) Count() int {
	return int(C.rocksdb_writebatch_count(wb.c))
}

// NewIterator returns a iterator to iterate over the records in the batch.
func (wb *WriteBatch) NewIterator() *WriteBatchIterator {
	data := wb.Data()
	if len(data) < 8+4 {
		return &WriteBatchIterator{}
	}
	return &WriteBatchIterator{data: data[12:]}
}

// Clear removes all the enqueued Put and Deletes.
func (wb *WriteBatch) Clear() {
	C.rocksdb_writebatch_clear(wb.c)
}

// Destroy deallocates the WriteBatch object.
func (wb *WriteBatch) Destroy() {
	C.rocksdb_writebatch_destroy(wb.c)
	wb.c = nil
}

// WriteBatchRecordType describes the type of a batch record.
type WriteBatchRecordType byte

// Types of batch records.
const (
	WriteBatchRecordTypeDeletion WriteBatchRecordType = 0x0
	WriteBatchRecordTypeValue    WriteBatchRecordType = 0x1
	WriteBatchRecordTypeMerge    WriteBatchRecordType = 0x2
	WriteBatchRecordTypeLogData  WriteBatchRecordType = 0x3
)

// WriteBatchRecord represents a record inside a WriteBatch.
type WriteBatchRecord struct {
	Key   []byte
	Value []byte
	Type  WriteBatchRecordType
}

// WriteBatchIterator represents a iterator to iterator over records.
type WriteBatchIterator struct {
	data   []byte
	record WriteBatchRecord
	err    error
}

// Next returns the next record.
// Returns false if no further record exists.
func (iter *WriteBatchIterator) Next() bool {
	if iter.err != nil || len(iter.data) == 0 {
		return false
	}
	// reset the current record
	iter.record.Key = nil
	iter.record.Value = nil

	// parse the record type
	recordType := WriteBatchRecordType(iter.data[0])
	iter.record.Type = recordType
	iter.data = iter.data[1:]

	// parse the key
	x, n := iter.decodeVarint(iter.data)
	if n == 0 {
		iter.err = io.ErrShortBuffer
		return false
	}
	k := n + int(x)
	iter.record.Key = iter.data[n:k]
	iter.data = iter.data[k:]

	// parse the data
	if recordType == WriteBatchRecordTypeValue || recordType == WriteBatchRecordTypeMerge {
		x, n := iter.decodeVarint(iter.data)
		if n == 0 {
			iter.err = io.ErrShortBuffer
			return false
		}
		k := n + int(x)
		iter.record.Value = iter.data[n:k]
		iter.data = iter.data[k:]
	}
	return true
}

// Record returns the current record.
func (iter *WriteBatchIterator) Record() *WriteBatchRecord {
	return &iter.record
}

// Error returns the error if the iteration is failed.
func (iter *WriteBatchIterator) Error() error {
	return iter.err
}

func (iter *WriteBatchIterator) decodeVarint(buf []byte) (x uint64, n int) {
	// x, n already 0
	for shift := uint(0); shift < 64; shift += 7 {
		if n >= len(buf) {
			return 0, 0
		}
		b := uint64(buf[n])
		n++
		x |= (b & 0x7F) << shift
		if (b & 0x80) == 0 {
			return x, n
		}
	}
	// The number is too large to represent in a 64-bit value.
	return 0, 0
}