From f496ba2589cb4dd3f65000d9dbdb87cf3fb954a5 Mon Sep 17 00:00:00 2001
From: Dmitry Kropachev <dmitry.kropachev@gmail.com>
Date: Wed, 15 Jan 2025 13:03:22 -0400
Subject: [PATCH] Make driver print warnings returned by server

---
 cluster.go         |  3 +++
 conn.go            | 29 +++++++++++++++++++++++++++--
 session.go         |  5 ++++-
 warning_handler.go | 28 ++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 warning_handler.go

diff --git a/cluster.go b/cluster.go
index 3d795322b..780621182 100644
--- a/cluster.go
+++ b/cluster.go
@@ -109,6 +109,8 @@ type ClusterConfig struct {
 	// Default: nil
 	Authenticator Authenticator
 
+	WarningsHandlerBuilder WarningHandlerBuilder
+
 	// An Authenticator factory. Can be used to create alternative authenticators.
 	// Default: nil
 	AuthProvider func(h *HostInfo) (Authenticator, error)
@@ -322,6 +324,7 @@ func NewCluster(hosts ...string) *ClusterConfig {
 		WriteCoalesceWaitTime:        200 * time.Microsecond,
 		MetadataSchemaRequestTimeout: 60 * time.Second,
 		DisableSkipMetadata:          true,
+		WarningsHandlerBuilder:       DefaultWarningHandlerBuilder,
 	}
 
 	return cfg
diff --git a/conn.go b/conn.go
index 670ff59be..c9e4f08ac 100644
--- a/conn.go
+++ b/conn.go
@@ -66,6 +66,12 @@ type Authenticator interface {
 	Success(data []byte) error
 }
 
+type WarningHandlerBuilder func(session *Session) WarningHandler
+
+type WarningHandler interface {
+	HandleWarnings(qry ExecutableQuery, host *HostInfo, warnings []string)
+}
+
 type PasswordAuthenticator struct {
 	Username              string
 	Password              string
@@ -1404,7 +1410,16 @@ func marshalQueryValue(typ TypeInfo, value interface{}, dst *queryValues) error
 	return nil
 }
 
-func (c *Conn) executeQuery(ctx context.Context, qry *Query) *Iter {
+func (c *Conn) executeQuery(ctx context.Context, qry *Query) (iter *Iter) {
+	defer func() {
+		if iter == nil || c.session == nil {
+			return
+		}
+		warnings := iter.Warnings()
+		if len(warnings) > 0 && c.session.warningHandler != nil {
+			c.session.warningHandler.HandleWarnings(qry, iter.host, warnings)
+		}
+	}()
 	params := queryParams{
 		consistency: qry.cons,
 	}
@@ -1670,7 +1685,17 @@ func (c *Conn) UseKeyspace(keyspace string) error {
 	return nil
 }
 
-func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
+func (c *Conn) executeBatch(ctx context.Context, batch *Batch) (iter *Iter) {
+	defer func() {
+		if iter == nil || c.session == nil {
+			return
+		}
+		warnings := iter.Warnings()
+		if len(warnings) > 0 && c.session.warningHandler != nil {
+			c.session.warningHandler.HandleWarnings(batch, iter.host, warnings)
+		}
+	}()
+
 	if c.version == protoVersion1 {
 		return &Iter{err: ErrUnsupported}
 	}
diff --git a/session.go b/session.go
index bca366282..1f029da5b 100644
--- a/session.go
+++ b/session.go
@@ -87,6 +87,7 @@ type Session struct {
 	tabletsRoutingV1 bool
 
 	usingTimeoutClause string
+	warningHandler     WarningHandler
 }
 
 var queryPool = &sync.Pool{
@@ -182,7 +183,9 @@ func newSessionCommon(cfg ClusterConfig) (*Session, error) {
 		return nil, fmt.Errorf("gocql: unable to create session: %v", err)
 	}
 	s.connCfg = connCfg
-
+	if cfg.WarningsHandlerBuilder != nil {
+		s.warningHandler = cfg.WarningsHandlerBuilder(s)
+	}
 	return s, nil
 }
 
diff --git a/warning_handler.go b/warning_handler.go
new file mode 100644
index 000000000..a913560ac
--- /dev/null
+++ b/warning_handler.go
@@ -0,0 +1,28 @@
+package gocql
+
+type DefaultWarningHandler struct {
+	logger StdLogger
+}
+
+func DefaultWarningHandlerBuilder(session *Session) WarningHandler {
+	return DefaultWarningHandler{
+		logger: session.logger,
+	}
+}
+
+func (d DefaultWarningHandler) HandleWarnings(qry ExecutableQuery, host *HostInfo, warnings []string) {
+	if d.logger == nil {
+		return
+	}
+	if host != nil && len(host.hostId) > 0 {
+		d.logger.Printf("[%s] warnings: %v", host.hostId, warnings)
+	} else {
+		d.logger.Printf("Cluster warnings: %v", warnings)
+	}
+}
+
+var _ WarningHandler = DefaultWarningHandler{}
+
+func NoopWarningHandlerBuilder(session *Session) WarningHandler {
+	return nil
+}