Add recordSize option to Webhooks (#107)

MaterializeInc · Sep 20, 2023 · 0dbb72e · 0dbb72e
1 parent 66466b4
commit 0dbb72e
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/src/kafkaDataGenerator.ts b/src/kafkaDataGenerator.ts
@@ -1,5 +1,5 @@
 import alert from 'cli-alerts';
-import cryptoRandomString from 'crypto-random-string';
+import recordSize from './utils/recordSize.js';
 import { KafkaProducer } from './kafka/producer.js';
 import { generateMegaRecord } from './schemas/generateMegaRecord.js';
 import { OutputFormat } from './formats/outputFormat.js';
@@ -22,7 +22,7 @@ export default async function kafkaDataGenerator({
 
     let payload: string;
     if (global.recordSize) {
-        payload = cryptoRandomString({ length: global.recordSize, type: 'base64' });
+        payload = await recordSize();
     }
 
     let producer: KafkaProducer | null = null;

diff --git a/src/utils/recordSize.ts b/src/utils/recordSize.ts
@@ -0,0 +1,15 @@
+import cryptoRandomString from 'crypto-random-string';
+
+export default async function recordSize(): Promise<string> {
+    let payload: string;
+    if (global.recordSize) {
+        if (global.recordSize <= 0) {
+            global.recordSize = 1;
+        }
+        payload = cryptoRandomString({
+            length: global.recordSize,
+            type: 'base64'
+        });
+    }
+    return payload;
+}
diff --git a/src/webhookDataGenerator.ts b/src/webhookDataGenerator.ts
@@ -1,9 +1,8 @@
 import alert from 'cli-alerts';
-import crypto from 'crypto';
-import * as pg from 'pg';
 import { generateMegaRecord } from './schemas/generateMegaRecord.js';
 import { OutputFormat } from './formats/outputFormat.js';
 import sleep from './utils/sleep.js';
+import recordSize from './utils/recordSize.js';
 import asyncGenerator from './utils/asyncGenerator.js';
 import webhookConfig from './webhook/webhookConfig.js';
 
@@ -38,6 +37,9 @@ export default async function webhookDataGenerator({
     }
 
     let payload: string;
+    if (global.recordSize) {
+        payload = await recordSize();
+    }
 
     for await (const iteration of asyncGenerator(iterations)) {
         global.iterationIndex = iteration;
@@ -56,6 +58,10 @@ export default async function webhookDataGenerator({
         const handler = async (megaRecord: any, iteration: number) => {
             for (const endpoint in megaRecord) {
                 for await (const record of megaRecord[endpoint].records) {
+                    if (global.recordSize) {
+                        record.recordSizePayload = payload;
+                    }
+
                     if (global.dryRun) {
                         alert({
                             type: `success`,