Java tutorial
/* * Password Management Servlets (PWM) * http://www.pwm-project.org * * Copyright (c) 2006-2009 Novell, Inc. * Copyright (c) 2009-2017 The PWM Project * * This program is free software; you can redistribute it and/or modify * it under the terms of the GNU General Public License as published by * the Free Software Foundation; either version 2 of the License, or * (at your option) any later version. * * This program is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU General Public License for more details. * * You should have received a copy of the GNU General Public License * along with this program; if not, write to the Free Software * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA */ package password.pwm.util.localdb; import org.apache.commons.csv.CSVPrinter; import org.apache.commons.csv.CSVRecord; import org.apache.commons.io.IOUtils; import org.apache.commons.io.input.CountingInputStream; import password.pwm.PwmApplication; import password.pwm.PwmConstants; import password.pwm.error.PwmError; import password.pwm.error.PwmOperationalException; import password.pwm.svc.stats.EventRateMeter; import password.pwm.util.ProgressInfo; import password.pwm.util.java.TimeDuration; import password.pwm.util.TransactionSizeCalculator; import password.pwm.util.java.JavaHelper; import password.pwm.util.logging.PwmLogger; import java.io.File; import java.io.FileInputStream; import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.io.OutputStream; import java.io.PrintStream; import java.io.Reader; import java.math.BigDecimal; import java.text.DecimalFormat; import java.time.Instant; import java.util.Date; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.Map; import java.util.Timer; import java.util.TimerTask; import java.util.TreeMap; import java.util.concurrent.TimeUnit; import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; public class LocalDBUtility { private static final PwmLogger LOGGER = PwmLogger.forClass(LocalDBUtility.class); private final LocalDB localDB; private int exportLineCounter; private int importLineCounter; private static final int GZIP_BUFFER_SIZE = 1024 * 512; public LocalDBUtility(final LocalDB localDB) { this.localDB = localDB; } public void exportLocalDB(final OutputStream outputStream, final Appendable debugOutput, final boolean showLineCount) throws PwmOperationalException, IOException { if (outputStream == null) { throw new PwmOperationalException(PwmError.ERROR_UNKNOWN, "outputFileStream for exportLocalDB cannot be null"); } final int totalLines; if (showLineCount) { writeStringToOut(debugOutput, "counting records in LocalDB..."); exportLineCounter = 0; for (final LocalDB.DB loopDB : LocalDB.DB.values()) { if (loopDB.isBackup()) { exportLineCounter += localDB.size(loopDB); } } totalLines = exportLineCounter; writeStringToOut(debugOutput, " total lines: " + totalLines); } else { totalLines = 0; } exportLineCounter = 0; writeStringToOut(debugOutput, "export beginning"); final long startTime = System.currentTimeMillis(); final Timer statTimer = new Timer(true); statTimer.schedule(new TimerTask() { @Override public void run() { if (showLineCount) { final float percentComplete = (float) exportLineCounter / (float) totalLines; final String percentStr = DecimalFormat.getPercentInstance().format(percentComplete); writeStringToOut(debugOutput, "exported " + exportLineCounter + " records, " + percentStr + " complete"); } else { writeStringToOut(debugOutput, "exported " + exportLineCounter + " records"); } } }, 30 * 1000, 30 * 1000); try (CSVPrinter csvPrinter = JavaHelper .makeCsvPrinter(new GZIPOutputStream(outputStream, GZIP_BUFFER_SIZE))) { csvPrinter.printComment(PwmConstants.PWM_APP_NAME + " " + PwmConstants.SERVLET_VERSION + " LocalDB export on " + JavaHelper.toIsoDate(new Date())); for (final LocalDB.DB loopDB : LocalDB.DB.values()) { if (loopDB.isBackup()) { csvPrinter.printComment("Export of " + loopDB.toString()); final LocalDB.LocalDBIterator<String> localDBIterator = localDB.iterator(loopDB); try { while (localDBIterator.hasNext()) { final String key = localDBIterator.next(); final String value = localDB.get(loopDB, key); csvPrinter.printRecord(loopDB.toString(), key, value); exportLineCounter++; } } finally { localDBIterator.close(); } csvPrinter.flush(); } } csvPrinter.printComment("export completed at " + JavaHelper.toIsoDate(new Date())); } catch (IOException e) { writeStringToOut(debugOutput, "IO error during localDB export: " + e.getMessage()); } finally { statTimer.cancel(); } writeStringToOut(debugOutput, "export complete, exported " + exportLineCounter + " records in " + TimeDuration.fromCurrent(startTime).asLongString()); } private static void writeStringToOut(final Appendable out, final String string) { if (out == null) { return; } final String msg = JavaHelper.toIsoDate(new Date()) + " " + string + "\n"; try { out.append(msg); } catch (IOException e) { LOGGER.error("error writing to output appender while performing operation: " + e.getMessage() + ", message:" + msg); } } public void importLocalDB(final File inputFile, final PrintStream out) throws PwmOperationalException, IOException { if (inputFile == null) { throw new PwmOperationalException(PwmError.ERROR_UNKNOWN, "inputFile for importLocalDB cannot be null"); } if (!inputFile.exists()) { throw new PwmOperationalException(PwmError.ERROR_UNKNOWN, "inputFile for importLocalDB does not exist"); } final long totalBytes = inputFile.length(); if (totalBytes <= 0) { throw new PwmOperationalException(PwmError.ERROR_UNKNOWN, "inputFile for importLocalDB is empty"); } final InputStream inputStream = new FileInputStream(inputFile); importLocalDB(inputStream, out, totalBytes); } public void importLocalDB(final InputStream inputStream, final Appendable out) throws PwmOperationalException, IOException { importLocalDB(inputStream, out, 0); } private void importLocalDB(final InputStream inputStream, final Appendable out, final long totalBytes) throws PwmOperationalException, IOException { this.prepareForImport(); importLineCounter = 0; if (totalBytes > 0) { writeStringToOut(out, "total bytes in localdb import source: " + totalBytes); } writeStringToOut(out, "beginning localdb import..."); final Instant startTime = Instant.now(); final TransactionSizeCalculator transactionCalculator = new TransactionSizeCalculator( new TransactionSizeCalculator.SettingsBuilder() .setDurationGoal(new TimeDuration(100, TimeUnit.MILLISECONDS)).setMinTransactions(50) .setMaxTransactions(5 * 1000).createSettings()); final Map<LocalDB.DB, Map<String, String>> transactionMap = new HashMap<>(); for (final LocalDB.DB loopDB : LocalDB.DB.values()) { transactionMap.put(loopDB, new TreeMap<>()); } final CountingInputStream countingInputStream = new CountingInputStream(inputStream); final EventRateMeter eventRateMeter = new EventRateMeter(TimeDuration.MINUTE); final Timer statTimer = new Timer(true); statTimer.scheduleAtFixedRate(new TimerTask() { @Override public void run() { String output = ""; if (totalBytes > 0) { final ProgressInfo progressInfo = new ProgressInfo(startTime, totalBytes, countingInputStream.getByteCount()); output += progressInfo.debugOutput(); } else { output += "recordsImported=" + importLineCounter; } output += ", avgTransactionSize=" + transactionCalculator.getTransactionSize() + ", recordsPerMinute=" + eventRateMeter.readEventRate().setScale(2, BigDecimal.ROUND_DOWN); writeStringToOut(out, output); } }, 30 * 1000, 30 * 1000); Reader csvReader = null; try { csvReader = new InputStreamReader(new GZIPInputStream(countingInputStream, GZIP_BUFFER_SIZE), PwmConstants.DEFAULT_CHARSET); for (final CSVRecord record : PwmConstants.DEFAULT_CSV_FORMAT.parse(csvReader)) { importLineCounter++; eventRateMeter.markEvents(1); final String dbName_recordStr = record.get(0); final LocalDB.DB db = JavaHelper.readEnumFromString(LocalDB.DB.class, null, dbName_recordStr); final String key = record.get(1); final String value = record.get(2); if (db == null) { writeStringToOut(out, "ignoring localdb import record #" + importLineCounter + ", invalid DB name '" + dbName_recordStr + "'"); } else { transactionMap.get(db).put(key, value); int cachedTransactions = 0; for (final LocalDB.DB loopDB : LocalDB.DB.values()) { cachedTransactions += transactionMap.get(loopDB).size(); } if (cachedTransactions >= transactionCalculator.getTransactionSize()) { final long startTxnTime = System.currentTimeMillis(); for (final LocalDB.DB loopDB : LocalDB.DB.values()) { localDB.putAll(loopDB, transactionMap.get(loopDB)); transactionMap.get(loopDB).clear(); } transactionCalculator.recordLastTransactionDuration(TimeDuration.fromCurrent(startTxnTime)); } } } } finally { LOGGER.trace("import process completed"); statTimer.cancel(); IOUtils.closeQuietly(csvReader); IOUtils.closeQuietly(countingInputStream); } for (final LocalDB.DB loopDB : LocalDB.DB.values()) { localDB.putAll(loopDB, transactionMap.get(loopDB)); transactionMap.get(loopDB).clear(); } this.markImportComplete(); writeStringToOut(out, "restore complete, restored " + importLineCounter + " records in " + TimeDuration.fromCurrent(startTime).asLongString()); statTimer.cancel(); } public static Map<StatsKey, Object> dbStats(final LocalDB localDB, final LocalDB.DB db) { int totalValues = 0; long storedChars = 0; final long totalChars = 0; LocalDB.LocalDBIterator<String> iter = null; try { iter = localDB.iterator(db); while (iter.hasNext()) { final String key = iter.next(); final String rawValue = localDB.get(db, key); if (rawValue != null) { totalValues++; storedChars += rawValue.length(); } } } catch (Exception e) { LOGGER.error("error while examining LocalDB: " + e.getMessage()); } finally { if (iter != null) { iter.close(); } } final int avgValueLength = totalValues == 0 ? 0 : (int) (totalChars / totalValues); final Map<StatsKey, Object> returnObj = new LinkedHashMap<>(); returnObj.put(StatsKey.TOTAL_VALUES, totalValues); returnObj.put(StatsKey.STORED_CHARS, storedChars); returnObj.put(StatsKey.AVG_VALUE_LENGTH, avgValueLength); return returnObj; } public enum StatsKey { TOTAL_VALUES, STORED_CHARS, AVG_VALUE_LENGTH, } public void prepareForImport() throws LocalDBException { LOGGER.info("preparing LocalDB for import procedure"); localDB.put(LocalDB.DB.PWM_META, PwmApplication.AppAttribute.LOCALDB_IMPORT_STATUS.getKey(), "inprogress"); for (final LocalDB.DB loopDB : LocalDB.DB.values()) { if (loopDB != LocalDB.DB.PWM_META) { localDB.truncate(loopDB); } } localDB.truncate(LocalDB.DB.PWM_META); // save meta for last so flag is cleared last. localDB.put(LocalDB.DB.PWM_META, PwmApplication.AppAttribute.LOCALDB_IMPORT_STATUS.getKey(), "inprogress"); } public void markImportComplete() throws LocalDBException { LOGGER.info("marking LocalDB import procedure completed"); localDB.remove(LocalDB.DB.PWM_META, PwmApplication.AppAttribute.LOCALDB_IMPORT_STATUS.getKey()); } public boolean readImportInprogressFlag() throws LocalDBException { return "inprogress".equals( localDB.get(LocalDB.DB.PWM_META, PwmApplication.AppAttribute.LOCALDB_IMPORT_STATUS.getKey())); } static boolean hasBooleanParameter(final LocalDBProvider.Parameter parameter, final Map<LocalDBProvider.Parameter, String> parameters) { return parameters != null && parameters.containsKey(parameter) && Boolean.parseBoolean(parameters.get(parameter)); } }