UNPKG

@electric-sql/d2ts

Version:

D2TS is a TypeScript implementation of Differential Dataflow.

241 lines 10.1 kB
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