bulltrackers-module 1.0.735 → 1.0.737

package/functions/computation-system-v2/framework/storage/StorageManager.js

@@ -1,11 +1,21 @@
 /**
  * @fileoverview StorageManager - Unified storage for computation results
- * Includes Optimized Checkpointing and Performance Monitoring.
+ * Includes Optimized Checkpointing, Heartbeating, and Zombie Detection with Dead Letter Logic.
+ * * V2.1 UPDATE: Replaced Streaming/DML with Batch Load Jobs.
+ * - commitResults: Now buffers data to GCS (NDJSON)
+ * - finalizeResults: Loads all GCS files to a temp table and performs a single MERGE
+ * * V2.2 UPDATE: Added saveCheckpoint for Append-Only Checkpointer support.
+ * * V2.3 UPDATE: Parallelized GCS and Firestore writes in commitResults.
+ * * FIX: Switched to bigquery.createJob for GCS imports to prevent local file path interpretation errors.
+ * * FIX: Improved error logging to catch swallowed BigQuery insert errors.
+ * * FIX: finalizeResults now checks for file existence to prevent "Not found" errors on empty results.
  */
 
 const { Firestore } = require('@google-cloud/firestore');
 const { BigQuery } = require('@google-cloud/bigquery');
+const { Storage } = require('@google-cloud/storage');
 const pLimit = require('p-limit');
+const crypto = require('crypto');
 
 class StorageManager {
     constructor(config, logger = null) {
@@ -17,6 +27,10 @@ class StorageManager {
             location: config.bigquery?.location || 'EU'
         });
         
+        this.storage = new Storage({
+            projectId: config.bigquery?.projectId
+        });
+        
         this._firestore = null;
         this.tableExists = new Map();
     }
@@ -30,71 +44,208 @@ class StorageManager {
         return this._firestore;
     }
     
+    // =========================================================================
+    // RESULT COMMITTING (Batch -> GCS Buffer)
+    // =========================================================================
+
+    /**
+     * Staging Phase: Write results to GCS in NDJSON format.
+     * Replaces the old direct-to-BigQuery INSERT/MERGE.
+     * NOW PARALLELIZED: GCS and Firestore writes happen concurrently.
+     */
     async commitResults(dateStr, entry, results, depResultHashes = {}) {
         const storageConfig = this._resolveStorageConfig(entry);
-        const writeResults = { bigquery: null, firestore: null };
         const startTime = Date.now();
         
-        if (storageConfig.bigquery !== false) {
+        // Define GCS Task (Fatal on error)
+        const gcsTask = async () => {
+            if (storageConfig.bigquery === false) return null;
             try {
-                writeResults.bigquery = await this._writeToBigQuery(
-                    dateStr, entry, results, depResultHashes
-                );
+                // Modified: Write to GCS instead of BQ directly
+                return await this._stageToGCS(dateStr, entry, results, depResultHashes);
             } catch (error) {
-                this._log('ERROR', `BigQuery write failed for ${entry.name}: ${error.message}`);
+                this._logError(`GCS staging failed for ${entry.name}`, error);
                 throw error;
             }
-        }
-        
-        if (storageConfig.firestore?.enabled) {
+        };
+
+        // Define Firestore Task (Non-fatal, logs warning)
+        const firestoreTask = async () => {
+            if (!storageConfig.firestore?.enabled) return null;
             try {
-                writeResults.firestore = await this._writeToFirestore(
-                    dateStr, entry, results, storageConfig.firestore
-                );
+                return await this._writeToFirestore(dateStr, entry, results, storageConfig.firestore);
             } catch (error) {
                 this._log('WARN', `Firestore write failed for ${entry.name}: ${error.message}`);
-                writeResults.firestore = { error: error.message };
+                return { error: error.message };
             }
-        }
+        };
+
+        // Execute IO operations in parallel
+        const [bigqueryResult, firestoreResult] = await Promise.all([
+            gcsTask(),
+            firestoreTask()
+        ]);
+
+        const writeResults = { 
+            bigquery: bigqueryResult, 
+            firestore: firestoreResult 
+        };
         
         const duration = Date.now() - startTime;
-        this._log('INFO', `Committed ${entry.name} results in ${duration}ms`);
+        this._log('INFO', `Committed (Staged) ${entry.name} results in ${duration}ms`);
         return writeResults;
     }
 
     // =========================================================================
