/**
* Copyright 2005-2012 Akiban Technologies, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.persistit;
import java.io.BufferedInputStream;
import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import com.persistit.CLI.Arg;
import com.persistit.CLI.Cmd;
import com.persistit.exception.CorruptImportStreamException;
import com.persistit.exception.PersistitException;
import com.persistit.policy.SplitPolicy;
import com.persistit.util.Util;
/**
* Loads Persistit records from a file or other stream in a format generated by
* a {@link StreamSaver}.
*
* @version 1.0
*/
public class StreamLoader extends Task {
/**
* Default for BufferedInputStream buffer size.
*/
public final static int DEFAULT_BUFFER_SIZE = 65536;
protected String _filePath;
protected DataInputStream _dis;
protected Key _key = new Key((Persistit) null);
protected Value _value = new Value((Persistit) null);
protected Volume _lastVolume;
protected Tree _lastTree;
protected int _dataRecordCount = 0;
protected int _otherRecordCount = 0;
protected boolean _stop;
protected Exception _lastException;
protected TreeSelector _treeSelector;
protected boolean _createMissingVolumes;
protected boolean _createMissingTrees;
protected ImportHandler _handler;
@Cmd("load")
static Task createStreamLoader(@Arg("file|string:|Load from file path") final String file,
@Arg("trees|string:|Tree selector - specify Volumes/Trees/Keys to save") final String treeSelectorString,
@Arg("_flag|r|Use regular expressions in tree selector") final boolean regex,
@Arg("_flag|n|Don't create missing Volumes (Default is to create them)") final boolean dontCreateVolumes,
@Arg("_flag|t|Don't create missing Trees (Default is to create them)") final boolean dontCreateTrees,
@Arg("_flag|v|verbose") final boolean verbose) throws Exception {
final StreamLoader task = new StreamLoader();
task._filePath = file;
task._treeSelector = TreeSelector.parseSelector(treeSelectorString, regex, '\\');
task._createMissingVolumes = !dontCreateVolumes;
task._createMissingTrees = !dontCreateTrees;
task.setMessageLogVerbosity(verbose ? LOG_VERBOSE : LOG_NORMAL);
return task;
}
/**
* Package-private constructor for use in a {@link Task}.
*
*/
StreamLoader() {
}
public StreamLoader(final Persistit persistit, final DataInputStream dis) {
super(persistit);
_dis = dis;
}
public StreamLoader(final Persistit persistit, final File file) throws IOException {
this(persistit, new DataInputStream(new BufferedInputStream(new FileInputStream(file))));
}
public StreamLoader(final Persistit persistit, final String fileName) throws IOException {
this(persistit, new DataInputStream(new BufferedInputStream(new FileInputStream(fileName))));
}
public void close() throws IOException {
_dis.close();
}
public void load() throws IOException, PersistitException {
load(new TreeSelector(), true, true);
}
public void load(final TreeSelector treeSelector, final boolean createMissingVolumes,
final boolean createMissingTrees) throws IOException, PersistitException {
_handler = new ImportHandler(_persistit, treeSelector, createMissingVolumes, createMissingTrees);
load(_handler);
close();
}
public void load(final ImportHandler handler) throws IOException, PersistitException {
while (next(handler)) {
}
postMessage(String.format("DONE - processed %,d data records and %,d other records", _dataRecordCount,
_otherRecordCount), Task.LOG_NORMAL);
}
public boolean next(final ImportHandler handler) throws IOException, PersistitException {
final int b1 = _dis.read();
if (b1 == -1) {
return false;
}
final int b2 = _dis.read();
final int recordType = ((b1 & 0xFF) << 8) + (b2 & 0xFF);
switch (recordType) {
case StreamSaver.RECORD_TYPE_FILL: {
handler.handleFillRecord();
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_DATA: {
final int keySize = _dis.readShort();
final int elisionCount = _dis.readShort();
final int valueSize = _dis.readInt();
_value.ensureFit(valueSize);
_dis.read(_key.getEncodedBytes(), elisionCount, keySize - elisionCount);
_key.setEncodedSize(keySize);
_dis.read(_value.getEncodedBytes(), 0, valueSize);
_value.setEncodedSize(valueSize);
handler.handleDataRecord(_key, _value);
_dataRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_KEY_FILTER: {
final String filterString = _dis.readUTF();
handler.handleKeyFilterRecord(filterString);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_VOLUME_ID: {
final long id = _dis.readLong();
final long initialPages = _dis.readLong();
final long extensionPages = _dis.readLong();
final long maximumPages = _dis.readLong();
final int bufferSize = _dis.readInt();
final String path = _dis.readUTF();
handler._volumeName = _dis.readUTF();
handler.handleVolumeIdRecord(id, initialPages, extensionPages, maximumPages, bufferSize, path,
handler._volumeName);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_TREE_ID: {
handler._treeName = _dis.readUTF();
handler.handleTreeIdRecord(handler._treeName);
_otherRecordCount++;
postMessage("Loading Tree " + handler._treeName, Task.LOG_VERBOSE);
break;
}
case StreamSaver.RECORD_TYPE_HOSTNAME: {
final String hostName = _dis.readUTF();
handler.handleHostNameRecord(hostName);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_USER: {
final String hostName = _dis.readUTF();
handler.handleUserRecord(hostName);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_COMMENT: {
final String comment = _dis.readUTF();
handler.handleCommentRecord(comment);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_COUNT: {
final long dataRecordCount = _dis.readLong();
final long otherRecordCount = _dis.readLong();
handler.handleCountRecord(dataRecordCount, otherRecordCount);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_START: {
handler.handleStartRecord();
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_END: {
handler.handleEndRecord();
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_TIMESTAMP: {
final long timeStamp = _dis.readLong();
handler.handleTimeStampRecord(timeStamp);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_EXCEPTION: {
final String exceptionString = _dis.readUTF();
handler.handleExceptionRecord(exceptionString);
_otherRecordCount++;
break;
}
case StreamSaver.RECORD_TYPE_COMPLETION: {
handler.handleCompletionRecord();
_otherRecordCount++;
break;
}
default: {
throw new CorruptImportStreamException("Invalid record type " + recordType + " ("
+ Util.bytesToHex(new byte[] { (byte) (recordType >>> 8), (byte) recordType }) + " after reading "
+ _dataRecordCount + " data records" + " and " + _otherRecordCount + " other records");
}
}
return true;
}
/**
* Handler for various record types in stream being loaded.
*
* @author peter
*
*/
protected static class ImportHandler {
protected Persistit _persistit;
protected TreeSelector _treeSelector;
protected Exchange _exchange;
protected Volume _volume;
protected Tree _tree;
protected KeyFilter _keyFilter;
protected boolean _createMissingVolumes;
protected boolean _createMissingTrees;
protected String _volumeName = null;
protected String _treeName = null;
protected ImportHandler(final Persistit persistit) {
this(persistit, new TreeSelector(), true, true);
}
protected ImportHandler(final Persistit persistit, final TreeSelector treeSelector,
final boolean createMissingVolumes, final boolean createMissingTrees) {
_persistit = persistit;
_treeSelector = treeSelector == null ? new TreeSelector() : treeSelector;
_createMissingTrees = createMissingTrees;
_createMissingVolumes = createMissingVolumes;
}
protected void handleFillRecord() throws PersistitException {
}
protected void handleDataRecord(final Key key, final Value value) throws PersistitException {
if (_keyFilter == null || _keyFilter.selected(key)) {
if (_volume == null || _tree == null)
return;
if (_exchange == null) {
_exchange = _persistit.getExchange(_volume, _tree.getName(), false);
}
key.copyTo(_exchange.getKey());
_exchange.setSplitPolicy(SplitPolicy.PACK_BIAS);
// Using this package-private method avoids copying
// the value field.
_exchange.store(_exchange.getKey(), value);
}
}
protected void handleKeyFilterRecord(final String keyFilterString) throws PersistitException {
}
protected void handleVolumeIdRecord(final long volumeId, final long initialPages, final long extensionPages,
final long maximumPages, final int bufferSize, final String path, final String name)
throws PersistitException {
final Exchange oldExchange = _exchange;
_exchange = null;
_volume = null;
_tree = null;
if (!_treeSelector.isVolumeNameSelected(name)) {
return;
}
_volume = _persistit.getVolume(name);
if (_volume != null) {
_volume.verifyId(volumeId);
} else if (_createMissingVolumes) {
_volume = new Volume(new VolumeSpecification(path, name, bufferSize, initialPages, maximumPages,
extensionPages, false, true, false));
_volume.setId(volumeId);
_volume.open(_persistit);
}
if (oldExchange != null && oldExchange.getVolume().equals(_volume)) {
_exchange = oldExchange;
}
}
protected void handleTreeIdRecord(final String treeName) throws PersistitException {
final Exchange oldExchange = _exchange;
_exchange = null;
_tree = null;
if (_volume == null) {
return;
}
if (!_treeSelector.isTreeNameSelected(_volume.getName(), treeName)) {
return;
}
_tree = _volume.getTree(treeName, _createMissingVolumes | _createMissingTrees);
if (oldExchange != null && oldExchange.getTree() == _tree) {
_exchange = oldExchange;
}
_keyFilter = _treeSelector.keyFilter(_volume.getName(), treeName);
}
protected void handleTimeStampRecord(final long timeStamp) throws PersistitException {
}
protected void handleHostNameRecord(final String hostName) throws PersistitException {
}
protected void handleUserRecord(final String userName) throws PersistitException {
}
protected void handleCommentRecord(final String comment) throws PersistitException {
}
protected void handleCountRecord(final long keyValueRecords, final long otherRecords) throws PersistitException {
}
protected void handleStartRecord() throws PersistitException {
}
protected void handleEndRecord() throws PersistitException {
}
protected void handleExceptionRecord(final String exceptionString) throws PersistitException {
}
protected void handleCompletionRecord() throws PersistitException {
}
}
@Override
public void runTask() throws Exception {
_dis = new DataInputStream(new BufferedInputStream(new FileInputStream(_filePath), DEFAULT_BUFFER_SIZE));
load(_treeSelector, _createMissingVolumes, _createMissingTrees);
}
@Override
public String getStatus() {
if (_handler == null || _handler._tree == null) {
return null;
}
final Tree tree = _handler._tree;
return tree.getName() + " in " + tree.getVolume().getPath() + " (" + _dataRecordCount + ")";
}
}