dbo11y: remove comments from normalized sql text in postgres

Strip away comments from `pg_stat_statements` normalized sql queries,
as postgres doesn't do that automatically.

While at it:
- move some stuff from `lexer.go` closer to postgres component
- keep a normalizer instance around in `query_details.go`
- polish some test cases formatting

Author: cristiangreco

internal/component/database_observability/lexer.go

@@ -7,20 +7,6 @@ import (
 	"github.com/DataDog/go-sqllexer"
 )
 
-// ExtractTableNames extracts the table names from a SQL query
-func ExtractTableNames(sql string) ([]string, error) {
-	normalizer := sqllexer.NewNormalizer(
-		sqllexer.WithCollectTables(true),
-	)
-	_, metadata, err := normalizer.Normalize(sql, sqllexer.WithDBMS(sqllexer.DBMSPostgres))
-	if err != nil {
-		return nil, fmt.Errorf("failed to normalize SQL: %w", err)
-	}
-
-	// Return all table names, including those that end with "..." for truncated queries, as we can't know if the table name was truncated or not
-	return metadata.Tables, nil
-}
-
 // RedactSql obfuscates a SQL query by replacing literals with ? placeholders
 func RedactSql(sql string) string {
 	obfuscatedSql := sqllexer.NewObfuscator().Obfuscate(sql)

internal/component/database_observability/lexer_test.go

@@ -49,26 +49,26 @@ func TestPgSqlParser_Redact(t *testing.T) {
 			sql: `WITH active_users AS (
 					SELECT * FROM users WHERE last_login > '2024-01-01'
 				), recent_orders AS (
-					SELECT o.* FROM orders o 
-					JOIN active_users u ON u.id = o.user_id 
+					SELECT o.* FROM orders o
+					JOIN active_users u ON u.id = o.user_id
 					WHERE o.created_at > '2024-03-01'
 				)
-				SELECT au.name, COUNT(ro.id) as order_count 
-				FROM active_users au 
-				LEFT JOIN recent_orders ro ON ro.user_id = au.id 
-				GROUP BY au.name 
+				SELECT au.name, COUNT(ro.id) as order_count
+				FROM active_users au
+				LEFT JOIN recent_orders ro ON ro.user_id = au.id
+				GROUP BY au.name
 				HAVING COUNT(ro.id) > 5`,
 			want: `WITH active_users AS (
 					SELECT * FROM users WHERE last_login > ?
 				), recent_orders AS (
-					SELECT o.* FROM orders o 
-					JOIN active_users u ON u.id = o.user_id 
+					SELECT o.* FROM orders o
+					JOIN active_users u ON u.id = o.user_id
 					WHERE o.created_at > ?
 				)
-				SELECT au.name, COUNT(ro.id) as order_count 
-				FROM active_users au 
-				LEFT JOIN recent_orders ro ON ro.user_id = au.id 
-				GROUP BY au.name 
+				SELECT au.name, COUNT(ro.id) as order_count
+				FROM active_users au
+				LEFT JOIN recent_orders ro ON ro.user_id = au.id
+				GROUP BY au.name
 				HAVING COUNT(ro.id) > ?`,
 		},
 		{
@@ -106,13 +106,13 @@ func TestPgSqlParser_Redact(t *testing.T) {
 		{
 			name: "WITH statement with UPDATE",
 			sql: `WITH inactive_users AS (
-					SELECT id FROM users 
+					SELECT id FROM users
 					WHERE last_login < '2023-01-01' AND status = 'active'
 				)
 				UPDATE users SET status = 'inactive', updated_at = '2024-03-20'
 				WHERE id IN (SELECT id FROM inactive_users)`,
 			want: `WITH inactive_users AS (
-					SELECT id FROM users 
+					SELECT id FROM users
 					WHERE last_login < ? AND status = ?
 				)
 				UPDATE users SET status = ?, updated_at = ?
@@ -121,16 +121,16 @@ func TestPgSqlParser_Redact(t *testing.T) {
 		{
 			name: "WITH statement with DELETE",
 			sql: `WITH old_orders AS (
-					SELECT id FROM orders 
+					SELECT id FROM orders
 					WHERE created_at < '2023-01-01' AND status = 'completed'
 				)
-				DELETE FROM order_items 
+				DELETE FROM order_items
 				WHERE order_id IN (SELECT id FROM old_orders)`,
 			want: `WITH old_orders AS (
-					SELECT id FROM orders 
+					SELECT id FROM orders
 					WHERE created_at < ? AND status = ?
 				)
-				DELETE FROM order_items 
+				DELETE FROM order_items
 				WHERE order_id IN (SELECT id FROM old_orders)`,
 		},
 		{
@@ -185,134 +185,6 @@ func TestPgSqlParser_Redact(t *testing.T) {
 	}
 }
 
-func TestPgSqlParser_ExtractTableNames(t *testing.T) {
-	tests := []struct {
-		name    string
-		sql     string
-		want    []string
-		wantErr bool
-	}{
-		{
-			name: "simple select",
-			sql:  "SELECT * FROM users",
-			want: []string{"users"},
-		},
-		{
-			name: "select with join",
-			sql:  "SELECT * FROM users u JOIN orders o ON u.id = o.user_id",
-			want: []string{"orders", "users"},
-		},
-		{
-			name: "select with schema qualified tables",
-			sql:  "SELECT * FROM public.users JOIN sales.orders ON users.id = orders.user_id",
-			want: []string{"public.users", "sales.orders"},
-		},
-		{
-			name: "insert statement",
-			sql:  "INSERT INTO users (name, email) VALUES ('John', '[email protected]')",
-			want: []string{"users"},
-		},
-		{
-			name: "update statement",
-			sql:  "UPDATE users SET last_login = NOW() WHERE id = 1",
-			want: []string{"users"},
-		},
-		{
-			name: "delete statement",
-			sql:  "DELETE FROM users WHERE id = 1",
-			want: []string{"users"},
-		},
-		{
-			name: "with clause",
-			sql: `WITH active_users AS (
-				SELECT * FROM users WHERE status = 'active'
-			)
-			SELECT * FROM active_users au
-			JOIN orders o ON o.user_id = au.id`,
-			want: []string{"orders", "users"},
-		},
-		{
-			name: "subquery in where clause",
-			sql: `SELECT * FROM orders 
-				WHERE user_id IN (SELECT id FROM users WHERE status = 'active')`,
-			want: []string{"orders", "users"},
-		},
-		{
-			name: "multiple schema qualified tables with aliases",
-			sql: `SELECT u.name, o.total, p.status 
-				FROM public.users u 
-				JOIN sales.orders o ON u.id = o.user_id
-				LEFT JOIN shipping.packages p ON o.id = p.order_id`,
-			want: []string{"public.users", "sales.orders", "shipping.packages"},
-		},
-		{
-			name: "truncated query with ...",
-			sql:  "SELECT * FROM users JOIN orders ON users.id = orders.user_id AND...",
-			want: []string{"users", "orders"},
-		},
-		{
-			name: "truncated query with incomplete comment",
-			sql:  "SELECT * FROM users JOIN orders ON users.id = orders.user_id /* some comment that gets truncated...",
-			want: []string{"users", "orders"},
-		},
-		{
-			name: "truncated query mid-table name",
-			sql:  "SELECT * FROM users JOIN ord...",
-			want: []string{"users", "ord..."},
-		},
-		{
-			name: "truncated query with schema qualified tables",
-			sql:  "SELECT * FROM public.users JOIN sales.orders ON users.id = orders.user_id AND...",
-			want: []string{"public.users", "sales.orders"},
-		},
-		{
-			name: "query with table.* expression",
-			sql:  "SELECT u.*, o.* FROM users u JOIN orders o ON u.id = o.user_id",
-			want: []string{"users", "orders"},
-		},
-		{
-			name: "query with type cast",
-			sql:  "SELECT u.id, '2024-03-20'::timestamp FROM users u",
-			want: []string{"users"},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := ExtractTableNames(tt.sql)
-			if (err != nil) != tt.wantErr {
-				t.Errorf("ExtractTableNames() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-			if !tt.wantErr {
-				if len(got) != len(tt.want) {
-					t.Errorf("ExtractTableNames()\nGOT = %v\nWANT = %v", got, tt.want)
-					return
-				}
-				// Compare slices ignoring order since table names might come in different order
-				gotMap := make(map[string]bool)
-				wantMap := make(map[string]bool)
-				for _, table := range got {
-					gotMap[table] = true
-				}
-				for _, table := range tt.want {
-					wantMap[table] = true
-				}
-				for table := range gotMap {
-					if !wantMap[table] {
-						t.Errorf("ExtractTableNames() got unexpected table = %v", table)
-					}
-				}
-				for table := range wantMap {
-					if !gotMap[table] {
-						t.Errorf("ExtractTableNames() missing expected table = %v", table)
-					}
-				}
-			}
-		})
-	}
-}
-
 func TestContainsReservedKeywords(t *testing.T) {
 	tests := []struct {
 		name          string

internal/component/database_observability/postgres/collector/query_details.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/DataDog/go-sqllexer"
 	"github.com/go-kit/log"
 	"go.uber.org/atomic"
 
@@ -55,6 +56,7 @@ type QueryDetails struct {
 	collectInterval time.Duration
 	entryHandler    loki.EntryHandler
 	tableRegistry   *TableRegistry
+	normalizer      *sqllexer.Normalizer
 
 	logger  log.Logger
 	running *atomic.Bool
@@ -68,6 +70,7 @@ func NewQueryDetails(args QueryDetailsArguments) (*QueryDetails, error) {
 		collectInterval: args.CollectInterval,
 		entryHandler:    args.EntryHandler,
 		tableRegistry:   args.TableRegistry,
+		normalizer:      sqllexer.NewNormalizer(sqllexer.WithCollectTables(true), sqllexer.WithCollectComments(true)),
 		logger:          log.With(args.Logger, "collector", QueryDetailsCollector),
 		running:         &atomic.Bool{},
 	}, nil
@@ -129,23 +132,25 @@ func (c QueryDetails) fetchAndAssociate(ctx context.Context) error {
 	for rs.Next() {
 		var queryID, queryText string
 		var databaseName database
-		err := rs.Scan(
-			&queryID,
-			&queryText,
-			&databaseName,
-		)
+		err := rs.Scan(&queryID, &queryText, &databaseName)
 		if err != nil {
 			level.Error(c.logger).Log("msg", "failed to scan result set for pg_stat_statements", "err", err)
 			continue
 		}
 
+		queryText, err = RemoveComments(c.normalizer, queryText)
+		if err != nil {
+			level.Error(c.logger).Log("msg", "failed to remove comments", "err", err)
+			continue
+		}
+
 		c.entryHandler.Chan() <- database_observability.BuildLokiEntry(
 			logging.LevelInfo,
 			OP_QUERY_ASSOCIATION,
-			fmt.Sprintf(`queryid="%s" querytext=%q datname="%s" engine="postgres"`, queryID, queryText, databaseName),
+			fmt.Sprintf(`queryid="%s" querytext=%q datname="%s"`, queryID, queryText, databaseName),
 		)
 
-		tables, err := c.tryTokenizeTableNames(queryText)
+		tables, err := TokenizeTableNames(c.normalizer, queryText)
 		if err != nil {
 			level.Error(c.logger).Log("msg", "failed to tokenize table names", "err", err)
 			continue
@@ -160,7 +165,7 @@ func (c QueryDetails) fetchAndAssociate(ctx context.Context) error {
 			c.entryHandler.Chan() <- database_observability.BuildLokiEntry(
 				logging.LevelInfo,
 				OP_QUERY_PARSED_TABLE_NAME,
-				fmt.Sprintf(`queryid="%s" datname="%s" table="%s" engine="postgres" validated="%t"`, queryID, databaseName, table, validated),
+				fmt.Sprintf(`queryid="%s" datname="%s" table="%s" validated="%t"`, queryID, databaseName, table, validated),
 			)
 		}
 	}
@@ -173,12 +178,29 @@ func (c QueryDetails) fetchAndAssociate(ctx context.Context) error {
 	return nil
 }
 
-func (c QueryDetails) tryTokenizeTableNames(sqlText string) ([]string, error) {
+func TokenizeTableNames(normalizer *sqllexer.Normalizer, sqlText string) ([]string, error) {
 	sqlText = strings.TrimSuffix(sqlText, "...")
-	tables, err := database_observability.ExtractTableNames(sqlText)
+	_, metadata, err := normalizer.Normalize(sqlText, sqllexer.WithDBMS(sqllexer.DBMSPostgres))
+	if err != nil {
+		return nil, fmt.Errorf("failed to tokenize table names: %w", err)
+	}
+
+	return metadata.Tables, nil
+}
+
+func RemoveComments(normalizer *sqllexer.Normalizer, sqlText string) (string, error) {
+	_, metadata, err := normalizer.Normalize(sqlText, sqllexer.WithDBMS(sqllexer.DBMSPostgres))
 	if err != nil {
-		return nil, fmt.Errorf("failed to extract table names: %w", err)
+		return sqlText, fmt.Errorf("failed to redact comments: %w", err)
+	}
+
+	if len(metadata.Comments) == 0 {
+		return sqlText, nil
+	}
+
+	for _, comment := range metadata.Comments {
+		sqlText = strings.ReplaceAll(sqlText, comment, "")
 	}
 
-	return tables, nil
+	return strings.TrimSpace(sqlText), nil
 }