Move row storage logic to separate package. (#68)

Row data message can be a simple byte slice.

Author: Michael

connection.go

@@ -355,12 +355,14 @@ func (v *connection) initializeSession() error {
 		return err
 	}
 
-	if len(result.Columns()) != 1 && result.Columns()[1] != "now" || len(result.resultData) != 1 {
+	firstRow := result.resultData.Peek()
+
+	if len(result.Columns()) != 1 && result.Columns()[1] != "now" || firstRow == nil {
 		return fmt.Errorf("unable to initialize session; functionality may be unreliable")
 	}
 
 	// Peek into the results manually.
-	colData := result.resultData[0].Columns()
+	colData := firstRow.Columns()
 	str := string(colData.Chunk())
 
 	if len(str) < 23 {

msgs/bedatarowmsg.go

@@ -33,15 +33,12 @@ package msgs
 // THE SOFTWARE.
 
 import (
-	"bytes"
 	"encoding/binary"
 	"fmt"
 )
 
 // BEDataRowMsg docs
-type BEDataRowMsg struct {
-	rowBuffer *bytes.Buffer
-}
+type BEDataRowMsg []byte
 
 // ColumnExtractor pulls columns out of a row
 type ColumnExtractor struct {
@@ -64,24 +61,22 @@ func (c *ColumnExtractor) Chunk() []byte {
 
 // CreateFromMsgBody docs
 func (b *BEDataRowMsg) CreateFromMsgBody(buf *msgBuffer) (BackEndMsg, error) {
-	newBuf := bytes.NewBuffer(buf.buf.Bytes())
+	res := BEDataRowMsg(buf.buf.Bytes())
 	buf.buf.Reset()
-	res := &BEDataRowMsg{rowBuffer: newBuf}
-	return res, nil
+	return &res, nil
 }
 
 // Columns provides an extractor to begin reading columns
 func (b *BEDataRowMsg) Columns() ColumnExtractor {
-	rowData := b.rowBuffer.Bytes()
 	return ColumnExtractor{
-		NumCols: binary.BigEndian.Uint16(rowData[0:2]),
-		data:    rowData[2:],
+		NumCols: binary.BigEndian.Uint16((*b)[0:2]),
+		data:    (*b)[2:],
 	}
 }
 
 // RevertToBytes dumps the message back into plain bytes
 func (b *BEDataRowMsg) RevertToBytes() []byte {
-	return b.rowBuffer.Bytes()
+	return *b
 }
 
 func (b *BEDataRowMsg) String() string {

rowcache/file.go

@@ -0,0 +1,148 @@
+package rowcache
+
+import (
+	"bufio"
+	"encoding/binary"
+	"io"
+	"io/ioutil"
+	"os"
+
+	"github.com/vertica/vertica-sql-go/msgs"
+)
+
+// FileCache stores rows from the wire and puts excess into a temporary file
+type FileCache struct {
+	maxInMemory int
+	rowCount    int
+	readIdx     int
+	resultData  []*msgs.BEDataRowMsg
+	file        *os.File
+	rwbuffer    *bufio.ReadWriter
+	scratch     [512]byte
+}
+
+// NewFileCache returns a file cache with a set row limit
+func NewFileCache(rowLimit int) (*FileCache, error) {
+	file, err := ioutil.TempFile("", ".vertica-sql-go.*.dat")
+	if err != nil {
+		return nil, err
+	}
+	return &FileCache{
+		maxInMemory: rowLimit,
+		resultData:  make([]*msgs.BEDataRowMsg, 0, rowLimit),
+		file:        file,
+		rwbuffer:    bufio.NewReadWriter(bufio.NewReader(file), bufio.NewWriterSize(file, 1<<16)),
+	}, nil
+}
+
+func (f *FileCache) writeCached(msg *msgs.BEDataRowMsg) {
+	sizeBuf := f.scratch[:4]
+	binary.LittleEndian.PutUint32(sizeBuf, uint32(len(*msg)))
+	f.rwbuffer.Write(sizeBuf)
+	f.rwbuffer.Write(*msg)
+}
+
+// AddRow adds a row to the cache
+func (f *FileCache) AddRow(msg *msgs.BEDataRowMsg) {
+	f.rowCount++
+	if len(f.resultData) >= f.maxInMemory {
+		f.writeCached(msg)
+		return
+	}
+	f.resultData = append(f.resultData, msg)
+}
+
+// Finalize signals the end of rows from the wire and readies the cache for reading
+func (f *FileCache) Finalize() error {
+	var err error
+	name := f.file.Name()
+	f.rwbuffer.Flush()
+	f.file.Close()
+	f.file, err = os.OpenFile(name, os.O_RDONLY|os.O_EXCL, 0600)
+	if err != nil {
+		return err
+	}
+	f.rwbuffer = bufio.NewReadWriter(bufio.NewReader(f.file), bufio.NewWriter(f.file))
+	return err
+}
+
+// GetRow pulls a row message out of the cache, returning nil of none remain
+func (f *FileCache) GetRow() *msgs.BEDataRowMsg {
+	if f.readIdx >= len(f.resultData) {
+		if !f.reloadFromCache() {
+			return nil
+		}
+	}
+	result := f.resultData[f.readIdx]
+	f.readIdx++
+	return result
+}
+
+// Peek returns the next row without changing the state
+func (f *FileCache) Peek() *msgs.BEDataRowMsg {
+	if len(f.resultData) == 0 {
+		return nil
+	}
+	return f.resultData[0]
+}
+
+// Close clears resources associated with the cache, deleting the temp file
+func (f *FileCache) Close() error {
+	name := f.file.Name()
+	f.rwbuffer.Flush()
+	f.file.Close()
+	return os.Remove(name)
+}
+
+func (f *FileCache) reloadFromCache() bool {
+	hadData := false
+
+	f.readIdx = 0
+	indexCount := 0
+
+	for {
+		sizeBuf := f.scratch[:4]
+
+		if _, err := io.ReadFull(f.rwbuffer, sizeBuf); err != nil {
+			if err == io.EOF {
+				if indexCount == 0 {
+					return false
+				}
+				f.resultData = f.resultData[0:indexCount]
+				return true
+			}
+			return false
+		}
+
+		rowDataSize := binary.LittleEndian.Uint32(sizeBuf)
+
+		var rowBuf []byte
+		rowBytes := f.scratch[4:]
+		if rowDataSize <= uint32(len(rowBytes)) {
+			rowBuf = rowBytes[:rowDataSize]
+		} else {
+			rowBuf = make([]byte, rowDataSize)
+		}
+		if _, err := io.ReadFull(f.rwbuffer, rowBuf); err != nil {
+			return false
+		}
+
+		msgBuf := msgs.NewMsgBufferFromBytes(rowBuf)
+
+		drm := &msgs.BEDataRowMsg{}
+
+		msg, _ := drm.CreateFromMsgBody(msgBuf)
+
+		f.resultData[indexCount] = msg.(*msgs.BEDataRowMsg)
+		indexCount++
+
+		hadData = true
+
+		// If we've reached the original capacity of the slice, we're done.
+		if indexCount == len(f.resultData) {
+			break
+		}
+	}
+
+	return hadData
+}

rowcache/file_test.go

@@ -0,0 +1,55 @@
+package rowcache
+
+import (
+	"testing"
+
+	"github.com/vertica/vertica-sql-go/msgs"
+)
+
+func TestFileCache(t *testing.T) {
+	t.Run("without enough rows to switch to a file", func(t *testing.T) {
+		rowCount := 100
+		cache, err := NewFileCache(1000)
+		if err != nil {
+			t.Fatalf("Unable to create temp file")
+		}
+		for i := 0; i < rowCount; i++ {
+			row := msgs.BEDataRowMsg([]byte("testRow"))
+			cache.AddRow(&row)
+		}
+		cache.Finalize()
+		if cache.Peek() == nil {
+			t.Error("Expected a row with Peek")
+		}
+		for i := 0; i < rowCount; i++ {
+			row := cache.GetRow()
+			if row == nil {
+				t.Errorf("Ran out of rows at %d", i)
+			}
+		}
+		cache.Close()
+	})
+	t.Run("with file writes", func(t *testing.T) {
+		rowCount := 10000
+		cache, err := NewFileCache(100)
+		if err != nil {
+			t.Fatalf("Unable to create temp file")
+		}
+		for i := 0; i < rowCount; i++ {
+			row := msgs.BEDataRowMsg([]byte("testRow"))
+			cache.AddRow(&row)
+		}
+		cache.Finalize()
+		if cache.Peek() == nil {
+			t.Error("Expected a row with Peek")
+		}
+		for i := 0; i < rowCount; i++ {
+			row := cache.GetRow()
+			if row == nil {
+				t.Errorf("Ran out of rows at %d", i)
+			}
+		}
+		cache.Close()
+	})
+
+}

rowcache/memory.go

@@ -0,0 +1,52 @@
+package rowcache
+
+import (
+	"github.com/vertica/vertica-sql-go/msgs"
+)
+
+// MemoryCache is a simple in memory row store
+type MemoryCache struct {
+	resultData []*msgs.BEDataRowMsg
+	readIdx    int
+}
+
+// NewMemoryCache initializes the memory store with a given size, but it can continue
+// to grow
+func NewMemoryCache(size int) *MemoryCache {
+	return &MemoryCache{
+		resultData: make([]*msgs.BEDataRowMsg, 0, size),
+	}
+}
+
+// AddRow adds a row to the store
+func (m *MemoryCache) AddRow(msg *msgs.BEDataRowMsg) {
+	m.resultData = append(m.resultData, msg)
+}
+
+// Finalize signals the end of new rows, a noop for the memory cache
+func (m *MemoryCache) Finalize() error {
+	return nil
+}
+
+// GetRow pulls a row from the cache, returning nil if none remain
+func (m *MemoryCache) GetRow() *msgs.BEDataRowMsg {
+	if m.readIdx >= len(m.resultData) {
+		return nil
+	}
+	result := m.resultData[m.readIdx]
+	m.readIdx++
+	return result
+}
+
+// Peek returns the next row without changing the state
+func (m *MemoryCache) Peek() *msgs.BEDataRowMsg {
+	if len(m.resultData) == 0 {
+		return nil
+	}
+	return m.resultData[0]
+}
+
+// Close provides an opportunity to free resources, a noop for the memory cache
+func (m *MemoryCache) Close() error {
+	return nil
+}

rowcache/memory_test.go

@@ -0,0 +1,27 @@
+package rowcache
+
+import (
+	"testing"
+
+	"github.com/vertica/vertica-sql-go/msgs"
+)
+
+func TestMemoryCache(t *testing.T) {
+	rowCount := 100
+	cache := NewMemoryCache(16)
+	for i := 0; i < rowCount; i++ {
+		row := msgs.BEDataRowMsg([]byte("testRow"))
+		cache.AddRow(&row)
+	}
+	cache.Finalize()
+	if cache.Peek() == nil {
+		t.Error("Expected a row with Peek")
+	}
+	for i := 0; i < rowCount; i++ {
+		row := cache.GetRow()
+		if row == nil {
+			t.Errorf("Ran out of rows at %d", i)
+		}
+	}
+	cache.Close()
+}