diff --git a/changes/35141-table-sizes b/changes/35141-table-sizes new file mode 100644 index 0000000000..4ff526c05d --- /dev/null +++ b/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. diff --git a/cmd/fleet/prepare.go b/cmd/fleet/prepare.go index 8af9aa63b2..170a87f52a 100644 --- a/cmd/fleet/prepare.go +++ b/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 diff --git a/server/datastore/mysql/statistics.go b/server/datastore/mysql/statistics.go index 8567e64ee1..13d9c2ed7b 100644 --- a/server/datastore/mysql/statistics.go +++ b/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 +} diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index 0cddfa988e..f33440fbbc 100644 --- a/server/fleet/datastore.go +++ b/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 diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index 61dde61421..ce5790786c 100644 --- a/server/mock/datastore_mock.go +++ b/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