Repository: usergrid Updated Branches: refs/heads/master 5fc3c506a -> abec1d9b6
Revamped Export code to be more efficient Export tool now works in parallel for multiple collections Exports can be resumed from last exported edge Export is substantially faster due to parallel processing Connection fetching is optimzed Export now creates multiple dirs for apps and collections Project: http://git-wip-us.apache.org/repos/asf/usergrid/repo Commit: http://git-wip-us.apache.org/repos/asf/usergrid/commit/abec1d9b Tree: http://git-wip-us.apache.org/repos/asf/usergrid/tree/abec1d9b Diff: http://git-wip-us.apache.org/repos/asf/usergrid/diff/abec1d9b Branch: refs/heads/master Commit: abec1d9b60ce39941aee5b363bd5cf0f03e18e11 Parents: 5fc3c50 Author: Keyur Karnik <keyur.karnik+git...@gmail.com> Authored: Tue Nov 6 16:42:40 2018 -0800 Committer: Keyur Karnik <keyur.karnik+git...@gmail.com> Committed: Tue Nov 6 16:42:55 2018 -0800 ---------------------------------------------------------------------- .../java/org/apache/usergrid/tools/Export.java | 942 ++++++++++++------- .../usergrid/tools/ExportingToolBase.java | 151 ++- .../org/apache/usergrid/tools/ToolBase.java | 48 +- 3 files changed, 803 insertions(+), 338 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/usergrid/blob/abec1d9b/stack/tools/src/main/java/org/apache/usergrid/tools/Export.java ---------------------------------------------------------------------- diff --git a/stack/tools/src/main/java/org/apache/usergrid/tools/Export.java b/stack/tools/src/main/java/org/apache/usergrid/tools/Export.java index 821d011..b07d09d 100644 --- a/stack/tools/src/main/java/org/apache/usergrid/tools/Export.java +++ b/stack/tools/src/main/java/org/apache/usergrid/tools/Export.java @@ -17,66 +17,201 @@ package org.apache.usergrid.tools; +import static org.apache.usergrid.management.AccountCreationProps.PROPERTIES_USERGRID_BINARY_UPLOADER; + +import java.io.BufferedOutputStream; +import java.io.BufferedWriter; import java.io.File; +import java.io.FileOutputStream; +import java.io.FileWriter; +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; import java.io.UnsupportedEncodingException; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Set; import java.util.UUID; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; -import com.fasterxml.jackson.core.JsonEncoding; -import com.fasterxml.jackson.core.JsonFactory; -import com.fasterxml.jackson.core.JsonGenerator; -import org.apache.usergrid.persistence.*; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; +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.corepersistence.ManagerCache; +import org.apache.usergrid.corepersistence.export.ExportRequestBuilder; +import org.apache.usergrid.corepersistence.export.ExportRequestBuilderImpl; +import org.apache.usergrid.corepersistence.rx.impl.AllEntityIdsObservable; +import org.apache.usergrid.corepersistence.rx.impl.EdgeScope; +import org.apache.usergrid.corepersistence.util.CpNamingUtils; import org.apache.usergrid.management.OrganizationInfo; import org.apache.usergrid.management.UserInfo; +import org.apache.usergrid.persistence.Entity; +import org.apache.usergrid.persistence.EntityManager; +import org.apache.usergrid.persistence.Results; +import org.apache.usergrid.persistence.SimpleEntityRef; +import org.apache.usergrid.persistence.collection.EntityCollectionManager; +import org.apache.usergrid.persistence.collection.EntityCollectionManagerFactory; +import org.apache.usergrid.persistence.core.scope.ApplicationScope; +import org.apache.usergrid.persistence.graph.Edge; +import org.apache.usergrid.persistence.graph.GraphManager; +import org.apache.usergrid.persistence.graph.SearchByEdgeType; +import org.apache.usergrid.persistence.graph.impl.SimpleEdge; +import org.apache.usergrid.persistence.graph.impl.SimpleSearchByEdgeType; +import org.apache.usergrid.persistence.model.entity.Id; +import org.apache.usergrid.persistence.model.entity.SimpleId; +import org.apache.usergrid.services.assets.BinaryStoreFactory; +import org.apache.usergrid.services.assets.data.BinaryStore; import org.apache.usergrid.tools.bean.ExportOrg; -import org.apache.usergrid.utils.JsonUtils; - -import org.apache.commons.cli.CommandLine; +import org.json.JSONObject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonGenerator; +import com.google.common.base.Optional; import com.google.common.collect.BiMap; +import com.google.gson.Gson; +import com.google.gson.GsonBuilder; + +import rx.Observable; +import rx.observables.ConnectableObservable; +import rx.schedulers.Schedulers; public class Export extends ExportingToolBase { static final Logger logger = LoggerFactory.getLogger( Export.class ); + public static final String LAST_ID = "lastId"; + + + @Autowired + private BinaryStoreFactory binaryStoreFactory; JsonFactory jsonFactory = new JsonFactory(); + + private AllEntityIdsObservable allEntityIdsObs; + private SimpleEdge lastEdge = null; + + //TODO : Add blocking queues for these executors where appropriate + private ExecutorService orgAppCollParallelizer = Executors.newFixedThreadPool(3); + private ExecutorService entityFetcher = Executors.newFixedThreadPool(10); + private ExecutorService enitityMemberFetcher = Executors.newFixedThreadPool(10); + private ExecutorService assetsFetcher = Executors.newFixedThreadPool(10); + + @Override + @SuppressWarnings("static-access") + public Options createOptions() { + + Options options = super.createOptions(); + + Option lastId = OptionBuilder.withArgName( LAST_ID ).hasArg() + .withDescription( "Last Entity Id to resume from" ).create( LAST_ID ); + options.addOption( lastId); + + return options; + } + @Override public void runTool( CommandLine line ) throws Exception { + startSpring(); - setVerbose( line ); - // ExportDataCreator dataCreator = new ExportDataCreator(emf, - // managementService); - // dataCreator.createTestData(); - - applyOrgId( line ); + Gson gson = new GsonBuilder().create(); + + this.allEntityIdsObs = injector.getInstance(AllEntityIdsObservable.class); + applyExportParams(line); prepareBaseOutputFileName( line ); + + + if(lastEdgeJson != null) { + JSONObject lastEdgeJsonObj = new JSONObject(lastEdgeJson); + UUID uuid = UUID.fromString(lastEdgeJsonObj.getJSONObject("sourceNode").getString("uuid")); + Id sourceId = new SimpleId(uuid, lastEdgeJsonObj.getJSONObject("sourceNode").getString("type")); + uuid = UUID.fromString(lastEdgeJsonObj.getJSONObject("targetNode").getString("uuid")); + Id targetId = new SimpleId(uuid, lastEdgeJsonObj.getJSONObject("targetNode").getString("type")); + lastEdge = new SimpleEdge(sourceId, lastEdgeJsonObj.getString("type"), targetId, lastEdgeJsonObj.getLong("timestamp")); + } + outputDir = createOutputParentDir(); logger.info( "Export directory: " + outputDir.getAbsolutePath() ); + // Export organizations separately. exportOrganizations(); + + logger.info("Finished export waiting for threads to end."); + + while(true) { + try { + //Spinning to prevent program execution from ending. + //Need to replace with some kind of countdown latch or task tracker + Thread.sleep(10000); + } catch (InterruptedException e) { + logger.error("Exception while waiting for export to complete.",e); + } + } - // Loop through the organizations - Map<UUID, String> organizations = getOrgs(); - for ( Entry<UUID, String> organization : organizations.entrySet() ) { + } + + + + private void exportOrganizations() throws Exception, UnsupportedEncodingException { - if ( organization.equals( properties.getProperty( "usergrid.test-account.organization" ) ) ) { - // Skip test data from being exported. - continue; - } + for (Entry<UUID, String> organizationName : getOrgs().entrySet()) { + + // Let's skip the test entities. + if (organizationName.equals(properties.getProperty("usergrid.test-account.organization"))) { + continue; + } + + OrganizationInfo orgInfo = managementService.getOrganizationByUuid(organizationName.getKey()); + logger.info("Exporting Organization: " + orgInfo.getName()); + + ExportOrg exportOrg = new ExportOrg(orgInfo); + + List<UserInfo> users = managementService.getAdminUsersForOrganization(organizationName.getKey()); + + for (UserInfo user : users) { + exportOrg.addAdmin(user.getUsername()); + } + + File orgDir = createOrgDir(orgInfo.getName()); + + // One file per Organization. + saveOrganizationMetadata(orgDir, exportOrg); - exportApplicationsForOrg( organization ); + exportApplicationsForOrg(orgDir, organizationName.getKey(), organizationName.getValue()); + } + } + + /** + * Serialize an Organization into a json file. + * @param orgDir + * + * @param acc OrganizationInfo + */ + private void saveOrganizationMetadata( File orgDir, ExportOrg acc ) { + + try { + + File outFile = createOutputFile( orgDir, "organization", acc.getName() ); + com.fasterxml.jackson.core.JsonGenerator jg = getJsonGenerator( outFile ); + jg.writeObject( acc ); + jg.close(); + } + catch ( Exception e ) { + throw new RuntimeException( e ); } } @@ -85,12 +220,18 @@ public class Export extends ExportingToolBase { // Loop through the organizations Map<UUID, String> organizationNames = null; - if ( orgId == null ) { + if ( orgId == null && (orgName == null || orgName.trim().equals(""))) { organizationNames = managementService.getOrganizations(); } - else { - OrganizationInfo info = managementService.getOrganizationByUuid( orgId ); + OrganizationInfo info = null; + + if( orgId != null ) { + info = managementService.getOrganizationByUuid( orgId ); + } + else { + info = managementService.getOrganizationByName( orgName ); + } if ( info == null ) { logger.error( "Organization info is null!" ); @@ -98,333 +239,480 @@ public class Export extends ExportingToolBase { } organizationNames = new HashMap<UUID, String>(); - organizationNames.put( orgId, info.getName() ); + organizationNames.put( info.getUuid(), info.getName() ); } - return organizationNames; } - - - private void exportApplicationsForOrg( Entry<UUID, String> organization ) throws Exception { - logger.info( "" + organization ); + + private void exportApplicationsForOrg(File orgDir, UUID orgId, String orgName ) throws Exception { + + logger.info("Exporting applications for {} : {} ",orgId, orgName); // Loop through the applications per organization - BiMap<UUID, String> applications = managementService.getApplicationsForOrganization( organization.getKey() ); - for ( Entry<UUID, String> application : applications.entrySet() ) { - - logger.info( application.getValue() + " : " + application.getKey() ); - - // Get the JSon serializer. - com.fasterxml.jackson.core.JsonGenerator jg = - getJsonGenerator( createOutputFile( "application", application.getValue() ) ); - - // load the dictionary - EntityManager rootEm = emf.getEntityManager( emf.getManagementAppId() ); - - Entity appEntity = rootEm.get( new SimpleEntityRef( "application", application.getKey())); - - Map<String, Object> dictionaries = new HashMap<String, Object>(); - - for ( String dictionary : rootEm.getDictionaries( appEntity ) ) { - Map<Object, Object> dict = rootEm.getDictionaryAsMap( appEntity, dictionary ); - - // nothing to do - if ( dict.isEmpty() ) { - continue; - } - - dictionaries.put( dictionary, dict ); - } - - EntityManager em = emf.getEntityManager( application.getKey() ); - - // Get application - Entity nsEntity = em.get( new SimpleEntityRef( "application", application.getKey())); - - Set<String> collections = em.getApplicationCollections(); - - // load app counters - - Map<String, Long> entityCounters = em.getApplicationCounters(); - - nsEntity.setMetadata( "organization", organization ); - nsEntity.setMetadata( "dictionaries", dictionaries ); - // counters for collections - nsEntity.setMetadata( "counters", entityCounters ); - nsEntity.setMetadata( "collections", collections ); - - jg.writeStartArray(); - jg.writeObject( nsEntity ); - - // Create a GENERATOR for the application collections. - JsonGenerator collectionsJg = - getJsonGenerator( createOutputFile( "collections", application.getValue() ) ); - collectionsJg.writeStartObject(); - - Map<String, Object> metadata = em.getApplicationCollectionMetadata(); - echo( JsonUtils.mapToFormattedJsonString( metadata ) ); - - // Loop through the collections. This is the only way to loop - // through the entities in the application (former namespace). - for ( String collectionName : metadata.keySet() ) { - - Query query = new Query(); - query.setLimit( MAX_ENTITY_FETCH ); - query.setResultsLevel( Query.Level.ALL_PROPERTIES ); - - Results entities = em.searchCollection( em.getApplicationRef(), collectionName, query ); - - while ( entities.size() > 0 ) { - - for ( Entity entity : entities ) { - // Export the entity first and later the collections for - // this entity. - jg.writeObject( entity ); - echo( entity ); - - saveCollectionMembers( collectionsJg, em, application.getValue(), entity ); - } - - //we're done - if ( entities.getCursor() == null ) { - break; - } - - - query.setCursor( entities.getCursor() ); - - entities = em.searchCollection( em.getApplicationRef(), collectionName, query ); - } - } - - // Close writer for the collections for this application. - collectionsJg.writeEndObject(); - collectionsJg.close(); - - // Close writer and file for this application. - jg.writeEndArray(); - jg.close(); - } - } - - - /** - * Serialize and save the collection members of this <code>entity</code> - * - * @param em Entity Manager - * @param application Application name - * @param entity entity - */ - private void saveCollectionMembers( JsonGenerator jg, EntityManager em, String application, Entity entity ) - throws Exception { - - Set<String> collections = em.getCollections( entity ); - - // Only create entry for Entities that have collections - if ( ( collections == null ) || collections.isEmpty() ) { - return; + BiMap<UUID, String> applications = managementService.getApplicationsForOrganization( orgId ); + + if ( applicationId == null && (applicationName == null || applicationName.trim().equals(""))) { + //export all apps as appId or name is not provided + + Observable.from(applications.entrySet()) + .subscribeOn(Schedulers.from(orgAppCollParallelizer)) + .subscribe(appEntry -> { + UUID appId = appEntry.getKey(); + String appName = appEntry.getValue().split("/")[1]; + try { + exportApplication(orgDir, appId, appName); + } catch (Exception e) { + logger.error("There was an exception exporting application {} : {}",appName, appId, e); + } + }); + } - - jg.writeFieldName( entity.getUuid().toString() ); - jg.writeStartObject(); - - for ( String collectionName : collections ) { - - jg.writeFieldName( collectionName ); - // Start collection array. - jg.writeStartArray(); - - Results collectionMembers = em.getCollection( entity, collectionName, null, 100000, Query.Level.IDS, false ); - - List<UUID> entityIds = collectionMembers.getIds(); - - if ( ( entityIds != null ) && !entityIds.isEmpty() ) { - for ( UUID childEntityUUID : entityIds ) { - jg.writeObject( childEntityUUID.toString() ); - } + else { + + UUID appId = applicationId; + String appName = applicationName; + + if( applicationId != null ) { + appName = applications.get(appId); } + else { + appId = applications.inverse().get(orgName+'/'+appName); + } + + try { + exportApplication(orgDir, appId, appName); + } catch (Exception e) { + logger.error("There was an exception exporting application {} : {}",appName, appId, e); + } - // End collection array. - jg.writeEndArray(); } - - // Write connections - saveConnections( entity, em, jg ); - - // Write dictionaries - saveDictionaries( entity, em, jg ); - - // End the object if it was Started - jg.writeEndObject(); } - - - /** Persists the connection for this entity. */ - private void saveDictionaries( Entity entity, EntityManager em, JsonGenerator jg ) throws Exception { - - jg.writeFieldName( "dictionaries" ); - jg.writeStartObject(); - - Set<String> dictionaries = em.getDictionaries( entity ); - for ( String dictionary : dictionaries ) { - - Map<Object, Object> dict = em.getDictionaryAsMap( entity, dictionary ); + + private void exportApplication(File orgDir, UUID appId, String appName) throws Exception { + + + logger.info( "Starting application export for {} : {} ",appName, appId ); + File appDir = createApplicationDir(orgDir, appName); + + JsonGenerator jg = + getJsonGenerator( createOutputFile( appDir, "application", appName) ); + + // load the dictionary + EntityManager rootEm = emf.getEntityManager( emf.getManagementAppId() ); + + Entity appEntity = rootEm.get( new SimpleEntityRef( "org_application", appId)); + + Map<String, Object> dictionaries = new HashMap<String, Object>(); + + for ( String dictionary : rootEm.getDictionaries( appEntity ) ) { + Map<Object, Object> dict = rootEm.getDictionaryAsMap( appEntity, dictionary ); // nothing to do if ( dict.isEmpty() ) { continue; } - jg.writeFieldName( dictionary ); - - jg.writeStartObject(); - - for ( Entry<Object, Object> entry : dict.entrySet() ) { - jg.writeFieldName( entry.getKey().toString() ); - jg.writeObject( entry.getValue() ); - } - - jg.writeEndObject(); - } - jg.writeEndObject(); - } - - - /** Persists the connection for this entity. */ - private void saveConnections( Entity entity, EntityManager em, JsonGenerator jg ) throws Exception { - - jg.writeFieldName( "connections" ); - jg.writeStartObject(); - - Set<String> connectionTypes = em.getConnectionTypes( entity ); - for ( String connectionType : connectionTypes ) { - - jg.writeFieldName( connectionType ); - jg.writeStartArray(); - - Results results = em.getTargetEntities( - entity, connectionType, null, Query.Level.IDS ); - - List<ConnectionRef> connections = results.getConnections(); - - for ( ConnectionRef connectionRef : connections ) { - jg.writeObject( connectionRef.getTargetRefs().getUuid() ); - } - - jg.writeEndArray(); + dictionaries.put( dictionary, dict ); } - jg.writeEndObject(); - } - - /*- - * Set<String> collections = em.getCollections(entity); - * for (String collection : collections) { - * Results collectionMembers = em.getCollection( - * entity, collection, null, - * MAX_ENTITY_FETCH, Level.IDS, false); - * write entity_id : { "collectionName" : [ids] - * } - * } - * - * - * { - * entity_id : - * { collection_name : - * [ - * collected_entity_id, - * collected_entity_id - * ] - * }, - * f47ac10b-58cc-4372-a567-0e02b2c3d479 : - * { "activtites" : - * [ - * f47ac10b-58cc-4372-a567-0e02b2c3d47A, - * f47ac10b-58cc-4372-a567-0e02b2c3d47B - * ] - * } - * } - * - * http://jackson.codehaus.org/1.8.0/javadoc/org/codehaus/jackson/JsonGenerator.html - * - * - *- - * List<ConnectedEntityRef> connections = em.getConnections(entityId, query); - */ - - - private void exportOrganizations() throws Exception, UnsupportedEncodingException { - - - for ( Entry<UUID, String> organizationName : getOrgs().entrySet() ) { - // Let's skip the test entities. - if ( organizationName.equals( properties.getProperty( "usergrid.test-account.organization" ) ) ) { - continue; - } - - OrganizationInfo acc = managementService.getOrganizationByUuid( organizationName.getKey() ); - logger.info( "Exporting Organization: " + acc.getName() ); - - ExportOrg exportOrg = new ExportOrg( acc ); - - List<UserInfo> users = managementService.getAdminUsersForOrganization( organizationName.getKey() ); - - for ( UserInfo user : users ) { - exportOrg.addAdmin( user.getUsername() ); - } - - // One file per Organization. - saveOrganizationInFile( exportOrg ); - } - } + EntityManager em = emf.getEntityManager( appId); + // Get application + Entity nsEntity = em.get( new SimpleEntityRef( "application", appId)); - /** - * Serialize an Organization into a json file. - * - * @param acc OrganizationInfo - */ - private void saveOrganizationInFile( ExportOrg acc ) { - try { + Set<String> collections = em.getApplicationCollections(); - File outFile = createOutputFile( "organization", acc.getName() ); - com.fasterxml.jackson.core.JsonGenerator jg = getJsonGenerator( outFile ); - jg.writeObject( acc ); - jg.close(); - } - catch ( Exception e ) { - throw new RuntimeException( e ); - } - } + // load app counters + Map<String, Long> entityCounters = em.getApplicationCounters(); - public void streamOutput( File file, List<Entity> entities ) throws Exception { - JsonFactory jsonFactory = new JsonFactory(); - // or, for data binding, - // org.codehaus.jackson.mapper.MappingJsonFactory - JsonGenerator jg = jsonFactory.createJsonGenerator( file, JsonEncoding.UTF8 ); - // or Stream, Reader + //nsEntity.setMetadata( "organization", orgName ); + nsEntity.setMetadata( "dictionaries", dictionaries ); + // counters for collections + nsEntity.setMetadata( "counters", entityCounters ); + nsEntity.setMetadata( "collections", collections ); jg.writeStartArray(); - for ( Entity entity : entities ) { - jg.writeObject( entity ); - } - jg.writeEndArray(); - + jg.writeObject( nsEntity ); jg.close(); - } + + if ( collNames == null || collNames.length <= 0) { + //export all collections as collection names are not provided + + Observable.from(collections) + .subscribeOn(Schedulers.from(orgAppCollParallelizer)) + .subscribe(collectionName -> { + exportCollection(appDir, appId, collectionName, em); + }); + + } + else { + Observable.from(collNames) + .subscribeOn(Schedulers.from(orgAppCollParallelizer)) + .subscribe(collectionName -> { + if(collections.contains(collectionName)) { + exportCollection(appDir, appId, collectionName, em); + } + }); + } + + } + + private void exportCollection(File appDir, UUID appId, String collectionName, EntityManager em) { + File collectionDir = createCollectionDir(appDir, collectionName); + extractEntityIdsForCollection(collectionDir, appId, collectionName); + } + + private void extractEntityIdsForCollection(File collectionDir, UUID applicationId, String collectionName) { + + AtomicInteger batch = new AtomicInteger(1); + + final EntityManager rootEm = emf.getEntityManager(applicationId); + final Gson gson = new GsonBuilder().create(); + ManagerCache managerCache = injector.getInstance(ManagerCache.class); + ExportRequestBuilder builder = new ExportRequestBuilderImpl().withApplicationId(applicationId); + final ApplicationScope appScope = builder.getApplicationScope().get(); + GraphManager gm = managerCache.getGraphManager(appScope); + EntityCollectionManagerFactory entityCollectionManagerFactory = injector + .getInstance(EntityCollectionManagerFactory.class); + final EntityCollectionManager ecm = entityCollectionManagerFactory.createCollectionManager(appScope); + + ExecutorService entityIdWriter = Executors.newFixedThreadPool(1); + allEntityIdsObs + .getEdgesToEntities(Observable.just(CpNamingUtils.getApplicationScope(applicationId)),Optional.fromNullable(CpNamingUtils.getEdgeTypeFromCollectionName(collectionName.toLowerCase())),(this.lastEdge == null ? Optional.absent() : Optional.fromNullable(lastEdge))) + .buffer(1000) + .finallyDo(()-> { + entityIdWriter.shutdown(); + logger.info("Finished fetching entity ids for {}. Shutting down entity id writer executor ", collectionName); + while(!entityIdWriter.isTerminated()) { + try { + entityIdWriter.awaitTermination(10, TimeUnit.SECONDS); + } catch (InterruptedException e) { + } + } + logger.info("Entity id writer executor terminated after shutdown for {}", collectionName); + }) + .subscribe(edges -> { + + logger.info("For collection {}" , collectionName); + Integer batchId = batch.getAndIncrement(); + logger.info("Started fetching details for collection {} batch {} ", collectionName, batchId); + Observable.just(edges) + .subscribeOn(Schedulers.from(entityIdWriter)) + .subscribe(edgeScopes -> { + + List<UUID> entityIds = new ArrayList<UUID>(1000); + + for (EdgeScope edgeScope : edgeScopes) { + // write to file + Id entityId = edgeScope.getEdge().getTargetNode(); + if (entityId != null) { + entityIds.add(entityId.getUuid()); + } else { + edgeScopes.remove(edgeScope); + } + } + // extract name for this batch + try { + writeEntityIdsBatch(collectionDir, edgeScopes, batchId, collectionName); + String type = edgeScopes.get(0).getEdge().getTargetNode().getType(); + + Observable.just(entityIds) + .subscribeOn(Schedulers.from(entityFetcher)) // change to + .subscribe(entIds -> { + + // get entities here + logger.info("entIds count {} for type {}", entIds.size(), type); + Results entities = rootEm.getEntities(entIds, type); + int size = entities.getEntities().size(); + logger.info("Got {} entities.", size); + + if(!skipConnections || !skipDictionaries || !skipAssets) { + + ConnectableObservable<Results> entityObs = Observable.just(entities) + .publish(); + entityObs.subscribeOn(Schedulers.from(enitityMemberFetcher)); + + + // fetch and write connections + if(!skipConnections) { + entityObs.subscribe(entity -> { + fetchConnections(gm, ecm, entity, collectionDir, collectionName, batchId, gson); + + }); + } + // fetch and write dictionaries + if(!skipDictionaries) { + entityObs.subscribe(entity -> { + fetchDictionaries(collectionDir, collectionName, rootEm, + entity, gson, batchId); + }); + } + + if(!skipAssets) { + File assetsDir = createDir(collectionDir.getAbsolutePath(), "files"); + entityObs.subscribe(entity -> { + try { + fetchAssets(assetsDir, applicationId, collectionName, batchId, entities); + } catch (Exception e) { + logger.error("Exception while trying to fetch assets for app {}, collection {}, batch {} ", + applicationId, collectionName, batchId, e); + } + }); + } + entityObs.connect(); + } + writeEntities(collectionDir, entities, batchId, collectionName, gson); + }); + + } catch (Exception e) { + logger.error("There was an error writing entity ids to file for " + + edgeScopes.get(0).getEdge(), e); + // since entity id writing has failed, we need to see how we can not exit the + // whole program + System.exit(0); + } + }); + + logger.info("Finished fetching details for collection {} for batch {}", collectionName, batchId); + }); + logger.info("Exiting extractEntityIdsForCollection() method."); + } + + private void fetchAssets(File assetsDir, UUID applicationId, String collectionName, Integer batchId, + Results entities) throws Exception { + + List<Entity> entitiesWithAssets = new ArrayList<>(); + + for (Entity e : entities.getEntities()) { + if (e.getProperty("file-metadata") != null) { + entitiesWithAssets.add(e); + } + } + + if (!entitiesWithAssets.isEmpty()) { + + writeAssets(assetsDir, collectionName, batchId, entitiesWithAssets); + + ConnectableObservable<Entity> entityAssets = Observable.from(entitiesWithAssets).publish(); + entityAssets.subscribeOn(Schedulers.from(assetsFetcher)); + entityAssets.subscribe(e -> { + // Write code to fetch these assets from entity store. + BinaryStore binaryStore = null; + try { + binaryStore = binaryStoreFactory + .getBinaryStore(properties.getProperty(PROPERTIES_USERGRID_BINARY_UPLOADER)); + } catch (Exception e2) { + logger.error("Except on while trying to get binary store for property {}, ", properties.getProperty(PROPERTIES_USERGRID_BINARY_UPLOADER), e2 ); + } + + File file = new File(assetsDir + "/" + collectionName + "_assets_" + e.getUuid()); + try (InputStream in = binaryStore.read(applicationId, e); + OutputStream out = new BufferedOutputStream(new FileOutputStream(file));) { + + int read = -1; + + while ((read = in.read()) != -1) { + out.write(read); + } + + } catch (Exception e1) { + logger.error("Exception while to write assets file for entity {}", e.getUuid(), e1); + } + + }); + entityAssets.connect(); + } + } + + private void writeAssets(final File collectionDir, final String collectionName, final Integer batchId, + List<Entity> entitiesWithAssets2) { + + try (BufferedWriter assetsWriter = new BufferedWriter( + new FileWriter(new File(collectionDir + "/" + collectionName + "_assets_" + batchId + ".json")));) { + for (Entity e : entitiesWithAssets2) { + JSONObject object = new JSONObject(); + object.put("uuid", e.getUuid()); + object.put("type", e.getType()); + object.put("file-metadata", e.getProperty("file-metadata")); + object.put("file", (e.getProperty("file") != null) ? e.getProperty("file") : null); + assetsWriter.write(object.toString()); + assetsWriter.newLine(); + } + } catch (Exception ex) { + logger.error("Exception while trying to write entities collection {} batch {}", collectionName, batchId, + ex); + } + } + + private void fetchDictionaries(File collectionDir, String collectionName, final EntityManager rootEm, + Results entity, Gson gson, Integer batchId) { + + //TODO : still using JsonGenerator + JsonGenerator jgDictionaries = null; + try { + jgDictionaries = getJsonGenerator(new File(collectionDir + "/" + collectionName + "_" + "dictionaries_" + batchId)); + + for (Entity et : entity.getEntities()) { + Set<String> dictionaries; + try { + dictionaries = rootEm.getDictionaries(et); + + jgDictionaries.writeStartArray(); + if (dictionaries != null && !dictionaries.isEmpty()) { + for (String dictionary : dictionaries) { + Map<Object, Object> dict = rootEm.getDictionaryAsMap(et, dictionary); + if (dict != null && dict.isEmpty()) { + continue; + } + + jgDictionaries.writeStartObject(); + jgDictionaries.writeObjectField(dictionary, dict); + jgDictionaries.writeEndObject(); + } + jgDictionaries.writeEndArray(); + } + } catch (Exception e) { + logger.error("Exception while trying to fetch dictionaries.", e); + } + } + } catch (Exception e) { + logger.error("Exception while trying to fetch dictionaries.", e); + } finally { + if (jgDictionaries != null) { + try { + jgDictionaries.close(); + } catch (IOException e) { + logger.error("Exception while trying to close dictionaries writer.", e); + } + } + } + } + + private void fetchConnections(GraphManager gm, final EntityCollectionManager ecm, Results entity, + File collectionDir, String collectionName, Integer batchId, Gson gson) { + + try(BufferedWriter bufferedWriter = new BufferedWriter( + new FileWriter(new File(collectionDir + "/" + collectionName + "_" + "connections_" + batchId)));){ + + for (Entity et : entity.getEntities()) { + + List<ConnectionPojo> connections = new ArrayList<>(); + + SimpleId id = new SimpleId(); + id.setType(et.getType()); + id.setUuid(et.getUuid()); + + gm.getEdgeTypesFromSource(CpNamingUtils.createConnectionTypeSearch(id)) + .flatMap(emittedEdgeType -> { + logger.debug("loading edges of type {} from node {}", emittedEdgeType, id); + return gm.loadEdgesFromSource(new SimpleSearchByEdgeType(id, emittedEdgeType, + Long.MAX_VALUE, SearchByEdgeType.Order.DESCENDING, Optional.absent())); + }).map(markedEdge -> { + + if (!markedEdge.isDeleted() && !markedEdge.isTargetNodeDeleted() + && markedEdge.getTargetNode() != null) { + + // doing the load to just again make sure bad + // connections are not exported + org.apache.usergrid.persistence.model.entity.Entity en = ecm + .load(markedEdge.getTargetNode()).toBlocking().lastOrDefault(null); + + if (en != null) { + + try { + + ConnectionPojo connectionPojo = new ConnectionPojo(); + connectionPojo.setRelationship(CpNamingUtils + .getConnectionNameFromEdgeName(markedEdge.getType())); + connectionPojo.setSourceNodeUUID(markedEdge.getSourceNode().getUuid().toString()); + connectionPojo.setTargetNodeUUID(markedEdge.getTargetNode().getUuid().toString()); + + connections.add(connectionPojo); + + } catch (Exception e) { + logger.error("Exception while trying process connection entity", e); + } + } else { + logger.warn( + "Exported connection has a missing target node, not creating connection in export. Edge: {}", + markedEdge); + } + } + return null; + + }).toBlocking().lastOrDefault(null); + + for(ConnectionPojo c : connections) { + bufferedWriter.write(gson.toJson(c)); + bufferedWriter.newLine(); + } + } + }catch (Exception e) { + logger.error("Exception while trying to write connection to file.", e); + } + logger.info("Finished fetching details for collection {} batch {}", collectionName, batchId); + } + + + class ConnectionPojo { + private String sourceNodeUUID; + private String relationship; + private String targetNodeUUID; + public String getSourceNodeUUID() { + return sourceNodeUUID; + } + public void setSourceNodeUUID(String sourceNodeUUID) { + this.sourceNodeUUID = sourceNodeUUID; + } + public String getRelationship() { + return relationship; + } + public void setRelationship(String relationship) { + this.relationship = relationship; + } + public String getTargetNodeUUID() { + return targetNodeUUID; + } + public void setTargetNodeUUID(String targetNodeUUID) { + this.targetNodeUUID = targetNodeUUID; + } + + } + + private void writeEntities(File collectionDir, Results entities, Integer batchId, String collectionName, Gson gson) { + logger.info("Started writing entities for collection {} batch {} ", collectionName, batchId); + + try(BufferedWriter bufferedWriter = new BufferedWriter( + new FileWriter(new File(collectionDir + "/" + collectionName + "_data_" + batchId + ".json")));) { + + logger.info("Got count {} entities for file writing", entities.getEntities().size()); + for(Entity e : entities.getEntities()) { + bufferedWriter.write(gson.toJson(e)); + bufferedWriter.newLine(); + } + + } catch (Exception e) { + logger.error("Exception while trying to write entities collection {} batch {}", collectionName, batchId, e); + } + + logger.info("Finised writing entities for collection {} batch {} ", collectionName, batchId); + } + + private void writeEntityIdsBatch(File collectionDir, List<EdgeScope> edgeScopes, Integer batchId, + String collectionName) throws Exception { + logger.info("Started writing ids for collection {} batch {} ", collectionName, batchId); + try (BufferedWriter bufferedWriter = new BufferedWriter( + new FileWriter(new File(collectionDir + "/" + collectionName + "_" + batchId)));) { + for (EdgeScope es : edgeScopes) { + bufferedWriter.write(es.getEdge().toString()); + bufferedWriter.newLine(); + } + } catch (Exception e) { + logger.error("Exception while tryign to write entity ids for collection {} batch {}", collectionName, batchId, e); + } + logger.info("Finished writing ids for collection {} batch {} ", collectionName, batchId); + } - // to generate the activities and user relationship, follow this: - - // write field name (id) - // write start object - // write field name (collection name) - // write start array - // write object/string - // write another object - // write end array - // write end object - // ...... more objects - // } http://git-wip-us.apache.org/repos/asf/usergrid/blob/abec1d9b/stack/tools/src/main/java/org/apache/usergrid/tools/ExportingToolBase.java ---------------------------------------------------------------------- diff --git a/stack/tools/src/main/java/org/apache/usergrid/tools/ExportingToolBase.java b/stack/tools/src/main/java/org/apache/usergrid/tools/ExportingToolBase.java index 6fb0911..5865d15 100644 --- a/stack/tools/src/main/java/org/apache/usergrid/tools/ExportingToolBase.java +++ b/stack/tools/src/main/java/org/apache/usergrid/tools/ExportingToolBase.java @@ -26,11 +26,15 @@ import com.fasterxml.jackson.core.JsonFactory; import com.fasterxml.jackson.core.JsonGenerator; import com.fasterxml.jackson.core.util.DefaultPrettyPrinter; import com.fasterxml.jackson.databind.ObjectMapper; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.apache.usergrid.persistence.graph.impl.SimpleEdge; +import org.apache.usergrid.services.assets.BinaryStoreFactory.Provider; import org.apache.usergrid.utils.ConversionUtils; import org.apache.commons.cli.CommandLine; +import org.apache.commons.cli.MissingOptionException; import org.apache.commons.cli.Option; import org.apache.commons.cli.OptionBuilder; import org.apache.commons.cli.Options; @@ -39,7 +43,6 @@ import org.apache.commons.cli.Options; /** * Base class for ToolBase implementations that write output to an output directory and file. * - * @author zznate */ public abstract class ExportingToolBase extends ToolBase { @@ -49,9 +52,31 @@ public abstract class ExportingToolBase extends ToolBase { /** Output dir option: -outputDir */ protected static final String OUTPUT_DIR = "outputDir"; + protected static final String ORG_ID = "orgid"; + protected static final String ORG_NAME = "orgName"; + protected static final String APP_ID = "appId"; + protected static final String APP_NAME = "appName"; + protected static final String COLL_NAMES = "collNames"; + protected static final String APPEND_TIMESTAMP = "appendTimestamp"; + protected static final String SKIP_CONN = "skipConnections"; + protected static final String SKIP_DICT = "skipDictionaries"; + protected static final String LAST_EDGE = "lastEdge"; + protected static final String SKIP_ASSETS = "skipAssets"; + protected static final String FIELD_TYPE = "fieldType"; + protected static final String COLLECTION_NAME = "collectionName"; protected String baseOutputDirName = "export"; protected UUID orgId; + protected String orgName; + protected UUID applicationId; + protected String applicationName; + protected String[] collNames; + protected String fieldType; + protected boolean skipConnections = false; + protected boolean skipDictionaries = false; + protected boolean skipAssets = false; + protected String lastEdgeJson = null; + JsonFactory jsonFactory = new JsonFactory(); protected long startTime = System.currentTimeMillis(); @@ -62,11 +87,41 @@ public abstract class ExportingToolBase extends ToolBase { Options options = super.createOptions(); - Option outputDir = OptionBuilder.hasArg().withDescription( "output file name -outputDir" ).create( OUTPUT_DIR ); - Option orgId = OptionBuilder.hasArg().withDescription( "Use a specific organization -orgId" ).create( "orgId" ); - + Option outputDir = OptionBuilder.hasArg().withDescription( "Output file name -outputDir" ).create( OUTPUT_DIR ); + Option orgId = OptionBuilder.hasArg().withDescription( "Use a specific organization -orgId" ).create( ORG_ID ); + Option appId = OptionBuilder.hasArg().withDescription( "Use a specific application -appId (Needs -orgId or -orgName)" ).create( APP_ID ); + Option orgName = OptionBuilder.hasArg().withDescription( "Use a specific organization name -orgName" ).create( ORG_NAME ); + Option appName = OptionBuilder.hasArg().withDescription( "Use a specific application name -appName (Needs -orgId or -orgName)" ).create( APP_NAME ); + Option collNames = OptionBuilder.hasArg().withDescription( "Export list of comma separated collections -collNames (Needs -orgId or -orgName and -appId or -appName)" ).create( COLL_NAMES ); + Option appendTimestamp = OptionBuilder.withDescription( "Attach timestamp to output directory -appendTimestamp" ).create( APPEND_TIMESTAMP ); + Option skipConns = OptionBuilder.withDescription( "Skip exporting connections for entities -skipConnections" ).create( SKIP_CONN ); + Option skipDicts = OptionBuilder.withDescription( "Skip exporting dictionaries for entities -skipDictionaries" ).create( SKIP_DICT ); + Option lastEdge = OptionBuilder.hasArg().withDescription( "Last Edge from previous run to resume export -lastEdge" ).create( LAST_EDGE ); + Option skipAssets = OptionBuilder.withDescription( "Skip exporting assets for entities -skipAssets" ).create( SKIP_ASSETS ); + Option awsKey = OptionBuilder.hasArg().withDescription( "AWS access key -awsKey" ).create( AWS_KEY ); + Option storeType = OptionBuilder.hasArg().withDescription( "Binary store type -storeType (aws, google, local)" ).create( STORE_TYPE ); + Option awsId = OptionBuilder.hasArg().withDescription( "AWS access id -awsId" ).create( AWS_ID ); + Option bucketName = OptionBuilder.hasArg().withDescription( "Binary storage bucket name -bucketName" ).create( BINARY_BUCKET_NAME ); + Option fieldType = OptionBuilder.hasArg().withDescription( "Field type for unique value check -fieldType" ).create( FIELD_TYPE ); + Option collectionName = OptionBuilder.hasArg().withDescription( "Collection name for unique value check -collectionName" ).create( COLLECTION_NAME ); + options.addOption( outputDir ); options.addOption( orgId ); + options.addOption( appId ); + options.addOption( collNames ); + options.addOption( appName ); + options.addOption( orgName ); + options.addOption( appendTimestamp ); + options.addOption( skipConns ); + options.addOption( skipDicts ); + options.addOption( lastEdge ); + options.addOption( skipAssets ); + options.addOption( awsKey ); + options.addOption( awsId ); + options.addOption( bucketName ); + options.addOption( storeType ); + options.addOption(fieldType); + options.addOption(collectionName); return options; } @@ -75,10 +130,15 @@ public abstract class ExportingToolBase extends ToolBase { protected void prepareBaseOutputFileName( CommandLine line ) { boolean hasOutputDir = line.hasOption( OUTPUT_DIR ); + boolean appendTimestamp = line.hasOption( APPEND_TIMESTAMP ); if ( hasOutputDir ) { baseOutputDirName = line.getOptionValue( OUTPUT_DIR ); } + + if(appendTimestamp) { + baseOutputDirName = baseOutputDirName + "_"+startTime; + } } @@ -87,6 +147,71 @@ public abstract class ExportingToolBase extends ToolBase { orgId = ConversionUtils.uuid( line.getOptionValue( "orgId" ) ); } } + + + protected void applyExportParams( CommandLine line ) { + + if ( line.hasOption( ORG_ID ) ) { + orgId = ConversionUtils.uuid( line.getOptionValue( ORG_ID ) ); + } + else if ( line.hasOption( ORG_NAME ) ) { + orgName = line.getOptionValue( ORG_NAME ) ; + } + + if ( line.hasOption( APP_ID ) ) { + applicationId = ConversionUtils.uuid( line.getOptionValue( APP_ID ) ); + } + else if ( line.hasOption( APP_NAME ) ) { + applicationName = line.getOptionValue( APP_NAME ) ; + } + if ( line.hasOption( COLL_NAMES ) ) { + collNames = line.getOptionValue( COLL_NAMES ).split(","); + } + if(line.hasOption( COLLECTION_NAME )) { + collNames = new String[] {line.getOptionValue( COLLECTION_NAME )}; + } + skipConnections = line.hasOption( SKIP_CONN ); + skipDictionaries = line.hasOption( SKIP_DICT ); + + if(line.hasOption(LAST_EDGE)) { + lastEdgeJson = line.getOptionValue(LAST_EDGE); + } + + skipAssets = line.hasOption( SKIP_ASSETS ); + + if(line.hasOption( FIELD_TYPE )) { + fieldType = line.getOptionValue( FIELD_TYPE ); + } + } + + protected void validateOptions(CommandLine line) throws MissingOptionException { + if ((line.hasOption(APP_ID) || line.hasOption(APP_NAME)) + && !(line.hasOption(ORG_ID) || line.hasOption(ORG_NAME))) { + throw new MissingOptionException("-orgId or -orgName is required if you pass -appId or -appName"); + } + if (line.hasOption(COLL_NAMES) && !(line.hasOption(APP_ID) || line.hasOption(APP_NAME))) { + throw new MissingOptionException( + "[-appId or -appName] and [-orgId or -orgName] are required if you pass -collNames"); + } + + if (!line.hasOption(SKIP_ASSETS)) { + if (line.hasOption(STORE_TYPE)) { + String storeType = line.getOptionValue(STORE_TYPE); + if (storeType.equals(Provider.aws.toString())) { + if (!line.hasOption(AWS_ID) || !line.hasOption(AWS_KEY) || !line.hasOption(BINARY_BUCKET_NAME)) { + throw new MissingOptionException( + "[-awsId and -awsKey and -bucketName] are required if you pass -storeType as aws"); + } + } else if (storeType.equals(Provider.google.toString())) { + if (!line.hasOption(BINARY_BUCKET_NAME)) { + throw new MissingOptionException("[-bucketName] is required if you pass -storeType as google"); + } + } + } else { + throw new MissingOptionException("[-storeType] is required if you do not pass -skipAssets"); + } + } + } /** @@ -116,10 +241,24 @@ public abstract class ExportingToolBase extends ToolBase { protected File createOutputFile( File parent, String type, String name ) { return new File( parent, prepareOutputFileName( type, name ) ); } + + protected File createOrgDir( String orgName ) { + return createDir( outputDir.getAbsolutePath(), orgName); + } + + protected File createApplicationDir( File orgDir, String applicationName ) { + return createDir( orgDir.getAbsolutePath(), applicationName); + } - protected File createCollectionsDir( String applicationName ) { - return createDir( String.format( "%s/%s.applicationName.collections", outputDir, applicationName ) ); + protected File createCollectionDir( File appDir, String collectionName ) { + return createDir( appDir.getAbsolutePath(), collectionName) ; + } + + + protected File createDir( String baseDir, String dirName ) { + + return createDir((baseDir!=null && !baseDir.trim().equals(""))?(baseDir + File.separator +dirName): dirName); } http://git-wip-us.apache.org/repos/asf/usergrid/blob/abec1d9b/stack/tools/src/main/java/org/apache/usergrid/tools/ToolBase.java ---------------------------------------------------------------------- diff --git a/stack/tools/src/main/java/org/apache/usergrid/tools/ToolBase.java b/stack/tools/src/main/java/org/apache/usergrid/tools/ToolBase.java index 97a7d5a..e2711bb 100644 --- a/stack/tools/src/main/java/org/apache/usergrid/tools/ToolBase.java +++ b/stack/tools/src/main/java/org/apache/usergrid/tools/ToolBase.java @@ -71,9 +71,11 @@ public abstract class ToolBase { protected CassandraService cass; protected Injector injector; - - - + + protected static final String AWS_KEY = "awsKey"; + protected static final String AWS_ID = "awsId"; + protected static final String BINARY_BUCKET_NAME = "bucketName"; + protected static final String STORE_TYPE = "storeType"; public void startTool( String[] args ) { startTool( args, true ); @@ -93,6 +95,15 @@ public abstract class ToolBase { if ( line == null ) { return; } + + + try { + validateOptions(line); + } + catch ( MissingOptionException exp ) { + printCliHelp( "Required or dependent options are missing. Reason: " + exp.getMessage() ); + } + // notification queue listener not needed for tools System.setProperty("usergrid.notifications.listener.run", "false"); @@ -127,6 +138,21 @@ public abstract class ToolBase { System.setProperty( "cassandra.lock.keyspace", line.getOptionValue( "lockskeyspace" ) ); } + if(line.hasOption(AWS_ID)) { + System.setProperty("AWS_ACCESS_KEY_ID", line.getOptionValue( AWS_ID )); + } + + if(line.hasOption(AWS_KEY)) { + System.setProperty("AWS_SECRET_KEY", line.getOptionValue( AWS_KEY )); + } + + if(line.hasOption(BINARY_BUCKET_NAME)) { + System.setProperty("usergrid.binary.bucketname", line.getOptionValue( BINARY_BUCKET_NAME )); + } + + if(line.hasOption( STORE_TYPE )) { + System.setProperty("usergrid.binary.uploader", line.getOptionValue( STORE_TYPE )); + } try { runTool( line ); @@ -193,6 +219,11 @@ public abstract class ToolBase { return options; } + + + protected void validateOptions(CommandLine line) throws MissingOptionException { + + } public void startEmbedded() throws Exception { @@ -304,7 +335,10 @@ public abstract class ToolBase { " cassandra.connections: {}\n" + " usergrid.notifications.listener.run: {}\n" + " usergrid.push.worker_count: {}\n" + - " usergrid.scheduler.enabled: {}\n", + " usergrid.scheduler.enabled: {}\n" + + " cassandra.readcl: {}\n" + + " usergrid.read.cl: {}\n" + + " usergrid.binary.uploader: {}\n", properties.get("cassandra.url"), properties.get("cassandra.datacenter.local"), properties.get("cassandra.username"), @@ -316,7 +350,11 @@ public abstract class ToolBase { properties.get("cassandra.connections"), properties.get("usergrid.notifications.listener.run"), properties.get("usergrid.push.worker_count"), - properties.get("usergrid.scheduler.enabled") + properties.get("usergrid.scheduler.enabled"), + properties.get("cassandra.readcl"), + properties.get("usergrid.read.cl"), + properties.get("usergrid.binary.uploader") + ); }
