From 210331ba1ea1104f437592f54572d0e2b22ceaec Mon Sep 17 00:00:00 2001 From: Rajendra kadam Date: Thu, 4 Jun 2026 13:09:49 +0530 Subject: [PATCH] Extract datastore initialization out of runServeCmd (#46742) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extracts the MySQL datastore initialization out of `runServeCmd` and into a new `cmd/fleet/datastore.go`. Same pattern as the prior extractions on this issue (#44929, #45343, #45583, #46166, #46421, #46517). Continues the path toward `serve.go` >60% coverage per the discussion on #33370. Three functions come out of the inline block: - `initDatastore` — builds the shared DB connections, the datastore, and the carve store (S3-backed when configured, otherwise the datastore itself). - `buildMySQLOpts` — assembles the DB options: base logger and config, plus the optional read replica, dev SQL interceptor, and tracing. - `evalMigrationStatus` — prints any operator guidance for the migration status and returns whether `runServeCmd` should exit. The `os.Exit` stays in `runServeCmd`, so the boot/refuse-to-boot decision becomes unit-testable without the function terminating the test binary. Behavior is preserved — `runServeCmd` calls these in the same order with the same arguments, the migration-exit conditions are unchanged, and the full `cmd/fleet` suite passes against MySQL + Redis. `initDatastore` returns early after `initFatal` so it's safe when the caller's `initFatal` doesn't terminate (the case in tests). On test scope: `TestEvalMigrationStatus` covers every migration status code across the dev-mode and allow-missing-migrations combinations — that's the real decision logic. I deliberately didn't add unit tests for `initDatastore`/`buildMySQLOpts`: their only failure paths are paranoid `initFatal` wrapping around constructors that don't dial at construction time, and the option builder returns opaque option closures. Those success paths are already exercised by booting the server, so a full datastore mock wasn't worth it for coverage's sake. Remaining slice per the broader plan: Redis init. **Related issue:** Refs #33370 # Checklist for submitter - [x] Added/updated automated tests - Changes file: not applicable — internal refactor with no user-visible behavior change ## Summary by CodeRabbit * **Refactor** * Reorganized database startup initialization and migration status evaluation for improved maintainability. * **Tests** * Added comprehensive test coverage for database migration status handling across various scenarios. --- cmd/fleet/datastore.go | 101 ++++++++++++++++++++++++++++ cmd/fleet/datastore_test.go | 130 ++++++++++++++++++++++++++++++++++++ cmd/fleet/serve.go | 74 +++----------------- 3 files changed, 241 insertions(+), 64 deletions(-) create mode 100644 cmd/fleet/datastore.go create mode 100644 cmd/fleet/datastore_test.go diff --git a/cmd/fleet/datastore.go b/cmd/fleet/datastore.go new file mode 100644 index 0000000000..4006b50264 --- /dev/null +++ b/cmd/fleet/datastore.go @@ -0,0 +1,101 @@ +package main + +import ( + "log/slog" + "os" + + "github.com/WatchBeam/clock" + "github.com/fleetdm/fleet/v4/server/config" + "github.com/fleetdm/fleet/v4/server/datastore/mysql" + "github.com/fleetdm/fleet/v4/server/datastore/s3" + "github.com/fleetdm/fleet/v4/server/dev_mode" + "github.com/fleetdm/fleet/v4/server/fleet" + common_mysql "github.com/fleetdm/fleet/v4/server/platform/mysql" +) + +// buildMySQLOpts assembles the DBOptions for the primary datastore connection +// from the Fleet config: the base logger and config, plus the optional read +// replica, dev SQL interceptor, and tracing. +func buildMySQLOpts(config config.FleetConfig, logger *slog.Logger) []mysql.DBOption { + opts := []mysql.DBOption{mysql.Logger(logger), mysql.WithFleetConfig(&config)} + if config.MysqlReadReplica.Address != "" { + opts = append(opts, mysql.Replica(&config.MysqlReadReplica)) + } + // NOTE this will disable OTEL/APM interceptor + if dev_mode.Env("FLEET_DEV_ENABLE_SQL_INTERCEPTOR") != "" { + opts = append(opts, mysql.WithInterceptor(&devSQLInterceptor{ + logger: logger.With("component", "sql-interceptor"), + })) + } + if config.Logging.TracingEnabled { + opts = append(opts, mysql.TracingEnabled(&config.Logging)) + } + return opts +} + +// initDatastore brings up the MySQL datastore: shared DB connections, the +// datastore itself, and the carve store (S3-backed when configured, otherwise +// the datastore). Failures go through initFatal. Returns nil values on the +// failure path so the function is safe when initFatal does not terminate +// (e.g., tests using a recorder). +func initDatastore(config config.FleetConfig, logger *slog.Logger, c clock.Clock, initFatal func(err error, msg string)) ( + *mysql.Datastore, + *common_mysql.DBConnections, + fleet.CarveStore, +) { + opts := buildMySQLOpts(config, logger) + + // Create database connections that can be shared across datastores + dbConns, err := mysql.NewDBConnections(config.Mysql, opts...) + if err != nil { + initFatal(err, "initializing database connections") + return nil, nil, nil + } + + mds, err := mysql.NewDatastore(dbConns, config.Mysql, c) + if err != nil { + initFatal(err, "initializing datastore") + return nil, nil, nil + } + + var carveStore fleet.CarveStore = mds + if config.S3.CarvesBucket != "" || config.S3.Bucket != "" { + carveStore, err = s3.NewCarveStore(config.S3, mds) + if err != nil { + initFatal(err, "initializing S3 carvestore") + return nil, nil, nil + } + } + + return mds, dbConns, carveStore +} + +// evalMigrationStatus prints any operator guidance for the current migration +// status and reports whether runServeCmd should exit instead of starting. It +// encodes the boot/refuse-to-boot decision: unknown migrations are only fatal +// in dev mode; the v4.73.2 and partial-migration states are fatal unless +// missing migrations are explicitly allowed; an uninitialized database is +// always fatal. +func evalMigrationStatus(status *fleet.MigrationStatus, devMode, allowMissing bool) (shouldExit bool) { + switch status.StatusCode { + case fleet.AllMigrationsCompleted: + // OK + return false + case fleet.UnknownMigrations: + printUnknownMigrationsMessage(status.UnknownTable, status.UnknownData) + return devMode + case fleet.NeedsFleetv4732Fix: + printFleetv4732FixNeededMessage() + return !allowMissing + case fleet.UnknownFleetv4732State: + printFleetv4732UnknownStateMessage(status.StatusCode) + return !allowMissing + case fleet.SomeMigrationsCompleted: + printMissingMigrationsWarning(os.Stdout, status.MissingTable, status.MissingData) + return !allowMissing + case fleet.NoMigrationsCompleted: + printDatabaseNotInitializedError() + return true + } + return false +} diff --git a/cmd/fleet/datastore_test.go b/cmd/fleet/datastore_test.go new file mode 100644 index 0000000000..6ca3f03975 --- /dev/null +++ b/cmd/fleet/datastore_test.go @@ -0,0 +1,130 @@ +package main + +import ( + "io" + "os" + "testing" + + "github.com/fleetdm/fleet/v4/server/fleet" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// captureStdout redirects os.Stdout for the duration of fn (the migration +// print helpers write there directly) and returns what was written, so test +// output stays clean and the banner wiring can be asserted. +func captureStdout(t *testing.T, fn func()) string { + t.Helper() + orig := os.Stdout + r, w, err := os.Pipe() + require.NoError(t, err) + os.Stdout = w + t.Cleanup(func() { os.Stdout = orig }) + + fn() + + require.NoError(t, w.Close()) + out, err := io.ReadAll(r) + require.NoError(t, err) + return string(out) +} + +func TestEvalMigrationStatus(t *testing.T) { + for _, tc := range []struct { + name string + status *fleet.MigrationStatus + devMode bool + allowMissing bool + wantExit bool + wantOut string // substring expected on stdout; empty means no output + }{ + { + name: "all completed never exits", + status: &fleet.MigrationStatus{StatusCode: fleet.AllMigrationsCompleted}, + wantExit: false, + wantOut: "", + }, + { + name: "all completed ignores dev and allow-missing", + status: &fleet.MigrationStatus{StatusCode: fleet.AllMigrationsCompleted}, + devMode: true, + wantExit: false, + wantOut: "", + }, + { + name: "unknown migrations fatal only in dev", + status: &fleet.MigrationStatus{StatusCode: fleet.UnknownMigrations, UnknownTable: []int64{1}}, + devMode: true, + wantExit: true, + wantOut: "unrecognized migrations", + }, + { + name: "unknown migrations tolerated outside dev still warns", + status: &fleet.MigrationStatus{StatusCode: fleet.UnknownMigrations, UnknownTable: []int64{1}}, + devMode: false, + wantExit: false, + wantOut: "unrecognized migrations", + }, + { + name: "needs v4732 fix exits unless allowed", + status: &fleet.MigrationStatus{StatusCode: fleet.NeedsFleetv4732Fix}, + allowMissing: false, + wantExit: true, + wantOut: "automatically perform this fix", + }, + { + name: "needs v4732 fix tolerated when missing allowed still warns", + status: &fleet.MigrationStatus{StatusCode: fleet.NeedsFleetv4732Fix}, + allowMissing: true, + wantExit: false, + wantOut: "automatically perform this fix", + }, + { + name: "unknown v4732 state exits unless allowed", + status: &fleet.MigrationStatus{StatusCode: fleet.UnknownFleetv4732State}, + wantExit: true, + wantOut: "contact Fleet support", + }, + { + name: "unknown v4732 state tolerated when missing allowed still warns", + status: &fleet.MigrationStatus{StatusCode: fleet.UnknownFleetv4732State}, + allowMissing: true, + wantExit: false, + wantOut: "contact Fleet support", + }, + { + name: "some migrations completed exits unless allowed", + status: &fleet.MigrationStatus{StatusCode: fleet.SomeMigrationsCompleted, MissingTable: []int64{7}}, + wantExit: true, + wantOut: "tables=[7]", + }, + { + name: "some migrations completed tolerated when missing allowed still warns", + status: &fleet.MigrationStatus{StatusCode: fleet.SomeMigrationsCompleted, MissingTable: []int64{7}}, + allowMissing: true, + wantExit: false, + wantOut: "tables=[7]", + }, + { + name: "no migrations always exits", + status: &fleet.MigrationStatus{StatusCode: fleet.NoMigrationsCompleted}, + allowMissing: true, + devMode: true, + wantExit: true, + wantOut: "not initialized", + }, + } { + t.Run(tc.name, func(t *testing.T) { + var got bool + out := captureStdout(t, func() { + got = evalMigrationStatus(tc.status, tc.devMode, tc.allowMissing) + }) + assert.Equal(t, tc.wantExit, got) + if tc.wantOut == "" { + assert.Empty(t, out) + } else { + assert.Contains(t, out, tc.wantOut) + } + }) + } +} diff --git a/cmd/fleet/serve.go b/cmd/fleet/serve.go index 42ef186223..b647979e35 100644 --- a/cmd/fleet/serve.go +++ b/cmd/fleet/serve.go @@ -248,79 +248,25 @@ func runServeCmd(cmd *cobra.Command, configManager configpkg.Manager, debug, dev config.MDM.CertificateProfilesLimit = 0 } - var ds fleet.Datastore - var carveStore fleet.CarveStore - - opts := []mysql.DBOption{mysql.Logger(logger), mysql.WithFleetConfig(&config)} - if config.MysqlReadReplica.Address != "" { - opts = append(opts, mysql.Replica(&config.MysqlReadReplica)) - } - // NOTE this will disable OTEL/APM interceptor - if dev_mode.Env("FLEET_DEV_ENABLE_SQL_INTERCEPTOR") != "" { - opts = append(opts, mysql.WithInterceptor(&devSQLInterceptor{ - logger: logger.With("component", "sql-interceptor"), - })) - } - - if config.Logging.TracingEnabled { - opts = append(opts, mysql.TracingEnabled(&config.Logging)) - } - // Configure default max request body size based on config platform_http.MaxRequestBodySize = config.Server.DefaultMaxRequestBodySize - // Create database connections that can be shared across datastores - dbConns, err := mysql.NewDBConnections(config.Mysql, opts...) - if err != nil { - initFatal(err, "initializing database connections") - } - - mds, err := mysql.NewDatastore(dbConns, config.Mysql, clock.C) - if err != nil { - initFatal(err, "initializing datastore") - } - ds = mds - - if config.S3.CarvesBucket != "" || config.S3.Bucket != "" { - carveStore, err = s3.NewCarveStore(config.S3, ds) - if err != nil { - initFatal(err, "initializing S3 carvestore") - } - } else { - carveStore = ds + mds, dbConns, carveStore := initDatastore(config, logger, clock.C, initFatal) + if mds == nil { + initFatal(errors.New("datastore was nil after initialization"), "initializing datastore") + return } + var ds fleet.Datastore = mds migrationStatus, err := ds.MigrationStatus(cmd.Context()) if err != nil { initFatal(err, "retrieving migration status") } - - switch migrationStatus.StatusCode { - case fleet.AllMigrationsCompleted: - // OK - case fleet.UnknownMigrations: - printUnknownMigrationsMessage(migrationStatus.UnknownTable, migrationStatus.UnknownData) - if dev_mode.IsEnabled { - os.Exit(1) - } - case fleet.NeedsFleetv4732Fix: - printFleetv4732FixNeededMessage() - if !config.Upgrades.AllowMissingMigrations { - os.Exit(1) - } - case fleet.UnknownFleetv4732State: - printFleetv4732UnknownStateMessage(migrationStatus.StatusCode) - if !config.Upgrades.AllowMissingMigrations { - os.Exit(1) - } - case fleet.SomeMigrationsCompleted: - tables, data := migrationStatus.MissingTable, migrationStatus.MissingData - printMissingMigrationsWarning(os.Stdout, tables, data) - if !config.Upgrades.AllowMissingMigrations { - os.Exit(1) - } - case fleet.NoMigrationsCompleted: - printDatabaseNotInitializedError() + if migrationStatus == nil { + initFatal(errors.New("migration status was nil"), "retrieving migration status") + return + } + if evalMigrationStatus(migrationStatus, dev_mode.IsEnabled, config.Upgrades.AllowMissingMigrations) { os.Exit(1) }