@tinybirdco/mockingbird/dist/client/generators/BaseGenerator.js

Mockingbird

import { z } from "zod";
import { extendedFaker } from "../extendedFaker";
import { validateSchema } from "../types";
// Helper function to safely get nested object properties
const getNestedValue = (obj, path) => {
    return path.split(".").reduce((current, key) => current?.[key], obj);
};
export const schemaSchema = z.record(z.object({
    type: z.string(),
    params: z.any().optional(),
    count: z.number().optional(),
}));
export const baseConfigSchema = z.object({
    schema: schemaSchema.refine((schemaSchema) => validateSchema(schemaSchema)),
    eps: z.number().optional().default(1),
    limit: z.number().optional().default(-1),
    logs: z.boolean().default(false).optional(),
});
export class BaseGenerator {
    config;
    state = {};
    constructor(config) {
        this.config = config;
    }
    log(level, message) {
        if (!this.config.logs)
            return;
        if (level === "info") {
            console.log(`INFO> ${message}`);
        }
        else if (level === "error") {
            console.error(`ERR > ${message}`);
        }
    }
    generateRow() {
        const generatedRow = Object.entries(this.config.schema).reduce((acc, [key, value]) => {
            const generator = getNestedValue(extendedFaker, value.type);
            const count = value.count ?? 1;
            // this.log(
            //   "info",
            //   `generateRow() - ${key} (${value.type}): ${JSON.stringify(value.params)}`
            // );
            const generatedValues = new Array(count)
                .fill(null)
                .map(() => generator(value.params, { state: this.state }));
            return {
                ...acc,
                [key]: count === 1 ? generatedValues[0] : generatedValues,
            };
        }, {});
        return generatedRow;
    }
    async generate(onMessage) {
        const minDelayPerBatch = 200;
        const maxBatchesPerSecond = 1000 / minDelayPerBatch;
        let batchSize, delayPerBatch;
        if (this.config.eps < 1000) {
            batchSize = this.config.eps;
            delayPerBatch = 1000;
        }
        else {
            batchSize = this.config.eps / maxBatchesPerSecond;
            delayPerBatch = minDelayPerBatch;
        }
        const rows = [];
        let limit = this.config.limit, sentRows = 0;
        while (true) {
            rows.push(this.generateRow());
            if (rows.length >= batchSize) {
                const data = rows.splice(0, batchSize);
                try {
                    await this.sendData(data);
                    if (onMessage)
                        onMessage(data);
                }
                catch (e) {
                    this.log("error", String(e));
                    break;
                }
                sentRows += data.length;
                this.log("info", `${sentRows} rows sent so far...`);
                if (limit != -1 && sentRows >= limit)
                    break;
                await new Promise((r) => setTimeout(r, delayPerBatch));
            }
        }
    }
}