From e0bcb2660051c6baec69d7439d8887475598e009 Mon Sep 17 00:00:00 2001
From: Sam Lord <sam@ghost.org>
Date: Mon, 6 Feb 2023 10:40:18 +0000
Subject: [PATCH] Improved data-generator tooling

no issue

This makes the data-generator more customisable.
---
 ghost/core/core/cli/generate-data.js          | 24 ++++++++++--
 ghost/data-generator/lib/data-generator.js    | 39 ++++++++++++++++++-
 .../lib/tables/members-created-events.js      |  2 +
 3 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/ghost/core/core/cli/generate-data.js b/ghost/core/core/cli/generate-data.js
index 3249c273ea..2057131b66 100644
--- a/ghost/core/core/cli/generate-data.js
+++ b/ghost/core/core/cli/generate-data.js
@@ -2,10 +2,13 @@ const Command = require('./command');
 const DataGenerator = require('@tryghost/data-generator');
 const config = require('../shared/config');
 
-module.exports = class REPL extends Command {
+module.exports = class DataGeneratorCommand extends Command {
     setup() {
         this.help('Generates random data to populate the database for development & testing');
         this.argument('--base-data-pack', {type: 'string', defaultValue: '', desc: 'Base data pack file location, imported instead of random content'});
+        this.argument('--scale', {type: 'string', defaultValue: 'small', desc: 'Scale of the data to generate. `small` for a quick run, `large` for more content'});
+        this.argument('--single-table', {type: 'string', desc: 'Import a single table'});
+        this.argument('--quantity', {type: 'number', desc: 'When importing a single table, the quantity to import'});
     }
 
     initializeContext(context) {
@@ -27,6 +30,17 @@ module.exports = class REPL extends Command {
     async handle(argv = {}) {
         const knex = require('../server/data/db/connection');
         const {tables: schema} = require('../server/data/schema/index');
+
+        const modelQuantities = {};
+        if (argv.scale) {
+            if (argv.scale === 'small') {
+                modelQuantities.members = 200;
+                modelQuantities.membersLoginEvents = 1;
+                modelQuantities.posts = 10;
+            }
+            // Defaults in data-generator package make a large set
+        }
+
         const dataGenerator = new DataGenerator({
             baseDataPack: argv['base-data-pack'],
             knex,
@@ -40,11 +54,15 @@ module.exports = class REPL extends Command {
                 fatal: this.fatal,
                 debug: this.debug
             },
-            modelQuantities: {},
+            modelQuantities,
             baseUrl: config.getSiteUrl()
         });
         try {
-            await dataGenerator.importData();
+            if (argv['single-table']) {
+                await dataGenerator.importSingleTable(argv['single-table'], argv.quantity ?? undefined);
+            } else {
+                await dataGenerator.importData();
+            }
         } catch (error) {
             this.fatal('Failed while generating data: ', error);
         }
diff --git a/ghost/data-generator/lib/data-generator.js b/ghost/data-generator/lib/data-generator.js
index 512e0d85ce..fe31b86a1e 100644
--- a/ghost/data-generator/lib/data-generator.js
+++ b/ghost/data-generator/lib/data-generator.js
@@ -115,7 +115,7 @@ class DataGenerator {
             newsletters = await jsonImporter.import({
                 name: 'newsletters',
                 data: baseData.newsletters,
-                rows: ['sort_order']
+                rows: ['name', 'sort_order']
             });
             newsletters.sort((a, b) => a.sort_order - b.sort_order);
 
@@ -132,7 +132,7 @@ class DataGenerator {
                 data: baseData.posts
             });
             await postsImporter.addNewsletters({posts});
-            posts = await transaction.select('id', 'newsletter_id').from('posts');
+            posts = await transaction.select('id', 'newsletter_id', 'published_at', 'slug').from('posts');
 
             await transaction('tags').delete();
             tags = await jsonImporter.import({
@@ -339,6 +339,41 @@ class DataGenerator {
         this.logger.info('Completed member data generation');
         this.logger.ok('Completed import process.');
     }
+
+    async importSingleTable(tableName, quantity) {
+        this.logger.info('Importing a single table');
+        const transaction = await this.knex.transaction();
+
+        const importMembers = async () => {
+            this.logger.info(`Importing ${quantity ?? this.modelQuantities.members} members`);
+            const membersImporter = new MembersImporter(transaction);
+            await membersImporter.import({amount: quantity ?? this.modelQuantities.members});
+        };
+
+        const importMentions = async () => {
+            const posts = await transaction.select('id', 'newsletter_id', 'published_at', 'slug').from('posts');
+            this.logger.info(`Importing up to ${posts.length * 4} mentions`);
+
+            const mentionsImporter = new MentionsImporter(transaction, {baseUrl: this.baseUrl});
+            await mentionsImporter.importForEach(posts, {amount: 4});
+        };
+
+        switch (tableName) {
+        case 'members':
+            await importMembers();
+            break;
+        case 'mentions':
+            await importMentions();
+            break;
+        default:
+            this.logger.warn(`Cannot import single table '${tableName}'`);
+            await transaction.commit(); // no-op, just close the transaction
+            return;
+        }
+
+        this.logger.ok('Completed import process.');
+        await transaction.commit();
+    }
 }
 
 module.exports = DataGenerator;
diff --git a/ghost/data-generator/lib/tables/members-created-events.js b/ghost/data-generator/lib/tables/members-created-events.js
index 8bf5bc3bfc..7121c45d01 100644
--- a/ghost/data-generator/lib/tables/members-created-events.js
+++ b/ghost/data-generator/lib/tables/members-created-events.js
@@ -17,6 +17,8 @@ class MembersCreatedEventsImporter extends TableImporter {
             source = 'admin';
         } else if (luck(5)) {
             source = 'api';
+        } else if (luck(5)) { // eslint-disable-line no-dupe-else-if
+            source = 'import';
         }
         return source;
     }