com.yahoo.ycsb.db.DynamoDBClient.java Source code

Java tutorial

Introduction

Here is the source code for com.yahoo.ycsb.db.DynamoDBClient.java

Source

/*
    
 * Copyright 2012 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.yahoo.ycsb.db;

import java.io.FileInputStream;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;
import java.util.Set;
import java.util.Vector;
import java.io.File;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.PropertiesCredentials;
import com.amazonaws.ClientConfiguration;
import com.amazonaws.services.dynamodb.AmazonDynamoDBClient;
import com.amazonaws.AmazonClientException;
import com.amazonaws.AmazonServiceException;
import com.amazonaws.services.dynamodb.model.AttributeValue;
import com.amazonaws.services.dynamodb.model.AttributeValueUpdate;
import com.amazonaws.services.dynamodb.model.DeleteItemRequest;
import com.amazonaws.services.dynamodb.model.DeleteItemResult;
import com.amazonaws.services.dynamodb.model.GetItemRequest;
import com.amazonaws.services.dynamodb.model.GetItemResult;
import com.amazonaws.services.dynamodb.model.Key;
import com.amazonaws.services.dynamodb.model.PutItemRequest;
import com.amazonaws.services.dynamodb.model.PutItemResult;
import com.amazonaws.services.dynamodb.model.ScanRequest;
import com.amazonaws.services.dynamodb.model.ScanResult;
import com.amazonaws.services.dynamodb.model.UpdateItemRequest;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.StringByteIterator;

/**
 * DynamoDB v1.3.14 client for YCSB
 */

public class DynamoDBClient extends DB {

    private static final int OK = 0;
    private static final int SERVER_ERROR = 1;
    private static final int CLIENT_ERROR = 2;
    private AmazonDynamoDBClient dynamoDB;
    private String primaryKeyName;
    private boolean debug = false;
    private boolean consistentRead = false;
    private String endpoint = "http://dynamodb.us-east-1.amazonaws.com";
    private int maxConnects = 50;
    private static Logger logger = Logger.getLogger(DynamoDBClient.class);

    public DynamoDBClient() {
    }

    /**
     * Initialize any state for this DB. Called once per DB instance; there is
     * one DB instance per client thread.
     */
    public void init() throws DBException {
        // initialize DynamoDb driver & table.
        String debug = getProperties().getProperty("dynamodb.debug", null);

        if (null != debug && "true".equalsIgnoreCase(debug)) {
            logger.setLevel(Level.DEBUG);
        }

        String endpoint = getProperties().getProperty("dynamodb.endpoint", null);
        String credentialsFile = getProperties().getProperty("dynamodb.awsCredentialsFile", null);
        String primaryKey = getProperties().getProperty("dynamodb.primaryKey", null);
        String consistentReads = getProperties().getProperty("dynamodb.consistentReads", null);
        String connectMax = getProperties().getProperty("dynamodb.connectMax", null);

        if (null != connectMax) {
            this.maxConnects = Integer.parseInt(connectMax);
        }

        if (null != consistentReads && "true".equalsIgnoreCase(consistentReads)) {
            this.consistentRead = true;
        }

        if (null != endpoint) {
            this.endpoint = endpoint;
        }

        if (null == primaryKey || primaryKey.length() < 1) {
            String errMsg = "Missing primary key attribute name, cannot continue";
            logger.error(errMsg);
        }

        try {
            AWSCredentials credentials = new PropertiesCredentials(new File(credentialsFile));
            ClientConfiguration cconfig = new ClientConfiguration();
            cconfig.setMaxConnections(maxConnects);
            dynamoDB = new AmazonDynamoDBClient(credentials, cconfig);
            dynamoDB.setEndpoint(this.endpoint);
            primaryKeyName = primaryKey;
            logger.info("dynamodb connection created with " + this.endpoint);
        } catch (Exception e1) {
            String errMsg = "DynamoDBClient.init(): Could not initialize DynamoDB client: " + e1.getMessage();
            logger.error(errMsg);
        }
    }

    @Override
    public int read(String table, String key, Set<String> fields, HashMap<String, ByteIterator> result) {

        logger.debug("readkey: " + key + " from table: " + table);
        GetItemRequest req = new GetItemRequest(table, createPrimaryKey(key));
        req.setAttributesToGet(fields);
        req.setConsistentRead(consistentRead);
        GetItemResult res = null;

        try {
            res = dynamoDB.getItem(req);
        } catch (AmazonServiceException ex) {
            logger.error(ex.getMessage());
            return SERVER_ERROR;
        } catch (AmazonClientException ex) {
            logger.error(ex.getMessage());
            return CLIENT_ERROR;
        }

        if (null != res.getItem()) {
            result.putAll(extractResult(res.getItem()));
            logger.debug("Result: " + res.toString());
        }
        return OK;
    }

