Added Batch insert for References and Vulnerabilities

Applies batch inserts for reference and vulnerability tables,
solves the slow one-by-one insert process, for Vulnerabilities
with several references/vulnerabilities associated.
This commit is contained in:
Ale Feltes
2017-11-17 10:36:15 -03:00
parent dea9fa1145
commit 38499898aa

View File

@@ -25,6 +25,7 @@ import java.sql.ResultSet;
import java.sql.SQLException; import java.sql.SQLException;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.Collections; import java.util.Collections;
import java.util.Date;
import java.util.EnumMap; import java.util.EnumMap;
import java.util.HashMap; import java.util.HashMap;
import java.util.HashSet; import java.util.HashSet;
@@ -70,6 +71,7 @@ public final class CveDB implements AutoCloseable {
* The logger. * The logger.
*/ */
private static final Logger LOGGER = LoggerFactory.getLogger(CveDB.class); private static final Logger LOGGER = LoggerFactory.getLogger(CveDB.class);
public static final int MAX_BATCH_SIZE = 1000;
/** /**
* The database connection factory. * The database connection factory.
*/ */
@@ -730,36 +732,51 @@ public final class CveDB implements AutoCloseable {
} }
} }
final PreparedStatement insertReference = getPreparedStatement(INSERT_REFERENCE); PreparedStatement insertReference = getPreparedStatement(INSERT_REFERENCE);
int countReferences = 0;
for (Reference r : vuln.getReferences()) { for (Reference r : vuln.getReferences()) {
insertReference.setInt(1, vulnerabilityId); insertReference.setInt(1, vulnerabilityId);
insertReference.setString(2, r.getName()); insertReference.setString(2, r.getName());
insertReference.setString(3, r.getUrl()); insertReference.setString(3, r.getUrl());
insertReference.setString(4, r.getSource()); insertReference.setString(4, r.getSource());
insertReference.execute(); insertReference.addBatch();
countReferences++;
if (countReferences % MAX_BATCH_SIZE == 0) {
insertReference.executeBatch();
insertReference = getPreparedStatement(INSERT_REFERENCE);
LOGGER.info(getLogForBatchInserts(countReferences, "Completed %s batch inserts to references table: %s"));
countReferences = 0;
} else if (countReferences == vuln.getReferences().size()) {
if (LOGGER.isTraceEnabled()) {
LOGGER.trace(getLogForBatchInserts(countReferences, "Completed %s batch inserts to reference table: %s"));
}
insertReference.executeBatch();
countReferences = 0;
}
} }
final PreparedStatement insertSoftware = getPreparedStatement(INSERT_SOFTWARE); PreparedStatement insertSoftware = getPreparedStatement(INSERT_SOFTWARE);
for (VulnerableSoftware s : vuln.getVulnerableSoftware()) { int countSoftware = 0;
for (VulnerableSoftware vulnerableSoftware : vuln.getVulnerableSoftware()) {
int cpeProductId = 0; int cpeProductId = 0;
final PreparedStatement selectCpeId = getPreparedStatement(SELECT_CPE_ID); final PreparedStatement selectCpeId = getPreparedStatement(SELECT_CPE_ID);
selectCpeId.setString(1, s.getName()); selectCpeId.setString(1, vulnerableSoftware.getName());
try { try {
rs = selectCpeId.executeQuery(); rs = selectCpeId.executeQuery();
if (rs.next()) { if (rs.next()) {
cpeProductId = rs.getInt(1); cpeProductId = rs.getInt(1);
} }
} catch (SQLException ex) { } catch (SQLException ex) {
throw new DatabaseException("Unable to get primary key for new cpe: " + s.getName(), ex); throw new DatabaseException("Unable to get primary key for new cpe: " + vulnerableSoftware.getName(), ex);
} finally { } finally {
DBUtils.closeResultSet(rs); DBUtils.closeResultSet(rs);
} }
if (cpeProductId == 0) { if (cpeProductId == 0) {
final PreparedStatement insertCpe = getPreparedStatement(INSERT_CPE); final PreparedStatement insertCpe = getPreparedStatement(INSERT_CPE);
insertCpe.setString(1, s.getName()); insertCpe.setString(1, vulnerableSoftware.getName());
insertCpe.setString(2, s.getVendor()); insertCpe.setString(2, vulnerableSoftware.getVendor());
insertCpe.setString(3, s.getProduct()); insertCpe.setString(3, vulnerableSoftware.getProduct());
insertCpe.executeUpdate(); insertCpe.executeUpdate();
cpeProductId = DBUtils.getGeneratedKey(insertCpe); cpeProductId = DBUtils.getGeneratedKey(insertCpe);
} }
@@ -770,20 +787,24 @@ public final class CveDB implements AutoCloseable {
insertSoftware.setInt(1, vulnerabilityId); insertSoftware.setInt(1, vulnerabilityId);
insertSoftware.setInt(2, cpeProductId); insertSoftware.setInt(2, cpeProductId);
if (s.getPreviousVersion() == null) { if (vulnerableSoftware.getPreviousVersion() == null) {
insertSoftware.setNull(3, java.sql.Types.VARCHAR); insertSoftware.setNull(3, java.sql.Types.VARCHAR);
} else { } else {
insertSoftware.setString(3, s.getPreviousVersion()); insertSoftware.setString(3, vulnerableSoftware.getPreviousVersion());
} }
try { insertSoftware.addBatch();
insertSoftware.execute(); countSoftware++;
} catch (SQLException ex) { if (countSoftware % MAX_BATCH_SIZE == 0) {
if (ex.getMessage().contains("Duplicate entry")) { executeBatch(vuln, vulnerableSoftware, insertSoftware);
final String msg = String.format("Duplicate software key identified in '%s:%s'", vuln.getName(), s.getName()); insertSoftware = getPreparedStatement(INSERT_SOFTWARE);
LOGGER.info(msg, ex); LOGGER.info(getLogForBatchInserts(countSoftware, "Completed %s batch inserts software table: %s"));
} else { countSoftware = 0;
throw ex; } else if (countSoftware == vuln.getVulnerableSoftware().size()) {
if (LOGGER.isTraceEnabled()) {
LOGGER.trace(getLogForBatchInserts(countSoftware, "Completed %s batch inserts software table: %s"));
countReferences = 0;
} }
executeBatch(vuln, vulnerableSoftware, insertSoftware);
} }
} }
@@ -796,6 +817,31 @@ public final class CveDB implements AutoCloseable {
} }
} }
private String getLogForBatchInserts(int pCountReferences, String pFormat) {
return String.format(pFormat, pCountReferences, new Date());
}
/**
* Executes batch inserts of vulnerabilities when MAX_BATCH_SIZE is reached
*
* @param pVulnerability
* @param pVulnerableSoftware
* @param pInsertSoftware
* @throws SQLException
*/
private void executeBatch(Vulnerability pVulnerability, VulnerableSoftware pVulnerableSoftware, PreparedStatement pInsertSoftware) throws SQLException {
try {
pInsertSoftware.executeBatch();
} catch (SQLException ex) {
if (ex.getMessage().contains("Duplicate entry")) {
final String msg = String.format("Duplicate software key identified in '%s:%s'", pVulnerability.getName(), pVulnerableSoftware.getName());
LOGGER.info(msg, ex);
} else {
throw ex;
}
}
}
/** /**
* Checks to see if data exists so that analysis can be performed. * Checks to see if data exists so that analysis can be performed.
* *