Merge pull request #99800 from ericharmeling/backport23.1-99674

release-23.1: sql: add explicit column family assignments to persisted stats upgrades

pkg/upgrade/upgrades/create_computed_indexes_sql_statistics.go

@@ -21,7 +21,7 @@ import (
 
 const addExecutionCountComputedColStmtStats = `
 ALTER TABLE system.statement_statistics
-ADD COLUMN IF NOT EXISTS "execution_count" INT8 
+ADD COLUMN IF NOT EXISTS "execution_count" INT8 FAMILY "primary" 
 AS ((statistics->'statistics'->'cnt')::INT8) STORED
 `
 
@@ -33,7 +33,7 @@ aggregated_ts, app_name, execution_count DESC) WHERE app_name NOT LIKE
 
 const addServiceLatencyComputedColStmtStats = `
 ALTER TABLE system.statement_statistics
-ADD COLUMN IF NOT EXISTS "service_latency" FLOAT 
+ADD COLUMN IF NOT EXISTS "service_latency" FLOAT FAMILY "primary" 
 AS ((statistics->'statistics'->'svcLat'->'mean')::FLOAT) STORED
 `
 
@@ -45,7 +45,7 @@ aggregated_ts, app_name, service_latency DESC) WHERE app_name NOT LIKE
 
 const addCpuSqlNanosComputedColStmtStats = `
 ALTER TABLE system.statement_statistics
-ADD COLUMN IF NOT EXISTS "cpu_sql_nanos" FLOAT 
+ADD COLUMN IF NOT EXISTS "cpu_sql_nanos" FLOAT FAMILY "primary" 
 AS ((statistics->'execution_statistics'->'cpuSQLNanos'->'mean')::FLOAT) STORED
 `
 
@@ -57,7 +57,7 @@ aggregated_ts, app_name, cpu_sql_nanos DESC) WHERE app_name NOT LIKE
 
 const addContentionTimeComputedColStmtStats = `
 ALTER TABLE system.statement_statistics
-ADD COLUMN IF NOT EXISTS "contention_time" FLOAT 
+ADD COLUMN IF NOT EXISTS "contention_time" FLOAT FAMILY "primary" 
 AS ((statistics->'execution_statistics'->'contentionTime'->'mean')::FLOAT) STORED
 `
 
@@ -69,7 +69,7 @@ aggregated_ts, app_name, contention_time DESC) WHERE app_name NOT LIKE
 
 const addTotalEstimatedExecutionTimeComputedColStmtStats = `
 ALTER TABLE system.statement_statistics
-ADD COLUMN IF NOT EXISTS "total_estimated_execution_time" FLOAT 
+ADD COLUMN IF NOT EXISTS "total_estimated_execution_time" FLOAT FAMILY "primary" 
 AS ((statistics->'statistics'->>'cnt')::FLOAT * (statistics->'statistics'->'svcLat'->>'mean')::FLOAT) STORED
 `
 
@@ -81,7 +81,7 @@ aggregated_ts, app_name, total_estimated_execution_time DESC) WHERE app_name NOT
 
 const addP99LatencyComputedColStmtStats = `
 ALTER TABLE system.statement_statistics
-ADD COLUMN IF NOT EXISTS "p99_latency" FLOAT 
+ADD COLUMN IF NOT EXISTS "p99_latency" FLOAT FAMILY "primary" 
 AS ((statistics->'statistics'->'latencyInfo'->'p99')::FLOAT) STORED
 `
 
@@ -94,7 +94,7 @@ aggregated_ts, app_name, p99_latency DESC) WHERE app_name NOT LIKE
 // transaction_statistics
 const addExecutionCountComputedColTxnStats = `
 ALTER TABLE system.transaction_statistics
-ADD COLUMN IF NOT EXISTS "execution_count" INT8 
+ADD COLUMN IF NOT EXISTS "execution_count" INT8 FAMILY "primary"
 AS ((statistics->'statistics'->'cnt')::INT8) STORED
 `
 
@@ -106,7 +106,7 @@ aggregated_ts, app_name, execution_count DESC) WHERE app_name NOT LIKE
 
 const addServiceLatencyComputedColTxnStats = `
 ALTER TABLE system.transaction_statistics
-ADD COLUMN IF NOT EXISTS "service_latency" FLOAT 
+ADD COLUMN IF NOT EXISTS "service_latency" FLOAT FAMILY "primary" 
 AS ((statistics->'statistics'->'svcLat'->'mean')::FLOAT) STORED
 `
 
@@ -118,7 +118,7 @@ aggregated_ts, app_name, service_latency DESC) WHERE app_name NOT LIKE
 
 const addCpuSqlNanosComputedColTxnStats = `
 ALTER TABLE system.transaction_statistics
-ADD COLUMN IF NOT EXISTS "cpu_sql_nanos" FLOAT 
+ADD COLUMN IF NOT EXISTS "cpu_sql_nanos" FLOAT FAMILY "primary" 
 AS ((statistics->'execution_statistics'->'cpuSQLNanos'->'mean')::FLOAT) STORED
 `
 
