added error handling when working with file cache #113

rowcache/file.go

@@ -67,29 +67,41 @@ func NewFileCache(rowLimit int) (*FileCache, error) {
 	}, nil
 }
 
-func (f *FileCache) writeCached(msg *msgs.BEDataRowMsg) {
+func (f *FileCache) writeCached(msg *msgs.BEDataRowMsg) error {
 	sizeBuf := f.scratch[:4]
 	binary.LittleEndian.PutUint32(sizeBuf, uint32(len(*msg)))
-	f.rwbuffer.Write(sizeBuf)
-	f.rwbuffer.Write(*msg)
+	if _, err := f.rwbuffer.Write(sizeBuf); err != nil {
+		return err
+	}
+	if _, err := f.rwbuffer.Write(*msg); err != nil {
+		return err
+	}
+	return nil
 }
 
 // AddRow adds a row to the cache
-func (f *FileCache) AddRow(msg *msgs.BEDataRowMsg) {
+func (f *FileCache) AddRow(msg *msgs.BEDataRowMsg) error {
 	f.rowCount++
 	if len(f.resultData) >= f.maxInMemory {
-		f.writeCached(msg)
-		return
+		if err := f.writeCached(msg); err != nil {
+			return err
+		}
+		return nil
 	}
 	f.resultData = append(f.resultData, msg)
+	return nil
 }
 
 // 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()
+	if err = f.rwbuffer.Flush(); err != nil {
+		return err
+	}
+	if err = f.file.Close(); err != nil {
+		return err
+	}
 	f.file, err = os.OpenFile(name, os.O_RDONLY|os.O_EXCL, 0600)
 	if err != nil {
 		return err

rowcache/file_test.go

@@ -47,9 +47,16 @@ func TestFileCache(t *testing.T) {
 		}
 		for i := 0; i < rowCount; i++ {
 			row := msgs.BEDataRowMsg([]byte("testRow"))
-			cache.AddRow(&row)
+			err = cache.AddRow(&row)
+			if err != nil {
+				t.Fatalf("Unable to add a row to the cache: %s", err)
+			}
+		}
+		err = cache.Finalize()
+		if err != nil {
+			t.Fatalf("Unable to finalize the cache: %s", err)
 		}
-		cache.Finalize()
+
 		if cache.Peek() == nil {
 			t.Error("Expected a row with Peek")
 		}
@@ -69,9 +76,15 @@ func TestFileCache(t *testing.T) {
 		}
 		for i := 0; i < rowCount; i++ {
 			row := msgs.BEDataRowMsg([]byte("testRow"))
-			cache.AddRow(&row)
+			err = cache.AddRow(&row)
+			if err != nil {
+				t.Fatalf("Unable to add a row to the cache: %s", err)
+			}
+		}
+		err = cache.Finalize()
+		if err != nil {
+			t.Fatalf("Unable to finalize the cache: %s", err)
 		}
-		cache.Finalize()
 		if cache.Peek() == nil {
 			t.Error("Expected a row with Peek")
 		}

rowcache/memory.go

@@ -51,8 +51,9 @@ func NewMemoryCache(size int) *MemoryCache {
 }
 
 // AddRow adds a row to the store
-func (m *MemoryCache) AddRow(msg *msgs.BEDataRowMsg) {
+func (m *MemoryCache) AddRow(msg *msgs.BEDataRowMsg) error {
 	m.resultData = append(m.resultData, msg)
+	return nil
 }
 
 // Finalize signals the end of new rows, a noop for the memory cache

rows.go

@@ -52,7 +52,7 @@ import (
 )
 
 type rowStore interface {
-	AddRow(msg *msgs.BEDataRowMsg)
+	AddRow(msg *msgs.BEDataRowMsg) error
 	GetRow() *msgs.BEDataRowMsg
 	Peek() *msgs.BEDataRowMsg
 	Close() error
@@ -196,12 +196,12 @@ func parseTimestampTZColumn(fullString string) (driver.Value, error) {
 	return result, err
 }
 
-func (r *rows) finalize() {
-	r.resultData.Finalize()
+func (r *rows) finalize() error {
+	return r.resultData.Finalize()
 }
 
-func (r *rows) addRow(rowData *msgs.BEDataRowMsg) {
-	r.resultData.AddRow(rowData)
+func (r *rows) addRow(rowData *msgs.BEDataRowMsg) error {
+	return r.resultData.AddRow(rowData)
 }
 
 func newRows(ctx context.Context, columnsDefsMsg *msgs.BERowDescMsg, tzOffset string) *rows {

stmt.go

@@ -304,7 +304,10 @@ func (s *stmt) QueryContextRaw(ctx context.Context, baseArgs []driver.NamedValue
 
 		switch msg := bMsg.(type) {
 		case *msgs.BEDataRowMsg:
-			rows.addRow(msg)
+			err = rows.addRow(msg)
+			if err != nil {
+				return rows, err
+			}
 		case *msgs.BERowDescMsg:
 			rows = newRows(ctx, msg, s.conn.serverTZOffset)
 		case *msgs.BECmdCompleteMsg:
@@ -314,7 +317,10 @@ func (s *stmt) QueryContextRaw(ctx context.Context, baseArgs []driver.NamedValue
 		case *msgs.BEEmptyQueryResponseMsg:
 			return newEmptyRows(), nil
 		case *msgs.BEReadyForQueryMsg, *msgs.BEPortalSuspendedMsg:
-			rows.finalize()
+			err = rows.finalize()
+			if err != nil {
+				return rows, err
+			}
 			return rows, ctx.Err()
 		case *msgs.BEInitSTDINLoadMsg:
 			s.copySTDIN(ctx)
@@ -528,7 +534,10 @@ func (s *stmt) collectResults(ctx context.Context) (*rows, error) {
 
 		switch msg := bMsg.(type) {
 		case *msgs.BEDataRowMsg:
-			rows.addRow(msg)
+			err = rows.addRow(msg)
+			if err != nil {
+				return rows, err
+			}
 		case *msgs.BERowDescMsg:
 			s.lastRowDesc = msg
 			rows = newRows(ctx, s.lastRowDesc, s.conn.serverTZOffset)
@@ -540,7 +549,10 @@ func (s *stmt) collectResults(ctx context.Context) (*rows, error) {
 		case *msgs.BEBindCompleteMsg, *msgs.BECmdDescriptionMsg:
 			continue
 		case *msgs.BEReadyForQueryMsg, *msgs.BEPortalSuspendedMsg, *msgs.BECmdCompleteMsg:
-			rows.finalize()
+			err = rows.finalize()
+			if err != nil {
+				return rows, err
+			}
 			return rows, ctx.Err()
 		case *msgs.BEInitSTDINLoadMsg:
 			s.copySTDIN(ctx)