From edf633e8fc8d864dcef4238eb45a91d894bbef3a Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Wed, 4 Dec 2024 15:27:36 +0200
Subject: [PATCH] fix: build

---
 db.go           | 25 ++++++++++++++++---------
 query_base.go   |  6 ++----
 query_select.go | 12 +++++++-----
 3 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/db.go b/db.go
index a19f4cc3..37ced956 100644
--- a/db.go
+++ b/db.go
@@ -40,29 +40,36 @@ func WithReadOnlyReplica(replica *sql.DB) DBOption {
 }
 
 type DB struct {
+	// Must be a pointer so we copy the state, not the state fields.
+	*noCopyState
+
+	queryHooks []QueryHook
+
+	fmter schema.Formatter
+	stats DBStats
+}
+
+type noCopyState struct {
 	*sql.DB
+	dialect schema.Dialect
 
 	replicas        []*sql.DB
 	healthyReplicas atomic.Pointer[[]*sql.DB]
 	nextReplica     atomic.Int64
 
-	dialect    schema.Dialect
-	queryHooks []QueryHook
-
-	fmter  schema.Formatter
 	flags  internal.Flag
 	closed atomic.Bool
-
-	stats DBStats
 }
 
 func NewDB(sqldb *sql.DB, dialect schema.Dialect, opts ...DBOption) *DB {
 	dialect.Init(sqldb)
 
 	db := &DB{
-		DB:      sqldb,
-		dialect: dialect,
-		fmter:   schema.NewFormatter(dialect),
+		noCopyState: &noCopyState{
+			DB:      sqldb,
+			dialect: dialect,
+		},
+		fmter: schema.NewFormatter(dialect),
 	}
 
 	for _, opt := range opts {
diff --git a/query_base.go b/query_base.go
index 70bedb99..27f55742 100644
--- a/query_base.go
+++ b/query_base.go
@@ -147,10 +147,8 @@ func (q *baseQuery) GetTableName() string {
 	}
 
 	for _, wq := range q.with {
-		if v, ok := wq.query.(Query); ok {
-			if model := v.GetModel(); model != nil {
-				return v.GetTableName()
-			}
+		if model := wq.query.GetModel(); model != nil {
+			return wq.query.GetTableName()
 		}
 	}
 
diff --git a/query_select.go b/query_select.go
index 70be52e2..d43424fd 100644
--- a/query_select.go
+++ b/query_select.go
@@ -748,7 +748,7 @@ func (q *SelectQuery) Rows(ctx context.Context) (*sql.Rows, error) {
 	query := internal.String(queryBytes)
 
 	ctx, event := q.db.beforeQuery(ctx, q, query, nil, query, q.model)
-	rows, err := q.conn.QueryContext(ctx, query)
+	rows, err := q.resolveConn(q).QueryContext(ctx, query)
 	q.db.afterQuery(ctx, event, nil, err)
 	return rows, err
 }
@@ -876,7 +876,7 @@ func (q *SelectQuery) Count(ctx context.Context) (int, error) {
 	ctx, event := q.db.beforeQuery(ctx, qq, query, nil, query, q.model)
 
 	var num int
-	err = q.conn.QueryRowContext(ctx, query).Scan(&num)
+	err = q.resolveConn(q).QueryRowContext(ctx, query).Scan(&num)
 
 	q.db.afterQuery(ctx, event, nil, err)
 
@@ -894,13 +894,15 @@ func (q *SelectQuery) ScanAndCount(ctx context.Context, dest ...interface{}) (in
 			return int(n), nil
 		}
 	}
-	if _, ok := q.conn.(*DB); ok {
-		return q.scanAndCountConc(ctx, dest...)
+	if _, ok := q.conn.(*DB); ok || q.conn == nil {
+		return q.scanAndCountConcurrently(ctx, dest...)
 	}
 	return q.scanAndCountSeq(ctx, dest...)
 }
 
-func (q *SelectQuery) scanAndCountConc(ctx context.Context, dest ...interface{}) (int, error) {
+func (q *SelectQuery) scanAndCountConcurrently(
+	ctx context.Context, dest ...interface{},
+) (int, error) {
 	var count int
 	var wg sync.WaitGroup
 	var mu sync.Mutex