@@ -130,7 +130,7 @@ aggregated_ts, app_name, cpu_sql_nanos DESC) WHERE app_name NOT LIKE
 
 const addContentionTimeComputedColTxnStats = `
 ALTER TABLE system.transaction_statistics
-ADD COLUMN IF NOT EXISTS "contention_time" FLOAT 
+ADD COLUMN IF NOT EXISTS "contention_time" FLOAT FAMILY "primary" 
 AS ((statistics->'execution_statistics'->'contentionTime'->'mean')::FLOAT) STORED
 `
 
@@ -141,7 +141,7 @@ aggregated_ts, app_name, contention_time DESC) WHERE app_name NOT LIKE
 `
 const addTotalEstimatedExecutionTimeComputedColTxnStats = `
 ALTER TABLE system.transaction_statistics
-ADD COLUMN IF NOT EXISTS "total_estimated_execution_time" FLOAT 
+ADD COLUMN IF NOT EXISTS "total_estimated_execution_time" FLOAT FAMILY "primary" 
 AS ((statistics->'statistics'->>'cnt')::FLOAT * (
 statistics->'statistics'->'svcLat'->>'mean')::FLOAT) STORED
 `
@@ -154,7 +154,7 @@ aggregated_ts, app_name, total_estimated_execution_time DESC) WHERE app_name NOT
 
 const addP99LatencyComputedColTxnStats = `
 ALTER TABLE system.transaction_statistics
-ADD COLUMN IF NOT EXISTS "p99_latency" FLOAT 
+ADD COLUMN IF NOT EXISTS "p99_latency" FLOAT FAMILY "primary" 
 AS ((statistics->'statistics'->'latencyInfo'->'p99')::FLOAT) STORED
 `
 

pkg/upgrade/upgrades/create_computed_indexes_sql_statistics_test.go

@@ -64,6 +64,7 @@ func TestCreateComputedIndexesOnSystemSQLStatistics(t *testing.T) {
 			{Name: "total_estimated_execution_time_idx", ValidationFn: upgrades.HasIndex},
 			{Name: "p99_latency", ValidationFn: upgrades.HasColumn},
 			{Name: "p99_latency_idx", ValidationFn: upgrades.HasIndex},
+			{Name: "primary", ValidationFn: upgrades.OnlyHasColumnFamily},
 		}
 	)
 

pkg/upgrade/upgrades/helpers_test.go

@@ -30,11 +30,12 @@ import (
 )
 
 var (
-	HasColumn         = hasColumn
-	HasIndex          = hasIndex
-	DoesNotHaveIndex  = doesNotHaveIndex
-	HasColumnFamily   = hasColumnFamily
-	CreateSystemTable = createSystemTable
+	HasColumn           = hasColumn
+	HasIndex            = hasIndex
+	DoesNotHaveIndex    = doesNotHaveIndex
+	HasColumnFamily     = hasColumnFamily
+	CreateSystemTable   = createSystemTable
+	OnlyHasColumnFamily = onlyHasColumnFamily
 )
 
 type Schema struct {

pkg/upgrade/upgrades/schema_changes.go

@@ -354,3 +354,48 @@ func hasColumnFamily(
 	}
 	return true, nil
 }
+
+// onlyHasColumnFamily returns true if storedTable has only the given column
+// family, comparing with expectedTable. storedTable descriptor must be read
+// from system storage as compared to reading from the systemschema package. On
+// the contrary, expectedTable must be accessed directly from systemschema
+// package. This function returns an error if there is more than one column
+// family, or if the only column family does not match the provided family name.
+func onlyHasColumnFamily(
+	storedTable, expectedTable catalog.TableDescriptor, colFamily string,
+) (bool, error) {
+	var storedFamily, expectedFamily *descpb.ColumnFamilyDescriptor
+	storedFamilies := storedTable.GetFamilies()
+	if len(storedFamilies) > 1 {
+		return false, nil
+	}
+	if storedFamilies[0].Name == colFamily {
+		storedFamily = &storedFamilies[0]
+	}
+
+	if storedFamily == nil {
+		return false, nil
+	}
+
+	expectedFamilies := expectedTable.GetFamilies()
+	if expectedFamilies[0].Name == colFamily {
+		expectedFamily = &expectedFamilies[0]
+	}
+
+	if expectedFamily == nil {
+		return false, errors.Errorf("column family %s does not exist", colFamily)
+	}
+
+	// Check that columns match.
+	storedFamilyCols := storedFamily.ColumnNames
+	expectedFamilyCols := expectedFamily.ColumnNames
+	if len(storedFamilyCols) != len(expectedFamilyCols) {
+		return false, nil
+	}
+	for i, storedCol := range storedFamilyCols {
+		if storedCol != expectedFamilyCols[i] {
+			return false, nil
+		}
+	}
+	return true, nil
+}