From 46cc21c5e231cab738f6a2f4a5129829b0f455f9 Mon Sep 17 00:00:00 2001
From: Aleksandr Razumov <ernado@ya.ru>
Date: Tue, 24 May 2022 08:16:56 +0300
Subject: [PATCH] feat(otelch): add more metrics

---
 otelch/keys.go   | 18 ++++++++++++++++++
 query.go         |  8 +++++++-
 query_metrics.go | 15 +++++++++++----
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/otelch/keys.go b/otelch/keys.go
index d427d471..5fa5bc95 100644
--- a/otelch/keys.go
+++ b/otelch/keys.go
@@ -13,6 +13,8 @@ const (
 	ErrorNameKey       = attribute.Key("ch.error.name")
 	BlocksSentKey      = attribute.Key("ch.blocks_sent")
 	BlocksReceivedKey  = attribute.Key("ch.blocks_received")
+	ColumnsReceivedKey = attribute.Key("ch.columns_received")
+	RowsReceivedKey    = attribute.Key("ch.rows_received")
 	RowsKey            = attribute.Key("ch.rows")
 	BytesKey           = attribute.Key("ch.bytes")
 )
@@ -33,6 +35,22 @@ func BlocksReceived(v int) attribute.KeyValue {
 	}
 }
 
+// RowsReceived is cumulative rows received count during query execution.
+func RowsReceived(v int) attribute.KeyValue {
+	return attribute.KeyValue{
+		Key:   RowsReceivedKey,
+		Value: attribute.IntValue(v),
+	}
+}
+
+// ColumnsReceived is count of columns in result.
+func ColumnsReceived(v int) attribute.KeyValue {
+	return attribute.KeyValue{
+		Key:   ColumnsReceivedKey,
+		Value: attribute.IntValue(v),
+	}
+}
+
 // Rows is cumulative rows processed count during query execution.
 func Rows(v int) attribute.KeyValue {
 	return attribute.KeyValue{
diff --git a/query.go b/query.go
index 8ddbfc09..fb86d268 100644
--- a/query.go
+++ b/query.go
@@ -227,7 +227,11 @@ func (c *Client) decodeBlock(ctx context.Context, opt decodeOptions) error {
 	if block.End() {
 		return nil
 	}
-	c.metricsInc(ctx, queryMetrics{BlocksReceived: 1})
+	c.metricsInc(ctx, queryMetrics{
+		BlocksReceived:  1,
+		RowsReceived:    block.Rows,
+		ColumnsReceived: block.Columns,
+	})
 	if err := opt.Handler(ctx, block); err != nil {
 		return errors.Wrap(err, "handler")
 	}
@@ -531,6 +535,8 @@ func (c *Client) Do(ctx context.Context, q Query) (err error) {
 			span.SetAttributes(
 				otelch.BlocksSent(m.BlocksSent),
 				otelch.BlocksReceived(m.BlocksReceived),
+				otelch.RowsReceived(m.RowsReceived),
+				otelch.ColumnsReceived(m.ColumnsReceived),
 				otelch.Rows(m.Rows),
 				otelch.Bytes(m.Bytes),
 			)
diff --git a/query_metrics.go b/query_metrics.go
index 13850198..8a632f17 100644
--- a/query_metrics.go
+++ b/query_metrics.go
@@ -5,10 +5,12 @@ import "context"
 type (
 	ctxQueryKey  struct{}
 	queryMetrics struct {
-		BlocksReceived int
-		BlocksSent     int
-		Rows           int
-		Bytes          int
+		ColumnsReceived int
+		RowsReceived    int
+		BlocksReceived  int
+		BlocksSent      int
+		Rows            int
+		Bytes           int
 	}
 )
 
@@ -23,6 +25,11 @@ func (c *Client) metricsInc(ctx context.Context, delta queryMetrics) {
 
 	v.Bytes += delta.Bytes
 	v.Rows += delta.Rows
+	v.RowsReceived += delta.RowsReceived
 	v.BlocksReceived += delta.BlocksReceived
 	v.BlocksSent += delta.BlocksSent
+
+	if delta.ColumnsReceived > 0 {
+		v.ColumnsReceived = delta.ColumnsReceived
+	}
 }