sql: update stats for descriptor table

This is needed in order to enforce the
sql.schema.approx_max_object_count cluster setting, which relies on
optimizer table statistics to find the count.

Since the schemachanger uses the KV API to write to the descriptor
table, we need to explicitly notify the stats refresher when it should
compute new stats for the table.

Release note: None

Author: rafiss

pkg/sql/catalog/descs/collection.go

@@ -273,21 +273,17 @@ func (tc *Collection) IsVersionBumpOfUncommittedDescriptor(id descpb.ID) bool {
 	return tc.uncommitted.versionBumpOnly[id]
 }
 
-// HasUncommittedNewOrDroppedDescriptors returns true if the collection contains
-// any uncommitted descriptors that are newly created or dropped.
-func (tc *Collection) HasUncommittedNewOrDroppedDescriptors() bool {
-	isNewDescriptor := false
-	err := tc.uncommitted.iterateUncommittedByID(func(desc catalog.Descriptor) error {
+// CountUncommittedNewOrDroppedDescriptors returns the number of uncommitted
+// descriptors that are newly created or dropped.
+func (tc *Collection) CountUncommittedNewOrDroppedDescriptors() int {
+	count := 0
+	_ = tc.uncommitted.iterateUncommittedByID(func(desc catalog.Descriptor) error {
 		if desc.GetVersion() == 1 || desc.Dropped() {
-			isNewDescriptor = true
-			return iterutil.StopIteration()
+			count++
 		}
 		return nil
 	})
-	if err != nil {
-		return false
-	}
-	return isNewDescriptor
+	return count
 }
 
 // HasUncommittedTypes returns true if the Collection contains uncommitted

pkg/sql/conn_executor.go

@@ -4176,7 +4176,7 @@ func (ex *connExecutor) txnStateTransitionsApplyWrapper(
 				return advanceInfo{}, err
 			}
 		}
-		if ex.extraTxnState.descCollection.HasUncommittedNewOrDroppedDescriptors() {
+		if ex.extraTxnState.descCollection.CountUncommittedNewOrDroppedDescriptors() > 0 {
 			execCfg := ex.planner.ExecCfg()
 			if err := UpdateDescriptorCount(ex.Ctx(), execCfg, execCfg.SchemaChangerMetrics); err != nil {
 				log.Dev.Warningf(ex.Ctx(), "failed to scan descriptor table: %v", err)
@@ -4578,6 +4578,16 @@ func (ex *connExecutor) notifyStatsRefresherOfNewTables(ctx context.Context) {
 			ex.planner.execCfg.StatsRefresher.NotifyMutation(desc, math.MaxInt32 /* rowsAffected */)
 		}
 	}
+	if cnt := ex.extraTxnState.descCollection.CountUncommittedNewOrDroppedDescriptors(); cnt > 0 {
+		// Notify the refresher of a mutation on the system.descriptor table.
+		// We conservatively assume that any transaction which creates or
+		desc, err := ex.extraTxnState.descCollection.ByIDWithLeased(ex.planner.txn).Get().Table(ctx, keys.DescriptorTableID)
+		if err != nil {
+			log.Dev.Warningf(ctx, "failed to fetch descriptor table to refresh stats: %v", err)
+			return
+		}
+		ex.planner.execCfg.StatsRefresher.NotifyMutation(desc, cnt)
+	}
 }
 
 func (ex *connExecutor) getDescIDGenerator() eval.DescIDGenerator {

pkg/sql/schemachanger/schemachanger_test.go

@@ -1165,6 +1165,9 @@ func TestApproxMaxSchemaObjects(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
+	skip.UnderDuress(t, "slow test, requires polling to wait for auto stats job")
+	skip.UnderShort(t, "slow test, requires polling to wait for auto stats job")
+
 	ctx := context.Background()
 
 	// Test with both declarative schema changer modes.
@@ -1174,40 +1177,41 @@ func TestApproxMaxSchemaObjects(t *testing.T) {
 			defer s.Stopper().Stop(ctx)
 
 			tdb := sqlutils.MakeSQLRunner(sqlDB)
+			tdb.Exec(t, "SET CLUSTER SETTING sql.stats.automatic_collection.min_stale_rows = 1")
 
 			// Configure the declarative schema changer mode.
 			if useDeclarative {
+				tdb.Exec(t, `SET sql.defaults.use_declarative_schema_changer = 'on'`)
 				tdb.Exec(t, `SET use_declarative_schema_changer = 'on'`)
 			} else {
+				tdb.Exec(t, `SET sql.defaults.use_declarative_schema_changer = 'off'`)
 				tdb.Exec(t, `SET use_declarative_schema_changer = 'off'`)
 			}
 
-			// Get the current count of descriptors to set a realistic limit.
-			var currentCount int
-			tdb.QueryRow(t, `SELECT count(*) FROM system.descriptor`).Scan(&currentCount)
+			var maxObjects int
+			updateMaxObjects := func() {
+				// Manually refresh stats.
+				tdb.Exec(t, "ANALYZE system.public.descriptor")
 
-			// Set the limit to be slightly more than current count.
-			maxObjects := currentCount + 5
-			tdb.Exec(t, fmt.Sprintf(`SET CLUSTER SETTING sql.schema.approx_max_object_count = %d`, maxObjects))
+				// Get the current count of descriptors to set a realistic limit.
+				var currentCount int
+				tdb.QueryRow(t, `SELECT count(*) FROM system.descriptor`).Scan(&currentCount)
 
-			// Create a test database and use it.
-			tdb.Exec(t, `CREATE DATABASE testdb`)
-			tdb.Exec(t, `USE testdb`)
+				// Set the limit to be slightly more than current count.
+				maxObjects = currentCount + 1
+				tdb.Exec(t, fmt.Sprintf(`SET CLUSTER SETTING sql.schema.approx_max_object_count = %d`, maxObjects))
+			}
+			updateMaxObjects()
 
 			// Test that different object types are subject to the limit.
 			objectTypes := []string{"table", "database", "schema", "type", "function"}
 			for _, objectType := range objectTypes {
 				t.Run(objectType, func(t *testing.T) {
 					// Increase the limit before each subtest to avoid interference.
-					maxObjects += 10
-					tdb.Exec(t, fmt.Sprintf(`SET CLUSTER SETTING sql.schema.approx_max_object_count = %d`, maxObjects))
-
-					// Try to create objects until we hit the limit.
-					// ANALYZE before starting to ensure stats are fresh.
-					tdb.Exec(t, `ANALYZE system.descriptor`)
+					updateMaxObjects()
 
 					objNum := 0
-					for {
+					testutils.SucceedsWithin(t, func() error {
 						var createStmt string
 						switch objectType {
 						case "table":
@@ -1224,30 +1228,22 @@ func TestApproxMaxSchemaObjects(t *testing.T) {
 
 						_, err := sqlDB.Exec(createStmt)
 						if err != nil {
-							// Check if we got the expected error.
+							// Check if we got the expected error and message.
 							if pqErr := (*pq.Error)(nil); errors.As(err, &pqErr) {
 								if string(pqErr.Code) == pgcode.ConfigurationLimitExceeded.String() {
-									// Verify the error message mentions "approximate maximum".
-									testutils.IsError(err, "would exceed approximate maximum")
-									break
+									if testutils.IsError(err, "would exceed approximate maximum") {
+										return nil
+									}
 								}
 							}
 							// Some other error occurred.
-							t.Fatal(err)
+							return err
 						}
 						objNum++
 
-						// Re-analyze periodically to update stats.
-						if objNum%2 == 0 {
-							tdb.Exec(t, `ANALYZE system.descriptor`)
-						}
-
-						// Safety check: if we created way more objects than expected,
-						// something is wrong.
-						if objNum > 30 {
-							t.Fatalf("created %d %ss without hitting limit (max=%d)", objNum, objectType, maxObjects)
-						}
-					}
+						// Haven't hit the limit yet, keep trying.
+						return errors.Errorf("created %d %ss without hitting limit (max=%d)", objNum, objectType, maxObjects)
+					}, 5*time.Minute)
 				})
 			}
 		})