Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 12 commits into from
Oct 23, 2024
6 changes: 1 addition & 5 deletions context.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,5 @@ func ContextWithSpan(ctx context.Context, sp ot.Span) context.Context {
// span, this method returns false.
func SpanFromContext(ctx context.Context) (ot.Span, bool) {
sp, ok := ctx.Value(activeSpanKey).(ot.Span)
if !ok {
sanojsubran marked this conversation as resolved.
Show resolved Hide resolved
return nil, false
}

return sp, true
return sp, ok
}
3 changes: 2 additions & 1 deletion context_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ func TestSpanFromContext_WithActiveSpan(t *testing.T) {
}

func TestSpanFromContext_NoActiveSpan(t *testing.T) {
_, ok := instana.SpanFromContext(context.Background())
sp, ok := instana.SpanFromContext(context.Background())
assert.Equal(t, sp, nil)
assert.False(t, ok)
}
65 changes: 47 additions & 18 deletions instrumentation_sql.go
Original file line number Diff line number Diff line change
Expand Up @@ -164,20 +164,17 @@ 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, cmd string, sensor TracerLogger) ot.Span {

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 + " "
}
// The command string will be empty if the database name is determined
// from the connection URL rather than the query.
// Therefore, the Redis command should be parsed from the query.
if cmd == "" {
sanojsubran marked this conversation as resolved.
Show resolved Hide resolved
cmd, _ = parseRedisQuery(query)
}

tags := ot.Tags{
"redis.command": strings.TrimSpace(q),
"redis.command": cmd,
}

if conn.Error != nil {
Expand Down Expand Up @@ -257,15 +254,45 @@ func genericSQLSpan(ctx context.Context, conn DbConnDetails, query string, senso
return sensor.StartSpan("sdk.database", opts...)
}

// dbNameByQuery attempts to guess what is the database based on the query.
func dbNameByQuery(q string) string {
qf := strings.Fields(q)
// retrieveDBNameAndCmd attempts to guess what is the database based on the query.
// It accepts a string that may be a database query
// And returns the database name and query command
func retrieveDBNameAndCmd(q string) (cmd string, dbName string) {

if len(qf) > 0 && redisCmds.MatchString(qf[0]) {
return "redis"
if cmd, ok := parseRedisQuery(q); ok {
return cmd, "redis"
}

return ""
return cmd, dbName
}

// 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 words 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
Angith marked this conversation as resolved.
Show resolved Hide resolved
}

// StartSQLSpan creates a span based on DbConnDetails and a query, and attempts to detect which kind of database it belongs.
Expand All @@ -277,15 +304,17 @@ func StartSQLSpan(ctx context.Context, conn DbConnDetails, query string, sensor
}

func startSQLSpan(ctx context.Context, conn DbConnDetails, query string, sensor TracerLogger) (sp ot.Span, dbKey string) {

var dbCmd string
if conn.DatabaseName == "" {
conn.DatabaseName = dbNameByQuery(query)
dbCmd, conn.DatabaseName = retrieveDBNameAndCmd(query)
}

switch conn.DatabaseName {
case "postgres":
return postgresSpan(ctx, conn, query, sensor), "pg"
case "redis":
return redisSpan(ctx, conn, query, sensor), "redis"
return redisSpan(ctx, conn, query, dbCmd, sensor), "redis"
case "mysql":
return mySQLSpan(ctx, conn, query, sensor), "mysql"
case "couchbase":
Expand Down
156 changes: 156 additions & 0 deletions instrumentation_sql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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: "redis-instrumentation-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",
Expand Down
85 changes: 85 additions & 0 deletions sql_data.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
// (c) Copyright IBM Corp. 2024

package instana

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": {},
}