Improve workload configuration logic and startup feedback

Author: daniel-percona

README.md

@@ -183,10 +183,9 @@ export PLGM_COLLECTIONS_PATH="./resources/custom_collections/"
 #### Default Workload Filtering
 When using **Directory Mode**, the behavior depends on the `PLGM_DEFAULT_WORKLOAD` setting:
 
-* **`true` (Default):** Loads **all** JSON files in the directory, including `default.json`.
+* **`true` (Default):** Loads **only** `default.json` (if present). It ignores all other files in the folder.
 * **`false` (Custom):** Loads all JSON files **except** `default.json`. 
-  * *Use Case:* You can keep the example `default.json` in your folder for reference but exclude it from your actual test run by setting `PLGM_DEFAULT_WORKLOAD=false`.
-
+  * *Use Case:* Set this to `false` to run your custom workload files while keeping `default.json` in the folder for reference (it will be ignored).
 
 ### 5. Docker & Kubernetes
 Prefer running in a container? We have a dedicated guide for building Docker images and running performance jobs directly inside Kubernetes (recommended for accurate network latency testing).

cmd/plgm/main.go

@@ -113,38 +113,31 @@ func main() {
 	}
 
 	// --- Secure Credentials Logic ---
-
-	// 1. Analyze the Base URI to see if it already has credentials
 	u, err := url.Parse(appCfg.URI)
 	if err != nil {
 		log.Fatalf("Invalid PLGM_URI: %v", err)
 	}
 	uriHasUser := u.User != nil && u.User.Username() != ""
 
-	// 2. Prompt for Username if missing (from both URI and Env/Config)
 	if !uriHasUser && appCfg.ConnectionParams.Username == "" {
 		fmt.Print("Enter MongoDB Username: ")
 		var inputUser string
 		if _, err := fmt.Scanln(&inputUser); err != nil {
-			// Handle case where user hits enter or pipe closes
 			if err.Error() != "unexpected newline" {
 				log.Fatal(err)
 			}
 		}
 		appCfg.ConnectionParams.Username = inputUser
 	}
 
-	// 3. Prompt for Password if missing
-	// Logic: If we have a username defined in config/prompt (overriding whatever is in URI),
-	// and no password is set for it, we must prompt.
 	if appCfg.ConnectionParams.Username != "" && appCfg.ConnectionParams.Password == "" {
 		fmt.Printf("Enter Password for user '%s': ", appCfg.ConnectionParams.Username)
 		bytePassword, err := term.ReadPassword(int(os.Stdin.Fd()))
 		if err != nil {
 			log.Fatal("\nError reading password:", err)
 		}
 		appCfg.ConnectionParams.Password = string(bytePassword)
-		fmt.Println() // Print newline after input
+		fmt.Println()
 	}
 
 	// --- Load Collections ---
@@ -185,7 +178,6 @@ func main() {
 	}
 	queriesCfg.Queries = filteredQueries
 
-	// Determine DB name from first collection for the banner
 	dbName := collectionsCfg.Collections[0].DatabaseName
 
 	// -----------------------------------------------------------------------------------
