refactor: optimized logic for identifying if the database is Redis from the query

context.go

@@ -9,9 +9,10 @@ import (
 	ot "github.com/opentracing/opentracing-go"
 )
 
-type contextKey struct{}
+type ContextKey string
 
-var activeSpanKey contextKey
+var activeSpanKey ContextKey = "active_span"
+var redisCommand ContextKey = "redis_command"
 
 // ContextWithSpan returns a new context.Context holding a reference to an active span
 func ContextWithSpan(ctx context.Context, sp ot.Span) context.Context {
@@ -28,3 +29,15 @@ func SpanFromContext(ctx context.Context) (ot.Span, bool) {
 
 	return sp, true
 }
+
+func AddToContext(ctx context.Context, key ContextKey, value string) context.Context {
+	return context.WithValue(ctx, key, value)
+}
+
+func GetValueFromContext(ctx context.Context, key ContextKey) string {
+	val, ok := ctx.Value(key).(string)
+	if !ok {
+		return ""
+	}
+	return val
+}

context_test.go

@@ -29,3 +29,17 @@ func TestSpanFromContext_NoActiveSpan(t *testing.T) {
 	_, ok := instana.SpanFromContext(context.Background())
 	assert.False(t, ok)
 }
+
+func TestAddToContext_WithAnExistingKey(t *testing.T) {
+	ctx := instana.AddToContext(context.Background(), instana.ContextKey("redis_command"), "GET")
+
+	val := instana.GetValueFromContext(ctx, instana.ContextKey("redis_command"))
+	assert.Equal(t, val, "GET")
+}
+
+func TestAddToContext_WithOutAnExistingKey(t *testing.T) {
+	ctx := context.Background()
+
+	val := instana.GetValueFromContext(ctx, instana.ContextKey("redis_command"))
+	assert.Equal(t, val, "")
+}

instrumentation_sql.go

@@ -22,6 +22,88 @@ var (
 	sqlDriverRegistrationMu sync.Mutex
 )
 
+var redisCommands = map[string]struct{}{
+	"SET":           {},
+	"GET":           {},
+	"DEL":           {},
+	"INCR":          {},
+	"DECR":          {},
+	"APPEND":        {},
+	"GETRANGE":      {},
+	"SETRANGE":      {},
+	"STRLEN":        {},
+	"HSET":          {},
+	"HGET":          {},
+	"HMSET":         {},
+	"HMGET":         {},
+	"HDEL":          {},
+	"HGETALL":       {},
+	"HKEYS":         {},
+	"HVALS":         {},
+	"HLEN":          {},
+	"HINCRBY":       {},
+	"LPUSH":         {},
+	"RPUSH":         {},
+	"LPOP":          {},
+	"RPOP":          {},
+	"LLEN":          {},
+	"LRANGE":        {},
+	"LREM":          {},
+	"LINDEX":        {},
+	"LSET":          {},
+	"SADD":          {},
+	"SREM":          {},
+	"SMEMBERS":      {},
+	"SISMEMBER":     {},
+	"SCARD":         {},
+	"SINTER":        {},
+	"SUNION":        {},
+	"SDIFF":         {},
+	"SRANDMEMBER":   {},
+	"SPOP":          {},
+	"ZADD":          {},
+	"ZREM":          {},
+	"ZRANGE":        {},
+	"ZREVRANGE":     {},
+	"ZRANK":         {},
+	"ZREVRANK":      {},
+	"ZRANGEBYSCORE": {},
+	"ZCARD":         {},
+	"ZSCORE":        {},
+	"PFADD":         {},
+	"PFCOUNT":       {},
+	"PFMERGE":       {},
+	"SUBSCRIBE":     {},
+	"UNSUBSCRIBE":   {},
+	"PUBLISH":       {},
+	"MULTI":         {},
+	"EXEC":          {},
+	"DISCARD":       {},
+	"WATCH":         {},
+	"UNWATCH":       {},
+	"KEYS":          {},
+	"EXISTS":        {},
+	"EXPIRE":        {},
+	"TTL":           {},
+	"PERSIST":       {},
+	"RENAME":        {},
+	"RENAMENX":      {},
+	"TYPE":          {},
+	"SCAN":          {},
+	"PING":          {},
+	"INFO":          {},
+	"CLIENT LIST":   {},
+	"CONFIG GET":    {},
+	"CONFIG SET":    {},
+	"FLUSHDB":       {},
+	"FLUSHALL":      {},
+	"DBSIZE":        {},
+	"SAVE":          {},
+	"BGSAVE":        {},
+	"BGREWRITEAOF":  {},
+	"SHUTDOWN":      {},
+}
+
 // InstrumentSQLDriver instruments provided database driver for  use with `sql.Open()`.
 // This method will ignore any attempt to register the driver with the same name again.
 //
@@ -164,20 +246,16 @@ func mySQLSpan(ctx context.Context, conn DbConnDetails, query string, sensor Tra
 	return sensor.StartSpan(string(MySQLSpanType), opts...)
 }
 
-var redisCmds = regexp.MustCompile(`(?i)SET|GET|DEL|INCR|DECR|APPEND|GETRANGE|SETRANGE|STRLEN|HSET|HGET|HMSET|HMGET|HDEL|HGETALL|HKEYS|HVALS|HLEN|HINCRBY|LPUSH|RPUSH|LPOP|RPOP|LLEN|LRANGE|LREM|LINDEX|LSET|SADD|SREM|SMEMBERS|SISMEMBER|SCARD|SINTER|SUNION|SDIFF|SRANDMEMBER|SPOP|ZADD|ZREM|ZRANGE|ZREVRANGE|ZRANK|ZREVRANK|ZRANGEBYSCORE|ZCARD|ZSCORE|PFADD|PFCOUNT|PFMERGE|SUBSCRIBE|UNSUBSCRIBE|PUBLISH|MULTI|EXEC|DISCARD|WATCH|UNWATCH|KEYS|EXISTS|EXPIRE|TTL|PERSIST|RENAME|RENAMENX|TYPE|SCAN|PING|INFO|CLIENT LIST|CONFIG GET|CONFIG SET|FLUSHDB|FLUSHALL|DBSIZE|SAVE|BGSAVE|BGREWRITEAOF|SHUTDOWN`)
-
 func redisSpan(ctx context.Context, conn DbConnDetails, query string, sensor TracerLogger) ot.Span {
-	qarr := strings.Fields(query)
-	var q string
 
-	for _, w := range qarr {
-		if redisCmds.MatchString(w) {
-			q += w + " "
-		}
+	var q string
+	q = GetValueFromContext(ctx, redisCommand)
+	if q == "" {
+		q, _ = parseRedisQuery(query)
 	}
 
 	tags := ot.Tags{
-		"redis.command": strings.TrimSpace(q),
+		"redis.command": q,
 	}
 
 	if conn.Error != nil {
@@ -258,14 +336,44 @@ func genericSQLSpan(ctx context.Context, conn DbConnDetails, query string, senso
 }
 
 // dbNameByQuery attempts to guess what is the database based on the query.
-func dbNameByQuery(q string) string {
-	qf := strings.Fields(q)
+func dbNameByQuery(ctx context.Context, q string) (context.Context, string) {
 
-	if len(qf) > 0 && redisCmds.MatchString(qf[0]) {
-		return "redis"
+	var command string
+	var ok bool
+
+	if command, ok = parseRedisQuery(q); ok {
+		return AddToContext(ctx, redisCommand, command), "redis"
 	}
 
-	return ""
+	return ctx, ""
+}
+
+// parseRedisQuery attempts to guess if the input string is a valid Redis query.
+// parameters:
+//   - query (string): a string that may be a redis query
+//
+// returns:
+//   - command (string): The Redis command if the input is identified as a Redis query.
+//     This would typically be the first word of the Redis command, such as "SET", "CONFIG GET" etc.
+//     If the input is not a Redis query, this value will be an empty string.
+//   - isRedis (bool): A boolean value, `true` if the input is recognized as a Redis query,
+//     otherwise `false`.
+func parseRedisQuery(query string) (command string, isRedis bool) {
+	query = strings.TrimSpace(query)
+	if len(query) == 0 {
+		return "", false
+	}
+
+	// getting first two word of the query
+	parts := strings.SplitN(query, " ", 3)
+	command = strings.ToUpper(parts[0])
+
+	_, isRedis = redisCommands[command]
+	if !isRedis && len(parts) > 1 {
+		command = strings.ToUpper(parts[0] + " " + parts[1])
+		_, isRedis = redisCommands[command]
+	}
+	return
 }
 
 // StartSQLSpan creates a span based on DbConnDetails and a query, and attempts to detect which kind of database it belongs.
@@ -276,9 +384,11 @@ func StartSQLSpan(ctx context.Context, conn DbConnDetails, query string, sensor
 	return startSQLSpan(ctx, conn, query, sensor)
 }
 
-func startSQLSpan(ctx context.Context, conn DbConnDetails, query string, sensor TracerLogger) (sp ot.Span, dbKey string) {
+func startSQLSpan(c context.Context, conn DbConnDetails, query string, sensor TracerLogger) (sp ot.Span, dbKey string) {
+
+	var ctx context.Context = c
 	if conn.DatabaseName == "" {
-		conn.DatabaseName = dbNameByQuery(query)
+		ctx, conn.DatabaseName = dbNameByQuery(c, query)
 	}
 
 	switch conn.DatabaseName {

instrumentation_sql_test.go

@@ -910,6 +910,162 @@ func TestOpenSQLDB_RedisKVConnString(t *testing.T) {
 	}, data.Tags)
 }
 
+func TestStmtExecContext_WithRedisCommands(t *testing.T) {
+
+	recorder := instana.NewTestRecorder()
+	s := instana.NewSensorWithTracer(instana.NewTracerWithEverything(&instana.Options{
+		Service:     "go-sensor-test",
+		AgentClient: alwaysReadyClient{},
+	}, recorder))
+	defer instana.ShutdownSensor()
+
+	span := s.Tracer().StartSpan("parent-span")
+	ctx := context.Background()
+	if span != nil {
+		ctx = instana.ContextWithSpan(ctx, span)
+	}
+
+	instana.InstrumentSQLDriver(s, "fake_redis_driver_2", sqlDriver{})
+	require.Contains(t, sql.Drivers(), "fake_redis_driver_2_with_instana")
+
+	db, err := instana.SQLOpen("fake_redis_driver_2", "192.168.2.10:6790")
+	require.NoError(t, err)
+
+	t.Run("valid redis command", func(t *testing.T) {
+
+		_, err = db.ExecContext(ctx, "GET key")
+		require.NoError(t, err)
+
+		spans := recorder.GetQueuedSpans()
+		require.Len(t, spans, 1)
+
+		require.IsType(t, instana.RedisSpanData{}, spans[0].Data)
+		data := spans[0].Data.(instana.RedisSpanData)
+
+		assert.Equal(t, instana.RedisSpanTags{
+			Connection: "192.168.2.10:6790",
+			Command:    "GET",
+			Error:      "",
+		}, data.Tags)
+	})
+
+	t.Run("With multi word command", func(t *testing.T) {
+
+		_, err = db.ExecContext(ctx, "CONFIG GET key")
+		require.NoError(t, err)
+
+		spans := recorder.GetQueuedSpans()
+		require.Len(t, spans, 1)
+
+		require.IsType(t, instana.RedisSpanData{}, spans[0].Data)
+		data := spans[0].Data.(instana.RedisSpanData)
+
+		assert.Equal(t, instana.RedisSpanTags{
+			Connection: "192.168.2.10:6790",
+			Command:    "CONFIG GET",
+			Error:      "",
+		}, data.Tags)
+	})
+
+	t.Run("wrong/unknown(to go sensor) redis command", func(t *testing.T) {
+
+		_, err = db.ExecContext(ctx, "SELECT key")
+		require.NoError(t, err)
+
+		spans := recorder.GetQueuedSpans()
+		require.Len(t, spans, 1)
+
+		require.IsType(t, instana.SDKSpanData{}, spans[0].Data)
+		data := spans[0].Data.(instana.SDKSpanData)
+
+		assert.Equal(t, instana.SDKSpanTags{
+			Name: "sdk.database",
+			Type: "exit",
+			Custom: map[string]interface{}{
+				"tags": ot.Tags{
+					"span.kind":    ext.SpanKindRPCClientEnum,
+					"db.instance":  "192.168.2.10:6790",
+					"db.statement": "SELECT key",
+					"db.type":      "sql",
+					"peer.address": "192.168.2.10:6790",
+				},
+			},
+		}, data.Tags)
+	})
+
+	t.Run("empty query", func(t *testing.T) {
+
+		_, err = db.ExecContext(ctx, "")
+		require.NoError(t, err)
+
+		spans := recorder.GetQueuedSpans()
+		require.Len(t, spans, 1)
+
+		require.IsType(t, instana.SDKSpanData{}, spans[0].Data)
+		data := spans[0].Data.(instana.SDKSpanData)
+
+		assert.Equal(t, instana.SDKSpanTags{
+			Name: "sdk.database",
+			Type: "exit",
+			Custom: map[string]interface{}{
+				"tags": ot.Tags{
+					"span.kind":    ext.SpanKindRPCClientEnum,
+					"db.instance":  "192.168.2.10:6790",
+					"db.statement": "",
+					"db.type":      "sql",
+					"peer.address": "192.168.2.10:6790",
+				},
+			},
+		}, data.Tags)
+	})
+
+	t.Run("transaction", func(t *testing.T) {
+
+		_, err = db.ExecContext(ctx, "MULTI")
+		require.NoError(t, err)
+
+		_, err = db.ExecContext(ctx, "SET", "key1", "value1")
+		require.NoError(t, err)
+
+		_, err = db.ExecContext(ctx, "INCR", "counter")
+		require.NoError(t, err)
+
+		_, err = db.ExecContext(ctx, "EXEC")
+		require.NoError(t, err)
+
+		spans := recorder.GetQueuedSpans()
+		require.Len(t, spans, 4)
+
+		testcases := []struct {
+			Command string
+		}{
+			{
+				Command: "MULTI",
+			},
+			{
+				Command: "SET",
+			},
+			{
+				Command: "INCR",
+			},
+			{
+				Command: "EXEC",
+			},
+		}
+
+		for i, tc := range testcases {
+			require.IsType(t, instana.RedisSpanData{}, spans[i].Data)
+			data := spans[i].Data.(instana.RedisSpanData)
+
+			assert.Equal(t, instana.RedisSpanTags{
+				Connection: "192.168.2.10:6790",
+				Command:    tc.Command,
+				Error:      "",
+			}, data.Tags)
+		}
+	})
+}
+
 func TestNoPanicWithNotParsableConnectionString(t *testing.T) {
 	s := instana.NewSensorWithTracer(instana.NewTracerWithEverything(&instana.Options{
 		Service:     "go-sensor-test",