@electric-sql/d2ts
Version:
D2TS is a TypeScript implementation of Differential Dataflow.
241 lines • 10.1 kB
JavaScript
import { MessageType, } from '../types.js';
import { DifferenceStreamWriter, UnaryOperator, } from '../graph.js';
import { StreamBuilder } from '../d2.js';
import { isChangeMessage, isControlMessage, } from '@electric-sql/client';
/**
* Connects an Electric ShapeStream to a D2 input stream
* IMPORTANT: Requires the ShapeStream to be configured with `replica: 'full'`
* @param options Configuration options
* @param options.stream The Electric ShapeStream to consume
* @param options.input The D2 input stream to send messages to
* @param options.lsnToVersion Optional function to convert LSN to version number/Version
* @param options.lsnToFrontier Optional function to convert LSN to frontier number/Antichain
* @returns The input stream for chaining
*
* @example
* ```ts
* // Create D2 graph
* const graph = new D2({ initialFrontier: 0 })
*
* // Create D2 input
* const input = graph.newInput<any>()
*
* // Configure the pipeline
* input
* .pipe(
* map(([key, data]) => data.value),
* filter(value => value > 10)
* )
*
* // Finalize graph
* graph.finalize()
*
* // Create Electric stream (example)
* const electricStream = new ShapeStream({
* url: 'http://localhost:3000/v1/shape',
* params: {
* table: 'items',
* replica: 'full',
* }
* })
*
* // Connect Electric stream to D2 input
* electricStreamToD2Input({
* stream: electricStream,
* input,
* })
* ```
*/
export function electricStreamToD2Input({ graph, stream, input, lsnToVersion = (lsn) => lsn, lsnToFrontier = (lsn) => lsn, initialLsn = 0, runOn = 'up-to-date', debug = false, }) {
let lastLsn = initialLsn;
let changes = [];
const sendChanges = (lsn) => {
const version = lsnToVersion(lsn);
log?.(`sending ${changes.length} changes at version ${version}`);
if (changes.length > 0) {
input.sendData(version, [...changes]);
}
changes = [];
};
const sendFrontier = (lsn) => {
const frontier = lsnToFrontier(lsn + 1); // +1 to account for the fact that the last LSN is the version of the last message
log?.(`sending frontier ${frontier}`);
input.sendFrontier(frontier);
};
const log = typeof debug === 'function'
? debug
: debug === true
? // eslint-disable-next-line no-console
console.log
: undefined;
log?.('subscribing to stream');
stream.subscribe((messages) => {
log?.(`received ${messages.length} messages`);
for (const message of messages) {
if (isControlMessage(message)) {
log?.(`- control message: ${message.headers.control}`);
// Handle control message
if (message.headers.control === 'up-to-date') {
log?.(`up-to-date ${JSON.stringify(message, null, 2)}`);
if (changes.length > 0) {
sendChanges(lastLsn);
}
if (typeof message.headers.global_last_seen_lsn !== `number`) {
throw new Error(`global_last_seen_lsn is not a number`);
}
const lsn = message.headers.global_last_seen_lsn;
sendFrontier(lsn);
if (runOn === 'up-to-date' || runOn === 'lsn-advance') {
log?.('running graph on up-to-date');
graph.run();
}
}
else if (message.headers.control === 'must-refetch') {
throw new Error('The server sent a "must-refetch" request, this is incompatible with a D2 pipeline and unresolvable. To handle this you will have to remove all state and start the pipeline again.');
}
}
else if (isChangeMessage(message)) {
log?.(`- change message: ${message.headers.operation}`);
// Handle change message
if (message.headers.lsn !== undefined &&
typeof message.headers.lsn !== `number`) {
throw new Error(`lsn is not a number`);
}
const lsn = message.headers.lsn ?? initialLsn; // The LSN is not present on the initial snapshot
const last = message.headers.last ?? false;
switch (message.headers.operation) {
case 'insert':
changes.push([[message.key, message.value], 1]);
break;
case 'update': {
// An update in D2 is a delete followed by an insert.
// `old_value` only holds the old values *that have changed*
// so we need to merge the old and new value to get a complete row
// that represents the row as it was before the update.
const oldValue = {
...message.value,
...(message.old_value ?? {}),
};
changes.push([[message.key, oldValue], -1]);
changes.push([[message.key, message.value], 1]);
break;
}
case 'delete':
changes.push([[message.key, message.value], -1]);
break;
}
if (last) {
sendChanges(lsn);
sendFrontier(lsn);
if (runOn === 'lsn-advance') {
log?.('running graph on lsn-advance');
graph.run();
}
}
if (lsn > lastLsn) {
lastLsn = lsn;
}
}
}
});
return input;
}
/**
* Operator that outputs the messages from the stream in ElectricSQL format
*
* TODO: Have two modes `replica=default` and `replica=full` to match the two modes
* of core Electric.
*/
export class OutputElectricMessagesOperator extends UnaryOperator {
#fn;
constructor(id, inputA, output, fn, initialFrontier) {
super(id, inputA, output, initialFrontier);
this.#fn = fn;
}
transformMessages(messages) {
const output = [];
let lastLSN = -Infinity;
for (const message of messages) {
if (message.type === MessageType.DATA) {
const { version, collection } = message.data;
// We use the first element of the D2version as the LSN as its what we used
// in the input
const lsn = version.getInner()[0];
if (lsn < lastLSN) {
throw new Error(`Invalid LSN ${lsn} less than lastLSN ${lastLSN}, you must consolidate your stream before passing it to the OutputElectricMessagesOperator`);
}
lastLSN = lsn;
const changesByKey = new Map();
for (const [[key, value], multiplicity] of collection.getInner()) {
let changes = changesByKey.get(key);
if (!changes) {
changes = { deletes: 0, inserts: 0, value: value };
changesByKey.set(key, changes);
}
if (multiplicity < 0) {
changes.deletes += Math.abs(multiplicity);
}
else if (multiplicity > 0) {
changes.inserts += multiplicity;
changes.value = value;
}
}
const newMessages = Array.from(changesByKey.entries()).map(([key, { deletes, inserts, value }]) => {
const operation = deletes > inserts
? 'delete'
: inserts > deletes
? 'insert'
: 'update';
return {
key: key.toString(),
value: value,
headers: { lsn, operation },
};
});
output.push(...newMessages);
}
}
return output;
}
run() {
const messages = this.inputMessages();
const outputMessages = this.transformMessages(messages);
if (outputMessages.length > 0) {
this.#fn(outputMessages);
}
for (const message of messages) {
if (message.type === MessageType.DATA) {
const { version, collection } = message.data;
this.output.sendData(version, collection);
}
else if (message.type === MessageType.FRONTIER) {
const frontier = message.data;
if (!this.inputFrontier().lessEqual(frontier)) {
throw new Error('Invalid frontier update');
}
this.setInputFrontier(frontier);
if (!this.outputFrontier.lessEqual(this.inputFrontier())) {
throw new Error('Invalid frontier state');
}
if (this.outputFrontier.lessThan(this.inputFrontier())) {
this.outputFrontier = this.inputFrontier();
this.output.sendFrontier(this.outputFrontier);
}
}
}
}
}
/**
* Outputs the messages from the stream in ElectricSQL format
* @param fn - The function to call with a batch of messages
*/
export function outputElectricMessages(fn) {
return (stream) => {
const output = new StreamBuilder(stream.graph, new DifferenceStreamWriter());
const operator = new OutputElectricMessagesOperator(stream.graph.getNextOperatorId(), stream.connectReader(), output.writer, fn, stream.graph.frontier());
stream.graph.addOperator(operator);
stream.graph.addStream(output.connectReader());
return output;
};
}
//# sourceMappingURL=index.js.map