Line data Source code
1 : import type * as arrow from 'apache-arrow';
2 :
3 : import type {
4 : EmbeddedClientContext,
5 : EmbeddedComputeDatabase,
6 : EmbeddedConnection,
7 : EmbeddedTableImportConnection,
8 : } from '../database/embedded_database.js';
9 : import { quoteIdentifier } from './hyperdb_wasm.js';
10 :
11 : /// A column of a query result or table.
12 : export interface HyperDBExtractColumn {
13 : /// The column name
14 : name: string;
15 : /// The SQL type of the column
16 : type: string;
17 : }
18 :
19 : /// A table of the extract, filled by a query over the new source rows.
20 : export interface HyperDBExtractTable {
21 : /// The table name in the schema of the extract
22 : name: string;
23 : /// A query over the relation of the source (e.g. `log`) that holds the new source rows.
24 : /// The columns of its result are the columns of the table, cast in the query to choose their types.
25 : query: string;
26 : /// The maximum number of rows that the table keeps, the rows with the lowest row numbers are deleted first
27 : maxRows?: number;
28 : /// The time in milliseconds that the table reaches back from its newest row, by the watermark column.
29 : /// Older rows are deleted, the table must have the watermark column.
30 : maxAgeMs?: number;
31 : }
32 :
33 : /// The column with the row numbers of every extract table, in insertion order and by the watermark within a transfer.
34 1 : export const HYPER_EXTRACT_ROW_COLUMN = '__row';
35 :
36 : /// A connection to the Hyper that an extract reads, e.g. an embedded connection or a channel to a remote Hyper
37 : export type HyperDBSourceConnection = Pick<EmbeddedConnection, 'query' | 'execute' | 'close'>;
38 :
39 : /// The Hyper that an extract reads
40 : export interface HyperDBSourceDatabase {
41 : /// Connect to the Hyper, the client context describes the queries of the connection
42 : connect(options?: { clientContext?: EmbeddedClientContext }): Promise<HyperDBSourceConnection>;
43 : }
44 :
45 : /// The rows that an incremental extract reads from a source Hyper.
46 : ///
47 : /// The rows are identified by all of their columns, and the newest extracted value of the watermark column bounds the
48 : /// backfills. Without an update query, every poll is a backfill from the watermark.
49 : export interface HyperDBIncrementalExtractSource {
50 : /// The name of the relation with the new source rows in the table queries
51 : relation: string;
52 : /// The columns of the source rows, the source queries return exactly these
53 : columns: readonly HyperDBExtractColumn[];
54 : /// The timestamp column whose newest extracted value is the watermark of backfills
55 : watermarkColumn: string;
56 : /// Check the source on the first source connection, before the first query
57 : open?(connection: HyperDBSourceConnection): Promise<void>;
58 : /// Did a read of the source rows fail in a way that the next read does not? It is then read again at once.
59 : isTransientReadError?(error: unknown): boolean;
60 : /// The query of the source rows that satisfy the watermark predicate, or of all rows when it is null
61 : backfillQuery(since: string | null): string;
62 : /// The query of the source rows since the connection that runs it was connected, cheaper than a backfill.
63 : /// E.g. the scope of the Hyper log of a session, which starts with the session.
64 : updateQuery?(): string;
65 : }
66 :
67 : /// Options of an incremental extract.
68 : export interface HyperDBIncrementalExtractOptions {
69 : /// The Hyper that is read
70 : sourceDatabase: HyperDBSourceDatabase;
71 : /// The rows that are read
72 : source: HyperDBIncrementalExtractSource;
73 : /// The Hyper that holds the extract
74 : sink: EmbeddedComputeDatabase;
75 : /// The path of an extract file in the sink, e.g. /mnt/fs/hyper_log.hyper.
76 : /// Without it, the extract lives in the compute database of the sink, where its compute connections read it.
77 : databasePath?: string;
78 : /// The alias of the extract file on the sink connection, readers that attach the file with this alias resolve
79 : /// the names of tableName(). hyper_extract when not set.
80 : databaseAlias?: string;
81 : /// The schema of the extract tables, public when not set
82 : schema?: string;
83 : /// The tables of the extract
84 : tables: readonly HyperDBExtractTable[];
85 : /// The client context of the source and sink connections, e.g. to attribute their queries in the Hyper log
86 : clientContext?: EmbeddedClientContext;
87 : /// The polling interval of start() in milliseconds
88 : intervalMs?: number;
89 : /// How far before the watermark a backfill reads again, in milliseconds
90 : backfillOverlapMs?: number;
91 : /// Receives the transfers of polls started by the timer
92 : onTransfer?: (transfer: HyperDBExtractTransfer) => void;
93 : /// Receives the errors of polls started by the timer
94 : onError?: (error: unknown) => void;
95 : /// Receives why open() rebuilt an extract whose tables no longer match the source or the table queries
96 : onRebuild?: (reason: string) => void;
97 : }
98 :
99 : /// The row counts of one transfer from the source into the extract.
100 : export interface HyperDBExtractTransfer {
101 : /// The rows read from the source
102 : fetched: number;
103 : /// The rows that were new to the extract
104 : appended: number;
105 : /// The rows appended to every extract table
106 : tables: Record<string, number>;
107 : /// The rows deleted from every extract table with a row limit
108 : deleted: Record<string, number>;
109 : }
110 :
111 : /// The default alias of an extract file on the sink connection.
112 1 : const EXTRACT_ALIAS = 'hyper_extract';
113 : /// The default database of sink connections, it holds the staging tables of transfers.
114 1 : const SINK_DEFAULT_DATABASE = '__dashql_compute';
115 :
116 : /// The qualified names of the tables of an extract
117 1 : class ExtractNames {
118 : /// The table with the recently extracted rows, the newest watermark value is the watermark of backfills.
119 : /// It keeps the rows of the backfill overlap, which suffices to skip rows that are read twice.
120 1 : readonly recentRows: string;
121 : /// The table with the new rows of the running transfer, the relation of the table queries.
122 1 : readonly newRows: string;
123 :
124 1 : constructor(
125 : /// The qualified schema of the extract tables
126 1 : readonly schema: string,
127 1 : ) {
128 1 : this.recentRows = this.table('__recent_rows');
129 1 : this.newRows = this.table('__new_rows');
130 1 : }
131 :
132 : /// The qualified name of a table of the extract
133 1 : table(name: string): string {
134 1 : return `${this.schema}.${quoteIdentifier(name)}`;
135 1 : }
136 1 : }
137 :
138 1 : const DEFAULT_INTERVAL_MS = 10_000;
139 1 : const DEFAULT_BACKFILL_OVERLAP_MS = 10_000;
140 : /// How often a read of the source rows is tried when it keeps failing with a transient error
141 1 : const SOURCE_READ_ATTEMPTS = 3;
142 :
143 : /// Incrementally copies rows of one Hyper (the source) into the tables of an extract file in another Hyper (the sink).
144 : ///
145 : /// Every poll reads the new source rows, with the update query of the source or else with a backfill from the
146 : /// watermark. The rows travel as Arrow to the sink, where every extract table receives the rows of its query over the
147 : /// new rows. open() always backfills from the watermark, so that the rows since the last extraction survive a
148 : /// disconnect. Backfills read again from a short overlap before the watermark: the sink skips the rows that it
149 : /// extracted before, comparing all columns. A transfer fills all extract tables in one transaction.
150 : ///
151 : /// An update query reads the rows since its connection was connected. Every poll connects the next reader before the
152 : /// current reader reads, and then replaces it: the rows of consecutive polls overlap and none fall between them, also
153 : /// not the rows of the read itself. The sink skips the rows of the overlap.
154 : ///
155 : /// open() checks the queries with `EXPLAIN (FORMAT SCHEMA)` before it reads any rows. The schema of every extract
156 : /// table is the result schema of its query. The extract is derived from the source, so when an existing table no
157 : /// longer matches the source columns or its query, open() drops all tables and backfills them from the source again.
158 : /// The tables share the watermark, a single rebuilt table would miss the rows before it. Every extract table also has
159 : /// the row number column `__row`, which its row limit deletes by.
160 : ///
161 : /// The extract is either a file that the sink attaches, or a schema of the compute database of the sink.
162 : ///
163 : /// Limits of this proof of concept:
164 : /// - Source rows older than the overlap before the watermark are never read again.
165 : /// - An extract holds the rows of one source. Changing the source or a table loses the rows that the source no longer
166 : /// holds.
167 1 : export class HyperDBIncrementalExtract {
168 : /// The pending timer of start(), null when not polling periodically
169 1 : private timer: ReturnType<typeof setTimeout> | null = null;
170 : /// The last poll, polls run one after another
171 1 : private polling: Promise<unknown> = Promise.resolve();
172 : /// Set by stop(), polls are rejected afterwards
173 1 : private stopped = false;
174 :
175 1 : private constructor(
176 : /// The connection to the source Hyper, the reader of the next poll
177 1 : private reader: HyperDBSourceConnection,
178 : /// The sink connection with the attached extract database
179 1 : private readonly sink: EmbeddedConnection,
180 : /// The sink connection that stages transfers, Arrow inserts need a connection with a single database
181 1 : private readonly stager: EmbeddedTableImportConnection,
182 : /// The name of the staging table of transfers in the sink
183 1 : private readonly stagingTable: string,
184 : /// The options of the extract
185 1 : private readonly options: HyperDBIncrementalExtractOptions,
186 : /// The qualified names of the extract tables
187 1 : private readonly names: ExtractNames,
188 : /// The columns of every extract table, by table name
189 1 : readonly schemas: Readonly<Record<string, readonly HyperDBExtractColumn[]>>,
190 : /// The transfer of the backfill in open()
191 1 : readonly backfill: HyperDBExtractTransfer,
192 1 : ) {}
193 :
194 : /// Open the extract, creating its tables when missing, and backfill since the watermark.
195 1 : static async open(options: HyperDBIncrementalExtractOptions): Promise<HyperDBIncrementalExtract> {
196 1 : const { source, clientContext } = options;
197 1 : let sink: EmbeddedConnection | null = null;
198 1 : let stager: EmbeddedTableImportConnection | null = null;
199 1 : let reader: HyperDBSourceConnection | null = null;
200 1 : try {
201 : // Open the source first, a source that cannot be read leaves no extract file behind
202 1 : reader = await options.sourceDatabase.connect({ clientContext });
203 1 : await source.open?.(reader);
204 1 : sink = await options.sink.connect({ defaultDatabase: SINK_DEFAULT_DATABASE, clientContext });
205 1 : stager = await options.sink.connect({ defaultDatabase: SINK_DEFAULT_DATABASE, clientContext });
206 1 : const database = options.databasePath == null ? SINK_DEFAULT_DATABASE : options.databaseAlias ?? EXTRACT_ALIAS;
207 1 : if (options.databasePath != null) {
208 1 : const path = quoteIdentifier(options.databasePath);
209 1 : await sink.execute(`CREATE DATABASE IF NOT EXISTS ${path}`);
210 1 : await sink.execute(`ATTACH DATABASE ${path} AS ${quoteIdentifier(database)}`);
211 1 : }
212 1 : const names = new ExtractNames(`${quoteIdentifier(database)}.${quoteIdentifier(options.schema ?? 'public')}`);
213 1 : if (options.schema != null) await sink.execute(`CREATE SCHEMA IF NOT EXISTS ${names.schema}`);
214 1 : let prepared = await prepareTables(sink, source, names, options.tables);
215 1 : if (prepared.outdated != null) {
216 1 : options.onRebuild?.(prepared.outdated);
217 1 : await dropTables(sink, names, options.tables);
218 1 : prepared = await prepareTables(sink, source, names, options.tables);
219 1 : if (prepared.outdated != null) throw new Error(`the rebuilt extract is outdated: ${prepared.outdated}`);
220 1 : }
221 1 : const { schemas } = prepared;
222 :
223 1 : await checkSourceQuery(reader, source, source.backfillQuery(null));
224 1 : const updateQuery = source.updateQuery?.();
225 1 : if (updateQuery != null) await checkSourceQuery(reader, source, updateQuery);
226 :
227 : // The first update reads the rows since the reader connected again, the sink skips them
228 1 : const stagingTable = `__hyper_extract_staging_${crypto.randomUUID().replace(/-/g, '_')}`;
229 1 : const rows = await readSource(reader, source, await backfillQuery(sink, names, options));
230 1 : const backfill = await transfer(sink, stager, stagingTable, options, names, schemas, rows);
231 1 : return new HyperDBIncrementalExtract(reader, sink, stager, stagingTable, options, names, schemas, backfill);
232 1 : } catch (error) {
233 : // Close every connection and keep the error that failed the open.
234 1 : for (const connection of [reader, stager, sink]) {
235 1 : await connection?.close().catch(() => {});
236 : }
237 1 : throw error;
238 1 : }
239 1 : }
240 :
241 : /// The qualified name of an extract table, e.g. to read it on a connection of the sink
242 1 : tableName(name: string): string {
243 1 : return this.names.table(name);
244 1 : }
245 :
246 : /// Transfer the source rows since the last poll into the extract.
247 1 : poll(): Promise<HyperDBExtractTransfer> {
248 1 : const poll = this.polling.catch(() => {}).then(async () => {
249 1 : if (this.stopped) throw new Error('extract is stopped');
250 1 : const { source, sourceDatabase, clientContext } = this.options;
251 1 : const updateQuery = source.updateQuery?.();
252 1 : if (updateQuery == null) {
253 1 : const rows = await readSource(this.reader, source, await backfillQuery(this.sink, this.names, this.options));
254 1 : return await transfer(this.sink, this.stager, this.stagingTable, this.options, this.names, this.schemas, rows);
255 1 : }
256 : // The next reader starts before the read, it replaces the current reader once the rows are transferred.
257 : // A failed poll keeps the current reader, the next poll reads its rows again.
258 1 : const next = await sourceDatabase.connect({ clientContext });
259 1 : try {
260 1 : const rows = await readSource(this.reader, source, projectSourceQuery(source, updateQuery));
261 1 : const transferred = await transfer(this.sink, this.stager, this.stagingTable, this.options, this.names, this.schemas, rows);
262 1 : const previous = this.reader;
263 1 : this.reader = next;
264 1 : await previous.close().catch(() => {});
265 1 : return transferred;
266 1 : } catch (error) {
267 1 : await next.close().catch(() => {});
268 1 : throw error;
269 1 : }
270 1 : });
271 1 : this.polling = poll;
272 1 : return poll;
273 1 : }
274 :
275 : /// Poll every intervalMs until stop().
276 1 : start(): void {
277 1 : if (this.stopped || this.timer != null) return;
278 1 : const schedule = () => {
279 1 : this.timer = setTimeout(() => {
280 : // A transfer that committed is reported, also when the extract stopped meanwhile
281 1 : void this.poll().then(transfer => {
282 1 : this.options.onTransfer?.(transfer);
283 1 : }, error => {
284 1 : if (!this.stopped) this.options.onError?.(error);
285 1 : }).finally(() => {
286 1 : if (!this.stopped) schedule();
287 1 : });
288 1 : }, this.options.intervalMs ?? DEFAULT_INTERVAL_MS);
289 1 : };
290 1 : schedule();
291 1 : }
292 :
293 : /// Stop polling, wait for a running poll and close all connections. The extract file is kept.
294 1 : async stop(): Promise<void> {
295 1 : if (this.stopped) return;
296 1 : this.stopped = true;
297 1 : if (this.timer != null) {
298 1 : clearTimeout(this.timer);
299 1 : this.timer = null;
300 1 : }
301 1 : await this.polling.catch(() => {});
302 1 : try {
303 1 : await this.reader.close();
304 1 : await this.stager.close();
305 1 : if (this.options.databasePath != null) {
306 1 : await this.sink.execute(`DETACH DATABASE ${quoteIdentifier(this.options.databaseAlias ?? EXTRACT_ALIAS)}`);
307 : }
308 1 : } finally {
309 1 : await this.sink.close();
310 1 : }
311 1 : }
312 1 : }
313 :
314 : /// Read source rows, again at once after a transient error. An error that stays fails the read.
315 1 : async function readSource(reader: HyperDBSourceConnection, source: HyperDBIncrementalExtractSource, query: string): Promise<arrow.Table> {
316 1 : for (let attempt = 1; ; ++attempt) {
317 1 : try {
318 1 : return await reader.query(query);
319 1 : } catch (error) {
320 1 : if (attempt >= SOURCE_READ_ATTEMPTS || !source.isTransientReadError?.(error)) throw error;
321 1 : }
322 : }
323 1 : }
324 :
325 : /// The backfill query from the overlap before the watermark, or of all rows for an empty extract.
326 1 : async function backfillQuery(sink: EmbeddedConnection, names: ExtractNames, options: HyperDBIncrementalExtractOptions): Promise<string> {
327 1 : const { source } = options;
328 1 : const column = quoteIdentifier(source.watermarkColumn);
329 1 : const type = source.columns.find(({ name }) => name === source.watermarkColumn)?.type;
330 1 : if (type == null) throw new Error(`the watermark column ${source.watermarkColumn} is not a source column`);
331 1 : const result = await sink.query(`SELECT max(${column})::text AS watermark FROM ${names.recentRows}`);
332 1 : const watermark = result.getChild('watermark')?.get(0);
333 1 : const since = typeof watermark === 'string'
334 1 : ? `${column} >= ${quoteLiteral(watermark)}::${type} - ${overlapInterval(options)}`
335 1 : : null;
336 1 : return projectSourceQuery(source, source.backfillQuery(since));
337 1 : }
338 :
339 : /// A source query with its columns cast to the source columns.
340 1 : function projectSourceQuery(source: HyperDBIncrementalExtractSource, query: string): string {
341 1 : return `WITH selected AS (${query}) SELECT ${castColumns(source.columns)} FROM selected`;
342 1 : }
343 :
344 : /// A table query over the relation of the source with the new rows of a transfer.
345 1 : function tableQuery(source: HyperDBIncrementalExtractSource, names: ExtractNames, table: HyperDBExtractTable): string {
346 1 : return `
347 1 : WITH ${quoteIdentifier(source.relation)} AS (SELECT * FROM ${names.newRows}), selected AS (${table.query})
348 : SELECT * FROM selected
349 : `;
350 1 : }
351 :
352 : /// Check that a source query returns exactly the source columns.
353 1 : async function checkSourceQuery(
354 1 : connection: HyperDBSourceConnection,
355 1 : source: HyperDBIncrementalExtractSource,
356 1 : query: string,
357 1 : ): Promise<void> {
358 1 : const columns = await describeQuery(connection, query, 'source query');
359 1 : const names = new Set(columns.map(({ name }) => name));
360 1 : const missing = source.columns.filter(({ name }) => !names.has(name)).map(({ name }) => name);
361 1 : const unexpected = [...names].filter(name => !source.columns.some(column => column.name === name));
362 1 : if (missing.length > 0 || unexpected.length > 0) {
363 1 : throw new Error(`the source query must return the source columns: ${[
364 1 : missing.length > 0 ? `missing ${missing.join(', ')}` : null,
365 1 : unexpected.length > 0 ? `unexpected ${unexpected.join(', ')}` : null,
366 1 : ].filter(part => part != null).join(', ')}`);
367 : }
368 1 : }
369 :
370 : /// The extract tables prepared by open()
371 : interface PreparedTables {
372 : /// The columns of every extract table, by table name
373 : schemas: Record<string, readonly HyperDBExtractColumn[]>;
374 : /// Why an existing table no longer matches the source columns or its query, null when all match
375 : outdated: string | null;
376 : }
377 :
378 : /// Create the missing extract tables with the result schemas of their queries, and check the existing ones.
379 1 : async function prepareTables(
380 1 : sink: EmbeddedConnection,
381 1 : source: HyperDBIncrementalExtractSource,
382 1 : names: ExtractNames,
383 1 : tables: readonly HyperDBExtractTable[],
384 1 : ): Promise<PreparedTables> {
385 : // The table queries read the new rows, which have the source columns
386 1 : for (const name of [names.recentRows, names.newRows]) {
387 1 : await sink.execute(`CREATE TABLE IF NOT EXISTS ${name} (${columnDefinitions(source.columns)})`);
388 1 : const existing = await describeQuery(sink, `SELECT * FROM ${name}`, `extract table ${name}`);
389 1 : if (formatColumns(existing) !== formatColumns(source.columns)) {
390 1 : return { schemas: {}, outdated: `the extract table ${name} has the columns (${formatColumns(existing)}), `
391 1 : + `but the source has (${formatColumns(source.columns)})` };
392 : }
393 1 : }
394 1 : const schemas: Record<string, readonly HyperDBExtractColumn[]> = {};
395 1 : for (const table of tables) {
396 1 : if (table.name in schemas) throw new Error(`the extract table ${table.name} is configured twice`);
397 1 : if (table.name.startsWith('__')) throw new Error(`the extract table name ${table.name} is reserved`);
398 1 : if (table.maxRows != null && !(Number.isInteger(table.maxRows) && table.maxRows > 0)) {
399 1 : throw new Error(`the row limit of the extract table ${table.name} must be a positive integer`);
400 : }
401 1 : if (table.maxAgeMs != null && !(table.maxAgeMs > 0)) {
402 1 : throw new Error(`the age limit of the extract table ${table.name} must be positive`);
403 : }
404 1 : const columns = await describeQuery(sink, tableQuery(source, names, table), `query of the extract table ${table.name}`);
405 1 : const reserved = columns.find(({ name }) => name.startsWith('__'));
406 1 : if (reserved != null) throw new Error(`the column name ${reserved.name} of the extract table ${table.name} is reserved`);
407 1 : if (table.maxAgeMs != null && !columns.some(({ name }) => name === source.watermarkColumn)) {
408 1 : throw new Error(`the extract table ${table.name} needs the watermark column ${source.watermarkColumn} for its age limit`);
409 : }
410 1 : await sink.execute(`CREATE TABLE IF NOT EXISTS ${names.table(table.name)} (${
411 1 : columnDefinitions([{ name: HYPER_EXTRACT_ROW_COLUMN, type: 'BIGINT' }, ...columns])
412 1 : })`);
413 1 : const existing = await describeQuery(sink, `SELECT * FROM ${names.table(table.name)}`, `extract table ${table.name}`)
414 1 : .then(existing => existing.filter(({ name }) => name !== HYPER_EXTRACT_ROW_COLUMN));
415 1 : if (formatColumns(existing) !== formatColumns(columns)) {
416 1 : return { schemas, outdated: `the extract table ${table.name} has the columns (${formatColumns(existing)}), `
417 1 : + `but its query returns (${formatColumns(columns)})` };
418 : }
419 1 : schemas[table.name] = columns;
420 1 : }
421 1 : return { schemas, outdated: null };
422 1 : }
423 :
424 : /// Drop all tables of an extract, the next backfill reads all source rows again
425 1 : async function dropTables(sink: EmbeddedConnection, names: ExtractNames, tables: readonly HyperDBExtractTable[]): Promise<void> {
426 1 : for (const name of [names.recentRows, names.newRows, ...tables.map(table => names.table(table.name))]) {
427 1 : await sink.execute(`DROP TABLE IF EXISTS ${name}`);
428 : }
429 1 : }
430 :
431 : /// The result columns of a query, from `EXPLAIN (FORMAT SCHEMA)`.
432 1 : async function describeQuery(connection: HyperDBSourceConnection, query: string, what: string): Promise<HyperDBExtractColumn[]> {
433 1 : let result: arrow.Table;
434 1 : try {
435 1 : result = await connection.query(`EXPLAIN (FORMAT SCHEMA) ${query}`);
436 1 : } catch (error) {
437 1 : throw new Error(`invalid ${what}: ${error instanceof Error ? error.message : String(error)}`);
438 1 : }
439 1 : const names = result.getChild('name');
440 1 : const types = result.getChild('type');
441 1 : const columns: HyperDBExtractColumn[] = [];
442 1 : for (let i = 0; i < result.numRows; ++i) {
443 1 : const name = names?.get(i) as string;
444 1 : if (columns.some(column => column.name === name)) throw new Error(`the ${what} returns the column ${name} twice`);
445 1 : columns.push({ name, type: types?.get(i) as string });
446 1 : }
447 1 : if (columns.length === 0) throw new Error(`the ${what} returns no columns`);
448 1 : return columns;
449 1 : }
450 :
451 : /// Append the fetched rows that are new to the extract to all extract tables, in one transaction of the extract.
452 1 : async function transfer(
453 1 : sink: EmbeddedConnection,
454 1 : stager: EmbeddedTableImportConnection,
455 1 : stagingTable: string,
456 1 : options: HyperDBIncrementalExtractOptions,
457 1 : names: ExtractNames,
458 1 : schemas: Readonly<Record<string, readonly HyperDBExtractColumn[]>>,
459 1 : rows: arrow.Table,
460 1 : ): Promise<HyperDBExtractTransfer> {
461 1 : const { source } = options;
462 1 : const tables = Object.fromEntries(options.tables.map(table => [table.name, 0]));
463 1 : const deleted = Object.fromEntries(options.tables
464 1 : .filter(table => table.maxRows != null || table.maxAgeMs != null)
465 1 : .map(table => [table.name, 0]));
466 1 : if (rows.numRows === 0) return { fetched: 0, appended: 0, tables, deleted };
467 1 : const staging = `${quoteIdentifier(SINK_DEFAULT_DATABASE)}.public.${quoteIdentifier(stagingTable)}`;
468 1 : const columns = columnList(source.columns);
469 1 : const watermark = quoteIdentifier(source.watermarkColumn);
470 1 : await stager.insertArrowTable(rows, { name: stagingTable, create: true });
471 1 : try {
472 : // A transaction may only modify one database and may not mix DDL with DML.
473 1 : await sink.execute('BEGIN');
474 1 : try {
475 1 : await sink.execute(`DELETE FROM ${names.newRows}`);
476 1 : const appended = await sink.execute(`
477 1 : INSERT INTO ${names.newRows} (${columns})
478 1 : SELECT DISTINCT ${source.columns.map(({ name }) => `s.${quoteIdentifier(name)}`).join(', ')}
479 1 : FROM ${staging} s
480 : WHERE NOT EXISTS (
481 1 : SELECT 1 FROM ${names.recentRows} e
482 1 : WHERE e.${watermark} = s.${watermark} AND ${source.columns.map(({ name }) =>
483 1 : `e.${quoteIdentifier(name)} IS NOT DISTINCT FROM s.${quoteIdentifier(name)}`).join(' AND ')}
484 : )
485 1 : `);
486 1 : await sink.execute(`INSERT INTO ${names.recentRows} (${columns}) SELECT ${columns} FROM ${names.newRows}`);
487 1 : const row = quoteIdentifier(HYPER_EXTRACT_ROW_COLUMN);
488 1 : for (const table of options.tables) {
489 1 : const name = names.table(table.name);
490 1 : const order = schemas[table.name].some(({ name }) => name === source.watermarkColumn) ? `ORDER BY ${watermark}` : '';
491 1 : tables[table.name] = Number(await sink.execute(`
492 1 : INSERT INTO ${name} (${row}, ${columnList(schemas[table.name])})
493 1 : WITH ${quoteIdentifier(source.relation)} AS (SELECT * FROM ${names.newRows}), selected AS (${table.query})
494 1 : SELECT (SELECT coalesce(max(${row}), 0) FROM ${name}) + row_number() OVER (${order}), *
495 : FROM selected
496 1 : `) ?? 0);
497 1 : if (table.maxRows != null) {
498 1 : deleted[table.name] += Number(await sink.execute(`
499 1 : DELETE FROM ${name} WHERE ${row} <= (SELECT max(${row}) FROM ${name}) - ${table.maxRows}
500 1 : `) ?? 0);
501 : }
502 1 : if (table.maxAgeMs != null) {
503 1 : deleted[table.name] += Number(await sink.execute(`
504 1 : DELETE FROM ${name} WHERE ${watermark} < (SELECT max(${watermark}) FROM ${name}) - ${interval(table.maxAgeMs)}
505 1 : `) ?? 0);
506 : }
507 1 : }
508 1 : await sink.execute(`
509 1 : DELETE FROM ${names.recentRows}
510 1 : WHERE ${watermark} < (SELECT max(${watermark}) FROM ${names.recentRows}) - ${overlapInterval(options)}
511 1 : `);
512 1 : await sink.execute('COMMIT');
513 1 : return { fetched: rows.numRows, appended: Number(appended ?? 0), tables, deleted };
514 1 : } catch (error) {
515 0 : await sink.execute('ROLLBACK');
516 0 : throw error;
517 0 : }
518 1 : } finally {
519 1 : await stager.execute(`DROP TABLE ${quoteIdentifier(stagingTable)}`);
520 1 : }
521 1 : }
522 :
523 1 : function formatColumns(columns: readonly HyperDBExtractColumn[]): string {
524 1 : return columns.map(({ name, type }) => `${name} ${type.toLowerCase()}`).join(', ');
525 1 : }
526 :
527 1 : function interval(ms: number): string {
528 1 : return `interval '${ms / 1000} seconds'`;
529 1 : }
530 :
531 1 : function overlapInterval(options: HyperDBIncrementalExtractOptions): string {
532 1 : return interval(options.backfillOverlapMs ?? DEFAULT_BACKFILL_OVERLAP_MS);
533 1 : }
534 :
535 1 : function columnDefinitions(columns: readonly HyperDBExtractColumn[]): string {
536 1 : return columns.map(({ name, type }) => `${quoteIdentifier(name)} ${type}`).join(', ');
537 1 : }
538 :
539 1 : function columnList(columns: readonly HyperDBExtractColumn[]): string {
540 1 : return columns.map(({ name }) => quoteIdentifier(name)).join(', ');
541 1 : }
542 :
543 1 : function castColumns(columns: readonly HyperDBExtractColumn[]): string {
544 1 : return columns.map(({ name, type }) => `${quoteIdentifier(name)}::${type} AS ${quoteIdentifier(name)}`).join(', ');
545 1 : }
546 :
547 1 : export function quoteLiteral(value: string): string {
548 1 : return `'${value.replace(/'/g, "''")}'`;
549 1 : }
|