/*
* 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 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.commons.io.filefilter.PrefixFileFilter;
import org.apache.usergrid.corepersistence.util.CpNamingUtils;
import org.apache.usergrid.management.OrganizationInfo;
import org.apache.usergrid.management.UserInfo;
import org.apache.usergrid.persistence.EntityManager;
import org.apache.usergrid.persistence.EntityRef;
import org.apache.usergrid.persistence.SimpleEntityRef;
import org.apache.usergrid.persistence.entities.User;
import org.apache.usergrid.persistence.exceptions.DuplicateUniquePropertyExistsException;
import org.codehaus.jackson.JsonFactory;
import org.codehaus.jackson.JsonParser;
import org.codehaus.jackson.JsonToken;
import org.codehaus.jackson.map.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import static org.apache.usergrid.persistence.Schema.PROPERTY_TYPE;
import static org.apache.usergrid.persistence.Schema.PROPERTY_UUID;
import org.apache.usergrid.persistence.index.query.Identifier;
/**
* Import Admin Users and metadata including organizations and passwords.
*
* Usage Example:
*
* java -Xmx8000m -Dlog4j.configuration=file:/home/me/log4j.properties -classpath . \
* -jar usergrid-tools-1.0.2.jar ImportAdmins -writeThreads 100 -auditThreads 100 \
* -host casshost -inputDir=/home/me/import-data
*
* If you want to provide any property overrides, put properties file named usergrid-custom-tools.properties
* in the same directory where you run the above command. For example, you might want to set the Cassandra
* client threads and import to a specific set of keyspaces:
*
* cassandra.connections=110
* cassandra.system.keyspace=My_Other_Usergrid
* cassandra.application.keyspace=My_Other_Usergrid_Applications
* cassandra.lock.keyspace=My_Other_Usergrid_Locks
*/
public class ImportAdmins extends ToolBase {
private static final Logger logger = LoggerFactory.getLogger(ImportAdmins.class);
/**
* Input directory where the .json export files are
*/
static final String INPUT_DIR = "inputDir";
static final String WRITE_THREAD_COUNT = "writeThreads";
static final String AUDIT_THREAD_COUNT = "auditThreads";
static File importDir;
static final String DEFAULT_INPUT_DIR = "export";
private Map<Stoppable, Thread> adminWriteThreads = new HashMap<Stoppable, Thread>();
private Map<Stoppable, Thread> adminAuditThreads = new HashMap<Stoppable, Thread>();
private Map<Stoppable, Thread> metadataWorkerThreadMap = new HashMap<Stoppable, Thread>();
Map<UUID, DuplicateUser> dupsByDupUuid = new HashMap<UUID, DuplicateUser>(200);
JsonFactory jsonFactory = new JsonFactory();
AtomicInteger userCount = new AtomicInteger( 0 );
AtomicInteger metadataCount = new AtomicInteger( 0 );
AtomicInteger writeEmptyCount = new AtomicInteger( 0 );
AtomicInteger auditEmptyCount = new AtomicInteger( 0 );
AtomicInteger metadataEmptyCount = new AtomicInteger( 0 );
static class DuplicateUser {
String email;
String username;
public DuplicateUser( String propName, Map<String, Object> user ) {
if ( "email".equals(propName)) {
email = user.get("email").toString();
} else {
username = user.get("username").toString();
}
}
}
@Override
@SuppressWarnings("static-access")
public Options createOptions() {
// inherit parent options
Options options = super.createOptions();
Option inputDir = OptionBuilder
.hasArg()
.withDescription("input directory -inputDir").create(INPUT_DIR);
Option writeThreads = OptionBuilder
.hasArg()
.withDescription("Write Threads -writeThreads").create(WRITE_THREAD_COUNT);
Option auditThreads = OptionBuilder
.hasArg()
.withDescription("Audit Threads -auditThreads").create(AUDIT_THREAD_COUNT);
Option verbose = OptionBuilder
.withDescription("Print on the console an echo of the content written to the file")
.create(VERBOSE);
options.addOption( writeThreads );
options.addOption( auditThreads );
options.addOption( inputDir );
options.addOption( verbose );
return options;
}
@Override
public void runTool(CommandLine line) throws Exception {
startSpring();
setVerbose(line);
openImportDirectory(line);
int auditThreadCount = 1;
int writeThreadCount = 1;
if (line.hasOption(AUDIT_THREAD_COUNT)) {
auditThreadCount = Integer.parseInt(line.getOptionValue(AUDIT_THREAD_COUNT));
}
if (line.hasOption(WRITE_THREAD_COUNT)) {
writeThreadCount = Integer.parseInt( line.getOptionValue(WRITE_THREAD_COUNT));
}
importAdminUsers( writeThreadCount, auditThreadCount );
importMetadata( writeThreadCount );
}
/**
* Import admin users.
*/
private void importAdminUsers(int writeThreadCount, int auditThreadCount) throws Exception {
String[] fileNames = importDir.list(new PrefixFileFilter(ExportAdmins.ADMIN_USERS_PREFIX + "."));
logger.info( "Applications to read: " + fileNames.length );
for (String fileName : fileNames) {
try {
importAdminUsers(fileName, writeThreadCount, auditThreadCount);
} catch (Exception e) {
logger.warn("Unable to import application: " + fileName, e);
}
}
}
/**
* Imports admin users.
*
* @param fileName Name of admin user data file.
*/
private void importAdminUsers(final String fileName,
final int writeThreadCount,
final int auditThreadCount) throws Exception {
int count = 0;
File adminUsersFile = new File(importDir, fileName);
logger.info("----- Loading file: " + adminUsersFile.getAbsolutePath());
JsonParser jp = getJsonParserForFile(adminUsersFile);
int loopCounter = 0;
BlockingQueue<Map<String, Object>> workQueue = new LinkedBlockingQueue<Map<String, Object>>();
BlockingQueue<Map<String, Object>> auditQueue = new LinkedBlockingQueue<Map<String, Object>>();
startAdminWorkers(workQueue, auditQueue, writeThreadCount);
startAdminAuditors(auditQueue, auditThreadCount);
JsonToken token = jp.nextToken();
validateStartArray(token);
while (jp.nextValue() != JsonToken.END_ARRAY) {
loopCounter += 1;
@SuppressWarnings("unchecked")
Map<String, Object> entityProps = jp.readValueAs(HashMap.class);
if (loopCounter % 1000 == 0) {
logger.debug( "Publishing to queue... counter=" + loopCounter );
}
workQueue.add( entityProps );
}
waitForQueueAndMeasure(workQueue, writeEmptyCount, adminWriteThreads, "Admin Write");
waitForQueueAndMeasure(auditQueue, auditEmptyCount, adminAuditThreads, "Admin Audit");
logger.info("----- End: Imported {} admin users from file {}",
count, adminUsersFile.getAbsolutePath());
jp.close();
}
private static void waitForQueueAndMeasure(final BlockingQueue workQueue,
final AtomicInteger emptyCounter,
final Map<Stoppable, Thread> threadMap,
final String identifier) throws InterruptedException {
double rateAverageSum = 0;
int iterations = 0;
while ( emptyCounter.get() < threadMap.size() ) {
iterations += 1;
int sizeLast = workQueue.size();
long lastTime = System.currentTimeMillis();
logger.info("Queue {} is not empty, remaining size={}, waiting...", identifier, sizeLast);
Thread.sleep(10000);
long timeNow = System.currentTimeMillis();
int sizeNow = workQueue.size();
int processed = sizeLast - sizeNow;
long timeDelta = timeNow - lastTime;
double rateLast = (double) processed / (timeDelta / 1000);
rateAverageSum += rateLast;
long timeRemaining = (long) ( sizeLast / (rateAverageSum / iterations) );
logger.info("++PROGRESS ({}): sizeLast={} nowSize={} processed={} rateLast={}/s rateAvg={}/s timeRemaining={}s",
new Object[] {
identifier, sizeLast, sizeNow, processed, rateLast, (rateAverageSum / iterations), timeRemaining } );
}
for (Stoppable worker : threadMap.keySet()) {
worker.setDone(true);
}
}
private void startAdminAuditors(BlockingQueue<Map<String, Object>> auditQueue, int workerCount) {
for (int x = 0; x < workerCount; x++) {
AuditWorker worker = new AuditWorker(auditQueue);
Thread workerThread = new Thread(worker, "AdminAuditor-" + x);
workerThread.start();
adminAuditThreads.put(worker, workerThread);
}
logger.info("Started {} admin auditors", workerCount);
}
private void startAdminWorkers(BlockingQueue<Map<String, Object>> workQueue,
BlockingQueue<Map<String, Object>> auditQueue,
int workerCount) {
for (int x = 0; x < workerCount; x++) {
ImportAdminWorker worker = new ImportAdminWorker(workQueue, auditQueue);
Thread workerThread = new Thread(worker, "AdminWriter-" + x);
workerThread.start();
adminWriteThreads.put(worker, workerThread);
}
logger.info("Started {} admin workers", workerCount);
}
private String getType(Map<String, Object> entityProps) {
return (String) entityProps.get(PROPERTY_TYPE);
}
private UUID getId(Map<String, Object> entityProps) {
return UUID.fromString((String) entityProps.get(PROPERTY_UUID));
}
private void validateStartArray(JsonToken token) {
if (token != JsonToken.START_ARRAY) {
throw new RuntimeException("Token should be START ARRAY but it is:" + token.asString());
}
}
private JsonParser getJsonParserForFile(File organizationFile) throws Exception {
JsonParser jp = jsonFactory.createJsonParser( organizationFile );
jp.setCodec( new ObjectMapper() );
return jp;
}
/**
* Import collections. Collections files are named: collections.<application_name>.Timestamp.json
*/
private void importMetadata(int writeThreadCount) throws Exception {
String[] fileNames = importDir.list(
new PrefixFileFilter( ExportAdmins.ADMIN_USER_METADATA_PREFIX + "." ) );
logger.info( "Metadata files to read: " + fileNames.length );
for (String fileName : fileNames) {
try {
importMetadata(fileName, writeThreadCount);
} catch (Exception e) {
logger.warn("Unable to import metadata file: " + fileName, e);
}
}
}
private void startMetadataWorkers(BlockingQueue<ImportMetadataTask> workQueue, int writeThreadCount) {
for (int x = 0; x < writeThreadCount; x++) {
ImportMetadataWorker worker = new ImportMetadataWorker(workQueue);
Thread workerThread = new Thread(worker, "ImportMetadataTask-" + x );
workerThread.start();
metadataWorkerThreadMap.put(worker, workerThread);
}
logger.info( "Started {} metadata workers", writeThreadCount );
}
@SuppressWarnings("unchecked")
private void importMetadata(String fileName, int writeThreads) throws Exception {
EntityManager em = emf.getEntityManager( CpNamingUtils.MANAGEMENT_APPLICATION_ID);
File metadataFile = new File(importDir, fileName);
logger.info("----- Loading metadata file: " + metadataFile.getAbsolutePath());
JsonParser jp = getJsonParserForFile(metadataFile);
JsonToken jsonToken = null; // jp.nextToken();// START_OBJECT this is the outer hashmap
int depth = 1;
BlockingQueue<ImportMetadataTask> workQueue = new LinkedBlockingQueue<ImportMetadataTask>();
startMetadataWorkers(workQueue, writeThreads);
while (depth > 0) {
jsonToken = jp.nextToken();
if (jsonToken == null) {
logger.info("token is null, breaking");
break;
}
if (jsonToken.equals(JsonToken.START_OBJECT)) {
depth++;
} else if (jsonToken.equals(JsonToken.END_OBJECT)) {
depth--;
}
if (jsonToken.equals(JsonToken.FIELD_NAME) && depth == 2) {
jp.nextToken();
String entityOwnerId = jp.getCurrentName();
try {
EntityRef entityRef = new SimpleEntityRef( "user", UUID.fromString( entityOwnerId ) );
Map<String, Object> metadata = (Map<String, Object>) jp.readValueAs( Map.class );
workQueue.put( new ImportMetadataTask( entityRef, metadata ) );
logger.debug( "Put user {} in metadata queue", entityRef.getUuid() );
} catch ( Exception e ) {
logger.debug( "Error with user {}, not putting in metadata queue", entityOwnerId );
}
}
}
waitForQueueAndMeasure(workQueue, metadataEmptyCount, metadataWorkerThreadMap, "Metadata Load");
logger.info("----- End of metadata -----");
jp.close();
}
/**
* Imports the entity's connecting references (collections and connections)
*/
@SuppressWarnings("unchecked")
private void importEntityMetadata(
EntityManager em, EntityRef entityRef, Map<String, Object> metadata) throws Exception {
DuplicateUser dup = dupsByDupUuid.get( entityRef.getUuid() );
if ( dup == null ) { // not a duplicate
User user = em.get( entityRef, User.class );
final UserInfo userInfo = managementService.getAdminUserByEmail( user.getEmail() );
if (user == null || userInfo == null) {
logger.error( "User {} does not exist, not processing metadata", entityRef.getUuid() );
return;
}
List<Object> organizationsList = (List<Object>) metadata.get("organizations");
if (organizationsList != null && !organizationsList.isEmpty()) {
for (Object orgObject : organizationsList) {
Map<String, Object> orgMap = (Map<String, Object>) orgObject;
UUID orgUuid = UUID.fromString( (String) orgMap.get( "uuid" ) );
String orgName = (String) orgMap.get( "name" );
OrganizationInfo orgInfo = managementService.getOrganizationByUuid( orgUuid );
if (orgInfo == null) { // org does not exist yet, create it and add user
try {
managementService.createOrganization( orgUuid, orgName, userInfo, false );
orgInfo = managementService.getOrganizationByUuid( orgUuid );
logger.debug( "Created new org {} for user {}",
new Object[]{orgInfo.getName(), user.getEmail()} );
} catch (DuplicateUniquePropertyExistsException dpee) {
logger.debug( "Org {} already exists", orgName );
}
} else { // org exists, add original user to it
try {
managementService.addAdminUserToOrganization( userInfo, orgInfo, false );
logger.debug( "Added to org user {}:{}:{}",
new Object[]{
orgInfo.getName(),
user.getUsername(),
user.getEmail(),
user.getUuid()
});
} catch (Exception e) {
logger.error( "Error Adding user {} to org {}", new Object[]{user.getEmail(), orgName} );
}
}
}
}
Map<String, Object> dictionariesMap = (Map<String, Object>) metadata.get("dictionaries");
if (dictionariesMap != null && !dictionariesMap.isEmpty()) {
for (String name : dictionariesMap.keySet()) {
try {
Map<String, Object> dictionary = (Map<String, Object>) dictionariesMap.get(name);
em.addMapToDictionary( entityRef, name, dictionary);
logger.debug( "Creating dictionary for {} name {}",
new Object[]{entityRef, name} );
} catch (Exception e) {
if (logger.isDebugEnabled()) {
logger.error("Error importing dictionary name "
+ name + " for user " + entityRef.getUuid(), e);
} else {
logger.error("Error importing dictionary name "
+ name + " for user " + entityRef.getUuid());
}
}
}
} else {
logger.warn("User {} has no dictionaries", entityRef.getUuid() );
}
} else { // let the DuplicateAdminRepair tool handle merging of admins
logger.info("Not processing duplicate username={} email={}", dup.email, dup.username );
}
}
/**
* Open up the import directory based on <code>importDir</code>
*/
private void openImportDirectory(CommandLine line) {
boolean hasInputDir = line.hasOption(INPUT_DIR);
if (hasInputDir) {
importDir = new File(line.getOptionValue(INPUT_DIR));
} else {
importDir = new File(DEFAULT_INPUT_DIR);
}
logger.info("Importing from:" + importDir.getAbsolutePath());
logger.info("Status. Exists: " + importDir.exists() + " - Readable: " + importDir.canRead());
}
interface Stoppable {
void setDone(boolean done);
}
class AuditWorker implements Runnable, Stoppable {
private BlockingQueue<Map<String, Object>> workQueue;
private boolean done;
public AuditWorker(BlockingQueue<Map<String, Object>> workQueue) {
this.workQueue = workQueue;
}
@Override
public void setDone(boolean done) {
this.done = done;
}
@Override
public void run() {
int count = 0;
EntityManager em = emf.getEntityManager(CpNamingUtils.MANAGEMENT_APPLICATION_ID);
long durationSum = 0;
while (!done) {
try {
Map<String, Object> entityProps = this.workQueue.poll(30, TimeUnit.SECONDS);
if (entityProps == null) {
logger.warn("Reading from AUDIT queue was null!");
auditEmptyCount.getAndIncrement();
Thread.sleep(1000);
continue;
}
auditEmptyCount.set(0);
count++;
long startTime = System.currentTimeMillis();
UUID uuid = (UUID) entityProps.get(PROPERTY_UUID);
String type = getType(entityProps);
if (em.get(uuid) == null) {
logger.error( "FATAL ERROR: wrote an entity {}:{} and it's missing", uuid, type );
System.exit(1);
}
echo(entityProps);
long stopTime = System.currentTimeMillis();
long duration = stopTime - startTime;
durationSum += duration;
//logger.debug( "Audited {}th admin", userCount );
if ( count % 100 == 0 ) {
logger.info( "Audited {}. Average Audit Rate: {}(ms)", count, durationSum / count );
}
} catch (InterruptedException e) {
e.printStackTrace();
} catch (Exception e) {
e.printStackTrace();
}
}
}
}
class ImportMetadataTask {
public EntityRef entityRef;
public Map<String, Object> metadata;
public ImportMetadataTask(EntityRef entityRef, Map<String, Object> metadata) {
this.entityRef = entityRef;
this.metadata = metadata;
}
}
class ImportMetadataWorker implements Runnable, Stoppable {
private BlockingQueue<ImportMetadataTask> workQueue;
private boolean done = false;
public ImportMetadataWorker(final BlockingQueue<ImportMetadataTask> workQueue) {
this.workQueue = workQueue;
}
@Override
public void setDone(boolean done) {
this.done = done;
}
@Override
public void run() {
int count = 0;
EntityManager em = emf.getEntityManager(CpNamingUtils.MANAGEMENT_APPLICATION_ID);
long durationSum = 0;
while (!done) {
try {
ImportMetadataTask task = this.workQueue.poll( 30, TimeUnit.SECONDS );
if (task == null) {
logger.warn("Reading from metadata queue was null!");
metadataEmptyCount.getAndIncrement();
Thread.sleep(1000);
continue;
}
metadataEmptyCount.set( 0 );
long startTime = System.currentTimeMillis();
importEntityMetadata( em, task.entityRef, task.metadata );
long stopTime = System.currentTimeMillis();
long duration = stopTime - startTime;
durationSum += duration;
metadataCount.getAndIncrement();
count++;
if ( count % 30 == 0 ) {
logger.info( "Imported {} metadata of total {} expected. " +
"Average metadata Imported Rate: {}(ms)",
new Object[] { metadataCount.get(), userCount.get(), durationSum / count });
}
} catch (Exception e) {
logger.debug("Error reading writing metadata", e);
}
}
}
}
class ImportAdminWorker implements Runnable, Stoppable {
private BlockingQueue<Map<String, Object>> workQueue;
private BlockingQueue<Map<String, Object>> auditQueue;
private boolean done = false;
public ImportAdminWorker(final BlockingQueue<Map<String, Object>> workQueue,
final BlockingQueue<Map<String, Object>> auditQueue) {
this.workQueue = workQueue;
this.auditQueue = auditQueue;
}
@Override
public void setDone(boolean done) {
this.done = done;
}
@Override
public void run() {
int count = 0;
EntityManager em = emf.getEntityManager(CpNamingUtils.MANAGEMENT_APPLICATION_ID);
long durationSum = 0;
while (!done) {
try {
Map<String, Object> entityProps = this.workQueue.poll(30, TimeUnit.SECONDS);
if (entityProps == null) {
logger.warn("Reading from admin import queue was null!");
writeEmptyCount.getAndIncrement();
Thread.sleep( 1000 );
continue;
}
writeEmptyCount.set(0);
// Import/create the entity
UUID uuid = getId(entityProps);
String type = getType( entityProps );
try {
long startTime = System.currentTimeMillis();
em.create(uuid, type, entityProps);
logger.debug( "Imported admin user {}:{}:{}",
new Object[] { entityProps.get( "username" ), entityProps.get("email"), uuid } );
userCount.getAndIncrement();
auditQueue.put(entityProps);
long stopTime = System.currentTimeMillis();
long duration = stopTime - startTime;
durationSum += duration;
count++;
if (count % 30 == 0) {
logger.info( "This worked has imported {} users of total {} imported so far. " +
"Average Creation Rate: {}ms",
new Object[] { count, userCount.get(), durationSum / count });
}
} catch (DuplicateUniquePropertyExistsException de) {
String dupProperty = de.getPropertyName();
handleDuplicateAccount( em, dupProperty, entityProps );
continue;
} catch (Exception e) {
logger.error("Error", e);
}
} catch (InterruptedException e) {
logger.error( "Error", e );
}
}
}
private void handleDuplicateAccount(EntityManager em, String dupProperty, Map<String, Object> entityProps ) {
logger.info( "Processing duplicate user {}:{}:{} with duplicate {}", new Object[]{
entityProps.get( "username" ),
entityProps.get( "email" ),
entityProps.get( "uuid" ),
dupProperty} );
UUID dupUuid = UUID.fromString( entityProps.get("uuid").toString() );
try {
dupsByDupUuid.put( dupUuid, new DuplicateUser( dupProperty, entityProps ) );
} catch (Exception e) {
logger.error("Error processing dup user {}:{}:{}",
new Object[] {entityProps.get( "username" ), entityProps.get("email"), dupUuid});
return;
}
}
}
}