Improving performance of load script, can generate thousands of users a second.

packages/backend-core/src/utils.js

@@ -197,11 +197,16 @@ exports.getBuildersCount = async () => {
   return builders.length
 }
 
-exports.saveUser = async (
+const DEFAULT_SAVE_USER = {
+  hashPassword: true,
+  requirePassword: true,
+  bulkCreate: false,
+}
+
+exports.internalSaveUser = async (
   user,
   tenantId,
-  hashPassword = true,
+  { hashPassword, requirePassword, bulkCreate } = DEFAULT_SAVE_USER
-  requirePassword = true
 ) => {
   if (!tenantId) {
     throw "No tenancy specified."
@@ -213,7 +218,10 @@ exports.saveUser = async (
     let { email, password, _id } = user
     // make sure another user isn't using the same email
     let dbUser
-    if (email) {
+    // user can't exist in bulk creation
+    if (bulkCreate) {
+      dbUser = null
+    } else if (email) {
       // check budibase users inside the tenant
       dbUser = await exports.getGlobalUserByEmail(email)
       if (dbUser != null && (dbUser._id !== _id || Array.isArray(dbUser))) {
@@ -267,11 +275,17 @@ exports.saveUser = async (
       user.status = UserStatus.ACTIVE
     }
     try {
-      const response = await db.put({
+      const putOpts = {
         password: hashedPassword,
         ...user,
-      })
+      }
-      await tryAddTenant(tenantId, _id, email)
+      if (bulkCreate) {
+        return putOpts
+      }
+      const response = await db.put(putOpts)
+      if (env.MULTI_TENANCY) {
+        await tryAddTenant(tenantId, _id, email)
+      }
       await userCache.invalidateUser(response.id)
       return {
         _id: response.id,
@@ -288,6 +302,19 @@ exports.saveUser = async (
   })
 }
 
+// maintained for api compat, don't want to change function signature
+exports.saveUser = async (
+  user,
+  tenantId,
+  hashPassword = true,
+  requirePassword = true
+) => {
+  return exports.internalSaveUser(user, tenantId, {
+    hashPassword,
+    requirePassword,
+  })
+}
+
 /**
  * Logs a user out from budibase. Re-used across account portal and builder.
  */

packages/server/scripts/load/users.js

@@ -1,13 +1,21 @@
-const fetch = require("node-fetch")
+// get the JWT secret etc
-const { getProdAppID } = require("@budibase/backend-core/db")
+require("../../src/environment")
+require("@budibase/backend-core").init()
+const {
+  getProdAppID,
+  generateGlobalUserID,
+} = require("@budibase/backend-core/db")
+const { doInTenant, getGlobalDB } = require("@budibase/backend-core/tenancy")
+const { internalSaveUser } = require("@budibase/backend-core/utils")
+const { publicApiUserFix } = require("../../src/utilities/users")
+const { hash } = require("@budibase/backend-core/utils")
 
 const USER_LOAD_NUMBER = 10000
-const BATCH_SIZE = 25
+const BATCH_SIZE = 200
-const SERVER_URL = "http://localhost:4001"
 const PASSWORD = "test"
+const TENANT_ID = "default"
 
 const APP_ID = process.argv[2]
-const API_KEY = process.argv[3]
 
 const words = [
   "test",
@@ -31,17 +39,15 @@ if (!APP_ID) {
   console.error("Must supply app ID as first CLI option!")
   process.exit(-1)
 }
-if (!API_KEY) {
-  console.error("Must supply API key as second CLI option!")
-  process.exit(-1)
-}
 
 const WORD_1 = words[Math.floor(Math.random() * words.length)]
 const WORD_2 = words[Math.floor(Math.random() * words.length)]
+let HASHED_PASSWORD
 
 function generateUser(count) {
   return {
-    password: PASSWORD,
+    _id: generateGlobalUserID(),
+    password: HASHED_PASSWORD,
     email: `${WORD_1}${count}@${WORD_2}.com`,
     roles: {
       [getProdAppID(APP_ID)]: "BASIC",
@@ -54,23 +60,31 @@ function generateUser(count) {
 }
 
 async function run() {
-  for (let i = 0; i < USER_LOAD_NUMBER; i += BATCH_SIZE) {
+  HASHED_PASSWORD = await hash(PASSWORD)
-    let promises = []
+  return doInTenant(TENANT_ID, async () => {
-    for (let j = 0; j < BATCH_SIZE; j++) {
+    const db = getGlobalDB()
-      promises.push(
+    for (let i = 0; i < USER_LOAD_NUMBER; i += BATCH_SIZE) {
-        fetch(`${SERVER_URL}/api/public/v1/users`, {
+      let userSavePromises = []
-          method: "POST",
+      for (let j = 0; j < BATCH_SIZE; j++) {
-          body: JSON.stringify(generateUser(i + j)),
+        // like the public API
-          headers: {
+        const ctx = publicApiUserFix({
-            "x-budibase-api-key": API_KEY,
+          request: {
-            "Content-Type": "application/json",
+            body: generateUser(i + j),
           },
         })
-      )
+        userSavePromises.push(
+          internalSaveUser(ctx.request.body, TENANT_ID, {
+            hashPassword: false,
+            requirePassword: true,
+            bulkCreate: true,
+          })
+        )
+      }
+      const users = await Promise.all(userSavePromises)
+      await db.bulkDocs(users)
+      console.log(`${i + BATCH_SIZE} users have been created.`)
     }
-    await Promise.all(promises)
+  })
-    console.log(`${i + BATCH_SIZE} users have been created.`)
-  }
 }
 
 run()

packages/server/src/api/controllers/public/users.ts

@@ -4,30 +4,9 @@ import {
   readGlobalUser,
   saveGlobalUser,
 } from "../../../utilities/workerRequests"
+import { publicApiUserFix } from "../../../utilities/users"
 import { search as stringSearch } from "./utils"
 
-const { getProdAppID } = require("@budibase/backend-core/db")
-
-function fixUser(ctx: any) {
-  if (!ctx.request.body) {
-    return ctx
-  }
-  if (!ctx.request.body._id && ctx.params.userId) {
-    ctx.request.body._id = ctx.params.userId
-  }
-  if (!ctx.request.body.roles) {
-    ctx.request.body.roles = {}
-  } else {
-    const newRoles: { [key: string]: string } = {}
-    for (let [appId, role] of Object.entries(ctx.request.body.roles)) {
-      // @ts-ignore
-      newRoles[getProdAppID(appId)] = role
-    }
-    ctx.request.body.roles = newRoles
-  }
-  return ctx
-}
-
 function getUser(ctx: any, userId?: string) {
   if (userId) {
     ctx.params = { userId }
@@ -45,7 +24,7 @@ export async function search(ctx: any, next: any) {
 }
 
 export async function create(ctx: any, next: any) {
-  const response = await saveGlobalUser(fixUser(ctx))
+  const response = await saveGlobalUser(publicApiUserFix(ctx))
   ctx.body = await getUser(ctx, response._id)
   await next()
 }
@@ -61,7 +40,7 @@ export async function update(ctx: any, next: any) {
     ...ctx.request.body,
     _rev: user._rev,
   }
-  const response = await saveGlobalUser(fixUser(ctx))
+  const response = await saveGlobalUser(publicApiUserFix(ctx))
   ctx.body = await getUser(ctx, response._id)
   await next()
 }

packages/server/src/environment.js

@@ -1,3 +1,5 @@
+const { join } = require("path")
+
 function isTest() {
   return (
     process.env.NODE_ENV === "jest" ||
@@ -20,7 +22,9 @@ function isCypress() {
 
 let LOADED = false
 if (!LOADED && isDev() && !isTest()) {
-  require("dotenv").config()
+  require("dotenv").config({
+    path: join(__dirname, "..", ".env"),
+  })
   LOADED = true
 }
 

packages/server/src/utilities/users.js

@@ -1,6 +1,7 @@
 const { InternalTables } = require("../db/utils")
 const { getGlobalUser } = require("../utilities/global")
 const { getAppDB } = require("@budibase/backend-core/context")
+const { getProdAppID } = require("@budibase/backend-core/db")
 
 exports.getFullUser = async (ctx, userId) => {
   const global = await getGlobalUser(userId)
@@ -22,3 +23,23 @@ exports.getFullUser = async (ctx, userId) => {
     _id: userId,
   }
 }
+
+exports.publicApiUserFix = ctx => {
+  if (!ctx.request.body) {
+    return ctx
+  }
+  if (!ctx.request.body._id && ctx.params.userId) {
+    ctx.request.body._id = ctx.params.userId
+  }
+  if (!ctx.request.body.roles) {
+    ctx.request.body.roles = {}
+  } else {
+    const newRoles = {}
+    for (let [appId, role] of Object.entries(ctx.request.body.roles)) {
+      // @ts-ignore
+      newRoles[getProdAppID(appId)] = role
+    }
+    ctx.request.body.roles = newRoles
+  }
+  return ctx
+}

packages/worker/src/environment.js

@@ -1,3 +1,5 @@
+const { join } = require("path")
+
 function isDev() {
   return process.env.NODE_ENV !== "production"
 }
@@ -12,7 +14,9 @@ function isTest() {
 
 let LOADED = false
 if (!LOADED && isDev() && !isTest()) {
-  require("dotenv").config()
+  require("dotenv").config({
+    path: join(__dirname, "..", ".env"),
+  })
   LOADED = true
 }