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

This commit is contained in:
mike12345567 2022-05-22 18:29:02 +01:00
parent 513d901ab4
commit aede23d44e
6 changed files with 106 additions and 57 deletions

View File

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

View File

@ -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 USER_LOAD_NUMBER = 10000
const BATCH_SIZE = 25 const BATCH_SIZE = 200
const SERVER_URL = "http://localhost:4001"
const PASSWORD = "test" const PASSWORD = "test"
const TENANT_ID = "default"
const APP_ID = process.argv[2] const APP_ID = process.argv[2]
const API_KEY = process.argv[3]
const words = [ const words = [
"test", "test",
@ -31,17 +39,15 @@ if (!APP_ID) {
console.error("Must supply app ID as first CLI option!") console.error("Must supply app ID as first CLI option!")
process.exit(-1) 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_1 = words[Math.floor(Math.random() * words.length)]
const WORD_2 = words[Math.floor(Math.random() * words.length)] const WORD_2 = words[Math.floor(Math.random() * words.length)]
let HASHED_PASSWORD
function generateUser(count) { function generateUser(count) {
return { return {
password: PASSWORD, _id: generateGlobalUserID(),
password: HASHED_PASSWORD,
email: `${WORD_1}${count}@${WORD_2}.com`, email: `${WORD_1}${count}@${WORD_2}.com`,
roles: { roles: {
[getProdAppID(APP_ID)]: "BASIC", [getProdAppID(APP_ID)]: "BASIC",
@ -54,23 +60,31 @@ function generateUser(count) {
} }
async function run() { 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() run()

View File

@ -4,30 +4,9 @@ import {
readGlobalUser, readGlobalUser,
saveGlobalUser, saveGlobalUser,
} from "../../../utilities/workerRequests" } from "../../../utilities/workerRequests"
import { publicApiUserFix } from "../../../utilities/users"
import { search as stringSearch } from "./utils" 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) { function getUser(ctx: any, userId?: string) {
if (userId) { if (userId) {
ctx.params = { userId } ctx.params = { userId }
@ -45,7 +24,7 @@ export async function search(ctx: any, next: any) {
} }
export async function create(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) ctx.body = await getUser(ctx, response._id)
await next() await next()
} }
@ -61,7 +40,7 @@ export async function update(ctx: any, next: any) {
...ctx.request.body, ...ctx.request.body,
_rev: user._rev, _rev: user._rev,
} }
const response = await saveGlobalUser(fixUser(ctx)) const response = await saveGlobalUser(publicApiUserFix(ctx))
ctx.body = await getUser(ctx, response._id) ctx.body = await getUser(ctx, response._id)
await next() await next()
} }

View File

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

View File

@ -1,6 +1,7 @@
const { InternalTables } = require("../db/utils") const { InternalTables } = require("../db/utils")
const { getGlobalUser } = require("../utilities/global") const { getGlobalUser } = require("../utilities/global")
const { getAppDB } = require("@budibase/backend-core/context") const { getAppDB } = require("@budibase/backend-core/context")
const { getProdAppID } = require("@budibase/backend-core/db")
exports.getFullUser = async (ctx, userId) => { exports.getFullUser = async (ctx, userId) => {
const global = await getGlobalUser(userId) const global = await getGlobalUser(userId)
@ -22,3 +23,23 @@ exports.getFullUser = async (ctx, userId) => {
_id: 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
}

View File

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