Optionally output database table sizes after migrations complete (#38620)

Resolves #35314.

# Checklist for submitter

If some of the following don't apply, delete the relevant line.

- [x] Changes file added for user-visible changes in `changes/`,
`orbit/changes/` or `ee/fleetd-chrome/changes`.
See [Changes
files](https://github.com/fleetdm/fleet/blob/main/docs/Contributing/guides/committing-changes.md#changes-files)
for more information.

- [x] Input data is properly validated, `SELECT *` is avoided, SQL
injection is prevented (using placeholders for values in statements)

## Testing

- [x] QA'd all new/changed functionality manually

changes/35141-table-sizes

@@ -0,0 +1 @@
+* Added `--with-table-sizes` option to `prepare` command to get approximate row counts of all database tables after a migration completes.

cmd/fleet/prepare.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"encoding/json"
 	"fmt"
 	"os"
 
@@ -29,6 +30,8 @@ To setup Fleet infrastructure, use one of the available commands.
 	noPrompt := false
 	// Whether to enable developer options
 	dev := false
+	// Whether to show table stats before and after the migration
+	showTableStats := false
 
 	dbCmd := &cobra.Command{
 		Use:   "db",
@@ -68,6 +71,20 @@ To setup Fleet infrastructure, use one of the available commands.
 				}
 			}
 
+			if showTableStats {
+				defer func() {
+					stats, err := ds.GetTableRowCounts(cmd.Context())
+					if err != nil {
+						initFatal(err, "getting table stats")
+					}
+					statsAsJSON, err := json.Marshal(stats)
+					if err != nil {
+						initFatal(err, "encoding table row counts to JSON")
+					}
+					fmt.Printf("Table Row Counts: %s\n", statsAsJSON)
+				}()
+			}
+
 			switch status.StatusCode {
 			case fleet.NoMigrationsCompleted:
 				// OK
@@ -102,6 +119,7 @@ To setup Fleet infrastructure, use one of the available commands.
 
 	dbCmd.PersistentFlags().BoolVar(&noPrompt, "no-prompt", false, "disable prompting before migrations (for use in scripts)")
 	dbCmd.PersistentFlags().BoolVar(&dev, "dev", false, "Enable developer options")
+	dbCmd.PersistentFlags().BoolVar(&showTableStats, "with-table-stats", false, "Show approximate table row counts after migrations")
 
 	prepareCmd.AddCommand(dbCmd)
 	return prepareCmd

server/datastore/mysql/statistics.go

@@ -233,3 +233,29 @@ func (ds *Datastore) CleanupStatistics(ctx context.Context) error {
 	}
 	return nil
 }
+
+func (ds *Datastore) GetTableRowCounts(ctx context.Context) (map[string]uint, error) {
+	return ds.getTableRowCountsViaInformationSchema(ctx)
+}
+
+func (ds *Datastore) getTableRowCountsViaInformationSchema(ctx context.Context) (map[string]uint, error) {
+	var results []struct {
+		Table string `db:"TABLE_NAME"`
+		Rows  uint   `db:"table_rows"`
+	}
+
+	if err := sqlx.SelectContext(
+		ctx,
+		ds.reader(ctx),
+		&results,
+		"SELECT table_name, COALESCE(table_rows, 0) table_rows FROM information_schema.tables WHERE table_schema = (SELECT DATABASE())",
+	); err != nil {
+		return nil, err
+	}
+
+	var byName = make(map[string]uint)
+	for _, row := range results {
+		byName[row.Table] = row.Rows
+	}
+	return byName, nil
+}

server/fleet/datastore.go

@@ -782,6 +782,8 @@ type Datastore interface {
 	// CleanupStatistics executes cleanup tasks to be performed upon successful transmission of
 	// statistics.
 	CleanupStatistics(ctx context.Context) error
+	// GetTableRowCounts returns approximate DB row counts for all tables in a map indexed by table name
+	GetTableRowCounts(ctx context.Context) (map[string]uint, error)
 
 	///////////////////////////////////////////////////////////////////////////////
 	// GlobalPoliciesStore

server/mock/datastore_mock.go

@@ -603,6 +603,8 @@ type RecordStatisticsSentFunc func(ctx context.Context) error
 
 type CleanupStatisticsFunc func(ctx context.Context) error
 
+type GetTableRowCountsFunc func(ctx context.Context) (map[string]uint, error)
+
 type ApplyPolicySpecsFunc func(ctx context.Context, authorID uint, specs []*fleet.PolicySpec) error
 
 type NewGlobalPolicyFunc func(ctx context.Context, authorID *uint, args fleet.PolicyPayload) (*fleet.Policy, error)
@@ -2628,6 +2630,9 @@ type DataStore struct {
 	CleanupStatisticsFunc        CleanupStatisticsFunc
 	CleanupStatisticsFuncInvoked bool
 
+	GetTableRowCountsFunc        GetTableRowCountsFunc
+	GetTableRowCountsFuncInvoked bool
+
 	ApplyPolicySpecsFunc        ApplyPolicySpecsFunc
 	ApplyPolicySpecsFuncInvoked bool
 
@@ -6392,6 +6397,13 @@ func (s *DataStore) CleanupStatistics(ctx context.Context) error {
 	return s.CleanupStatisticsFunc(ctx)
 }
 
+func (s *DataStore) GetTableRowCounts(ctx context.Context) (map[string]uint, error) {
+	s.mu.Lock()
+	s.GetTableRowCountsFuncInvoked = true
+	s.mu.Unlock()
+	return s.GetTableRowCountsFunc(ctx)
+}
+
 func (s *DataStore) ApplyPolicySpecs(ctx context.Context, authorID uint, specs []*fleet.PolicySpec) error {
 	s.mu.Lock()
 	s.ApplyPolicySpecsFuncInvoked = true