@@ -17,7 +17,6 @@ import (
1717 "math"
1818 "regexp"
1919 "sort"
20- "strconv"
2120 "strings"
2221 "testing"
2322 "time"
@@ -106,7 +105,7 @@ func TestTelemetryLogging(t *testing.T) {
106105 execTimestampsSeconds []float64 // Execute the query with the following timestamps.
107106 expectedLogStatement string
108107 stubMaxEventFrequency int64
109- expectedSkipped []int // Expected skipped query count per expected log line.
108+ expectedSkipped []uint64 // Expected skipped query count per expected log line.
110109 expectedUnredactedTags []string
111110 expectedApplicationName string
112111 expectedFullScan bool
@@ -131,7 +130,7 @@ func TestTelemetryLogging(t *testing.T) {
131130 execTimestampsSeconds : []float64 {1 , 1.1 , 1.2 , 2 },
132131 expectedLogStatement : `TRUNCATE TABLE defaultdb.public.t` ,
133132 stubMaxEventFrequency : 1 ,
134- expectedSkipped : []int {0 , 0 , 0 , 0 },
133+ expectedSkipped : []uint64 {0 , 0 , 0 , 0 },
135134 expectedUnredactedTags : []string {"client" },
136135 expectedApplicationName : "telemetry-logging-test" ,
137136 expectedFullScan : false ,
@@ -176,7 +175,7 @@ func TestTelemetryLogging(t *testing.T) {
176175 execTimestampsSeconds : []float64 {3 },
177176 expectedLogStatement : `SELECT * FROM \"\".\"\".t LIMIT ‹1›` ,
178177 stubMaxEventFrequency : 1 ,
179- expectedSkipped : []int {0 },
178+ expectedSkipped : []uint64 {0 },
180179 expectedUnredactedTags : []string {"client" },
181180 expectedApplicationName : "telemetry-logging-test" ,
182181 expectedFullScan : false ,
@@ -199,7 +198,7 @@ func TestTelemetryLogging(t *testing.T) {
199198 execTimestampsSeconds : []float64 {4 , 4.1 , 4.2 , 5 },
200199 expectedLogStatement : `SELECT * FROM \"\".\"\".u LIMIT ‹2›` ,
201200 stubMaxEventFrequency : 1 ,
202- expectedSkipped : []int {0 , 2 },
201+ expectedSkipped : []uint64 {0 , 2 },
203202 expectedUnredactedTags : []string {"client" },
204203 expectedApplicationName : "telemetry-logging-test" ,
205204 expectedFullScan : false ,
@@ -224,7 +223,7 @@ func TestTelemetryLogging(t *testing.T) {
224223 execTimestampsSeconds : []float64 {6 , 6.01 , 6.05 , 6.06 , 6.1 , 6.2 },
225224 expectedLogStatement : `SELECT * FROM \"\".\"\".u LIMIT ‹3›` ,
226225 stubMaxEventFrequency : 10 ,
227- expectedSkipped : []int {0 , 3 , 0 },
226+ expectedSkipped : []uint64 {0 , 3 , 0 },
228227 expectedUnredactedTags : []string {"client" },
229228 expectedApplicationName : "telemetry-logging-test" ,
230229 expectedFullScan : false ,
@@ -252,7 +251,7 @@ func TestTelemetryLogging(t *testing.T) {
252251 execTimestampsSeconds : []float64 {7 },
253252 expectedLogStatement : `SELECT x FROM \"\".\"\".u` ,
254253 stubMaxEventFrequency : 10 ,
255- expectedSkipped : []int {0 },
254+ expectedSkipped : []uint64 {0 },
256255 expectedUnredactedTags : []string {"client" },
257256 expectedApplicationName : "telemetry-logging-test" ,
258257 expectedFullScan : true ,
@@ -279,7 +278,7 @@ func TestTelemetryLogging(t *testing.T) {
279278 execTimestampsSeconds : []float64 {8 },
280279 expectedLogStatement : `UPDATE \"\".\"\".u SET x = ‹5› WHERE x > ‹50› RETURNING x` ,
281280 stubMaxEventFrequency : 10 ,
282- expectedSkipped : []int {0 },
281+ expectedSkipped : []uint64 {0 },
283282 expectedUnredactedTags : []string {"client" },
284283 expectedApplicationName : "telemetry-logging-test" ,
285284 expectedFullScan : true ,
@@ -304,7 +303,7 @@ func TestTelemetryLogging(t *testing.T) {
304303 execTimestampsSeconds : []float64 {9 },
305304 expectedLogStatement : `CREATE USER root` ,
306305 stubMaxEventFrequency : 1 ,
307- expectedSkipped : []int {0 },
306+ expectedSkipped : []uint64 {0 },
308307 expectedUnredactedTags : []string {"client" },
309308 expectedApplicationName : "telemetry-logging-test" ,
310309 expectedFullScan : false ,
@@ -325,7 +324,7 @@ func TestTelemetryLogging(t *testing.T) {
325324 execTimestampsSeconds : []float64 {10 , 10.01 , 10.02 , 10.03 , 10.04 , 10.05 },
326325 expectedLogStatement : `SELECT * FROM \"\".\"\".u LIMIT ‹4›` ,
327326 stubMaxEventFrequency : 10 ,
328- expectedSkipped : []int {0 , 0 , 0 , 0 , 0 , 0 },
327+ expectedSkipped : []uint64 {0 , 0 , 0 , 0 , 0 , 0 },
329328 expectedUnredactedTags : []string {"client" },
330329 expectedApplicationName : "telemetry-logging-test" ,
331330 expectedFullScan : false ,
@@ -351,7 +350,7 @@ func TestTelemetryLogging(t *testing.T) {
351350 execTimestampsSeconds : []float64 {11 , 11.01 , 11.02 , 11.03 , 11.04 , 11.05 },
352351 expectedLogStatement : `SELECT * FROM \"\".\"\".u WHERE x > ‹10› LIMIT ‹3›` ,
353352 stubMaxEventFrequency : 10 ,
354- expectedSkipped : []int {0 },
353+ expectedSkipped : []uint64 {0 },
355354 expectedUnredactedTags : []string {"client" },
356355 expectedApplicationName : "telemetry-logging-test" ,
357356 expectedFullScan : true ,
@@ -371,7 +370,7 @@ func TestTelemetryLogging(t *testing.T) {
371370 execTimestampsSeconds : []float64 {20 },
372371 expectedLogStatement : `SELECT x FROM \"\".\"\".u ORDER BY x DESC` ,
373372 stubMaxEventFrequency : 1 ,
374- expectedSkipped : []int {24 },
373+ expectedSkipped : []uint64 {24 },
375374 expectedUnredactedTags : []string {"client" },
376375 expectedApplicationName : "telemetry-logging-test" ,
377376 expectedFullScan : true ,
@@ -479,17 +478,10 @@ func TestTelemetryLogging(t *testing.T) {
479478 err = json .Unmarshal ([]byte (e .Message ), & sampledQueryFromLog )
480479 require .NoError (t , err )
481480
482- expectedSkipped := tc .expectedSkipped [logCount ]
481+ require .Equal (t , tc .expectedSkipped [logCount ], sampledQueryFromLog .SkippedQueries )
482+
483483 logCount ++
484- if expectedSkipped == 0 {
485- if strings .Contains (e .Message , "SkippedQueries" ) {
486- t .Errorf ("%s: expected no skipped queries, found:\n %s" , tc .name , e .Message )
487- }
488- } else {
489- if expected := fmt .Sprintf (`"SkippedQueries":%d` , expectedSkipped ); ! strings .Contains (e .Message , expected ) {
490- t .Errorf ("%s: expected %s found:\n %s" , tc .name , expected , e .Message )
491- }
492- }
484+
493485 costRe := regexp .MustCompile ("\" CostEstimate\" :[0-9]*\\ .[0-9]*" )
494486 if ! costRe .MatchString (e .Message ) {
495487 t .Errorf ("expected to find CostEstimate but none was found" )
@@ -521,19 +513,18 @@ func TestTelemetryLogging(t *testing.T) {
521513 }
522514 }
523515 }
524- if ! strings .Contains (e .Message , "\" ApplicationName\" :\" " + tc .expectedApplicationName + "\" " ) {
525- t .Errorf ("expected to find unredacted Application Name: %s" , tc .expectedApplicationName )
526- }
527- if ! strings .Contains (e .Message , "\" SessionID\" :\" " + sessionID + "\" " ) {
528- t .Errorf ("expected to find sessionID: %s" , sessionID )
529- }
530- if ! strings .Contains (e .Message , "\" Database\" :\" " + databaseName + "\" " ) {
531- t .Errorf ("expected to find Database: %s" , databaseName )
532- }
516+
517+ require .Equal (t , tc .expectedApplicationName , sampledQueryFromLog .ApplicationName )
518+ require .Equal (t , sessionID , sampledQueryFromLog .SessionID )
519+ require .Equal (t , databaseName , sampledQueryFromLog .Database )
520+
521+ // All expected logs in this test are single stmt txns.
522+ require .Equal (t , uint32 (1 ), sampledQueryFromLog .StmtPosInTxn )
523+
533524 stmtFingerprintID := appstatspb .ConstructStatementFingerprintID (tc .queryNoConstants , tc .expectedErr != "" , true , databaseName )
534- if ! strings . Contains ( e . Message , " \" StatementFingerprintID \" :" + strconv . FormatUint ( uint64 ( stmtFingerprintID ), 10 )) {
535- t . Errorf ( "expected to find StatementFingerprintID: %v" , stmtFingerprintID )
536- }
525+
526+ require . Equal ( t , uint64 ( stmtFingerprintID ), sampledQueryFromLog . StatementFingerprintID )
527+
537528 maxFullScanRowsRe := regexp .MustCompile ("\" MaxFullScanRowsEstimate\" :[0-9]*" )
538529 foundFullScan := maxFullScanRowsRe .MatchString (e .Message )
539530 if tc .expectedFullScan && ! foundFullScan {
@@ -1476,3 +1467,86 @@ $$`
14761467 t .Errorf ("expected 1 log entries, found %d" , numLogsFound )
14771468 }
14781469}
1470+
1471+ // TestTelemetryLoggingStmtPosInTxn verifies that StmtCount is logged correctly to
1472+ // telemetry. The structure of the other test cases makes it difficult to verify
1473+ // this property, so we do it explicitly here.
1474+ func TestTelemetryLoggingStmtPosInTxn (t * testing.T ) {
1475+ defer leaktest .AfterTest (t )()
1476+ sc := log .ScopeWithoutShowLogs (t )
1477+ defer sc .Close (t )
1478+
1479+ cleanup := logtestutils .InstallLogFileSink (sc , t , logpb .Channel_TELEMETRY )
1480+ defer cleanup ()
1481+
1482+ st := logtestutils.StubTime {}
1483+ sts := logtestutils.StubTracingStatus {}
1484+
1485+ s , sqlDB , _ := serverutils .StartServer (t , base.TestServerArgs {
1486+ Knobs : base.TestingKnobs {
1487+ EventLog : & EventLogTestingKnobs {
1488+ // The sampling checks below need to have a deterministic
1489+ // number of statements run by internal executor.
1490+ SyncWrites : true ,
1491+ },
1492+ TelemetryLoggingKnobs : & TelemetryLoggingTestingKnobs {
1493+ getTimeNow : st .TimeNow ,
1494+ getTracingStatus : sts .TracingStatus ,
1495+ },
1496+ },
1497+ })
1498+
1499+ defer s .Stopper ().Stop (context .Background ())
1500+
1501+ db := sqlutils .MakeSQLRunner (sqlDB )
1502+ db .Exec (t , `SET application_name = 'telemetry-stmt=count-logging-test'` )
1503+
1504+ db .Exec (t , `SET CLUSTER SETTING sql.telemetry.query_sampling.enabled = true;` )
1505+
1506+ st .SetTime (timeutil .FromUnixMicros (int64 (1e6 )))
1507+ db .Exec (t , `BEGIN;` )
1508+ db .Exec (t , `SELECT 1` )
1509+ st .SetTime (timeutil .FromUnixMicros (int64 (2 * 1e6 )))
1510+ db .Exec (t , `SELECT 2` )
1511+ st .SetTime (timeutil .FromUnixMicros (int64 (3 * 1e6 )))
1512+ db .Exec (t , `SELECT 3` )
1513+ db .Exec (t , `COMMIT;` )
1514+
1515+ expectedQueries := []string {
1516+ `SELECT ‹1›` , `SELECT ‹2›` , `SELECT ‹3›` ,
1517+ }
1518+
1519+ log .Flush ()
1520+
1521+ entries , err := log .FetchEntriesFromFiles (
1522+ 0 ,
1523+ math .MaxInt64 ,
1524+ 10000 ,
1525+ regexp .MustCompile (`"EventType":"sampled_query"` ),
1526+ log .WithMarkedSensitiveData ,
1527+ )
1528+
1529+ if err != nil {
1530+ t .Fatal (err )
1531+ }
1532+
1533+ require .NotEmpty (t , entries )
1534+
1535+ // Attempt to find all expected logs.
1536+ for i , expected := range expectedQueries {
1537+ found := false
1538+ for _ , e := range entries {
1539+ if strings .Contains (e .Message , expected ) {
1540+ var sq eventpb.SampledQuery
1541+ require .NoError (t , json .Unmarshal ([]byte (e .Message ), & sq ))
1542+ require .Equal (t , uint32 (i + 1 ), sq .StmtPosInTxn , "%s" , entries )
1543+ found = true
1544+ break
1545+ }
1546+ }
1547+
1548+ if ! found {
1549+ t .Errorf ("did not find expected query log in log entries: %s" , expected )
1550+ }
1551+ }
1552+ }
0 commit comments