Skip to content

Commit 0fdd301

Browse files
themaroltthemarolt
committed
fix: comments
Signed-off-by: Uroš Marolt <uros@marolt.me>
1 parent c7a6dd0 commit 0fdd301

3 files changed

Lines changed: 164 additions & 134 deletions

File tree

‎services/apps/packages_worker/src/deps-dev/activities/bqExportToGcs.ts‎

Lines changed: 54 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -217,60 +217,71 @@ export async function bqExportToGcs(input: BqExportToGcsInput): Promise<BqExport
217217
...(ecosystems ? { tableRowCounts: { 'meta:ecosystems': ecosystems } } : {}),
218218
})
219219

220-
// Both modes finish by exporting from a `_export_data` TEMP table — materializing first forces
221-
// BQ to shard the parquet output properly (direct EXPORT DATA from a subquery emits O(rows)
222-
// ~1 KB micro-files).
223-
// - Single-SELECT: wrap `sql` in SELECT * FROM (...) so QUALIFY / top-level set ops don't break
224-
// the CREATE TEMP TABLE, then export.
225-
// - Script (isScript): the script already builds `_export_data` as its final statement, so we
226-
// append only the EXPORT DATA — its temp tables run inline and auto-drop when the session ends.
227-
const exportTail = `
220+
// From here the row is 'exporting'; any BQ failure (incl. script-mode maximumBytesBilled aborts)
221+
// must flip it to 'failed' with the reason, else it stays stuck 'exporting' forever and the
222+
// failure reason is lost to monitoring. Mirrors the rankPackages pattern (criticality/activities).
223+
try {
224+
// Both modes finish by exporting from a `_export_data` TEMP table — materializing first forces
225+
// BQ to shard the parquet output properly (direct EXPORT DATA from a subquery emits O(rows)
226+
// ~1 KB micro-files).
227+
// - Single-SELECT: wrap `sql` in SELECT * FROM (...) so QUALIFY / top-level set ops don't break
228+
// the CREATE TEMP TABLE, then export.
229+
// - Script (isScript): the script already builds `_export_data` as its final statement, so we
230+
// append only the EXPORT DATA — its temp tables run inline and auto-drop when the session ends.
231+
const exportTail = `
228232
EXPORT DATA OPTIONS(
229233
uri='${gcsPrefix}*.parquet',
230234
format='PARQUET',
231235
compression='SNAPPY',
232236
overwrite=true
233237
) AS SELECT * FROM _export_data;
234238
`
235-
const exportSql = isScript
236-
? `${sql.replace(/;\s*$/, '')};\n${exportTail}`
237-
: `\nCREATE TEMP TABLE _export_data AS SELECT * FROM (${sql.replace(/;\s*$/, '')});\n${exportTail}`
239+
const exportSql = isScript
240+
? `${sql.replace(/;\s*$/, '')};\n${exportTail}`
241+
: `\nCREATE TEMP TABLE _export_data AS SELECT * FROM (${sql.replace(/;\s*$/, '')});\n${exportTail}`
238242

239-
log.info({ jobKind, jobId, gcsPrefix, isScript: Boolean(isScript) }, 'Starting BQ export')
243+
log.info({ jobKind, jobId, gcsPrefix, isScript: Boolean(isScript) }, 'Starting BQ export')
240244

241-
const [job] = await bigquery.createQueryJob({
242-
query: exportSql,
243-
location: 'US',
244-
// Server-side runaway guard for script mode — aborts the job if a statement scans beyond the
245-
// ceiling. Single-SELECT mode is already gated by the dry-run check above.
246-
...(isScript ? { maximumBytesBilled: String(Math.floor(ceiling)) } : {}),
247-
})
248-
await job.promise()
249-
const bqStats = await extractBqStats(job, bigquery)
250-
251-
const rowCount = bqStats.outputRows ?? 0
245+
const [job] = await bigquery.createQueryJob({
246+
query: exportSql,
247+
location: 'US',
248+
// Server-side runaway guard for script mode — aborts the job if a statement scans beyond the
249+
// ceiling. Single-SELECT mode is already gated by the dry-run check above.
250+
...(isScript ? { maximumBytesBilled: String(Math.floor(ceiling)) } : {}),
251+
})
252+
await job.promise()
253+
const bqStats = await extractBqStats(job, bigquery)
252254

253-
await markJobStatus(qx, jobId, 'exported', {
254-
gcsPrefix,
255-
rowCountBq: rowCount,
256-
bqBytesBilled: bqStats.bqBytesBilled,
257-
bqJobId: bqStats.bqJobId,
258-
bqStats,
259-
tableRowCounts: { 'bq:export': rowCount },
260-
})
255+
const rowCount = bqStats.outputRows ?? 0
261256

262-
log.info(
263-
{
264-
jobKind,
265-
jobId,
266-
rowCount,
257+
await markJobStatus(qx, jobId, 'exported', {
258+
gcsPrefix,
259+
rowCountBq: rowCount,
260+
bqBytesBilled: bqStats.bqBytesBilled,
267261
bqJobId: bqStats.bqJobId,
268-
totalBytesProcessed: bqStats.totalBytesProcessed,
269-
totalSlotMs: bqStats.totalSlotMs,
270-
durationMs: bqStats.durationMs,
271-
},
272-
'BQ export complete',
273-
)
262+
bqStats,
263+
tableRowCounts: { 'bq:export': rowCount },
264+
})
274265

275-
return { gcsPrefix, rowCount, bqBytesBilled: bqStats.bqBytesBilled, jobId }
266+
log.info(
267+
{
268+
jobKind,
269+
jobId,
270+
rowCount,
271+
bqJobId: bqStats.bqJobId,
272+
totalBytesProcessed: bqStats.totalBytesProcessed,
273+
totalSlotMs: bqStats.totalSlotMs,
274+
durationMs: bqStats.durationMs,
275+
},
276+
'BQ export complete',
277+
)
278+
279+
return { gcsPrefix, rowCount, bqBytesBilled: bqStats.bqBytesBilled, jobId }
280+
} catch (err) {
281+
await markJobStatus(qx, jobId, 'failed', {
282+
errorMessage: (err as Error).message,
283+
finishedAt: new Date(),
284+
})
285+
throw err
286+
}
276287
}

‎services/apps/packages_worker/src/deps-dev/activities/gcsParquetToStaging.ts‎

Lines changed: 63 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -114,66 +114,76 @@ export async function gcsParquetToStaging(input: GcsToStagingInput): Promise<Gcs
114114
const qx = await getPackagesDb()
115115

116116
await markJobStatus(qx, jobId, 'loading')
117-
if (stagingDdl) {
118-
const stmts = Array.isArray(stagingDdl) ? stagingDdl : [stagingDdl]
119-
for (const stmt of stmts) await qx.result(stmt)
120-
}
121-
await qx.result(`TRUNCATE ${stagingTable}`)
122-
// Reset progress at the start of a fresh run (first chunk). This handles job ID reuse
123-
// via --export-name: a previous run's 193/193 would otherwise persist due to GREATEST.
124-
if (filesOffset === 0 && input.totalFiles != null) {
125-
await updateLoadingProgress(qx, jobId, 0, input.totalFiles, true)
126-
}
127-
128-
let parquetFileNames: string[]
129-
if (input.fileNames) {
130-
parquetFileNames = input.fileNames
131-
} else if (input.gcsPrefix) {
132-
const objectPrefix = gcsPrefixToObjectPrefix(input.gcsPrefix)
133-
const [files] = await bucket.getFiles({ prefix: objectPrefix })
134-
parquetFileNames = files.filter((f) => f.name.endsWith('.parquet')).map((f) => f.name)
135-
} else {
136-
throw new Error('gcsParquetToStaging: must provide either fileNames or gcsPrefix')
137-
}
138-
139-
const totalFiles = input.totalFiles ?? parquetFileNames.length
117+
// Row is 'loading' from here; a staging-load failure must flip it to 'failed' with the reason
118+
// rather than leave it stuck 'loading'. Mirrors the rankPackages pattern (criticality/activities).
119+
try {
120+
if (stagingDdl) {
121+
const stmts = Array.isArray(stagingDdl) ? stagingDdl : [stagingDdl]
122+
for (const stmt of stmts) await qx.result(stmt)
123+
}
124+
await qx.result(`TRUNCATE ${stagingTable}`)
125+
// Reset progress at the start of a fresh run (first chunk). This handles job ID reuse
126+
// via --export-name: a previous run's 193/193 would otherwise persist due to GREATEST.
127+
if (filesOffset === 0 && input.totalFiles != null) {
128+
await updateLoadingProgress(qx, jobId, 0, input.totalFiles, true)
129+
}
140130

141-
log.info(
142-
{ jobId, stagingTable, fileCount: parquetFileNames.length, filesOffset, totalFiles },
143-
'Loading parquet files into staging',
144-
)
131+
let parquetFileNames: string[]
132+
if (input.fileNames) {
133+
parquetFileNames = input.fileNames
134+
} else if (input.gcsPrefix) {
135+
const objectPrefix = gcsPrefixToObjectPrefix(input.gcsPrefix)
136+
const [files] = await bucket.getFiles({ prefix: objectPrefix })
137+
parquetFileNames = files.filter((f) => f.name.endsWith('.parquet')).map((f) => f.name)
138+
} else {
139+
throw new Error('gcsParquetToStaging: must provide either fileNames or gcsPrefix')
140+
}
145141

146-
let totalLoaded = 0
142+
const totalFiles = input.totalFiles ?? parquetFileNames.length
147143

148-
for (let i = 0; i < parquetFileNames.length; i += MAX_CONCURRENT) {
149-
const chunk = parquetFileNames.slice(i, i + MAX_CONCURRENT)
150-
const counts = await Promise.all(
151-
chunk.map((name) => loadParquetFile(qx, stagingTable, pgColumns, name, tsCols, decCols)),
152-
)
153-
totalLoaded += counts.reduce((a, b) => a + b, 0)
154-
const doneInBatch = i + chunk.length
155-
const doneGlobal = filesOffset + doneInBatch
156144
log.info(
157-
{
158-
jobId,
159-
totalLoaded,
160-
progress: `${doneGlobal}/${totalFiles} (${Math.round((doneGlobal / totalFiles) * 100)}%)`,
161-
},
162-
'Staging load progress',
145+
{ jobId, stagingTable, fileCount: parquetFileNames.length, filesOffset, totalFiles },
146+
'Loading parquet files into staging',
163147
)
164-
Context.current().heartbeat({ done: doneGlobal, total: totalFiles })
165-
await updateLoadingProgress(qx, jobId, doneGlobal, totalFiles)
166-
}
167148

168-
const cumulativeStagingRows = (input.priorStagingRows ?? 0) + totalLoaded
169-
await markJobStatus(qx, jobId, 'loading', {
170-
rowCountStaging: cumulativeStagingRows,
171-
tableRowCounts: { [`staging:${stagingTable}`]: totalLoaded },
172-
})
149+
let totalLoaded = 0
150+
151+
for (let i = 0; i < parquetFileNames.length; i += MAX_CONCURRENT) {
152+
const chunk = parquetFileNames.slice(i, i + MAX_CONCURRENT)
153+
const counts = await Promise.all(
154+
chunk.map((name) => loadParquetFile(qx, stagingTable, pgColumns, name, tsCols, decCols)),
155+
)
156+
totalLoaded += counts.reduce((a, b) => a + b, 0)
157+
const doneInBatch = i + chunk.length
158+
const doneGlobal = filesOffset + doneInBatch
159+
log.info(
160+
{
161+
jobId,
162+
totalLoaded,
163+
progress: `${doneGlobal}/${totalFiles} (${Math.round((doneGlobal / totalFiles) * 100)}%)`,
164+
},
165+
'Staging load progress',
166+
)
167+
Context.current().heartbeat({ done: doneGlobal, total: totalFiles })
168+
await updateLoadingProgress(qx, jobId, doneGlobal, totalFiles)
169+
}
170+
171+
const cumulativeStagingRows = (input.priorStagingRows ?? 0) + totalLoaded
172+
await markJobStatus(qx, jobId, 'loading', {
173+
rowCountStaging: cumulativeStagingRows,
174+
tableRowCounts: { [`staging:${stagingTable}`]: totalLoaded },
175+
})
173176

174-
await qx.result(`ANALYZE ${stagingTable}`)
177+
await qx.result(`ANALYZE ${stagingTable}`)
175178

176-
log.info({ jobId, stagingTable, totalLoaded }, 'Staging load complete')
179+
log.info({ jobId, stagingTable, totalLoaded }, 'Staging load complete')
177180

178-
return { rowsLoaded: totalLoaded }
181+
return { rowsLoaded: totalLoaded }
182+
} catch (err) {
183+
await markJobStatus(qx, jobId, 'failed', {
184+
errorMessage: (err as Error).message,
185+
finishedAt: new Date(),
186+
})
187+
throw err
188+
}
179189
}

‎services/apps/packages_worker/src/deps-dev/activities/mergeStagingToTable.ts‎

Lines changed: 47 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -50,51 +50,60 @@ export async function mergeStagingToTable(input: MergeStagingInput): Promise<Mer
5050
const qx = await getPackagesDb()
5151

5252
await markJobStatus(qx, jobId, 'merging')
53+
// Row is 'merging' from here; a merge failure must flip it to 'failed' with the reason rather
54+
// than leave it stuck 'merging'. Mirrors the rankPackages pattern (criticality/activities).
55+
try {
56+
let rowsAffected = 0
57+
const tableRowCounts: Record<string, number> = {}
5358

54-
let rowsAffected = 0
55-
const tableRowCounts: Record<string, number> = {}
59+
await qx.tx(async (tx) => {
60+
for (const sql of prepareStatements) {
61+
await tx.result(sql)
62+
}
63+
for (let i = 0; i < statements.length; i++) {
64+
const count = await tx.result(statements[i])
65+
rowsAffected += count
66+
const table = names[i] ?? `table_${i}`
67+
tableRowCounts[table] = (tableRowCounts[table] ?? 0) + count
68+
}
69+
})
5670

57-
await qx.tx(async (tx) => {
58-
for (const sql of prepareStatements) {
59-
await tx.result(sql)
60-
}
61-
for (let i = 0; i < statements.length; i++) {
62-
const count = await tx.result(statements[i])
63-
rowsAffected += count
64-
const table = names[i] ?? `table_${i}`
65-
tableRowCounts[table] = (tableRowCounts[table] ?? 0) + count
71+
if (chunkInfo) {
72+
log.info(
73+
{ jobId, rowsAffected, tableRowCounts, chunk: `${chunkInfo.index + 1}/${chunkInfo.total}` },
74+
'Chunk merge complete',
75+
)
6676
}
67-
})
6877

69-
if (chunkInfo) {
70-
log.info(
71-
{ jobId, rowsAffected, tableRowCounts, chunk: `${chunkInfo.index + 1}/${chunkInfo.total}` },
72-
'Chunk merge complete',
73-
)
74-
}
75-
76-
if (!isFinal) {
77-
await markJobStatus(qx, jobId, 'merging', {
78-
rowCountPg: priorRowsAffected + rowsAffected,
79-
})
80-
}
78+
if (!isFinal) {
79+
await markJobStatus(qx, jobId, 'merging', {
80+
rowCountPg: priorRowsAffected + rowsAffected,
81+
})
82+
}
8183

82-
if (isFinal) {
83-
const totalRowsAffected = priorRowsAffected + rowsAffected
84-
const totalTableRowCounts: Record<string, number> = { ...priorTableRowCounts }
85-
for (const [table, count] of Object.entries(tableRowCounts)) {
86-
totalTableRowCounts[table] = (totalTableRowCounts[table] ?? 0) + count
84+
if (isFinal) {
85+
const totalRowsAffected = priorRowsAffected + rowsAffected
86+
const totalTableRowCounts: Record<string, number> = { ...priorTableRowCounts }
87+
for (const [table, count] of Object.entries(tableRowCounts)) {
88+
totalTableRowCounts[table] = (totalTableRowCounts[table] ?? 0) + count
89+
}
90+
await markJobStatus(qx, jobId, 'done', {
91+
finishedAt: new Date(),
92+
rowCountPg: totalRowsAffected,
93+
tableRowCounts: totalTableRowCounts,
94+
})
95+
log.info(
96+
{ jobId, rowsAffected: totalRowsAffected, tableRowCounts: totalTableRowCounts },
97+
'Merge complete',
98+
)
8799
}
88-
await markJobStatus(qx, jobId, 'done', {
100+
101+
return { rowsAffected, tableRowCounts }
102+
} catch (err) {
103+
await markJobStatus(qx, jobId, 'failed', {
104+
errorMessage: (err as Error).message,
89105
finishedAt: new Date(),
90-
rowCountPg: totalRowsAffected,
91-
tableRowCounts: totalTableRowCounts,
92106
})
93-
log.info(
94-
{ jobId, rowsAffected: totalRowsAffected, tableRowCounts: totalTableRowCounts },
95-
'Merge complete',
96-
)
107+
throw err
97108
}
98-
99-
return { rowsAffected, tableRowCounts }
100109
}

0 commit comments

Comments
 (0)