From 8b0cb9e44c719306f55484ba180f6a6844e20298 Mon Sep 17 00:00:00 2001
From: Cristian Greco <cristian@regolo.cc>
Date: Fri, 17 Jan 2025 16:00:40 +0100
Subject: [PATCH] database_observability: attempt to parse queries truncated in
 a comment (#2439)

The `query_sample` collector for now attempts to detect if a query is
truncated in a multi-line comment (those between `/*` and `*/`).

This is a best-effort attempt to avoid skipping queries that might be
truncated in the middle of a trailing comment that is not closed, but
the first part of the query is still valid.
---
 CHANGELOG.md                                  |  2 ++
 .../mysql/collector/query_sample.go           | 12 +++++++--
 .../mysql/collector/query_sample_test.go      | 26 +++++++++++++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 232df6b1bd..819ad1005c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,8 @@ Main (unreleased)
 
 - (_Experimental_) Log instance label key in `database_observability.mysql` (@cristiangreco)
 
+- (_Experimental_) Improve parsing of truncated queries in `database_observability.mysql` (@cristiangreco)
+
 - Add json format support for log export via faro receiver (@ravishankar15)
 
 - Add livedebugging support for `prometheus.remote_write` (@ravishankar15)
diff --git a/internal/component/database_observability/mysql/collector/query_sample.go b/internal/component/database_observability/mysql/collector/query_sample.go
index 28422196c2..5252ed2a76 100644
--- a/internal/component/database_observability/mysql/collector/query_sample.go
+++ b/internal/component/database_observability/mysql/collector/query_sample.go
@@ -135,8 +135,16 @@ func (c *QuerySample) fetchQuerySamples(ctx context.Context) error {
 		}
 
 		if strings.HasSuffix(sampleText, "...") {
-			level.Debug(c.logger).Log("msg", "skipping parsing truncated query", "digest", digest)
-			continue
+			// best-effort attempt to detect truncated trailing comment
+			if idx := strings.LastIndex(sampleText, "/*"); idx >= 0 {
+				trailingPart := sampleText[idx:]
+				if strings.LastIndex(trailingPart, "*/") < 0 {
+					sampleText = sampleText[:idx]
+				}
+			} else {
+				level.Debug(c.logger).Log("msg", "skipping parsing truncated query", "digest", digest)
+				continue
+			}
 		}
 
 		stmt, err := sqlparser.Parse(sampleText)
diff --git a/internal/component/database_observability/mysql/collector/query_sample_test.go b/internal/component/database_observability/mysql/collector/query_sample_test.go
index 3bd50675d0..96dc30eff1 100644
--- a/internal/component/database_observability/mysql/collector/query_sample_test.go
+++ b/internal/component/database_observability/mysql/collector/query_sample_test.go
@@ -111,6 +111,19 @@ func TestQuerySample(t *testing.T) {
 				`level=info msg="table name parsed" op="query_parsed_table_name" instance="mysql-db" digest="abc123" table="performance_schema.events_statements_summary_by_digest"`,
 			},
 		},
+		{
+			name: "with comment",
+			rows: [][]driver.Value{{
+				"abc123",
+				"select val1, /* val2,*/ val3 from some_table where id = 1",
+				"2024-01-01T00:00:00.000Z",
+				"1000",
+			}},
+			logs: []string{
+				`level=info msg="query samples fetched" op="query_sample" instance="mysql-db" digest="abc123" query_type="select" query_sample_seen="2024-01-01T00:00:00.000Z" query_sample_timer_wait="1000" query_sample_redacted="select val1, val3 from some_table where id = :redacted1"`,
+				`level=info msg="table name parsed" op="query_parsed_table_name" instance="mysql-db" digest="abc123" table="some_table"`,
+			},
+		},
 		{
 			name: "truncated query",
 			rows: [][]driver.Value{{
@@ -129,6 +142,19 @@ func TestQuerySample(t *testing.T) {
 				`level=info msg="table name parsed" op="query_parsed_table_name" instance="mysql-db" digest="abc123" table="some_table"`,
 			},
 		},
+		{
+			name: "truncated in multi-line comment",
+			rows: [][]driver.Value{{
+				"abc123",
+				"select * from some_table where id = 1 /*traceparent='00-abc...",
+				"2024-01-01T00:00:00.000Z",
+				"1000",
+			}},
+			logs: []string{
+				`level=info msg="query samples fetched" op="query_sample" instance="mysql-db" digest="abc123" query_type="select" query_sample_seen="2024-01-01T00:00:00.000Z" query_sample_timer_wait="1000" query_sample_redacted="select * from some_table where id = :redacted1"`,
+				`level=info msg="table name parsed" op="query_parsed_table_name" instance="mysql-db" digest="abc123" table="some_table"`,
+			},
+		},
 		{
 			name: "start transaction",
 			rows: [][]driver.Value{{