com.wordnik.system.mongodb.RestoreUtil Maven / Gradle / Ivy
The newest version!
// Copyright (C) 2010 Wordnik, Inc.
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 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 Lesser
// General Public License for more details. You should have received a copy
// of the GNU Lesser General Public License along with this program. If not,
// see collections = getCollections();
// restore in single-threaded manner
for(CollectionInfo info : collections){
restore(info);
}
}
static long REPORT_DURATION = 10000;
protected void restore(CollectionInfo info) {
try{
System.out.println("restoring collection " + info.getName());
File[] files = new File(INPUT_DIR).listFiles();
if(files != null){
List filesToProcess = new ArrayList();
for(File file : files){
if(file.getName().startsWith(info.getName()) && file.getName().indexOf(".bson") > 0){
filesToProcess.add(file);
}
}
Collections.sort(filesToProcess, new FilenameComparator());
if(DROP_EXISTING){
try{
System.out.println("Dropping collection " + info.getName());
DB db = MongoDBConnectionManager.getConnection("DB", DATABASE_HOST, "local", DATABASE_USER_NAME, DATABASE_PASSWORD, SchemaType.READ_WRITE());
DBCollection coll = db.getCollection(info.getName());
coll.drop();
}
catch(Exception e){
e.printStackTrace();
}
}
long count = 0;
long startTime = System.currentTimeMillis();
long lastOutput = System.currentTimeMillis();
for(File file : filesToProcess){
System.out.println("restoring file " + file.getName() + " to collection " + info.getName());
InputStream inputStream = null;
try{
if(file.getName().endsWith(".gz")){
inputStream = new GZIPInputStream(new FileInputStream(file));
}
else{
inputStream = new BufferedInputStream(new FileInputStream(file));
}
BSONDecoder decoder = new DefaultDBDecoder();
while(true){
if(inputStream.available() == 0){
break;
}
BSONObject obj = decoder.readObject(inputStream);
if(obj == null){
break;
}
write(info, new BasicDBObject((BasicBSONObject)obj));
count++;
long duration = System.currentTimeMillis() - lastOutput;
if(duration > REPORT_DURATION){
report(info.getName(), count, System.currentTimeMillis() - startTime);
lastOutput = System.currentTimeMillis();
}
}
}
catch(java.io.EOFException e){
break;
}
finally{
inputStream.close();
}
}
}
}
catch(Exception e){
e.printStackTrace();
}
}
void report(String collectionName, long count, long duration){
double brate = (double)count / ((duration) / 1000.0);
System.out.println(collectionName + ": " + PrintFormat.LONG_FORMAT.format(count) + " records, " + PrintFormat.LONG_FORMAT.format(brate) + " req/sec");
}
protected void write(CollectionInfo info, DBObject dbo) throws Exception {
DB db = MongoDBConnectionManager.getConnection("TARGET", DATABASE_HOST, DATABASE_NAME, DATABASE_USER_NAME, DATABASE_PASSWORD, SchemaType.READ_WRITE());
DBCollection coll = db.getCollection(info.getName());
coll.save(dbo);
}
private List getCollections() {
List collections = new ArrayList();
try{
Collection collectionsFromFiles = getCollectionNamesFromFiles(INPUT_DIR);
List collectionsToAdd = new ArrayList();
List collectionsToSkip = new ArrayList();
selectCollections(COLLECTION_STRING, collectionsToAdd, collectionsToSkip);
boolean exclusionsOnly = collectionsToAdd.contains("*");
if(exclusionsOnly){
for(String collectionName : collectionsFromFiles){
if(!collectionsToSkip.contains(collectionName)){
collectionsToAdd.add(collectionName);
}
}
}
else{
if(collectionsToAdd.size() == 0){
// add everything
collectionsToAdd.addAll(collectionsFromFiles);
}
}
if(exclusionsOnly){
for(String collectionName : collectionsFromFiles){
if(!collectionsToSkip.contains(collectionName)){
collectionsToAdd.add(collectionName);
}
}
}
else{
if(collectionsToAdd.size() == 0){
// add everything
collectionsToAdd.addAll(collectionsFromFiles);
}
}
for(String collection : collectionsToAdd){
if(!"*".equals(collection))
collections.add(new CollectionInfo(collection, 0));
}
}
catch(Exception e){
throw new RuntimeException(e);
}
return collections;
}
private Collection getCollectionNamesFromFiles(String inputDir) {
File[] files = new File(inputDir).listFiles();
Set collectionNames = new HashSet();
for(File file : files){
if(file.getName().contains(".bson")){
int pos = file.getName().indexOf(".bson");
String collectionName = file.getName().substring(0, pos);
if(collectionName.indexOf('.') > 0){
StringTokenizer tk = new StringTokenizer(collectionName, ".");
if(tk.countTokens() > 1){
String base = null;
try{
base = tk.nextToken();
Integer.parseInt(tk.nextToken());
collectionName = base;
}
catch(NumberFormatException e){
// continue
}
}
}
collectionNames.add(collectionName);
}
}
return collectionNames;
}
public static boolean parseArgs(String...args){
for (int i = 0; i < args.length; i++) {
switch (args[i].charAt(1)) {
case 'i':
INPUT_DIR = args[++i];
break;
case 'c':
COLLECTION_STRING = args[++i];
break;
case 'D':
DROP_EXISTING = true;
break;
case 'd':
DATABASE_NAME = args[++i];
break;
case 'u':
DATABASE_USER_NAME = args[++i];
break;
case 'p':
DATABASE_PASSWORD = args[++i];
break;
case 'h':
DATABASE_HOST = args[++i];
break;
default:
return false;
}
}
return true;
}
public static void usage(){
System.out.println("usage: RestoreUtil");
System.out.println(" -i : input directory");
System.out.println(" -c : CSV collection string (prefix with ! to exclude)");
System.out.println(" -D : drop existing collections");
System.out.println(" -h : target database host[:port]");
System.out.println(" -d : target database name");
System.out.println(" [-u : target database username]");
System.out.println(" [-p : target database password]");
}
class FilenameComparator implements Comparator {
@Override
public int compare(File o1, File o2) {
return o1.getName().compareTo(o2.getName());
}
}
} © 2015 - 2025 Weber Informatics LLC | Privacy Policy