com.amazon.janusgraph.TestGraphUtil.java Source code

Java tutorial

Introduction

Here is the source code for com.amazon.janusgraph.TestGraphUtil.java

Source

/*
 * Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file 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 com.amazon.janusgraph;

import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

import org.apache.commons.configuration.Configuration;
import org.apache.commons.configuration.ConfigurationException;
import org.apache.commons.configuration.PropertiesConfiguration;
import org.janusgraph.core.JanusGraph;
import org.janusgraph.core.JanusGraphFactory;
import org.janusgraph.diskstorage.BackendException;
import org.janusgraph.diskstorage.configuration.BasicConfiguration;
import org.janusgraph.diskstorage.configuration.BasicConfiguration.Restriction;
import org.janusgraph.diskstorage.configuration.WriteConfiguration;
import org.janusgraph.diskstorage.configuration.backend.CommonsConfiguration;
import org.janusgraph.graphdb.configuration.GraphDatabaseConfiguration;

import com.amazon.janusgraph.diskstorage.dynamodb.BackendDataModel;
import com.amazon.janusgraph.diskstorage.dynamodb.Client;
import com.amazon.janusgraph.diskstorage.dynamodb.Constants;
import com.amazon.janusgraph.diskstorage.dynamodb.DynamoDBDelegate;
import com.amazonaws.services.dynamodbv2.model.DeleteTableRequest;
import com.amazonaws.services.dynamodbv2.model.ListTablesResult;
import com.amazonaws.services.dynamodbv2.model.ResourceNotFoundException;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;

/**
 *
 * @author Alexander Patrikalakis
 * @author Michael Rodaitis
 *
 */
public enum TestGraphUtil {
    instance;

    private final int dynamoDBPartitions;
    private final int controlPlaneRate;
    private final boolean unlimitedIops;
    private final int provisionedReadAndWriteTps;
    private final File propertiesFile;

    TestGraphUtil() {
        dynamoDBPartitions = Integer.valueOf(System.getProperty("dynamodb-partitions", String.valueOf(1)));
        Preconditions.checkArgument(dynamoDBPartitions > 0);
        provisionedReadAndWriteTps = 750 * dynamoDBPartitions;
        unlimitedIops = Boolean
                .valueOf(System.getProperty("dynamodb-unlimited-iops", String.valueOf(Boolean.TRUE)));
        controlPlaneRate = Integer.valueOf(System.getProperty("dynamodb-control-plane-rate", String.valueOf(100)));
        Preconditions.checkArgument(controlPlaneRate > 0);
        //This is a configuration file for test code. not part of production applications.
        //no validation necessary.
        propertiesFile = new File(
                System.getProperty("properties-file", "src/test/resources/dynamodb-local.properties"));
        Preconditions.checkArgument(propertiesFile.exists());
        Preconditions.checkArgument(propertiesFile.isFile());
        Preconditions.checkArgument(propertiesFile.canRead());
    }

    public boolean isUnlimitedIops() {
        return unlimitedIops;
    }

    public JanusGraph openGraph(final BackendDataModel backendDataModel) {
        return JanusGraphFactory.open(createTestGraphConfig(backendDataModel));
    }

    public void tearDownGraph(final JanusGraph graph) throws BackendException {
        if (null != graph) {
            graph.close();
        }
        cleanUpTables();
    }

    @VisibleForTesting // used in ClientTest.java
    Client createClient() {
        return new Client(new BasicConfiguration(GraphDatabaseConfiguration.ROOT_NS,
                new CommonsConfiguration(loadProperties()), Restriction.NONE));
    }

    public Configuration loadProperties() {
        PropertiesConfiguration storageConfig;
        try {
            storageConfig = new PropertiesConfiguration(propertiesFile);
        } catch (ConfigurationException e) {
            throw new RuntimeException(e);
        }
        return storageConfig;
    }

    public Configuration createTestGraphConfig(final BackendDataModel backendDataModel) {
        final String dataModelName = backendDataModel.name();

        final Configuration properties = createTestConfig(backendDataModel);

        final String storesNsPrefix = "storage.dynamodb.stores.";
        final List<String> storeList = new ArrayList<>(Constants.REQUIRED_BACKEND_STORES);
        storeList.add(GraphDatabaseConfiguration.IDS_STORE_NAME.getDefaultValue());
        for (String store : storeList) {
            configureStore(dataModelName, provisionedReadAndWriteTps, properties, unlimitedIops,
                    storesNsPrefix + store);
        }

        return properties;
    }

    private Configuration createTestConfig(final BackendDataModel backendDataModel) {
        final Configuration properties = loadProperties();
        final Configuration dynamodb = properties.subset("storage").subset("dynamodb");
        dynamodb.setProperty(Constants.DYNAMODB_TABLE_PREFIX.getName(), backendDataModel.name() /*prefix*/);
        dynamodb.setProperty(Constants.DYNAMODB_CONTROL_PLANE_RATE.getName(), controlPlaneRate);
        return properties;
    }

