/*
* 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.google.common.collect.BiMap;
import com.google.common.collect.HashBiMap;
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.util.CpNamingUtils;
import org.apache.usergrid.management.UserInfo;
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.utils.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
/**
* Export 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/export-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 export from a specific set of keyspaces:
*
* cassandra.connections=110
* cassandra.system.keyspace=My_Usergrid
* cassandra.application.keyspace=My_Usergrid_Applications
* cassandra.lock.keyspace=My_Usergrid_Locks
*/
public class ExportAdmins extends ExportingToolBase {
static final Logger logger = LoggerFactory.getLogger( ExportAdmins.class );
public static final String ADMIN_USERS_PREFIX = "admin-users";
public static final String ADMIN_USER_METADATA_PREFIX = "admin-user-metadata";
// map admin user UUID to list of organizations to which user belongs
private Map<UUID, List<Org>> userToOrgsMap = new HashMap<UUID, List<Org>>(50000);
private Map<String, UUID> orgNameToUUID = new HashMap<String, UUID>(50000);
private Set<UUID> orgsWritten = new HashSet<UUID>(50000);
private Set<UUID> duplicateOrgs = new HashSet<UUID>();
private static final String READ_THREAD_COUNT = "readThreads";
private int readThreadCount;
AtomicInteger userCount = new AtomicInteger( 0 );
boolean ignoreInvalidUsers = false; // true to ignore users with no credentials or orgs
/**
* Represents an AdminUser that has been read and is ready for export.
*/
class AdminUserWriteTask {
Entity adminUser;
Map<String, Map<Object, Object>> dictionariesByName;
BiMap<UUID, String> orgNamesByUuid;
}
/**
* Represents an organization associated with a user.
*/
private class Org {
UUID orgId;
String orgName;
public Org( UUID orgId, String orgName ) {
this.orgId = orgId;
this.orgName = orgName;
}
}
/**
* Export admin users using multiple threads.
* <p/>
* How it works:
* In main thread we query for IDs of all admin users, add each ID to read queue.
* Read-queue workers read admin user data, add data to write queue.
* One write-queue worker reads data writes to file.
*/
@Override
public void runTool(CommandLine line) throws Exception {
startSpring();
setVerbose( line );
applyOrgId( line );
prepareBaseOutputFileName( line );
outputDir = createOutputParentDir();
logger.info( "Export directory: " + outputDir.getAbsolutePath() );
if (StringUtils.isNotEmpty( line.getOptionValue( READ_THREAD_COUNT ) )) {
try {
readThreadCount = Integer.parseInt( line.getOptionValue( READ_THREAD_COUNT ) );
} catch (NumberFormatException nfe) {
logger.error( "-" + READ_THREAD_COUNT + " must be specified as an integer. Aborting..." );
return;
}
} else {
readThreadCount = 20;
}
buildOrgMap();
// start write queue worker
BlockingQueue<AdminUserWriteTask> writeQueue = new LinkedBlockingQueue<AdminUserWriteTask>();
AdminUserWriter adminUserWriter = new AdminUserWriter( writeQueue );
Thread writeThread = new Thread( adminUserWriter );
writeThread.start();
logger.debug( "Write thread started" );
// start read queue workers
BlockingQueue<UUID> readQueue = new LinkedBlockingQueue<UUID>();
for (int i = 0; i < readThreadCount; i++) {
AdminUserReader worker = new AdminUserReader( readQueue, writeQueue );
Thread readerThread = new Thread( worker, "AdminUserReader-" + i );
readerThread.start();
}
logger.debug( readThreadCount + " read worker threads started" );
// query for IDs, add each to read queue
Query query = new Query();
query.setLimit( MAX_ENTITY_FETCH );
query.setResultsLevel( Query.Level.IDS );
EntityManager em = emf.getEntityManager( CpNamingUtils.MANAGEMENT_APPLICATION_ID );
Results ids = em.searchCollection( em.getApplicationRef(), "users", query );
while (ids.size() > 0) {
for (UUID uuid : ids.getIds()) {
readQueue.add( uuid );
//logger.debug( "Added uuid to readQueue: " + uuid );
}
if (ids.getCursor() == null) {
break;
}
query.setCursor( ids.getCursor() );
ids = em.searchCollection( em.getApplicationRef(), "users", query );
}
logger.debug( "Waiting for write thread to complete" );
boolean done = false;
while ( !done ) {
writeThread.join( 10000, 0 );
done = !writeThread.isAlive();
logger.info( "Wrote {} users", userCount.get() );
}
}
@Override
@SuppressWarnings("static-access")
public Options createOptions() {
Options options = super.createOptions();
Option readThreads = OptionBuilder
.hasArg().withType(0).withDescription("Read Threads -" + READ_THREAD_COUNT).create(READ_THREAD_COUNT);
options.addOption( readThreads );
return options;
}
/**
* Shouldn't have to do this but getOrganizationsForAdminUser() is not 100% reliable in some Usergrid installations.
*/
private void buildOrgMap() throws Exception {
logger.info( "Building org map" );
ExecutorService execService = Executors.newFixedThreadPool( readThreadCount );
EntityManager em = emf.getEntityManager( CpNamingUtils.MANAGEMENT_APPLICATION_ID );
String queryString = "select *";
Query query = Query.fromQL( queryString );
query.withLimit( 1000 );
Results organizations = null;
int count = 0;
do {
organizations = em.searchCollection( em.getApplicationRef(), "groups", query );
for ( Entity organization : organizations.getEntities() ) {
execService.submit( new OrgMapWorker( organization ) );
count++;
}
if ( count % 1000 == 0 ) {
logger.info("Queued {} org map workers", count);
}
query.setCursor( organizations.getCursor() );
}
while ( organizations != null && organizations.hasCursor() );
execService.shutdown();
while ( !execService.awaitTermination( 10, TimeUnit.SECONDS ) ) {
logger.info( "Processed {} orgs for map", userToOrgsMap.size() );
}
logger.info("Org map complete, counted {} organizations", count);
}
public class OrgMapWorker implements Runnable {
private final Entity orgEntity;
public OrgMapWorker( Entity orgEntity ) {
this.orgEntity = orgEntity;
}
@Override
public void run() {
try {
final String orgName = orgEntity.getProperty( "path" ).toString();
final UUID orgId = orgEntity.getUuid();
for (UserInfo user : managementService.getAdminUsersForOrganization( orgEntity.getUuid() )) {
try {
Entity admin = managementService.getAdminUserEntityByUuid( user.getUuid() );
Org org = new Org( orgId, orgName );
synchronized (userToOrgsMap) {
List<Org> userOrgs = userToOrgsMap.get( admin.getUuid() );
if (userOrgs == null) {
userOrgs = new ArrayList<Org>();
userToOrgsMap.put( admin.getUuid(), userOrgs );
}
userOrgs.add( org );
}
synchronized (orgNameToUUID) {
UUID existingOrgId = orgNameToUUID.get( orgName );
;
if (existingOrgId != null && !orgId.equals( existingOrgId )) {
if ( !duplicateOrgs.contains( orgId )) {
logger.info( "Org {}:{} is a duplicate", orgId, orgName );
duplicateOrgs.add(orgId);
}
} else {
orgNameToUUID.put( orgName, orgId );
}
}
} catch (Exception e) {
logger.warn( "Cannot get orgs for userId {}", user.getUuid() );
}
}
} catch ( Exception e ) {
logger.error("Error getting users for org {}:{}", orgEntity.getName(), orgEntity.getUuid());
}
}
}
public class AdminUserReader implements Runnable {
private final BlockingQueue<UUID> readQueue;
private final BlockingQueue<AdminUserWriteTask> writeQueue;
public AdminUserReader( BlockingQueue<UUID> readQueue, BlockingQueue<AdminUserWriteTask> writeQueue ) {
this.readQueue = readQueue;
this.writeQueue = writeQueue;
}
@Override
public void run() {
try {
readAndQueueAdminUsers();
} catch (Exception e) {
logger.error("Error reading data for export", e);
}
}
private void readAndQueueAdminUsers() throws Exception {
EntityManager em = emf.getEntityManager( CpNamingUtils.MANAGEMENT_APPLICATION_ID );
while ( true ) {
UUID uuid = null;
try {
uuid = readQueue.poll( 30, TimeUnit.SECONDS );
if ( uuid == null ) {
break;
}
Entity entity = em.get( uuid );
AdminUserWriteTask task = new AdminUserWriteTask();
task.adminUser = entity;
addDictionariesToTask( task, entity );
addOrganizationsToTask( task );
String actionTaken = "Processed";
if (ignoreInvalidUsers && (task.orgNamesByUuid.isEmpty()
|| task.dictionariesByName.isEmpty()
|| task.dictionariesByName.get( "credentials" ).isEmpty())) {
actionTaken = "Ignored";
} else {
writeQueue.add( task );
}
Map<String, Object> creds = (Map<String, Object>) (task.dictionariesByName.isEmpty() ?
0 : task.dictionariesByName.get( "credentials" ));
logger.error( "{} admin user {}:{}:{} has organizations={} dictionaries={} credentials={}",
new Object[]{
actionTaken,
task.adminUser.getProperty( "username" ),
task.adminUser.getProperty( "email" ),
task.adminUser.getUuid(),
task.orgNamesByUuid.size(),
task.dictionariesByName.size(),
creds == null ? 0 : creds.size()
} );
} catch ( Exception e ) {
logger.error("Error reading data for user " + uuid, e );
}
}
}
private void addDictionariesToTask(AdminUserWriteTask task, Entity entity) throws Exception {
EntityManager em = emf.getEntityManager( CpNamingUtils.MANAGEMENT_APPLICATION_ID );
task.dictionariesByName = new HashMap<String, Map<Object, Object>>();
Set<String> dictionaries = em.getDictionaries( entity );
if ( dictionaries.isEmpty() ) {
logger.error("User {}:{} has no dictionaries", task.adminUser.getName(), task.adminUser.getUuid() );
return;
}
Map<Object, Object> credentialsDictionary = em.getDictionaryAsMap( entity, "credentials" );
if ( credentialsDictionary != null ) {
task.dictionariesByName.put( "credentials", credentialsDictionary );
}
}
private void addOrganizationsToTask(AdminUserWriteTask task) throws Exception {
task.orgNamesByUuid = managementService.getOrganizationsForAdminUser( task.adminUser.getUuid() );
List<Org> orgs = userToOrgsMap.get( task.adminUser.getUuid() );
if ( orgs != null && task.orgNamesByUuid.size() < orgs.size() ) {
// list of orgs from getOrganizationsForAdminUser() is less than expected, use userToOrgsMap
BiMap<UUID, String> bimap = HashBiMap.create();
for (Org org : orgs) {
bimap.put( org.orgId, org.orgName );
}
task.orgNamesByUuid = bimap;
}
}
}
class AdminUserWriter implements Runnable {
private final BlockingQueue<AdminUserWriteTask> taskQueue;
public AdminUserWriter( BlockingQueue<AdminUserWriteTask> taskQueue ) {
this.taskQueue = taskQueue;
}
@Override
public void run() {
try {
writeEntities();
} catch (Exception e) {
logger.error("Error writing export data", e);
}
}
private void writeEntities() throws Exception {
EntityManager em = emf.getEntityManager( CpNamingUtils.MANAGEMENT_APPLICATION_ID );
// write one JSON file for management application users
JsonGenerator usersFile =
getJsonGenerator( createOutputFile( ADMIN_USERS_PREFIX, em.getApplication().getName() ) );
usersFile.writeStartArray();
// write one JSON file for metadata: collections, connections and dictionaries of those users
JsonGenerator metadataFile =
getJsonGenerator( createOutputFile( ADMIN_USER_METADATA_PREFIX, em.getApplication().getName() ) );
metadataFile.writeStartObject();
while ( true ) {
try {
AdminUserWriteTask task = taskQueue.poll( 30, TimeUnit.SECONDS );
if ( task == null ) {
break;
}
// write user to application file
usersFile.writeObject( task.adminUser );
echo( task.adminUser );
// write metadata to metadata file
metadataFile.writeFieldName( task.adminUser.getUuid().toString() );
metadataFile.writeStartObject();
saveOrganizations( metadataFile, task );
saveDictionaries( metadataFile, task );
metadataFile.writeEndObject();
logger.debug( "Exported user {}:{}:{}", new Object[] {
task.adminUser.getProperty("username"),
task.adminUser.getProperty("email"),
task.adminUser.getUuid() } );
userCount.addAndGet( 1 );
} catch (InterruptedException e) {
throw new Exception("Interrupted", e);
}
}
metadataFile.writeEndObject();
metadataFile.close();
usersFile.writeEndArray();
usersFile.close();
logger.info( "Exported TOTAL {} admin users and {} organizations", userCount.get(), orgsWritten.size() );
}
private void saveDictionaries( JsonGenerator jg, AdminUserWriteTask task ) throws Exception {
jg.writeFieldName( "dictionaries" );
jg.writeStartObject();
for (String dictionary : task.dictionariesByName.keySet() ) {
Map<Object, Object> dict = task.dictionariesByName.get( dictionary );
if (dict.isEmpty()) {
continue;
}
jg.writeFieldName( dictionary );
jg.writeStartObject();
for (Map.Entry<Object, Object> entry : dict.entrySet()) {
jg.writeFieldName( entry.getKey().toString() );
jg.writeObject( entry.getValue() );
}
jg.writeEndObject();
}
jg.writeEndObject();
}
private void saveOrganizations( JsonGenerator jg, AdminUserWriteTask task ) throws Exception {
final BiMap<UUID, String> orgs = task.orgNamesByUuid;
jg.writeFieldName( "organizations" );
jg.writeStartArray();
for (UUID uuid : orgs.keySet()) {
jg.writeStartObject();
jg.writeFieldName( "uuid" );
jg.writeObject( uuid );
jg.writeFieldName( "name" );
jg.writeObject( orgs.get( uuid ) );
jg.writeEndObject();
synchronized (orgsWritten) {
orgsWritten.add( uuid );
}
}
jg.writeEndArray();
}
}
}