-    // OPTIMIZED CHECKPOINTING (Lightweight State)
+    // FINALIZATION (GCS -> Load Job -> Merge)
     // =========================================================================
 
-    async initCheckpoint(dateStr, computationName, checkpointId, totalEntities) {
+    /**
+     * Finalize Phase: Bulk Load GCS files -> Temp Table -> Merge to Final.
+     * This is called ONCE per computation run.
+     */
+    async finalizeResults(dateStr, entry) {
+        const table = this.config.resultStore?.table || 'computation_results';
+        const bucketName = this.config.gcs?.bucket;
+        const prefix = this.config.gcs?.prefix || 'staging';
+        
+        // 1. Define GCS path pattern: gs://bucket/prefix/date/computation/*.json
+        const filePrefix = `${prefix}/${dateStr}/${entry.name}/`;
+        const gcsPath = `gs://${bucketName}/${filePrefix}*.json`;
+        
+        this._log('INFO', `Finalizing ${entry.name}...`);
+
+        try {
+            // FIX: Check if files actually exist before trying to load them
+            // If the computation produced 0 results, no files exist, and BQ will throw "Not Found".
+            const [files] = await this.storage.bucket(bucketName).getFiles({ prefix: filePrefix });
+            
+            if (!files || files.length === 0) {
+                this._log('INFO', `No staged files found for ${entry.name}. Skipping finalization (Empty Result).`);
+                return;
+            }
+
+            const tempTableId = `${table}_temp_${crypto.randomUUID().replace(/-/g, '')}`;
+            this._log('INFO', `Loading ${files.length} files from ${gcsPath} into ${tempTableId}...`);
+
+            // 2. Load GCS files into a Temporary Table
+            // We create the temp table with the exact schema we expect first
+            await this._createTempTableForLoad(tempTableId);
+            
+            // FIX: Use bigquery.createJob directly.
+            const [job] = await this.bigquery.createJob({
+                configuration: {
+                    load: {
+                        destinationTable: {
+                            projectId: this.config.bigquery.projectId,
+                            datasetId: this.config.bigquery.dataset,
+                            tableId: tempTableId
+                        },
+                        sourceUris: [gcsPath],
+                        sourceFormat: 'NEWLINE_DELIMITED_JSON',
+                        writeDisposition: 'WRITE_TRUNCATE',
+                        autodetect: false // Use schema from table creation
+                    }
+                }
+            });
+
+            // Wait for job completion
+            await new Promise((resolve, reject) => {
+                job.on('complete', (metadata) => resolve(metadata));
+                job.on('error', (err) => reject(err));
+            });
+
+            this._log('INFO', `Load Job completed for ${entry.name}`);
+
+            // 3. Perform the bulk MERGE from Temp -> Final
+            await this._mergeStagedData(table, tempTableId);
+
+            // 4. Cleanup: Delete Temp Table & GCS Files
+            try {
+                await this.bigquery.dataset(this.config.bigquery.dataset).table(tempTableId).delete();
+                await this._cleanupGCSFiles(bucketName, filePrefix);
+            } catch (cleanupErr) {
+                this._log('WARN', `Cleanup warning for ${entry.name}: ${cleanupErr.message}`);
+            }
+
+            this._log('INFO', `Finalization complete for ${entry.name}`);
+        } catch (e) {
+            this._logError(`Finalization failed for ${entry.name}`, e);
+            throw e;
+        }
+    }
+
+    // =========================================================================
+    // CHECKPOINTING & RESILIENCE
+    // =========================================================================
+
+    async saveCheckpoint(state) {
+        const table = 'computation_checkpoints';
+        const dataset = this.config.bigquery.dataset;
+        await this._ensureCheckpointTable(table);
+
+        const row = {
+            date: state.date,
+            computation_name: state.computation_name,
+            checkpoint_id: state.checkpoint_id,
+            worker_instance_id: state.worker_instance_id,
+            code_hash: state.code_hash || null,
+            status: state.status,
+            processed_count: state.processed_count,
+            total_entities: state.total_entities,
+            last_entity_id: state.last_entity_id,
+            completed_batches: state.completed_batches || [],
+            attempts: state.attempts,
+            started_at: this.bigquery.timestamp(new Date(state.started_at)),
+            last_updated: this.bigquery.timestamp(new Date(state.last_updated))
+        };
+
+        try {
+            await this.bigquery.dataset(dataset).table(table).insert([row]);
+        } catch (error) {
+            if (error.name === 'PartialFailureError' || error.errors) {
+                 this._log('ERROR', `Checkpoint insert failed: ${JSON.stringify(error.errors)}`);
+            } else {
+                 this._log('ERROR', `Checkpoint insert failed: ${error.message}`);
+            }
+            throw error;
+        }
+    }
+
+    // Legacy/Alternative methods
+    async initCheckpoint(dateStr, computationName, checkpointId, totalEntities, workerId, codeHash) {
         const table = 'computation_checkpoints';
         const dataset = this.config.bigquery.dataset;
-        
         await this._ensureCheckpointTable(table);
 
         const row = {
             date: dateStr,
             computation_name: computationName,
             checkpoint_id: checkpointId,
+            worker_instance_id: workerId,
+            code_hash: codeHash || null,
             status: 'running',
             processed_count: 0,
             total_entities: totalEntities,
             last_entity_id: null,
             completed_batches: [],
-            // FIX: Use instance method, remove 'new', pass Date object
+            attempts: 1,
             started_at: this.bigquery.timestamp(new Date()),
             last_updated: this.bigquery.timestamp(new Date())
         };
 
-        // Standard insert for initialization
-        await this.bigquery.dataset(dataset).table(table).insert([row]);
+        try {
+            await this.bigquery.dataset(dataset).table(table).insert([row]);
+        } catch (e) {
+            this._logError('initCheckpoint failed', e);
+            throw e;
+        }
     }
 
     async updateCheckpoint(dateStr, computationName, checkpointId, state) {
         const table = 'computation_checkpoints';
         const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
 
-        // Efficient UPDATE using structured fields
         const query = `
             UPDATE ${fullTable}
             SET 
@@ -115,131 +266,131 @@ class StorageManager {
                 checkpointId,
                 processedCount: state.processedCount,
                 lastEntityId: state.lastEntityId,
-                batchIndex: state.batchIndex // Assuming we add one batch at a time
+                batchIndex: state.batchIndex
             },
             location: this.config.bigquery.location
         });
     }
 
-    async completeCheckpoint(dateStr, computationName, checkpointId) {
+    async updateHeartbeat(checkpointId, workerId) {
+        const table = 'computation_checkpoints';
+        const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
+        const query = `
+            UPDATE ${fullTable}
+            SET last_updated = CURRENT_TIMESTAMP()
+            WHERE checkpoint_id = @checkpointId AND worker_instance_id = @workerId
+        `;
+        this.bigquery.query({ query, params: { checkpointId, workerId }, location: this.config.bigquery.location })
+            .catch(e => this._log('WARN', `Heartbeat failed: ${e.message}`));
+    }
+
+    async claimCheckpoint(checkpointId, newWorkerId, minutesThreshold = 15) {
+        const table = 'computation_checkpoints';
+        const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
+        const query = `
+            UPDATE ${fullTable}
+            SET worker_instance_id = @newWorkerId, attempts = IFNULL(attempts, 1) + 1, last_updated = CURRENT_TIMESTAMP()
+            WHERE checkpoint_id = @checkpointId AND status = 'running'
+              AND last_updated < TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL @minutes MINUTE)
+        `;
+        try {
+            await this.bigquery.query({
+                query,
+                params: { checkpointId, newWorkerId, minutes: minutesThreshold },
+                location: this.config.bigquery.location
+            });
+            const verify = await this.getLatestCheckpointById(checkpointId);
+            return verify && verify.worker_instance_id === newWorkerId;
+        } catch (e) {
+            this._log('ERROR', `Failed to claim checkpoint: ${e.message}`);
+            return false;
+        }
+    }
+
+    async failCheckpoint(checkpointId, reason) {
         const table = 'computation_checkpoints';
         const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
+        const query = `UPDATE ${fullTable} SET status = 'failed', last_updated = CURRENT_TIMESTAMP() WHERE checkpoint_id = @checkpointId`;
+        await this.bigquery.query({ query, params: { checkpointId }, location: this.config.bigquery.location });
+    }
 
+    async findZombies(minutesThreshold = 15) {
+        const table = 'computation_checkpoints';
+        const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
+        try {
+            const query = `
+                SELECT computation_name, date, checkpoint_id, last_updated
+                FROM ${fullTable}
+                WHERE status = 'running' AND last_updated < TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL @minutes MINUTE)
+                LIMIT 50
+            `;
+            const [rows] = await this.bigquery.query({ query, params: { minutes: minutesThreshold }, location: this.config.bigquery.location });
+            return rows.map(r => ({ name: r.computation_name, date: r.date.value || r.date, checkpointId: r.checkpoint_id }));
+        } catch (e) { return []; }
+    }
+
+    async completeCheckpoint(dateStr, computationName, checkpointId) {
+        const table = 'computation_checkpoints';
+        const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
         const query = `
             UPDATE ${fullTable}
             SET status = 'completed', last_updated = CURRENT_TIMESTAMP()
             WHERE date = @date AND computation_name = @computationName AND checkpoint_id = @checkpointId
         `;
-
-        await this.bigquery.query({
-            query, 
-            params: { date: dateStr, computationName, checkpointId },
-            location: this.config.bigquery.location
-        });
+        await this.bigquery.query({ query, params: { date: dateStr, computationName, checkpointId }, location: this.config.bigquery.location });
     }
 
     async getLatestCheckpoint(dateStr, computationName) {
         const table = 'computation_checkpoints';
         const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
-        
         try {
-            // Get the most recent attempt
             const query = `