    private static void configureStore(final String dataModelName, final int tps, final Configuration config,
            final boolean unlimitedIops, final String prefix) {
        final String prefixPeriod = prefix + ".";
        config.setProperty(prefixPeriod + Constants.STORES_DATA_MODEL.getName(), dataModelName);
        config.setProperty(prefixPeriod + Constants.STORES_SCAN_LIMIT.getName(), 10000);
        config.setProperty(prefixPeriod + Constants.STORES_INITIAL_CAPACITY_READ.getName(), tps);
        config.setProperty(prefixPeriod + Constants.STORES_READ_RATE_LIMIT.getName(),
                unlimitedIops ? Integer.MAX_VALUE : tps);
        config.setProperty(prefixPeriod + Constants.STORES_INITIAL_CAPACITY_WRITE.getName(), tps);
        config.setProperty(prefixPeriod + Constants.STORES_WRITE_RATE_LIMIT.getName(),
                unlimitedIops ? Integer.MAX_VALUE : tps);
    }

    private static void configureStore(final String dataModelName, final int tps, final WriteConfiguration config,
            final boolean unlimitedIops, final String prefix) {
        final String prefixPeriod = prefix + ".";
        config.set(prefixPeriod + Constants.STORES_DATA_MODEL.getName(), dataModelName);
        config.set(prefixPeriod + Constants.STORES_SCAN_LIMIT.getName(), 10000);
        config.set(prefixPeriod + Constants.STORES_INITIAL_CAPACITY_READ.getName(), tps);
        config.set(prefixPeriod + Constants.STORES_READ_RATE_LIMIT.getName(),
                unlimitedIops ? Integer.MAX_VALUE : tps);
        config.set(prefixPeriod + Constants.STORES_INITIAL_CAPACITY_WRITE.getName(), tps);
        config.set(prefixPeriod + Constants.STORES_WRITE_RATE_LIMIT.getName(),
                unlimitedIops ? Integer.MAX_VALUE : tps);
    }

    public WriteConfiguration getStoreConfig(final BackendDataModel model, final List<String> storeNames) {
        return appendClusterPartitionsAndStores(model,
                new CommonsConfiguration(TestGraphUtil.instance.createTestConfig(model)), storeNames,
                1 /*partitions*/);
    }

    public WriteConfiguration appendStoreConfig(final BackendDataModel model, final WriteConfiguration config,
            final List<String> storeNames) {
        final Configuration baseconfig = createTestConfig(model);
        final Iterator<String> it = baseconfig.getKeys();
        while (it.hasNext()) {
            final String key = it.next();
            config.set(key, baseconfig.getProperty(key));
        }
        return appendClusterPartitionsAndStores(model, config, storeNames, 1 /*titanClusterPartitions*/);
    }

    public WriteConfiguration graphConfigWithClusterPartitionsAndExtraStores(final BackendDataModel model,
            final List<String> extraStoreNames, final int janusGraphClusterPartitions) {
        return appendClusterPartitionsAndStores(model,
                new CommonsConfiguration(TestGraphUtil.instance.createTestGraphConfig(model)), extraStoreNames,
                janusGraphClusterPartitions);
    }

    public WriteConfiguration graphConfig(BackendDataModel model) {
        return graphConfigWithClusterPartitionsAndExtraStores(model, Collections.emptyList(), 1);
    }

    public WriteConfiguration appendClusterPartitionsAndStores(final BackendDataModel model,
            final WriteConfiguration config, final List<String> storeNames, final int janusGraphClusterPartitions) {
        final Configuration baseconfig = createTestConfig(model);
        final Iterator<String> it = baseconfig.getKeys();
        while (it.hasNext()) {
            final String key = it.next();
            config.set(key, baseconfig.getProperty(key));
        }

        Preconditions.checkArgument(janusGraphClusterPartitions > 0);
        if (janusGraphClusterPartitions > 1) {
            config.set("cluster.max-partitions", Integer.toString(janusGraphClusterPartitions));
        }

        for (String store : storeNames) {
            final String prefix = "storage.dynamodb.stores." + store;
            TestGraphUtil.configureStore(model.name(), provisionedReadAndWriteTps, config, isUnlimitedIops(),
                    prefix);
        }

        return config;
    }

    private static void deleteAllTables(final String prefix, final DynamoDBDelegate delegate)
            throws BackendException {
        final ListTablesResult result = delegate.listAllTables();
        for (String tableName : result.getTableNames()) {
            if (prefix != null && !tableName.startsWith(prefix)) {
                continue;
            }
            try {
                delegate.deleteTable(new DeleteTableRequest().withTableName(tableName));
            } catch (ResourceNotFoundException e) {
            }
        }
    }

    public void cleanUpTables() throws BackendException {
        final Client client = instance.createClient();
        deleteAllTables(null /*prefix - delete all tables*/, client.delegate());
        client.delegate().shutdown();
    }
}