    @Override
    public int scan(String table, String startkey, int recordcount, Set<String> fields,
            Vector<HashMap<String, ByteIterator>> result) {
        logger.debug("scan " + recordcount + " records from key: " + startkey + " on table: " + table);
        /*
         * on DynamoDB's scan, startkey is *exclusive* so we need to
         * getItem(startKey) and then use scan for the res
         */
        GetItemRequest greq = new GetItemRequest(table, createPrimaryKey(startkey));
        greq.setAttributesToGet(fields);

        GetItemResult gres = null;

        try {
            gres = dynamoDB.getItem(greq);
        } catch (AmazonServiceException ex) {
            logger.error(ex.getMessage());
            return SERVER_ERROR;
        } catch (AmazonClientException ex) {
            logger.error(ex.getMessage());
            return CLIENT_ERROR;
        }

        if (null != gres.getItem()) {
            result.add(extractResult(gres.getItem()));
        }

        int count = 1; // startKey is done, rest to go.

        Key startKey = createPrimaryKey(startkey);
        ScanRequest req = new ScanRequest(table);
        req.setAttributesToGet(fields);
        while (count < recordcount) {
            req.setExclusiveStartKey(startKey);
            req.setLimit(recordcount - count);
            ScanResult res = null;
            try {
                res = dynamoDB.scan(req);
            } catch (AmazonServiceException ex) {
                logger.error(ex.getMessage());
                ex.printStackTrace();
                return SERVER_ERROR;
            } catch (AmazonClientException ex) {
                logger.error(ex.getMessage());
                ex.printStackTrace();
                return CLIENT_ERROR;
            }

            count += res.getCount();
            for (Map<String, AttributeValue> items : res.getItems()) {
                result.add(extractResult(items));
            }
            startKey = res.getLastEvaluatedKey();

        }

        return OK;
    }

    @Override
    public int update(String table, String key, HashMap<String, ByteIterator> values) {
        logger.debug("updatekey: " + key + " from table: " + table);

        Map<String, AttributeValueUpdate> attributes = new HashMap<String, AttributeValueUpdate>(values.size());
        for (Entry<String, ByteIterator> val : values.entrySet()) {
            AttributeValue v = new AttributeValue(val.getValue().toString());
            attributes.put(val.getKey(), new AttributeValueUpdate().withValue(v).withAction("PUT"));
        }

        UpdateItemRequest req = new UpdateItemRequest(table, createPrimaryKey(key), attributes);

        try {
            dynamoDB.updateItem(req);
        } catch (AmazonServiceException ex) {
            logger.error(ex.getMessage());
            return SERVER_ERROR;
        } catch (AmazonClientException ex) {
            logger.error(ex.getMessage());
            return CLIENT_ERROR;
        }
        return OK;
    }

    @Override
    public int insert(String table, String key, HashMap<String, ByteIterator> values) {
        logger.debug("insertkey: " + primaryKeyName + "-" + key + " from table: " + table);
        Map<String, AttributeValue> attributes = createAttributes(values);
        // adding primary key
        attributes.put(primaryKeyName, new AttributeValue(key));

        PutItemRequest putItemRequest = new PutItemRequest(table, attributes);
        PutItemResult res = null;
        try {
            res = dynamoDB.putItem(putItemRequest);
        } catch (AmazonServiceException ex) {
            logger.error(ex.getMessage());
            return SERVER_ERROR;
        } catch (AmazonClientException ex) {
            logger.error(ex.getMessage());
            return CLIENT_ERROR;
        }
        return OK;
    }

    @Override
    public int delete(String table, String key) {
        logger.debug("deletekey: " + key + " from table: " + table);
        DeleteItemRequest req = new DeleteItemRequest(table, createPrimaryKey(key));
        DeleteItemResult res = null;

        try {
            res = dynamoDB.deleteItem(req);
        } catch (AmazonServiceException ex) {
            logger.error(ex.getMessage());
            return SERVER_ERROR;
        } catch (AmazonClientException ex) {
            logger.error(ex.getMessage());
            return CLIENT_ERROR;
        }
        return OK;
    }

    private static Map<String, AttributeValue> createAttributes(HashMap<String, ByteIterator> values) {
        Map<String, AttributeValue> attributes = new HashMap<String, AttributeValue>(values.size() + 1); //leave space for the PrimaryKey
        for (Entry<String, ByteIterator> val : values.entrySet()) {
            attributes.put(val.getKey(), new AttributeValue(val.getValue().toString()));
        }
        return attributes;
    }

    private HashMap<String, ByteIterator> extractResult(Map<String, AttributeValue> item) {
        if (null == item)
            return null;
        HashMap<String, ByteIterator> rItems = new HashMap<String, ByteIterator>(item.size());

        for (Entry<String, AttributeValue> attr : item.entrySet()) {
            logger.debug(String.format("Result- key: %s, value: %s", attr.getKey(), attr.getValue()));
            rItems.put(attr.getKey(), new StringByteIterator(attr.getValue().getS()));
        }
        return rItems;
    }

    private static Key createPrimaryKey(String key) {
        Key k = new Key().withHashKeyElement(new AttributeValue().withS(key));
        return k;
    }
}