From 50c29772c94d6436188aceb9e8619c6399987ada Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <dmitry.kropachev@gmail.com>
Date: Fri, 14 Jun 2024 11:21:57 -0400
Subject: [PATCH] Introduce `Unsafe` method on `Queryx`

It enables local control over `unsafe` mode for .Bind methods of `Queryx` and iterators
spawn by it.
---
 iterx.go      |  8 ++++----
 iterx_test.go |  6 +-----
 queryx.go     | 18 ++++++++++++++----
 session.go    |  5 ++++-
 4 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/iterx.go b/iterx.go
index 54f37dd..da72085 100644
--- a/iterx.go
+++ b/iterx.go
@@ -13,7 +13,7 @@ import (
 	"github.com/scylladb/go-reflectx"
 )
 
-// DefaultUnsafe enables the behavior of forcing the iterator to ignore
+// DefaultUnsafe enables the behavior of forcing queries and iterators to ignore
 // missing fields for all queries. See Unsafe below for more information.
 var DefaultUnsafe bool
 
@@ -206,9 +206,9 @@ func (iter *Iterx) scanAll(dest interface{}) bool {
 
 // 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, gocql.UDTUnmarshaler or UDT
-//   * it is not a struct
-//   * it has no exported fields
+//   - ptr to t implements gocql.Unmarshaler, gocql.UDTUnmarshaler or UDT
+//   - it is not a struct
+//   - it has no exported fields
 func (iter *Iterx) isScannable(t reflect.Type) bool {
 	ptr := reflect.PtrTo(t)
 	switch {
diff --git a/iterx_test.go b/iterx_test.go
index 2a1304f..757f20c 100644
--- a/iterx_test.go
+++ b/iterx_test.go
@@ -480,12 +480,8 @@ func TestIterxUnsafe(t *testing.T) {
 	})
 
 	t.Run("select default unsafe", func(t *testing.T) {
-		gocqlx.DefaultUnsafe = true
-		defer func() {
-			gocqlx.DefaultUnsafe = false
-		}()
 		var v []UnsafeTable
-		err := session.Query(stmt, nil).Iter().Select(&v)
+		err := session.Query(stmt, nil).Unsafe().Iter().Select(&v)
 		if err != nil {
 			t.Fatal("Select() failed:", err)
 		}
diff --git a/queryx.go b/queryx.go
index 36e8acc..f5e8736 100644
--- a/queryx.go
+++ b/queryx.go
@@ -94,8 +94,9 @@ type Queryx struct {
 	Names  []string
 	Mapper *reflectx.Mapper
 
-	tr  Transformer
-	err error
+	unsafe bool
+	tr     Transformer
+	err    error
 }
 
 // Query creates a new Queryx from gocql.Query using a default mapper.
@@ -107,6 +108,7 @@ func Query(q *gocql.Query, names []string) *Queryx {
 		Names:  names,
 		Mapper: DefaultMapper,
 		tr:     DefaultBindTransformer,
+		unsafe: DefaultUnsafe,
 	}
 }
 
@@ -210,7 +212,7 @@ func (q *Queryx) bindMapArgs(arg map[string]interface{}) ([]interface{}, error)
 // Bind sets query arguments of query. This can also be used to rebind new query arguments
 // to an existing query instance.
 func (q *Queryx) Bind(v ...interface{}) *Queryx {
-	q.Query.Bind(udtWrapSlice(q.Mapper, DefaultUnsafe, v)...)
+	q.Query.Bind(udtWrapSlice(q.Mapper, q.unsafe, v)...)
 	return q
 }
 
@@ -343,6 +345,14 @@ func (q *Queryx) Iter() *Iterx {
 	return &Iterx{
 		Iter:   q.Query.Iter(),
 		Mapper: q.Mapper,
-		unsafe: DefaultUnsafe,
+		unsafe: q.unsafe,
 	}
 }
+
+// Unsafe forces the query and iterators to ignore missing fields. By default when scanning
+// a struct if result row has a column that cannot be mapped to any destination
+// field an error is reported. With unsafe such columns are ignored.
+func (q *Queryx) Unsafe() *Queryx {
+	q.unsafe = true
+	return q
+}
diff --git a/session.go b/session.go
index 608a3ae..cb82c3c 100644
--- a/session.go
+++ b/session.go
@@ -34,7 +34,8 @@ func NewSession(session *gocql.Session) Session {
 // the created session to gocqlx.Session.
 //
 // Example:
-//     session, err := gocqlx.WrapSession(cluster.CreateSession())
+//
+//	session, err := gocqlx.WrapSession(cluster.CreateSession())
 func WrapSession(session *gocql.Session, err error) (Session, error) {
 	return Session{
 		Session: session,
@@ -50,6 +51,7 @@ func (s Session) ContextQuery(ctx context.Context, stmt string, names []string)
 		Names:  names,
 		Mapper: s.Mapper,
 		tr:     DefaultBindTransformer,
+		unsafe: DefaultUnsafe,
 	}
 }
 
@@ -64,6 +66,7 @@ func (s Session) Query(stmt string, names []string) *Queryx {
 		Names:  names,
 		Mapper: s.Mapper,
 		tr:     DefaultBindTransformer,
+		unsafe: DefaultUnsafe,
 	}
 }