se.gawell.fakeriak.servlet.InMemoryHandler Maven / Gradle / Ivy

Go to download
package se.gawell.fakeriak.servlet;

import java.io.DataInputStream;
import java.io.IOException;
import java.util.Enumeration;
import java.util.StringTokenizer;
import java.util.concurrent.atomic.AtomicLong;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import se.gawell.fakeriak.buckets.Bucket;
import se.gawell.fakeriak.buckets.Buckets;
import se.gawell.fakeriak.buckets.InMemoryBucket;
import se.gawell.fakeriak.index.Indexes;

public abstract class InMemoryHandler {
	private static final String RIAK_INDEX_PREFIX = "x-riak-index-";
	private static final String NOT_FOUND = "not found";

	private final Buckets buckets;
	protected final Indexes indexes;
	private AtomicLong idGenerator = new AtomicLong();

	public InMemoryHandler(Buckets buckets, Indexes indexes) {
		super();
		this.buckets = buckets;
		this.indexes = indexes;
	}

	protected String getBucketName(HttpServletRequest request) {
		String path = request.getPathInfo();
		String bucketName = getPathElement(path, 1);
		return bucketName;
	}

	protected String getPathElement(String path, int elementNumber) {
		try {
			StringTokenizer tokenizer = new StringTokenizer(path, "/");
			for (int i = 0; i < elementNumber; i++) {
				tokenizer.nextToken();
			}
			String bucketName = tokenizer.nextToken();
			return bucketName;
		} catch (RuntimeException e) {
			return null;
		}
	}

	protected String getKey(HttpServletRequest request) {
		String path = request.getPathInfo();
		String key = getPathElement(path, 2);
		return key;
	}

	protected Bucket getBucket(String bucketName) {
		Bucket bucket = buckets.get(bucketName);
		if (bucket == null) {
			bucket = new InMemoryBucket();
			buckets.put(bucketName, bucket);
		}
		return bucket;
	}

	protected void writeResultToResponse(HttpServletResponse response, String result) throws IOException {
		byte[] bytes = new byte[0];

		if (result != null) {
			bytes = result.getBytes();
		} else {
			bytes = NOT_FOUND.getBytes();
			response.setStatus(404);
		}
		response.getOutputStream().write(bytes);
	}

	protected synchronized void addOrUpdate(HttpServletRequest request, HttpServletResponse response) throws IOException {
		String bucketName = getBucketName(request);
		String key = getKey(request);
		String data = readDataFromRequest(request);
		Bucket bucket = getBucket(bucketName);
		if (key == null) {
			key = generateUniqueKey(bucket);
		}
		bucket.put(key, data);

		Enumeration headerNames = request.getHeaderNames();
		while (headerNames.hasMoreElements()) {
			String headerName = headerNames.nextElement();
			if (headerName.toLowerCase().startsWith(RIAK_INDEX_PREFIX)) {
				String indexName = headerName.substring(RIAK_INDEX_PREFIX.length());
				String indexValue = request.getHeader(headerName);
				indexes.put(bucketName, indexName, indexValue, key);
			}
		}

		String result = bucket.get(key);
		
		writeLocationToResponse(bucket,key, response);
		writeResultToResponse(response, result);
	}

	private void writeLocationToResponse(Bucket bucket, String key, HttpServletResponse response) {
		response.addHeader("Location", "/riak/" + bucket + "/" + key);
	}

	private String generateUniqueKey(Bucket bucket) {
		String key = Long.toString(idGenerator.addAndGet(1));
		return key;
	}

	private String readDataFromRequest(HttpServletRequest request) throws IOException {
		DataInputStream dataInputStream = new DataInputStream(request.getInputStream());
		StringBuffer buffer = new StringBuffer();
		boolean moreDataToRead = true;
		while (moreDataToRead) {
			int c = dataInputStream.read();
			if (c < 0) {
				moreDataToRead = false;
			} else {
				buffer.append((char) c);
			}
		}
		return buffer.toString();
	}

}