-                SELECT checkpoint_id, status, processed_count, last_entity_id, completed_batches
+                SELECT checkpoint_id, status, processed_count, last_entity_id, completed_batches, worker_instance_id, last_updated, attempts, code_hash, started_at
                 FROM ${fullTable}
                 WHERE date = @date AND computation_name = @computationName
-                ORDER BY started_at DESC
-                LIMIT 1
+                ORDER BY last_updated DESC LIMIT 1
             `;
-
-            const [rows] = await this.bigquery.query({
-                query,
-                params: { date: dateStr, computationName },
-                location: this.config.bigquery.location
-            });
-
+            const [rows] = await this.bigquery.query({ query, params: { date: dateStr, computationName }, location: this.config.bigquery.location });
             if (rows.length === 0) return null;
-            return rows[0];
-        } catch (e) {
-            return null; // Table might not exist yet
-        }
+            
+            const r = rows[0];
+            return {
+                ...r,
+                started_at: r.started_at?.value || r.started_at,
+                last_updated: r.last_updated?.value || r.last_updated
+            };
+        } catch (e) { return null; }
     }
 
-    async _ensureCheckpointTable(tableName) {
-        if (this.tableExists.get(tableName)) return;
-        const dataset = this.bigquery.dataset(this.config.bigquery.dataset);
-        const table = dataset.table(tableName);
-        const [exists] = await table.exists();
-        
-        if (!exists) {
-            await dataset.createTable(tableName, {
-                schema: [
-                    { name: 'date', type: 'DATE', mode: 'REQUIRED' },
-                    { name: 'computation_name', type: 'STRING', mode: 'REQUIRED' },
-                    { name: 'checkpoint_id', type: 'STRING', mode: 'REQUIRED' },
-                    { name: 'status', type: 'STRING', mode: 'REQUIRED' }, // running, completed, failed
-                    { name: 'processed_count', type: 'INTEGER', mode: 'NULLABLE' },
-                    { name: 'total_entities', type: 'INTEGER', mode: 'NULLABLE' },
-                    { name: 'last_entity_id', type: 'STRING', mode: 'NULLABLE' },
-                    { name: 'completed_batches', type: 'INTEGER', mode: 'REPEATED' }, // ARRAY<INT64>
-                    { name: 'started_at', type: 'TIMESTAMP', mode: 'REQUIRED' },
-                    { name: 'last_updated', type: 'TIMESTAMP', mode: 'REQUIRED' }
-                ],
-                timePartitioning: { type: 'DAY', field: 'date' },
-                clustering: { fields: ['computation_name', 'status'] }
-            });
-        }
-        this.tableExists.set(tableName, true);
+    async getLatestCheckpointById(checkpointId) {
+        const table = 'computation_checkpoints';
+        const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
+        const query = `SELECT worker_instance_id, status FROM ${fullTable} WHERE checkpoint_id = @id LIMIT 1`;
+        const [rows] = await this.bigquery.query({ query, params: { id: checkpointId }, location: this.config.bigquery.location });
+        return rows[0];
     }
 
-    // =========================================================================
-    // PERFORMANCE REPORTING
-    // =========================================================================
-
-    async savePerformanceReport(dateStr, report) {
+    async savePerformanceReport(report) {
         const table = 'computation_performance';
-        
-        await this._ensurePerformanceTable(table);
-        
-        const rows = report.computations.map(c => ({
-            date: dateStr,
-            computation_name: c.name,
-            entity_count: c.entityCount,
-            avg_duration_ms: c.avgDuration,
-            max_duration_ms: c.maxDuration,
-            p95_duration_ms: c.p95Duration,
-            avg_memory_delta: c.avgMemory,
-            total_duration_ms: c.totalDuration,
-            recorded_at: new Date().toISOString()
-        }));
-
-        if (rows.length === 0) return;
-        await this.bigquery.dataset(this.config.bigquery.dataset).table(table).insert(rows);
-    }
+        const dataset = this.config.bigquery.dataset;
+        await this._ensurePerformanceTable(table);  
 
-    async _ensurePerformanceTable(tableName) {
-        if (this.tableExists.get(tableName)) return;
-        const dataset = this.bigquery.dataset(this.config.bigquery.dataset);
-        const table = dataset.table(tableName);
-        const [exists] = await table.exists();
-        
-        if (!exists) {
-            await dataset.createTable(tableName, {
-                schema: [
-                    { name: 'date', type: 'DATE', mode: 'REQUIRED' },
-                    { name: 'computation_name', type: 'STRING', mode: 'REQUIRED' },
-                    { name: 'entity_count', type: 'INTEGER', mode: 'NULLABLE' },
-                    { name: 'avg_duration_ms', type: 'FLOAT', mode: 'NULLABLE' },
-                    { name: 'max_duration_ms', type: 'FLOAT', mode: 'NULLABLE' },
-                    { name: 'p95_duration_ms', type: 'FLOAT', mode: 'NULLABLE' },
-                    { name: 'avg_memory_delta', type: 'FLOAT', mode: 'NULLABLE' },
-                    { name: 'total_duration_ms', type: 'FLOAT', mode: 'NULLABLE' },
-                    { name: 'recorded_at', type: 'TIMESTAMP', mode: 'REQUIRED' }
-                ],
-                timePartitioning: { type: 'DAY', field: 'date' }
-            });
+        const row = {
+            run_id : report.runId || 'unknown',
+            computation_name : report.computationName,
+            date: report.date,
+            duration_ms: report.durationMs,
+            metrics: JSON.stringify (report.metrics || {}),
+            entity_count: report.entityCount || 0,
+            status : report.status || 'completed',
+            created_at : this.bigquery.timestamp(new Date())
+        };
+        try {
+            await this.bigquery.dataset(dataset).table(table).insert([row]);
+            this._log('INFO', `Performance report saved for ${report.computationName}`);
+        } catch (e) {
+            this._logError('Failed to save performance report', e);
         }
-        this.tableExists.set(tableName, true);
     }
 
     // =========================================================================
@@ -259,21 +410,85 @@ class StorageManager {
         };
     }
 
-    async _writeToBigQuery(dateStr, entry, results, depResultHashes) {
-        const table = this.config.resultStore?.table || 'computation_results';
-        const fullTable = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${table}\``;
-        
+    async _stageToGCS(dateStr, entry, results, depResultHashes) {
         const rows = this._buildBigQueryRows(dateStr, entry, results, depResultHashes);
         if (rows.length === 0) return { rowCount: 0 };
         
-        await this._ensureBigQueryTable(table);
+        const bucketName = this.config.gcs?.bucket;
+        const prefix = this.config.gcs?.prefix || 'staging';
+        const filename = `${prefix}/${dateStr}/${entry.name}/${crypto.randomUUID()}.json`;
+        
+        const file = this.storage.bucket(bucketName).file(filename);
+        
+        const ndjson = rows.map(r => JSON.stringify(r)).join('\n');
+        
+        await file.save(ndjson, {
+            contentType: 'application/json',
+            resumable: false 
+        });
+        
+        return { rowCount: rows.length, gcsUri: `gs://${bucketName}/${filename}` };
+    }
+
+    async _createTempTableForLoad(tableName) {
+        const dataset = this.bigquery.dataset(this.config.bigquery.dataset);
+        const table = dataset.table(tableName);
+        
+        const schema = [
+            { name: 'date', type: 'DATE', mode: 'REQUIRED' },
+            { name: 'computation_name', type: 'STRING', mode: 'REQUIRED' },
+            { name: 'category', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'entity_id', type: 'STRING', mode: 'REQUIRED' },
+            { name: 'code_hash', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'result_hash', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'dependency_result_hashes', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'entity_count', type: 'INTEGER', mode: 'NULLABLE' },
+            { name: 'result_data', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'updated_at', type: 'TIMESTAMP', mode: 'NULLABLE' }
+        ];
+
+        await table.create({ schema, timePartitioning: { type: 'DAY', field: 'date' } });
+    }
+
+    async _mergeStagedData(targetTable, tempTable) {
+        const fullTarget = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${targetTable}\``;
+        const fullTemp = `\`${this.config.bigquery.projectId}.${this.config.bigquery.dataset}.${tempTable}\``;
+        
+        await this._ensureBigQueryTable(targetTable);
+
+        const mergeQuery = `
+            MERGE INTO ${fullTarget} T
+            USING (
+                SELECT * FROM (
+                    SELECT *, ROW_NUMBER() OVER(PARTITION BY date, computation_name, entity_id ORDER BY updated_at DESC) as rn
+                    FROM ${fullTemp}
+                ) WHERE rn = 1
+            ) S
+            ON T.date = S.date AND T.computation_name = S.computation_name AND T.entity_id = S.entity_id
+            WHEN MATCHED THEN
+                UPDATE SET 
+                    code_hash = S.code_hash, 
+                    result_hash = S.result_hash, 
+                    dependency_result_hashes = S.dependency_result_hashes, 
+                    entity_count = S.entity_count, 
+                    result_data = S.result_data, 
+                    updated_at = S.updated_at
+            WHEN NOT MATCHED THEN
+                INSERT (date, computation_name, category, entity_id, code_hash, result_hash, 
+                        dependency_result_hashes, entity_count, result_data, updated_at)
+                VALUES (S.date, S.computation_name, S.category, S.entity_id, S.code_hash, S.result_hash, 
+                        S.dependency_result_hashes, S.entity_count, S.result_data, S.updated_at)
+        `;
         
-        const batchSize = this.config.execution?.insertBatchSize || 500;
-        for (let i = 0; i < rows.length; i += batchSize) {
-            const batch = rows.slice(i, i + batchSize);
-            await this._upsertBatch(fullTable, batch);
+        await this.bigquery.query({ query: mergeQuery, location: this.config.bigquery.location });
+    }
+
+    async _cleanupGCSFiles(bucketName, prefix) {
+        const [files] = await this.storage.bucket(bucketName).getFiles({ prefix });
+        if (files.length > 0) {
+            const limit = pLimit(10);
+            await Promise.all(files.map(f => limit(() => f.delete())));
         }
-        return { rowCount: rows.length };
     }
 
     _buildBigQueryRows(dateStr, entry, results, depResultHashes) {
@@ -313,60 +528,6 @@ class StorageManager {
         return rows;
     }
 
-    async _upsertBatch(fullTable, rows) {
-        const dates = rows.map(r => r.date);
-        const names = rows.map(r => r.computation_name);
-        const entityIds = rows.map(r => r.entity_id);
-        const codeHashes = rows.map(r => r.code_hash);
-        const resultHashes = rows.map(r => r.result_hash);
-        const depHashes = rows.map(r => r.dependency_result_hashes);
-        const entityCounts = rows.map(r => r.entity_count);
-        const resultData = rows.map(r => r.result_data);
-        const updatedAts = rows.map(r => r.updated_at);
-        
-        const mergeQuery = `
-            MERGE INTO ${fullTable} T
-            USING (
-                SELECT 
-                    PARSE_DATE('%Y-%m-%d', date) as date,
-                    computation_name, entity_id, code_hash, result_hash, 
-                    dependency_result_hashes, entity_count, result_data, 
-                    PARSE_TIMESTAMP('%Y-%m-%dT%H:%M:%E*SZ', updated_at) as updated_at
-                FROM UNNEST(@dates) AS date WITH OFFSET pos
-                JOIN UNNEST(@names) AS computation_name WITH OFFSET USING(pos)
-                JOIN UNNEST(@entity_ids) AS entity_id WITH OFFSET USING(pos)
-                JOIN UNNEST(@code_hashes) AS code_hash WITH OFFSET USING(pos)
-                JOIN UNNEST(@result_hashes) AS result_hash WITH OFFSET USING(pos)
-                JOIN UNNEST(@dep_hashes) AS dependency_result_hashes WITH OFFSET USING(pos)
-                JOIN UNNEST(@entity_counts) AS entity_count WITH OFFSET USING(pos)
-                JOIN UNNEST(@result_data) AS result_data WITH OFFSET USING(pos)
-                JOIN UNNEST(@updated_ats) AS updated_at WITH OFFSET USING(pos)
-            ) S
-            ON T.date = S.date AND T.computation_name = S.computation_name AND T.entity_id = S.entity_id
-            WHEN MATCHED THEN
-                UPDATE SET code_hash = S.code_hash, result_hash = S.result_hash, 
-                           dependency_result_hashes = S.dependency_result_hashes, 
-                           entity_count = S.entity_count, result_data = S.result_data, 
-                           updated_at = S.updated_at
-            WHEN NOT MATCHED THEN
-                INSERT (date, computation_name, category, entity_id, code_hash, result_hash, 
-                        dependency_result_hashes, entity_count, result_data, updated_at)
-                VALUES (S.date, S.computation_name, @category, S.entity_id, S.code_hash, S.result_hash, 
-                        S.dependency_result_hashes, S.entity_count, S.result_data, S.updated_at)
-        `;
-        
-        await this.bigquery.query({
-            query: mergeQuery,
-            params: {
-                dates, names, entity_ids: entityIds, code_hashes: codeHashes, 
-                result_hashes: resultHashes, dep_hashes: depHashes, 
-                entity_counts: entityCounts, result_data: resultData, 
-                updated_ats: updatedAts, category: rows[0].category
-            },
-            location: this.config.bigquery.location
-        });
-    }
-
     async _ensureBigQueryTable(tableName) {
         if (this.tableExists.get(tableName)) return;
         const dataset = this.bigquery.dataset(this.config.bigquery.dataset);
@@ -393,7 +554,69 @@ class StorageManager {
         this.tableExists.set(tableName, true);
     }
 
-    async _writeToFirestore(dateStr, entry, results, firestoreConfig) {
+    async _ensurePerformanceTable(tableName) {
+        if (this.tableExists.get(tableName)) return;
+
+        const dataset = this.bigquery.dataset(this.config.bigquery.dataset);
+        const table = dataset.table(tableName);
+        const [exists] = await table.exists();
+
+        if (!exists) {
+            const schema = [
+                { name: 'run_id', type: 'STRING', mode: 'REQUIRED' },
+                { name: 'computation_name', type: 'STRING', mode: 'REQUIRED' },
+                { name: 'date', type: 'DATE', mode: 'REQUIRED' },
+                { name: 'duration_ms', type: 'INTEGER', mode: 'NULLABLE' },
+                { name: 'metrics', type: 'STRING', mode: 'NULLABLE' }, // JSON string
+                { name: 'entity_count', type: 'INTEGER', mode: 'NULLABLE' },
+                { name: 'status', type: 'STRING', mode: 'NULLABLE' },
+                { name: 'created_at', type: 'TIMESTAMP', mode: 'REQUIRED' }
+            ];
+
+            await dataset.createTable(tableName, {
+                schema,
+                timePartitioning: { type: 'DAY', field: 'date' },
+                clustering: { fields: ['computation_name', 'run_id'] }
+            });
+            this._log('INFO', `Created table ${tableName}`);
+        }
+        
+        this.tableExists.set(tableName, true);
+    }
+    
+    async _ensureCheckpointTable(tableName) {
+        if (this.tableExists.get(tableName)) return;
+        const dataset = this.bigquery.dataset(this.config.bigquery.dataset);
+        const table = dataset.table(tableName);
+        const [exists] = await table.exists();
+        
+        const schema = [
+            { name: 'date', type: 'DATE', mode: 'REQUIRED' },
+            { name: 'computation_name', type: 'STRING', mode: 'REQUIRED' },
+            { name: 'checkpoint_id', type: 'STRING', mode: 'REQUIRED' },
+            { name: 'worker_instance_id', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'code_hash', type: 'STRING', mode: 'NULLABLE' }, 
+            { name: 'status', type: 'STRING', mode: 'REQUIRED' },
+            { name: 'processed_count', type: 'INTEGER', mode: 'NULLABLE' },
+            { name: 'total_entities', type: 'INTEGER', mode: 'NULLABLE' },
+            { name: 'last_entity_id', type: 'STRING', mode: 'NULLABLE' },
+            { name: 'completed_batches', type: 'INTEGER', mode: 'REPEATED' },
+            { name: 'attempts', type: 'INTEGER', mode: 'NULLABLE' }, 
+            { name: 'started_at', type: 'TIMESTAMP', mode: 'REQUIRED' },
+            { name: 'last_updated', type: 'TIMESTAMP', mode: 'REQUIRED' }
+        ];
+
+        if (!exists) {
+            await dataset.createTable(tableName, {
+                schema,
+                timePartitioning: { type: 'DAY', field: 'date' },
+                clustering: { fields: ['computation_name', 'status'] }
+            });
+        } 
+        this.tableExists.set(tableName, true);
+    }
+
+    _writeToFirestore(dateStr, entry, results, firestoreConfig) {
         const { path, merge, includeMetadata } = firestoreConfig;
         if (!path) throw new Error(`Firestore path not configured for ${entry.name}`);
         
@@ -431,7 +654,7 @@ class StorageManager {
             if (batchCount > 0) batches.push(currentBatch);
             
             const limit = pLimit(10);
-            await Promise.all(batches.map(b => limit(() => b.commit())));
+            return Promise.all(batches.map(b => limit(() => b.commit()))).then(() => ({ docCount }));
             
         } else {
             const docPath = this._resolvePath(path, {
@@ -439,10 +662,8 @@ class StorageManager {
             });
             const docRef = this.firestore.doc(docPath);
             const docData = { ...results, ...metadata };
-            merge ? await docRef.set(docData, { merge: true }) : await docRef.set(docData);
-            docCount = 1;
+            return (merge ? docRef.set(docData, { merge: true }) : docRef.set(docData)).then(() => ({ docCount: 1 }));
         }
-        return { docCount };
     }
 
     _resolvePath(pathTemplate, values) {
@@ -456,7 +677,6 @@ class StorageManager {
     }
 
     _hashResult(data) {
-        const crypto = require('crypto');
         const str = typeof data === 'string' ? data : JSON.stringify(data);
         return crypto.createHash('md5').update(str).digest('hex').substring(0, 16);
     }
@@ -464,6 +684,18 @@ class StorageManager {
     _log(level, message) {
         this.logger?.log ? this.logger.log(level, `[StorageManager] ${message}`) : console.log(`[${level}] [StorageManager] ${message}`);
     }
+
+    _logError(context, error) {
+        // Safe logging for BigQuery PartialFailureError which hides details in .errors
+        let details = error.message;
+        if (error.errors && Array.isArray(error.errors)) {
+            details = JSON.stringify(error.errors, null, 2);
+        } else if (error.response && error.response.insertErrors) {
+            details = JSON.stringify(error.response.insertErrors, null, 2);
+        }
+        
+        this._log('ERROR', `${context}: ${details}`);
+    }
 }
 
 module.exports = { StorageManager };