Table Output doesn't commit all records in Beam pipeline
@@ -584,18 +584,28 @@ public class TableOutput extends BaseTransform<TableOutputMeta, TableOutputData>
public void dispose() {
if (data.db != null) {
+ try {
+ emptyAndCommitBatchBuffers();
+ } finally {
+ data.db.disconnect();
+ }
+ super.dispose();
+ }
+ }
+
+ private void emptyAndCommitBatchBuffers() {
try {
for (String schemaTable : data.preparedStatements.keySet()) {
// Get a commit counter per prepared statement to keep track of separate tables, etc.
//
Integer batchCounter = data.commitCounterMap.get(schemaTable);
- if (batchCounter == null) {
- batchCounter = 0;
+ if (batchCounter == null || batchCounter == 0) {
+ continue; // Skip this one, no work required
}
PreparedStatement insertStatement = data.preparedStatements.get(schemaTable);
-
data.db.emptyAndCommit(insertStatement, data.batchMode, batchCounter);
+ data.commitCounterMap.put(schemaTable, 0);
}
for (int i = 0; i < data.batchBuffer.size(); i++) {
Object[] row = data.batchBuffer.get(i);
@@ -635,10 +645,14 @@ public class TableOutput extends BaseTransform<TableOutputMeta, TableOutputData>
logError("Unexpected error rolling back the database connection.", e);
}
}
-
- data.db.disconnect();
}
- super.dispose();
}
+
+ // Force the batched up rows to the database in a single-threaded scenario
+ // (Beam as well)
+ //
+ @Override
+ public void batchComplete() throws HopException {
+ emptyAndCommitBatchBuffers();
}
}
