Fix streaming splits in realtime streams v2

This commit is contained in:
Eric Allam
2024-12-11 11:00:46 +00:00
parent 30ea5eb13a
commit 9336397fba
3 changed files with 79 additions and 28 deletions
@@ -37,13 +37,14 @@ export class DatabaseRealtimeStreams implements StreamIngestor, StreamResponder
): Promise<Response> {
try {
const textStream = stream.pipeThrough(new TextDecoderStream());
const reader = textStream.getReader();
let sequence = 0;
while (true) {
const { done, value } = await reader.read();
if (done) {
if (done || !value) {
break;
}
@@ -53,25 +54,13 @@ export class DatabaseRealtimeStreams implements StreamIngestor, StreamResponder
value,
});
const chunks = value
.split("\n")
.filter((chunk) => chunk) // Remove empty lines
.map((line) => {
return {
sequence: sequence++,
value: line,
};
});
await this.options.prisma.realtimeStreamChunk.createMany({
data: chunks.map((chunk) => {
return {
runId,
key: streamId,
sequence: chunk.sequence,
value: chunk.value,
};
}),
await this.options.prisma.realtimeStreamChunk.create({
data: {
runId,
key: streamId,
sequence: sequence++,
value,
},
});
}
+28 -8
View File
@@ -16,7 +16,12 @@ import {
} from "../utils/ioSerialization.js";
import { ApiError } from "./errors.js";
import { ApiClient } from "./index.js";
import { AsyncIterableStream, createAsyncIterableReadable, zodShapeStream } from "./stream.js";
import {
AsyncIterableStream,
createAsyncIterableReadable,
LineTransformStream,
zodShapeStream,
} from "./stream.js";
export type RunShape<TRunTypes extends AnyRunTypes> = TRunTypes extends AnyRunTypes
? {
@@ -209,13 +214,28 @@ export class ElectricStreamSubscription implements StreamSubscription {
) {}
async subscribe(): Promise<ReadableStream<unknown>> {
return zodShapeStream(SubscribeRealtimeStreamChunkRawShape, this.url, this.options).pipeThrough(
new TransformStream({
transform(chunk, controller) {
controller.enqueue(safeParseJSON(chunk.value));
},
})
);
return zodShapeStream(SubscribeRealtimeStreamChunkRawShape, this.url, this.options)
.pipeThrough(
new TransformStream({
transform(chunk, controller) {
console.log("ElectricStreamSubscription chunk.value", chunk.value);
controller.enqueue(chunk.value);
},
})
)
.pipeThrough(new LineTransformStream(this.url))
.pipeThrough(
new TransformStream({
transform(chunk, controller) {
for (const line of chunk) {
console.log("ElectricStreamSubscription line", line);
controller.enqueue(safeParseJSON(line));
}
},
})
);
}
}
+42
View File
@@ -203,3 +203,45 @@ class ReadableShapeStream<T extends Row<unknown> = Row> {
}
}
}
export class LineTransformStream extends TransformStream<string, string[]> {
private buffer = "";
constructor(streamId: string) {
super({
transform: (chunk, controller) => {
// Append the chunk to the buffer
this.buffer += chunk;
// Split on newlines
const lines = this.buffer.split("\n");
// The last element might be incomplete, hold it back in buffer
this.buffer = lines.pop() || "";
// Filter out empty or whitespace-only lines
const fullLines = lines.filter((line) => line.trim().length > 0);
console.log("LineTransformStream", {
chunk,
lines,
fullLines,
buffer: this.buffer,
streamId,
});
// If we got any complete lines, emit them as an array
if (fullLines.length > 0) {
controller.enqueue(fullLines);
}
},
flush: (controller) => {
// On stream end, if there's leftover text, emit it as a single-element array
const trimmed = this.buffer.trim();
if (trimmed.length > 0) {
controller.enqueue([trimmed]);
}
},
});
}
}