iterx: Allow forcing scanning as struct

We have a structure type that implements UnmarshalCQL method.
We use it to unmarshal a user defined type. We also want to use the same struct for scanning an entire row.

There is StructScan method available in gocqlx for this purpose when iterating over rows, but no equivalent when doing a Select.
This commit introduces the possibility when doing select/get as well.

Co-authored-by: Michał Matczuk <michal@scylladb.com>

gocqlx.go

@@ -14,22 +14,29 @@ import (
 )
 
 // structOnlyError returns an error appropriate for type when a non-scannable
-// struct is expected but something else is given
+// struct is expected but something else is given.
 func structOnlyError(t reflect.Type) error {
-	isStruct := t.Kind() == reflect.Struct
+	if isStruct := t.Kind() == reflect.Struct; !isStruct {
-	isScanner := reflect.PtrTo(t).Implements(_unmarshallerInterface)
+		return fmt.Errorf("expected a struct but got %s", t.Kind())
-	if !isStruct {
-		return fmt.Errorf("expected %s but got %s", reflect.Struct, t.Kind())
 	}
-	if isScanner {
+
-		return fmt.Errorf("structscan expects a struct dest but the provided struct type %s implements unmarshaler", t.Name())
+	if isUnmarshaller := reflect.PtrTo(t).Implements(unmarshallerInterface); isUnmarshaller {
+		return fmt.Errorf("expected a struct but the provided struct type %s implements gocql.Unmarshaler", t.Name())
 	}
+
+	if isUDTUnmarshaller := reflect.PtrTo(t).Implements(udtUnmarshallerInterface); isUDTUnmarshaller {
+		return fmt.Errorf("expected a struct but the provided struct type %s implements gocql.UDTUnmarshaler", t.Name())
+	}
+
 	return fmt.Errorf("expected a struct, but struct %s has no exported fields", t.Name())
 }
 
 // reflect helpers
 
-var _unmarshallerInterface = reflect.TypeOf((*gocql.Unmarshaler)(nil)).Elem()
+var (
+	unmarshallerInterface    = reflect.TypeOf((*gocql.Unmarshaler)(nil)).Elem()
+	udtUnmarshallerInterface = reflect.TypeOf((*gocql.UDTUnmarshaler)(nil)).Elem()
+)
 
 func baseType(t reflect.Type, expected reflect.Kind) (reflect.Type, error) {
 	t = reflectx.Deref(t)

iterx.go

@@ -22,9 +22,10 @@ type Iterx struct {
 	*gocql.Iter
 	Mapper *reflectx.Mapper
 
-	unsafe  bool
+	unsafe     bool
-	started bool
+	structOnly bool
-	err     error
+	started    bool
+	err        error
 
 	// Cache memory for a rows during iteration in StructScan.
 	fields [][]int
@@ -48,14 +49,29 @@ func (iter *Iterx) Unsafe() *Iterx {
 	return iter
 }
 
-// Get scans first row into a destination and closes the iterator. If the
+// StructOnly forces the iterator to treat a single-argument struct as
-// destination type is a struct pointer, then StructScan will be used.
+// non-scannable. This is is useful if you need to scan a row into a struct
+// that also implements gocql.UDTUnmarshaler or in rare cases gocql.Unmarshaler.
+func (iter *Iterx) StructOnly() *Iterx {
+	iter.structOnly = true
+	return iter
+}
+
+// Get scans first row into a destination and closes the iterator.
+//
+// If the destination type is a struct pointer, then StructScan will be
+// used.
 // If the destination is some other type, then the row must only have one column
 // which can scan into that type.
+// This includes types that implement gocql.Unmarshaler and gocql.UDTUnmarshaler.
+//
+// If you'd like to treat a type that implements gocql.Unmarshaler or
+// gocql.UDTUnmarshaler as an ordinary struct you should call
+// StructOnly().Get(dest) instead.
 //
 // If no rows were selected, ErrNotFound is returned.
 func (iter *Iterx) Get(dest interface{}) error {
-	iter.scanAny(dest, false)
+	iter.scanAny(dest)
 	iter.Close()
 
 	return iter.checkErrAndNotFound()
@@ -63,11 +79,11 @@ func (iter *Iterx) Get(dest interface{}) error {
 
 // isScannable takes the reflect.Type and the actual dest value and returns
 // whether or not it's Scannable. t is scannable if:
-//   * ptr to t implements gocql.Unmarshaler
+//   * ptr to t implements gocql.Unmarshaler or gocql.UDTUnmarshaler
 //   * it is not a struct
 //   * it has no exported fields
 func (iter *Iterx) isScannable(t reflect.Type) bool {
-	if reflect.PtrTo(t).Implements(_unmarshallerInterface) {
+	if ptr := reflect.PtrTo(t); ptr.Implements(unmarshallerInterface) || ptr.Implements(udtUnmarshallerInterface) {
 		return true
 	}
 	if t.Kind() != reflect.Struct {
@@ -77,7 +93,7 @@ func (iter *Iterx) isScannable(t reflect.Type) bool {
 	return len(iter.Mapper.TypeMap(t).Index) == 0
 }
 
-func (iter *Iterx) scanAny(dest interface{}, structOnly bool) bool {
+func (iter *Iterx) scanAny(dest interface{}) bool {
 	value := reflect.ValueOf(dest)
 	if value.Kind() != reflect.Ptr {
 		iter.err = fmt.Errorf("expected a pointer but got %T", dest)
@@ -91,13 +107,17 @@ func (iter *Iterx) scanAny(dest interface{}, structOnly bool) bool {
 	base := reflectx.Deref(value.Type())
 	scannable := iter.isScannable(base)
 
-	if structOnly && scannable {
+	if iter.structOnly && scannable {
-		iter.err = structOnlyError(base)
+		if base.Kind() == reflect.Struct {
-		return false
+			scannable = false
+		} else {
+			iter.err = structOnlyError(base)
+			return false
+		}
 	}
 
 	if scannable && len(iter.Columns()) > 1 {
-		iter.err = fmt.Errorf("scannable dest type %s with >1 columns (%d) in result", base.Kind(), len(iter.Columns()))
+		iter.err = fmt.Errorf("expected 1 column in result while scanning scannable type %s but got %d", base.Kind(), len(iter.Columns()))
 		return false
 	}
 
@@ -109,20 +129,27 @@ func (iter *Iterx) scanAny(dest interface{}, structOnly bool) bool {
 }
 
 // Select scans all rows into a destination, which must be a pointer to slice
-// of any type and closes the iterator. If the destination slice type is
+// of any type, and closes the iterator.
-// a struct, then StructScan will be used on each row. If the destination is
+//
-// some other type, then each row must only have one column which can scan into
+// If the destination slice type is a struct, then StructScan will be used
-// that type.
+// on each row.
+// If the destination is some other type, then each row must only have one
+// column which can scan into that type.
+// This includes types that implement gocql.Unmarshaler and gocql.UDTUnmarshaler.
+//
+// If you'd like to treat a type that implements gocql.Unmarshaler or
+// gocql.UDTUnmarshaler as an ordinary struct you should call
+// StructOnly().Select(dest) instead.
 //
 // If no rows were selected, ErrNotFound is NOT returned.
 func (iter *Iterx) Select(dest interface{}) error {
-	iter.scanAll(dest, false)
+	iter.scanAll(dest)
 	iter.Close()
 
 	return iter.err
 }
 
-func (iter *Iterx) scanAll(dest interface{}, structOnly bool) bool {
+func (iter *Iterx) scanAll(dest interface{}) bool {
 	value := reflect.ValueOf(dest)
 
 	// json.Unmarshal returns errors for these
@@ -145,14 +172,18 @@ func (iter *Iterx) scanAll(dest interface{}, structOnly bool) bool {
 	base := reflectx.Deref(slice.Elem())
 	scannable := iter.isScannable(base)
 
-	if structOnly && scannable {
+	if iter.structOnly && scannable {
-		iter.err = structOnlyError(base)
+		if base.Kind() == reflect.Struct {
-		return false
+			scannable = false
+		} else {
+			iter.err = structOnlyError(base)
+			return false
+		}
 	}
 
 	// if it's a base type make sure it only has 1 column;  if not return an error
 	if scannable && len(iter.Columns()) > 1 {
-		iter.err = fmt.Errorf("non-struct dest type %s with >1 columns (%d)", base.Kind(), len(iter.Columns()))
+		iter.err = fmt.Errorf("expected 1 column in result while scanning scannable type %s but got %d", base.Kind(), len(iter.Columns()))
 		return false
 	}
 

iterx_test.go

@@ -35,9 +35,29 @@ func (n *FullName) UnmarshalCQL(info gocql.TypeInfo, data []byte) error {
 	return nil
 }
 
+type FullNameUDT struct {
+	FirstName string
+	LastName  string
+}
+
+func (n FullNameUDT) MarshalUDT(name string, info gocql.TypeInfo) ([]byte, error) {
+	f := gocqlx.DefaultMapper.FieldByName(reflect.ValueOf(n), name)
+	return gocql.Marshal(info, f.Interface())
+}
+
+func (n *FullNameUDT) UnmarshalUDT(name string, info gocql.TypeInfo, data []byte) error {
+	f := gocqlx.DefaultMapper.FieldByName(reflect.ValueOf(n), name)
+	return gocql.Unmarshal(info, data, f.Addr().Interface())
+}
+
 func TestStruct(t *testing.T) {
 	session := CreateSession(t)
 	defer session.Close()
+
+	if err := ExecStmt(session, `CREATE TYPE gocqlx_test.FullName (first_Name text, last_name text)`); err != nil {
+		t.Fatal("create type:", err)
+	}
+
 	if err := ExecStmt(session, `CREATE TABLE gocqlx_test.struct_table (
 			testuuid       timeuuid PRIMARY KEY,
 			testtimestamp  timestamp,
@@ -54,8 +74,8 @@ func TestStruct(t *testing.T) {
 			testmap        map<varchar, varchar>,
 			testvarint     varint,
 			testinet       inet,
-			testcustom     text
+			testcustom     text,
-
+			testudt        gocqlx_test.FullName
 		)`); err != nil {
 		t.Fatal("create table:", err)
 	}
@@ -77,6 +97,7 @@ func TestStruct(t *testing.T) {
 		Testvarint    *big.Int
 		Testinet      string
 		Testcustom    FullName
+		Testudt       FullNameUDT
 	}
 
 	bigInt := new(big.Int)
@@ -101,9 +122,10 @@ func TestStruct(t *testing.T) {
 		Testvarint:    bigInt,
 		Testinet:      "213.212.2.19",
 		Testcustom:    FullName{FirstName: "John", LastName: "Doe"},
+		Testudt:       FullNameUDT{FirstName: "John", LastName: "Doe"},
 	}
 
-	if err := session.Query(`INSERT INTO struct_table (testuuid, testtimestamp, testvarchar, testbigint, testblob, testbool, testfloat,testdouble, testint, testdecimal, testlist, testset, testmap, testvarint, testinet, testcustom) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
+	if err := session.Query(`INSERT INTO struct_table (testuuid, testtimestamp, testvarchar, testbigint, testblob, testbool, testfloat,testdouble, testint, testdecimal, testlist, testset, testmap, testvarint, testinet, testcustom, testudt) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
 		m.Testuuid,
 		m.Testtimestamp,
 		m.Testvarchar,
@@ -119,7 +141,8 @@ func TestStruct(t *testing.T) {
 		m.Testmap,
 		m.Testvarint,
 		m.Testinet,
-		m.Testcustom).Exec(); err != nil {
+		m.Testcustom,
+		m.Testudt).Exec(); err != nil {
 		t.Fatal("insert:", err)
 	}
 
@@ -219,6 +242,148 @@ func TestScannable(t *testing.T) {
 	})
 }
 
+func TestStructOnly(t *testing.T) {
+	session := CreateSession(t)
+	defer session.Close()
+	if err := ExecStmt(session, `CREATE TABLE gocqlx_test.struct_only_table (first_name text, last_name text, PRIMARY KEY (first_name, last_name))`); err != nil {
+		t.Fatal("create table:", err)
+	}
+
+	m := FullName{"John", "Doe"}
+
+	if err := session.Query(`INSERT INTO struct_only_table (first_name, last_name) values (?, ?)`, m.FirstName, m.LastName).Exec(); err != nil {
+		t.Fatal("insert:", err)
+	}
+
+	t.Run("get", func(t *testing.T) {
+		var v FullName
+		if err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_table`)).StructOnly().Get(&v); err != nil {
+			t.Fatal("get failed", err)
+		}
+
+		if !reflect.DeepEqual(m, v) {
+			t.Fatal("not equals")
+		}
+	})
+
+	t.Run("select", func(t *testing.T) {
+		var v []FullName
+		if err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_table`)).StructOnly().Select(&v); err != nil {
+			t.Fatal("select failed", err)
+		}
+
+		if len(v) != 1 {
+			t.Fatal("select unexpected number of rows", len(v))
+		}
+
+		if !reflect.DeepEqual(m, v[0]) {
+			t.Fatal("not equals")
+		}
+	})
+
+	t.Run("select ptr", func(t *testing.T) {
+		var v []*FullName
+		if err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_table`)).StructOnly().Select(&v); err != nil {
+			t.Fatal("select failed", err)
+		}
+
+		if len(v) != 1 {
+			t.Fatal("select unexpected number of rows", len(v))
+		}
+
+		if !reflect.DeepEqual(&m, v[0]) {
+			t.Fatal("not equals")
+		}
+	})
+
+	t.Run("get error", func(t *testing.T) {
+		var v FullName
+		err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_table`)).Get(&v)
+		if err == nil || !strings.HasPrefix(err.Error(), "expected 1 column in result") {
+			t.Fatal("get expected validation error got", err)
+		}
+	})
+
+	t.Run("select error", func(t *testing.T) {
+		var v []FullName
+		err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_table`)).Select(&v)
+		if err == nil || !strings.HasPrefix(err.Error(), "expected 1 column in result") {
+			t.Fatal("select expected validation error got", err)
+		}
+	})
+}
+
+func TestStructOnlyUDT(t *testing.T) {
+	session := CreateSession(t)
+	defer session.Close()
+	if err := ExecStmt(session, `CREATE TABLE gocqlx_test.struct_only_udt_table (first_name text, last_name text, PRIMARY KEY (first_name, last_name))`); err != nil {
+		t.Fatal("create table:", err)
+	}
+
+	m := FullNameUDT{"John", "Doe"}
+
+	if err := session.Query(`INSERT INTO struct_only_udt_table (first_name, last_name) values (?, ?)`, m.FirstName, m.LastName).Exec(); err != nil {
+		t.Fatal("insert:", err)
+	}
+
+	t.Run("get", func(t *testing.T) {
+		var v FullNameUDT
+		if err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_udt_table`)).StructOnly().Get(&v); err != nil {
+			t.Fatal("get failed", err)
+		}
+
+		if !reflect.DeepEqual(m, v) {
+			t.Fatal("not equals")
+		}
+	})
+
+	t.Run("select", func(t *testing.T) {
+		var v []FullNameUDT
+		if err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_udt_table`)).StructOnly().Select(&v); err != nil {
+			t.Fatal("select failed", err)
+		}
+
+		if len(v) != 1 {
+			t.Fatal("select unexpected number of rows", len(v))
+		}
+
+		if !reflect.DeepEqual(m, v[0]) {
+			t.Fatal("not equals")
+		}
+	})
+
+	t.Run("select ptr", func(t *testing.T) {
+		var v []*FullNameUDT
+		if err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_udt_table`)).StructOnly().Select(&v); err != nil {
+			t.Fatal("select failed", err)
+		}
+
+		if len(v) != 1 {
+			t.Fatal("select unexpected number of rows", len(v))
+		}
+
+		if !reflect.DeepEqual(&m, v[0]) {
+			t.Fatal("not equals")
+		}
+	})
+
+	t.Run("get error", func(t *testing.T) {
+		var v FullNameUDT
+		err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_udt_table`)).Get(&v)
+		if err == nil || !strings.HasPrefix(err.Error(), "expected 1 column in result") {
+			t.Fatal("get expected validation error got", err)
+		}
+	})
+
+	t.Run("select error", func(t *testing.T) {
+		var v []FullNameUDT
+		err := gocqlx.Iter(session.Query(`SELECT first_name, last_name FROM struct_only_udt_table`)).Select(&v)
+		if err == nil || !strings.HasPrefix(err.Error(), "expected 1 column in result") {
+			t.Fatal("select expected validation error got", err)
+		}
+	})
+}
+
 func TestUnsafe(t *testing.T) {
 	session := CreateSession(t)
 	defer session.Close()

queryx.go

@@ -200,10 +200,17 @@ func (q *Queryx) ExecRelease() error {
 	return q.Exec()
 }
 
-// Get scans first row into a destination. If the destination type is a struct
+// Get scans first row into a destination and closes the iterator.
-// pointer, then Iter.StructScan will be used. If the destination is some
+//
-// other type, then the row must only have one column which can scan into that
+// If the destination type is a struct pointer, then Iter.StructScan will be
-// type.
+// used.
+// If the destination is some other type, then the row must only have one column
+// which can scan into that type.
+// This includes types that implement gocql.Unmarshaler and gocql.UDTUnmarshaler.
+//
+// If you'd like to treat a type that implements gocql.Unmarshaler or
+// gocql.UDTUnmarshaler as an ordinary struct you should call
+// Iter().StructOnly().Get(dest) instead.
 //
 // If no rows were selected, ErrNotFound is returned.
 func (q *Queryx) Get(dest interface{}) error {
@@ -221,9 +228,17 @@ func (q *Queryx) GetRelease(dest interface{}) error {
 }
 
 // Select scans all rows into a destination, which must be a pointer to slice
-// of any type. If the destination slice type is a struct, then Iter.StructScan
+// of any type, and closes the iterator.
-// will be used on each row. If the destination is some other type, then each
+//
-// row must only have one column which can scan into that type.
+// If the destination slice type is a struct, then Iter.StructScan will be used
+// on each row.
+// If the destination is some other type, then each row must only have one
+// column which can scan into that type.
+// This includes types that implement gocql.Unmarshaler and gocql.UDTUnmarshaler.
+//
+// If you'd like to treat a type that implements gocql.Unmarshaler or
+// gocql.UDTUnmarshaler as an ordinary struct you should call
+// Iter().StructOnly().Select(dest) instead.
 //
 // If no rows were selected, ErrNotFound is NOT returned.
 func (q *Queryx) Select(dest interface{}) error {