org.apache.usergrid.tools.ImportAdmins.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.usergrid.tools.ImportAdmins.java

Source

/*
 * 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;
            }

        }
    }
}