From 3ce82982bbbff3483873b591c468d3e552a72d02 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal@scylladb.com>
Date: Thu, 28 Sep 2017 15:10:47 +0200
Subject: [PATCH] iterx: ErrNotFound for Get only

---
 integration_test.go |  4 ++--
 iterx.go            | 40 +++++++++++++++++++++++-----------------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/integration_test.go b/integration_test.go
index 79a256e..2e52a9c 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -319,8 +319,8 @@ func TestNotFound(t *testing.T) {
 	t.Run("select", func(t *testing.T) {
 		var v []NotFoundTable
 		i := gocqlx.Iter(session.Query(`SELECT * FROM not_found_table`))
-		if err := i.Select(&v); err != gocql.ErrNotFound {
-			t.Fatal("expected ErrNotFound", "got", err)
+		if err := i.Select(&v); err != nil {
+			t.Fatal(err)
 		}
 		if cap(v) > 0 {
 			t.Fatal("side effect alloc")
diff --git a/iterx.go b/iterx.go
index d19c2f8..1fb7121 100644
--- a/iterx.go
+++ b/iterx.go
@@ -13,12 +13,12 @@ import (
 	"github.com/jmoiron/sqlx/reflectx"
 )
 
-// Get is a convenience function for creating iterator and calling Get on it.
+// Get is a convenience function for creating iterator and calling Get.
 func Get(dest interface{}, q *gocql.Query) error {
 	return Iter(q).Get(dest)
 }
 
-// Select is a convenience function for creating iterator and calling Select on it.
+// Select is a convenience function for creating iterator and calling Select.
 func Select(dest interface{}, q *gocql.Query) error {
 	return Iter(q).Select(dest)
 }
@@ -67,7 +67,7 @@ func (iter *Iterx) Get(dest interface{}) error {
 	iter.Close()
 	iter.ReleaseQuery()
 
-	return iter.err
+	return iter.checkErrAndNotFound()
 }
 
 func (iter *Iterx) scanAny(dest interface{}, structOnly bool) bool {
@@ -80,8 +80,9 @@ func (iter *Iterx) scanAny(dest interface{}, structOnly bool) bool {
 		iter.err = errors.New("nil pointer passed to StructScan destination")
 		return false
 	}
+
+	// no results or query error
 	if iter.Iter.NumRows() == 0 {
-		// no results or query error
 		return false
 	}
 
@@ -134,8 +135,9 @@ func (iter *Iterx) scanAll(dest interface{}, structOnly bool) bool {
 		iter.err = errors.New("nil pointer passed to StructScan destination")
 		return false
 	}
+
+	// no results or query error
 	if iter.Iter.NumRows() == 0 {
-		// no results or query error
 		return false
 	}
 
@@ -201,12 +203,12 @@ func (iter *Iterx) scanAll(dest interface{}, structOnly bool) bool {
 	return true
 }
 
-// StructScan is like gocql.Scan, but scans a single row into a single Struct.
-// Use this and iterate manually when the memory load of Select() might be
-// prohibitive. StructScan caches the reflect work of matching up column
+// StructScan is like gocql.Iter.Scan, but scans a single row into a single
+// Struct. Use this and iterate manually when the memory load of Select() might
+// be prohibitive. StructScan caches the reflect work of matching up column
 // positions to fields to avoid that overhead per scan, which means it is not
 // safe to run StructScan on the same Iterx instance with different struct
-// types. If no rows were selected, ErrNotFound is returned.
+// types.
 func (iter *Iterx) StructScan(dest interface{}) bool {
 	if iter.query == nil {
 		iter.err = errors.New("using released query")
@@ -219,8 +221,8 @@ func (iter *Iterx) StructScan(dest interface{}) bool {
 		return false
 	}
 
+	// no results or query error
 	if iter.Iter.NumRows() == 0 {
-		// no results or query error
 		return false
 	}
 
@@ -261,15 +263,9 @@ func columnNames(ci []gocql.ColumnInfo) []string {
 // the query or the iteration.
 func (iter *Iterx) Close() error {
 	err := iter.Iter.Close()
-
 	if iter.err == nil {
-		if err != nil {
-			iter.err = err
-		} else if iter.Iter.NumRows() == 0 {
-			iter.err = gocql.ErrNotFound
-		}
+		iter.err = err
 	}
-
 	return iter.err
 }
 
@@ -281,3 +277,13 @@ func (iter *Iterx) ReleaseQuery() {
 		iter.query = nil
 	}
 }
+
+// checkErrAndNotFound handle error and NotFound in one method.
+func (iter *Iterx) checkErrAndNotFound() error {
+	if iter.err != nil {
+		return iter.err
+	} else if iter.Iter.NumRows() == 0 {
+		return gocql.ErrNotFound
+	}
+	return nil
+}