@rockset/cli
Version:
Official Rockset CLI
141 lines (139 loc) • 5.64 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.loadTest = exports.shouldLoadTest = exports.runApiCall = exports.showTable = void 0;
/* eslint-disable unicorn/no-abusive-eslint-disable */
const _ = require("lodash");
const pathutil_1 = require("@rockset/core/dist/filesystem/pathutil");
const perf_hooks_1 = require("perf_hooks");
const helper_1 = require("@rockset/core/dist/helper");
const cli_ux_1 = require("cli-ux");
const YAML = require("yaml");
const prompts = require("prompts");
function showTable(data, flags) {
var _a;
const columns = Object.getOwnPropertyNames((_a = data === null || data === void 0 ? void 0 : data[0]) !== null && _a !== void 0 ? _a : {});
const col = columns.reduce((obj, cur) => (Object.assign(Object.assign({}, obj), { [cur]: { header: cur } })), {});
cli_ux_1.cli.table(data, col, Object.assign({}, flags));
}
exports.showTable = showTable;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
async function runApiCall({ args, flags, namedArgs, apicall, method, endpoint, }) {
var _a, _b, _c;
const log = (_a = this.log.bind(this)) !== null && _a !== void 0 ? _a : console.log;
// This function shows a regular object in a table by wrapping it in an array
// Special case when the output is set to JSON or YAML (otherwise there will be an extra bracket in the output)
function showObjectAsTable(data, flags) {
if (flags.output === 'json') {
log(helper_1.prettyPrint(data));
}
else if (flags.output === 'yaml') {
log(YAML.stringify(data));
}
else {
showTable([data], flags);
}
}
// Pass the URL arguments directly into the api call
const allArgs = namedArgs
.filter((arg) => arg.name !== 'body')
.map((arg) => args[arg.name]);
// This path is only taken for POST requests
if (flags.body) {
// Load the specified file
const configRaw = (await pathutil_1.readConfigFromPath(pathutil_1.makeAbsolute(flags.body), 'YAML'));
allArgs.push(configRaw);
}
this.info(`${method}: ${endpoint}`);
const argObj = _.zipObject(namedArgs.map((arg) => arg.name), allArgs);
if (!_.isEmpty(argObj)) {
this.info(`Arguments: \n${JSON.stringify(argObj, null, 2)}`);
}
const loadTestRps = flags.loadTestRps;
if (loadTestRps) {
await loadTest.bind(this)(() => apicall(...allArgs), loadTestRps, flags.yes);
}
else {
const data = (await apicall(...allArgs));
const uData = (_c = (_b = data === null || data === void 0 ? void 0 : data.results) !== null && _b !== void 0 ? _b : data === null || data === void 0 ? void 0 : data.data) !== null && _c !== void 0 ? _c : data;
const unwrapData = flags.raw ? data : uData;
if (_.isArray(unwrapData)) {
showTable(unwrapData, Object.assign({}, flags));
}
else if (_.isObject(unwrapData)) {
showObjectAsTable(unwrapData, flags);
}
else {
log(helper_1.prettyPrint(unwrapData));
}
}
}
exports.runApiCall = runApiCall;
async function shouldLoadTest(rps, skipConfirmation = false) {
if (skipConfirmation) {
return true;
}
const { c } = (await prompts({
type: 'confirm',
name: 'c',
initial: false,
message: `Please confirm that you would like to send ${rps} API requests per second to the endpoint show above.
Sending huge amounts of requests may cause performance issues for the rest of your organization. Please be careful`,
}));
return c;
}
exports.shouldLoadTest = shouldLoadTest;
async function loadTest(apiCall, rps, skipConfirmation = false) {
const proceed = await shouldLoadTest(rps, skipConfirmation);
if (!proceed) {
this.error('Load test aborted.');
}
const exceptions = [];
const data = [];
const successLatency = [];
const errLatency = [];
let total = 0;
this.info('Starting load test');
const execute = async () => {
const startTime = perf_hooks_1.performance.now();
try {
total += 1;
const d = await apiCall();
const endTime = perf_hooks_1.performance.now();
const latency = endTime - startTime;
// Update the log
data.push(d);
successLatency.push(latency);
}
catch (error) {
const endTime = perf_hooks_1.performance.now();
const latency = endTime - startTime;
exceptions.push(error);
errLatency.push(latency);
}
};
// Execute rps api calls
const arr = [...new Array(rps).keys()];
const executeLoop = async () => arr.map(async () => {
// Wait a random amount of time less than 1 second (to distribute the load)
await helper_1.wait(Math.random() * 1000);
return execute();
});
// Run once per second
setInterval(() => {
// Silence errors: they should be handled above
executeLoop().catch(() => null);
const avgLatency = successLatency.reduce((a, b) => a + b, 0) / successLatency.length;
const avgFailLatency = errLatency.reduce((a, b) => a + b, 0) / errLatency.length;
const pending = total - successLatency.length - errLatency.length;
this.log(`
****
Sent: ${total}
Success: ${successLatency.length}
Failure: ${errLatency.length}
Pending: ${pending}
Average Success Latency: ${Math.round(avgLatency)} ms
Average Failure Latency: ${Math.round(avgFailLatency)} ms
`);
}, 1000);
}
exports.loadTest = loadTest;