Repository: snowflakedb/gosnowflake
Pull request #443 (2021-06-21 22:58:58)

SNOW-81453 Inband Telemetry
### Description
Create inband telemetry.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/connection.go b/connection.go
index 307866d..dbdf796 100644
--- a/connection.go
+++ b/connection.go
@@ -15,6 +15,7 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -65,6 +66,7 @@ type snowflakeConn struct {
 	SequenceCounter uint64
 	QueryID         string
 	SQLState        string
+	telemetry       *snowflakeTelemetry
 	internal        InternalClient
 }
 
@@ -250,6 +252,7 @@ func (sc *snowflakeConn) cleanup() {
 
 func (sc *snowflakeConn) Close() (err error) {
 	logger.WithContext(sc.ctx).Infoln("Close")
+	sc.telemetry.sendBatch()
 	sc.stopHeartBeat()
 
 	if !sc.cfg.KeepSessionAlive {
@@ -992,6 +995,10 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 	} else {
 		tokenAccessor = getSimpleTokenAccessor()
 	}
+	if sc.cfg.DisableTelemetry {
+		sc.telemetry.enabled = false
+	}
+
 	// authenticate
 	sc.rest = &snowflakeRestful{
 		Host:     sc.cfg.Host,
@@ -1017,5 +1024,11 @@ func buildSnowflakeConn(ctx context.Context, config Config) (*snowflakeConn, err
 		FuncPostAuthOKTA:    postAuthOKTA,
 		FuncGetSSO:          getSSO,
 	}
+	sc.telemetry = &snowflakeTelemetry{
+		flushSize: defaultFlushSize,
+		sr:        sc.rest,
+		mutex:     &sync.Mutex{},
+		enabled:   true,
+	}
 	return sc, nil
 }
diff --git a/dsn.go b/dsn.go
index 100ed72..07643af 100644
--- a/dsn.go
+++ b/dsn.go
@@ -78,6 +78,8 @@ type Config struct {
 	PrivateKey *rsa.PrivateKey // Private key used to sign JWT
 
 	Transporter http.RoundTripper // RoundTripper to intercept HTTP requests and responses
+
+	DisableTelemetry bool // indicates whether to disable telemetry
 }
 
 // ocspMode returns the OCSP mode in string INSECURE, FAIL_OPEN, FAIL_CLOSED
diff --git a/restful.go b/restful.go
index 0c64a6d..19e2706 100644
--- a/restful.go
+++ b/restful.go
@@ -132,6 +132,14 @@ type cancelQueryResponse struct {
 	Success bool        `json:"success"`
 }
 
+type telemetryResponse struct {
+	Data    interface{}       `json:"data,omitempty"`
+	Message string            `json:"message"`
+	Code    string            `json:"code"`
+	Success bool              `json:"success"`
+	Headers map[string]string `json:"headers,omitempty"`
+}
+
 func postRestful(
 	ctx context.Context,
 	sr *snowflakeRestful,
diff --git a/telemetry.go b/telemetry.go
new file mode 100644
index 0000000..ed6b7cb
--- /dev/null
+++ b/telemetry.go
@@ -0,0 +1,91 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"sync"
+)
+
+const (
+	telemetryPath    = "/telemetry/send"
+	defaultFlushSize = 100
+)
+
+type telemetryMessage struct {
+	Type    string `json:"type,omitempty"`
+	QueryID string `json:"QueryID,omitempty"`
+}
+
+type telemetryData struct {
+	Timestamp int64             `json:"timestamp,omitempty"`
+	Message   *telemetryMessage `json:"message,omitempty"`
+}
+
+type snowflakeTelemetry struct {
+	logs      []*telemetryData
+	flushSize int
+	sr        *snowflakeRestful
+	mutex     *sync.Mutex
+	enabled   bool
+}
+
+func (st *snowflakeTelemetry) addLog(data *telemetryData) error {
+	if !st.enabled {
+		return fmt.Errorf("telemetry disabled; not adding log")
+	}
+	st.mutex.Lock()
+	st.logs = append(st.logs, data)
+	st.mutex.Unlock()
+	if len(st.logs) >= st.flushSize {
+		if err := st.sendBatch(); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (st *snowflakeTelemetry) sendBatch() error {
+	if !st.enabled {
+		return fmt.Errorf("telemetry disabled; not sending log")
+	}
+	type telemetry struct {
+		Logs []*telemetryData `json:"logs"`
+	}
+
+	st.mutex.Lock()
+	logsToSend := st.logs
+	st.logs = make([]*telemetryData, 0)
+	st.mutex.Unlock()
+
+	s := &telemetry{logsToSend}
+	body, err := json.Marshal(s)
+	if err != nil {
+		return err
+	}
+	headers := getHeaders()
+	if token, _, _ := st.sr.TokenAccessor.GetTokens(); token != "" {
+		headers[headerAuthorizationKey] = fmt.Sprintf(headerSnowflakeToken, token)
+	}
+	resp, err := st.sr.FuncPost(context.Background(), st.sr, st.sr.getFullURL(telemetryPath, nil), headers, body, 5, true)
+	if err != nil {
+		return err
+	}
+	if resp.StatusCode != http.StatusOK {
+		st.enabled = false
+		return fmt.Errorf("failed to upload metrics to telemetry")
+	}
+	var respd telemetryResponse
+	if err = json.NewDecoder(resp.Body).Decode(&respd); err != nil {
+		st.enabled = false
+		return err
+	}
+	if !respd.Success {
+		st.enabled = false
+		return fmt.Errorf("telemetry send failed with error code: %v, message: %v", respd.Code, respd.Message)
+	}
+	return nil
+}


Test patch:
diff --git a/connection_test.go b/connection_test.go
index f0e92f2..9d32326 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -173,6 +174,10 @@ func TestServiceName(t *testing.T) {
 
 var closedSessionCount = 0
 
+var testTelemetry = &snowflakeTelemetry{
+	mutex: &sync.Mutex{},
+}
+
 func closeSessionMock(_ context.Context, _ *snowflakeRestful, _ time.Duration) error {
 	closedSessionCount++
 	return &SnowflakeError{
@@ -187,6 +192,7 @@ func TestCloseIgnoreSessionGone(t *testing.T) {
 	sc := &snowflakeConn{
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
+		telemetry: testTelemetry,
 	}
 
 	if sc.Close() != nil {
@@ -201,6 +207,7 @@ func TestClientSessionPersist(t *testing.T) {
 	sc := &snowflakeConn{
 		cfg:  &Config{Params: map[string]*string{}},
 		rest: sr,
+		telemetry: testTelemetry,
 	}
 	sc.cfg.KeepSessionAlive = true
 	count := closedSessionCount
diff --git a/telemetry_test.go b/telemetry_test.go
new file mode 100644
index 0000000..d69e3b7
--- /dev/null
+++ b/telemetry_test.go
@@ -0,0 +1,49 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"context"
+	"math/rand"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestTelemetryAddLog(t *testing.T) {
+	config, _ := ParseDSN(dsn)
+	sc, err := buildSnowflakeConn(context.Background(), *config)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = authenticateWithConfig(sc); err != nil {
+		t.Fatal(err)
+	}
+
+	st := &snowflakeTelemetry{
+		sr:        sc.rest,
+		mutex:     &sync.Mutex{},
+		enabled:   true,
+		flushSize: defaultFlushSize,
+	}
+	rand.Seed(time.Now().UnixNano())
+	randNum := rand.Int() % 10000
+	for i := 0; i < randNum; i++ {
+		if err = st.addLog(&telemetryData{
+			Message: &telemetryMessage{
+				"client_telemetry_type",
+				"123",
+			},
+			Timestamp: time.Now().UnixNano(),
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if len(st.logs) != randNum%defaultFlushSize {
+		t.Errorf("length of remaining logs does not match. expected: %v, got: %v",
+			randNum%defaultFlushSize, len(st.logs))
+	}
+	if err := st.sendBatch(); err != nil {
+		t.Fatal(err)
+	}
+}