@@ -201,7 +193,6 @@ func main() {
 	defer conn.Disconnect(ctx)
 
 	// --- Log Setup Details (Execution Phase) ---
-	// Now we print the logs, so they appear below the plan
 	logger.Info("Loaded %d collection definition(s)", len(collectionsCfg.Collections))
 	logger.Info("Loaded %d query templates(s)", len(queriesCfg.Queries))
 
@@ -218,11 +209,9 @@ func main() {
 		log.Fatal(err)
 	}
 
-	// --- Seed documents (initial dataset) ---
+	// --- Seed documents ---
 	if !appCfg.SkipSeed {
 		if appCfg.DocumentsCount > 0 {
-			// NOTE: Logging is handled inside InsertRandomDocuments,
-			// but we can add a high-level log here if we want.
 			for _, col := range collectionsCfg.Collections {
 				if err := mongo.InsertRandomDocuments(ctx, conn.Database, col, appCfg.DocumentsCount, appCfg); err != nil {
 					log.Fatal(err)

internal/config/collections.go

@@ -43,13 +43,10 @@ type CollectionsFile struct {
 	Collections []CollectionDefinition `json:"collections"`
 }
 
-// LoadCollections loads files from a path.
-// If path is a folder:
-//   - If loadDefault is TRUE: Loads ALL .json files (including default.json).
-//   - If loadDefault is FALSE: Loads ALL .json files EXCEPT default.json.
-//
-// If path is a file:
-//   - Loads the file unconditionally.
+// LoadCollections filters files based on the 'loadDefault' flag.
+// - If loadDefault is TRUE: Load ONLY 'default.json'.
+// - If loadDefault is FALSE: Load ALL files EXCEPT 'default.json'.
+// - Single file paths are always loaded.
 func LoadCollections(path string, loadDefault bool) (*CollectionsFile, error) {
 	if path == "" {
 		return &CollectionsFile{}, nil
@@ -74,10 +71,16 @@ func LoadCollections(path string, loadDefault bool) (*CollectionsFile, error) {
 
 			isDefault := strings.EqualFold(entry.Name(), "default.json")
 
-			// If default_workload is FALSE, we explicitly ignore 'default.json' in folder mode.
-			// If default_workload is TRUE, we load everything (no continue).
-			if !loadDefault && isDefault {
-				continue
+			if loadDefault {
+				// Mode: Default Workload -> Load ONLY default.json
+				if !isDefault {
+					continue
+				}
+			} else {
+				// Mode: Custom Workload -> Load EVERYTHING ELSE
+				if isDefault {
+					continue
+				}
 			}
 
 			fullPath := filepath.Join(path, entry.Name())
@@ -88,7 +91,7 @@ func LoadCollections(path string, loadDefault bool) (*CollectionsFile, error) {
 			allCollections = append(allCollections, loaded.Collections...)
 		}
 	} else {
-		// Single file: Always load it if the user specifically pointed to it.
+		// Single file: Always load it (explicit user choice).
 		loaded, err := loadCollectionsFromFile(path)
 		if err != nil {
 			return nil, err

internal/config/config.go

@@ -111,17 +111,23 @@ func applyDefaults(cfg *AppConfig) {
 }
 
 func applyEnvOverrides(cfg *AppConfig) {
+	// 1. Credentials
 	if v := os.Getenv("PLGM_USERNAME"); v != "" {
 		cfg.ConnectionParams.Username = v
 	}
 	if v := os.Getenv("PLGM_PASSWORD"); v != "" {
 		cfg.ConnectionParams.Password = v
 	}
+
+	// 2. Default Workload (Explicit Override)
+	// If the user sets this Env Var, it takes precedence over everything.
 	if v := os.Getenv("PLGM_DEFAULT_WORKLOAD"); v != "" {
 		if b, err := strconv.ParseBool(v); err == nil {
 			cfg.DefaultWorkload = b
 		}
 	}
+
+	// 3. Other Settings
 	if envDebug := os.Getenv("PLGM_DEBUG_MODE"); envDebug != "" {
 		if b, err := strconv.ParseBool(envDebug); err == nil {
 			cfg.DebugMode = b
@@ -142,17 +148,12 @@ func applyEnvOverrides(cfg *AppConfig) {
 		cfg.ConnectionParams.ReadPreference = v
 	}
 
-	customWorkloadEnv := false
+	// 4. Custom Paths
 	if envCollectionsPath := os.Getenv("PLGM_COLLECTIONS_PATH"); envCollectionsPath != "" {
 		cfg.CollectionsPath = envCollectionsPath
-		customWorkloadEnv = true
 	}
 	if envQueriesPath := os.Getenv("PLGM_QUERIES_PATH"); envQueriesPath != "" {
 		cfg.QueriesPath = envQueriesPath
-		customWorkloadEnv = true
-	}
-	if customWorkloadEnv {
-		cfg.DefaultWorkload = false
 	}
 
 	if envDrop := os.Getenv("PLGM_DROP_COLLECTIONS"); envDrop != "" {

internal/config/queries.go

@@ -24,13 +24,10 @@ type QueriesFile struct {
 	Queries []QueryDefinition
 }
 
-// LoadQueries loads files from a path.
-// If path is a folder:
-//   - If loadDefault is TRUE: Loads ALL .json files (including default.json).
-//   - If loadDefault is FALSE: Loads ALL .json files EXCEPT default.json.
-//
-// If path is a file:
-//   - Loads the file unconditionally.
+// LoadQueries filters files based on the 'loadDefault' flag.
+// - If loadDefault is TRUE: Load ONLY 'default.json'.
+// - If loadDefault is FALSE: Load ALL files EXCEPT 'default.json'.
+// - Single file paths are always loaded.
 func LoadQueries(path string, loadDefault bool) (*QueriesFile, error) {
 	if path == "" {
 		return &QueriesFile{}, nil
@@ -55,10 +52,16 @@ func LoadQueries(path string, loadDefault bool) (*QueriesFile, error) {
 
 			isDefault := strings.EqualFold(entry.Name(), "default.json")
 
-			// If default_workload is FALSE, we explicitly ignore 'default.json' in folder mode.
-			// If default_workload is TRUE, we load everything.
-			if !loadDefault && isDefault {
-				continue
+			if loadDefault {
+				// Mode: Default Workload -> Load ONLY default.json
+				if !isDefault {
+					continue
+				}
+			} else {
+				// Mode: Custom Workload -> Load EVERYTHING ELSE
+				if isDefault {
+					continue
+				}
 			}
 
 			fullPath := filepath.Join(path, entry.Name())

internal/mongo/runner.go

@@ -177,6 +177,11 @@ func insertDocumentProducer(ctx context.Context, col config.CollectionDefinition
 	}
 }
 
+// Helper to get collection handle
+func getCollectionHandle(db *mongo.Database, col config.CollectionDefinition) *mongo.Collection {
+	return db.Client().Database(col.DatabaseName).Collection(col.Name)
+}
+
 func runTransaction(ctx context.Context, id int, wCfg workloadConfig, rng *rand.Rand) {
 	session, err := wCfg.database.Client().StartSession()
 	if err != nil {
@@ -215,7 +220,8 @@ func runTransaction(ctx context.Context, id int, wCfg workloadConfig, rng *rand.
 				continue
 			}
 
-			coll := wCfg.database.Collection(currentCol.Name)
+			coll := getCollectionHandle(wCfg.database, currentCol)
+
 			filter := cloneMap(q.Filter)
 			processRecursive(filter, rng)
 
@@ -304,7 +310,8 @@ func independentWorker(ctx context.Context, id int, wg *sync.WaitGroup, wCfg wor
 			continue
 		}
 
-		coll := wCfg.database.Collection(currentCol.Name)
+		coll := getCollectionHandle(wCfg.database, currentCol)
+
 		var filter map[string]interface{}
 		var pipeline []interface{}
 
@@ -532,7 +539,12 @@ func queryWorkerOnce(ctx context.Context, id int, tasks <-chan *queryTask, wg *s
 	dbOpCtx := context.Background()
 	for task := range tasks {
 		q := task.definition
-		coll := task.database.Collection(q.Collection)
+		coll := task.database.Client().Database(q.Database).Collection(q.Collection)
+		if q.Database == "" {
+			// Fallback if not specified in query
+			coll = task.database.Collection(q.Collection)
+		}
+
 		filter := cloneMap(q.Filter)
 		processRecursive(filter, task.rng)
 		switch q.Operation {

internal/mongo/setup.go

@@ -15,17 +15,19 @@ import (
 )
 
 func InsertRandomDocuments(ctx context.Context, db *mongo.Database, col config.CollectionDefinition, count int, cfg *config.AppConfig) error {
-	logger.Info("Seeding %d documents into '%s'...", count, col.Name)
+	logger.Info("Seeding %d documents into '%s.%s'...", count, col.DatabaseName, col.Name)
 	docs := make([]interface{}, count)
 
 	for i := 0; i < count; i++ {
 		docs[i] = workloads.GenerateDocument(col, cfg)
 	}
 
 	if len(docs) > 0 {
-		_, err := db.Collection(col.Name).InsertMany(ctx, docs)
+		// Use database from definition
+		targetDB := db.Client().Database(col.DatabaseName)
+		_, err := targetDB.Collection(col.Name).InsertMany(ctx, docs)
 		if err != nil {
-			return fmt.Errorf("insert documents into %s: %w", col.Name, err)
+			return fmt.Errorf("insert documents into %s.%s: %w", col.DatabaseName, col.Name, err)
 		}
 	}
 	return nil
@@ -35,7 +37,7 @@ func InsertRandomDocuments(ctx context.Context, db *mongo.Database, col config.C
 func CreateCollectionsFromConfig(ctx context.Context, db *mongo.Database, cfg *config.CollectionsFile, drop bool) error {
 	adminDB := db.Client().Database("admin")
 
-	// 1. Check if the cluster is actually sharded, if the output has isdbgrid this means it is a sharded cluster)
+	// 1. Check if the cluster is actually sharded
 	var helloResult bson.M
 	isShardedCluster := false
 	if err := adminDB.RunCommand(ctx, bson.D{{Key: "hello", Value: 1}}).Decode(&helloResult); err == nil {
@@ -45,16 +47,18 @@ func CreateCollectionsFromConfig(ctx context.Context, db *mongo.Database, cfg *c
 	}
 
 	for _, col := range cfg.Collections {
+		// Derive database handle
+		targetDB := db.Client().Database(col.DatabaseName)
+
 		// 2. Drop if requested
 		if drop {
-			_ = db.Collection(col.Name).Drop(ctx)
+			_ = targetDB.Collection(col.Name).Drop(ctx)
 		}
 
 		// 3. Create Collection
-		// We only log if we actually create it or if it's significant
-		if err := db.CreateCollection(ctx, col.Name); err != nil {
+		if err := targetDB.CreateCollection(ctx, col.Name); err != nil {
 			if drop {
-				return fmt.Errorf("create collection %s: %w", col.Name, err)
+				return fmt.Errorf("create collection %s.%s: %w", col.DatabaseName, col.Name, err)
 			}
 		}
 
@@ -90,7 +94,9 @@ func CreateIndexesFromConfig(ctx context.Context, db *mongo.Database, cfg *confi
 		if len(col.Indexes) == 0 {
 			continue
 		}
-		collection := db.Collection(col.Name)
+
+		targetDB := db.Client().Database(col.DatabaseName)
+		collection := targetDB.Collection(col.Name)
 		models := make([]mongo.IndexModel, 0, len(col.Indexes))
 
 		for _, idx := range col.Indexes {
@@ -112,11 +118,10 @@ func CreateIndexesFromConfig(ctx context.Context, db *mongo.Database, cfg *confi
 		opts := options.CreateIndexes()
 
 		if _, err := collection.Indexes().CreateMany(ctxCreate, models, opts); err != nil {
-			return fmt.Errorf("create indexes on %s: %w", col.Name, err)
+			return fmt.Errorf("create indexes on %s.%s: %w", col.DatabaseName, col.Name, err)
 		}
 
-		// Consolidated Log:
-		logger.Info("Created %d indexes on '%s'", len(models), col.Name)
+		logger.Info("Created %d indexes on '%s.%s'", len(models), col.DatabaseName, col.Name)
 	}
 	return nil
 }