/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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 org.apache.usergrid.tools;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.util.MinimalPrettyPrinter;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.Options;
import org.apache.usergrid.persistence.Entity;
import org.apache.usergrid.persistence.EntityManager;
import org.apache.usergrid.persistence.Query;
import org.apache.usergrid.persistence.Results;
import org.apache.usergrid.tools.export.ExportConnection;
import org.apache.usergrid.tools.export.ExportEntity;
import org.apache.usergrid.utils.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import rx.Observable;
import rx.Scheduler;
import rx.Subscriber;
import rx.functions.Action0;
import rx.functions.Action1;
import rx.schedulers.Schedulers;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
/**
* Export all entities and connections of a Usergrid app.
*
* Exports data files to specified directory.
*
* Will create as many output files as there are writeThreads (by default: 10).
*
* Will create two types of files: *.entities for Usegrird entities and *.collections for entity to entity connections.
*
* Every line of the data files is a complete JSON object.
*/
public class ExportApp extends ExportingToolBase {
static final Logger logger = LoggerFactory.getLogger( ExportApp.class );
static final String APPLICATION_NAME = "application";
private static final String WRITE_THREAD_COUNT = "writeThreads";
String applicationName;
String organizationName;
AtomicInteger entitiesWritten = new AtomicInteger(0);
AtomicInteger connectionsWritten = new AtomicInteger(0);
Scheduler writeScheduler;
ObjectMapper mapper = new ObjectMapper();
Map<Thread, JsonGenerator> entityGeneratorsByThread = new HashMap<Thread, JsonGenerator>();
Map<Thread, JsonGenerator> connectionGeneratorsByThread = new HashMap<Thread, JsonGenerator>();
int writeThreadCount = 10; // set via CLI option; limiting write will limit output files
@Override
@SuppressWarnings("static-access")
public Options createOptions() {
Options options = super.createOptions();
Option appNameOption = OptionBuilder.hasArg().withType("")
.withDescription( "Application Name -" + APPLICATION_NAME ).create( APPLICATION_NAME );
options.addOption( appNameOption );
Option writeThreadsOption = OptionBuilder.hasArg().withType(0)
.withDescription( "Write Threads -" + WRITE_THREAD_COUNT ).create(WRITE_THREAD_COUNT);
options.addOption( writeThreadsOption );
return options;
}
/**
* Tool entry point.
*/
@Override
public void runTool(CommandLine line) throws Exception {
applicationName = line.getOptionValue( APPLICATION_NAME );
if (StringUtils.isNotEmpty( line.getOptionValue( WRITE_THREAD_COUNT ) )) {
try {
writeThreadCount = Integer.parseInt( line.getOptionValue( WRITE_THREAD_COUNT ) );
} catch (NumberFormatException nfe) {
logger.error( "-" + WRITE_THREAD_COUNT + " must be specified as an integer. Aborting..." );
return;
}
}
setVerbose( line );
applyOrgId( line );
prepareBaseOutputFileName( line );
outputDir = createOutputParentDir();
logger.info( "Export directory: " + outputDir.getAbsolutePath() );
startSpring();
UUID applicationId = emf.lookupApplication( applicationName );
if (applicationId == null) {
throw new RuntimeException( "Cannot find application " + applicationName );
}
final EntityManager em = emf.getEntityManager( applicationId );
organizationName = em.getApplication().getOrganizationName();
ExecutorService writeThreadPoolExecutor = Executors.newFixedThreadPool( writeThreadCount );
writeScheduler = Schedulers.from( writeThreadPoolExecutor );
Observable<String> collectionsObservable = Observable.create( new CollectionsObservable( em ) );
logger.debug( "Starting export" );
collectionsObservable.flatMap( collection -> {
return Observable.create( new EntityObservable( em, collection ) )
.doOnNext( new EntityWriteAction() ).subscribeOn( writeScheduler );
} ).flatMap( exportEntity -> {
return Observable.create( new ConnectionsObservable( em, exportEntity ) )
.doOnNext( new ConnectionWriteAction() ).subscribeOn( writeScheduler );
} ).doOnCompleted( new FileWrapUpAction() ).toBlocking().lastOrDefault(null);
}
// ----------------------------------------------------------------------------------------
// reading data
/**
* Emits collection names found in application.
*/
class CollectionsObservable implements rx.Observable.OnSubscribe<String> {
EntityManager em;
public CollectionsObservable(EntityManager em) {
this.em = em;
}
public void call(Subscriber<? super String> subscriber) {
int count = 0;
try {
Map<String, Object> collectionMetadata = em.getApplicationCollectionMetadata();
logger.debug( "Emitting {} collection names for application {}",
collectionMetadata.size(), em.getApplication().getName() );
for ( String collection : collectionMetadata.keySet() ) {
subscriber.onNext( collection );
count++;
}
} catch (Exception e) {
subscriber.onError( e );
}
subscriber.onCompleted();
logger.info( "Completed. Read {} collection names", count );
}
}
/**
* Emits entities of collection.
*/
private class EntityObservable implements rx.Observable.OnSubscribe<ExportEntity> {
EntityManager em;
String collection;
public EntityObservable(EntityManager em, String collection) {
this.em = em;
this.collection = collection;
}
public void call(Subscriber<? super ExportEntity> subscriber) {
logger.info("Starting to read entities of collection {}", collection);
//subscriber.onStart();
try {
int count = 0;
Query query = new Query();
query.setLimit( MAX_ENTITY_FETCH );
Results results = em.searchCollection( em.getApplicationRef(), collection, query );
while (results.size() > 0) {
for (Entity entity : results.getEntities()) {
try {
Set<String> dictionaries = em.getDictionaries( entity );
Map dictionariesByName = new HashMap<String, Map<Object, Object>>();
for (String dictionary : dictionaries) {
Map<Object, Object> dict = em.getDictionaryAsMap( entity, dictionary );
if (dict.isEmpty()) {
continue;
}
dictionariesByName.put( dictionary, dict );
}
ExportEntity exportEntity = new ExportEntity(
organizationName,
applicationName,
entity,
dictionariesByName );
subscriber.onNext( exportEntity );
count++;
} catch (Exception e) {
logger.error("Error reading entity " + entity.getUuid() +" from collection " + collection);
}
}
if (results.getCursor() == null) {
break;
}
query.setCursor( results.getCursor() );
results = em.searchCollection( em.getApplicationRef(), collection, query );
}
subscriber.onCompleted();
logger.info("Completed collection {}. Read {} entities", collection, count);
} catch ( Exception e ) {
subscriber.onError(e);
}
}
}
/**
* Emits connections of an entity.
*/
private class ConnectionsObservable implements rx.Observable.OnSubscribe<ExportConnection> {
EntityManager em;
ExportEntity exportEntity;
public ConnectionsObservable(EntityManager em, ExportEntity exportEntity) {
this.em = em;
this.exportEntity = exportEntity;
}
public void call(Subscriber<? super ExportConnection> subscriber) {
// logger.debug( "Starting to read connections for entity {} type {}",
// exportEntity.getEntity().getName(), exportEntity.getEntity().getType() );
int count = 0;
try {
Set<String> connectionTypes = em.getConnectionTypes( exportEntity.getEntity() );
for (String connectionType : connectionTypes) {
Results results = em.getTargetEntities(
exportEntity.getEntity(), connectionType, null, Query.Level.CORE_PROPERTIES );
for (Entity connectedEntity : results.getEntities()) {
try {
ExportConnection connection = new ExportConnection(
applicationName,
organizationName,
connectionType,
exportEntity.getEntity().getUuid(),
connectedEntity.getUuid());
subscriber.onNext( connection );
count++;
} catch (Exception e) {
logger.error( "Error reading connection entity "
+ exportEntity.getEntity().getUuid() + " -> " + connectedEntity.getType());
}
}
}
} catch (Exception e) {
subscriber.onError( e );
}
subscriber.onCompleted();
if ( count == 0 ) {
logger.debug("Completed entity {} type {} no connections",
new Object[] { exportEntity.getEntity().getUuid(), exportEntity.getEntity().getType() });
}
// logger.debug("Completed entity {} type {} connections count {}",
// new Object[] { exportEntity.getEntity().getUuid(), exportEntity.getEntity().getType(), count });
}
}
// ----------------------------------------------------------------------------------------
// writing data
/**
* Writes entities to JSON file.
*/
private class EntityWriteAction implements Action1<ExportEntity> {
public void call(ExportEntity entity) {
String [] parts = Thread.currentThread().getName().split("-");
String fileName = outputDir.getAbsolutePath() + File.separator
+ applicationName.replace('/','-') + "-" + parts[3] + ".entities";
JsonGenerator gen = entityGeneratorsByThread.get( Thread.currentThread() );
if ( gen == null ) {
// no generator so we are opening new file and writing the start of an array
try {
gen = jsonFactory.createJsonGenerator( new FileOutputStream( fileName ) );
logger.info("Opened output file {}", fileName);
} catch (IOException e) {
throw new RuntimeException("Error opening output file: " + fileName, e);
}
gen.setPrettyPrinter( new MinimalPrettyPrinter(""));
gen.setCodec( mapper );
entityGeneratorsByThread.put( Thread.currentThread(), gen );
}
try {
gen.writeObject( entity );
gen.writeRaw('\n');
entitiesWritten.getAndIncrement();
} catch (IOException e) {
throw new RuntimeException("Error writing to output file: " + fileName, e);
}
}
}
/**
* Writes connection to JSON file.
*/
private class ConnectionWriteAction implements Action1<ExportConnection> {
public void call(ExportConnection conn) {
String [] parts = Thread.currentThread().getName().split("-");
String fileName = outputDir.getAbsolutePath() + File.separator
+ applicationName.replace('/','-') + "-" + parts[3] + ".connections";
JsonGenerator gen = connectionGeneratorsByThread.get( Thread.currentThread() );
if ( gen == null ) {
// no generator so we are opening new file and writing the start of an array
try {
gen = jsonFactory.createJsonGenerator( new FileOutputStream( fileName ) );
logger.info("Opened output file {}", fileName);
} catch (IOException e) {
throw new RuntimeException("Error opening output file: " + fileName, e);
}
gen.setPrettyPrinter( new MinimalPrettyPrinter(""));
gen.setCodec( mapper );
connectionGeneratorsByThread.put( Thread.currentThread(), gen );
}
try {
gen.writeObject( conn );
gen.writeRaw('\n');
connectionsWritten.getAndIncrement();
} catch (IOException e) {
throw new RuntimeException("Error writing to output file: " + fileName, e);
}
}
}
private class FileWrapUpAction implements Action0 {
@Override
public void call() {
logger.info("-------------------------------------------------------------------");
logger.info("DONE! Entities: {} Connections: {}", entitiesWritten.get(), connectionsWritten.get());
logger.info("-------------------------------------------------------------------");
for ( JsonGenerator gen : entityGeneratorsByThread.values() ) {
try {
//gen.writeEndArray();
gen.flush();
gen.close();
} catch (IOException e) {
logger.error("Error closing output file", e);
}
}
for ( JsonGenerator gen : connectionGeneratorsByThread.values() ) {
try {
//gen.writeEndArray();
gen.flush();
gen.close();
} catch (IOException e) {
logger.error("Error closing output file", e);
}
}
}
}
}