From c2575cf2da4e2dd0a47d168e2c805d5a94af8a5a Mon Sep 17 00:00:00 2001
From: Raphael 'kena' Poss <knz@cockroachlabs.com>
Date: Fri, 30 Nov 2018 21:22:53 +0100
Subject: [PATCH] sql: proprely record stmt stats for client-specified app
 names

The initialization of `(*connExecutor).appStats` in `newConnExecutor`
was based off a temporary local variable (`sd`) unrelated to the object
actually holding the application name (`ex.sessionData`). This caused
the initial appStats to always be wrong when the initial application
name was not empty.

Meanwhile, the changes in #30702 have ensured that `appStats` is
initialized via `resetSessionVars` for normal client connections. The
extra initialization in `newConnExecutor` was not just wrong; it was
not needed any more except for internal executors.

This patch corrects/simplifies the situation.

Release note (bug fix): CockroachDB now properly records statistics
for sessions where the value of `application_name` is given by the
client during initialization instead of `SET`.
---
 pkg/sql/conn_executor.go      |  8 +++++--
 pkg/sql/conn_executor_test.go | 43 +++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/pkg/sql/conn_executor.go b/pkg/sql/conn_executor.go
index 331fd48d7e58..da1156c40560 100644
--- a/pkg/sql/conn_executor.go
+++ b/pkg/sql/conn_executor.go
@@ -503,10 +503,14 @@ func (s *Server) newConnExecutor(
 			log.Errorf(ctx, "error setting up client session: %v", err)
 			return nil, err
 		}
+	} else {
+		// It's possible there were no defaults, for example when the
+		// connEx is serving an internal executor. In that case we still
+		// need to populate appStats according to the configured
+		// application name.
+		ex.appStats = s.sqlStats.getStatsForApplication(ex.sessionData.ApplicationName)
 	}
 
-	ex.appStats = s.sqlStats.getStatsForApplication(sd.ApplicationName)
-
 	ex.phaseTimes[sessionInit] = timeutil.Now()
 	ex.extraTxnState.tables = TableCollection{
 		leaseMgr:          s.cfg.LeaseManager,
diff --git a/pkg/sql/conn_executor_test.go b/pkg/sql/conn_executor_test.go
index 03891225c827..f60c90a9669c 100644
--- a/pkg/sql/conn_executor_test.go
+++ b/pkg/sql/conn_executor_test.go
@@ -16,6 +16,7 @@ package sql_test
 
 import (
 	"context"
+	gosql "database/sql"
 	"database/sql/driver"
 	"fmt"
 	"net/url"
@@ -353,3 +354,45 @@ CREATE TABLE t.test (k INT PRIMARY KEY, v TEXT);
 			"the EndTransaction with the expected key")
 	}
 }
+
+func TestAppNameStatisticsInitialization(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	params, _ := tests.CreateTestServerParams()
+	params.Insecure = true
+	s, _, _ := serverutils.StartServer(t, params)
+	defer s.Stopper().Stop(context.TODO())
+
+	// Prepare a session with a custom application name.
+	pgURL := url.URL{
+		Scheme:   "postgres",
+		User:     url.User(security.RootUser),
+		Host:     s.ServingAddr(),
+		RawQuery: "sslmode=disable&application_name=mytest",
+	}
+	rawSQL, err := gosql.Open("postgres", pgURL.String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rawSQL.Close()
+	sqlDB := sqlutils.MakeSQLRunner(rawSQL)
+
+	// Issue a query to be registered in stats.
+	sqlDB.Exec(t, "SELECT version()")
+
+	// Verify the query shows up in stats.
+	rows := sqlDB.Query(t, "SELECT application_name, key FROM crdb_internal.node_statement_statistics")
+	defer rows.Close()
+
+	counts := map[string]int{}
+	for rows.Next() {
+		var appName, key string
+		if err := rows.Scan(&appName, &key); err != nil {
+			t.Fatal(err)
+		}
+		counts[appName+":"+key]++
+	}
+	if counts["mytest:SELECT version()"] == 0 {
+		t.Fatalf("query was not counted properly: %+v", counts)
+	}
+}