package avrobase.redis;
import avrobase.AvroBaseException;
import avrobase.AvroBaseImpl;
import avrobase.AvroFormat;
import avrobase.Row;
import com.google.common.base.Supplier;
import org.apache.avro.Schema;
import org.apache.avro.specific.SpecificRecord;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisException;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.TransactionBlock;
import sun.reflect.generics.reflectiveObjects.NotImplementedException;
import java.util.List;
import java.util.concurrent.TimeoutException;
/**
* AvroBase on top of Redis
* <p/>
* User: sam
* Date: Oct 3, 2010
* Time: 11:32:10 AM
*/
@SuppressWarnings({"unchecked"})
public class RAB<T extends SpecificRecord> extends AvroBaseImpl<T, String> {
private JedisPool pool;
private int db;
private Supplier<String> kg;
private static final String d = "_d";
private static final String s = "_s";
private static final String v = "_v";
private static final String z = "_z";
public RAB(JedisPool pool, int db, Supplier<String> kg, Schema actualSchema) {
super(actualSchema, AvroFormat.JSON);
this.pool = pool;
this.db = db;
this.kg = kg;
}
@Override
public Row<T, String> get(final String row) throws AvroBaseException {
try {
boolean returned = false;
final Jedis j = pool.getResource();
try {
j.select(db);
List<Object> results;
do {
results = j.multi(new TransactionBlock() {
@Override
public void execute() throws JedisException {
mget(row + s, row + v, row + d);
}
});
} while (results == null);
if (results.size() != 1 || (results = (List<Object>) results.get(0)).size() != 3) {
throw new AvroBaseException("Incorrect number of results from redis transaction: " + results);
}
String schemaId = (String) results.get(0);
String versionStr = (String) results.get(1);
String data = (String) results.get(2);
if (versionStr == null || schemaId == null || data == null) {
return null;
}
Schema schema = schemaCache.get(schemaId);
if (schema == null) {
schema = loadSchema(j.get(schemaId + z).getBytes(), schemaId);
}
return new Row<T, String>(readValue(data.getBytes(), schema, format), row, Long.parseLong(versionStr));
} catch (Exception e) {
pool.returnBrokenResource(j);
returned = true;
throw new AvroBaseException(e);
} finally {
if (!returned) pool.returnResource(j);
}
} catch (TimeoutException e) {
throw new AvroBaseException("Timed out", e);
}
}
@Override
public String create(T value) throws AvroBaseException {
String row;
do {
row = kg.get();
} while(!put(row, value, 0));
return row;
}
@Override
public void put(final String row, final T value) throws AvroBaseException {
try {
boolean returned = false;
Jedis j = pool.getResource();
try {
j.select(db);
Schema schema = value.getSchema();
String schemaKey = hashCache.get(schema);
if (schemaKey == null) {
final String doc = schema.toString();
schemaKey = createSchemaKey(schema, doc);
j.set(schemaKey + z, doc);
}
final String finalSchemaKey = schemaKey;
List<Object> results;
do {
results = j.multi(new TransactionBlock() {
@Override
public void execute() throws JedisException {
incr(row + v);
mset(row + s, finalSchemaKey, row + d, new String(serialize(value), UTF8));
}
});
} while (results == null);
} catch (Exception e) {
pool.returnBrokenResource(j);
returned = true;
throw new AvroBaseException(e);
} finally {
if (!returned) pool.returnResource(j);
}
} catch (TimeoutException e) {
throw new AvroBaseException("Timed out", e);
}
}
@Override
public boolean put(final String row, final T value, final long version) throws AvroBaseException {
try {
boolean returned = false;
final Jedis j = pool.getResource();
try {
j.select(db);
Schema schema = value.getSchema();
String schemaKey = hashCache.get(schema);
if (schemaKey == null) {
final String doc = schema.toString();
schemaKey = createSchemaKey(schema, doc);
j.set(schemaKey + z, doc);
}
String watch = j.watch(row + v);
if (!watch.equals("OK")) {
return false;
}
String versionStr = j.get(row + v);
if ((versionStr == null && version != 0) || // row missing but version set
(versionStr != null && version == 0) || // row exists but insert requested
(versionStr != null && !versionStr.equals(String.valueOf(version))) // version in db doesn't match
) {
return false;
}
final String finalSchemaKey = schemaKey;
List<Object> results = j.multi(new TransactionBlock() {
@Override
public void execute() throws JedisException {
mset(row + v, String.valueOf(version + 1), row + s, finalSchemaKey, row + d, new String(serialize(value), UTF8));
}
});
return results != null;
} catch (Exception e) {
pool.returnBrokenResource(j);
returned = true;
throw new AvroBaseException(e);
} finally {
if (!returned) pool.returnResource(j);
}
} catch (TimeoutException e) {
throw new AvroBaseException("Timed out", e);
}
}
@Override
public void delete(final String row) throws AvroBaseException {
try {
boolean returned = false;
Jedis j = pool.getResource();
try {
j.select(db);
List<Object> results;
do {
results = j.multi(new TransactionBlock() {
@Override
public void execute() throws JedisException {
del(row + v); // Delete the version first and it is deleted
del(row + d);
del(row + s);
}
});
} while (results == null);
} catch (Exception e) {
pool.returnBrokenResource(j);
returned = true;
throw new AvroBaseException(e);
} finally {
if (!returned) pool.returnResource(j);
}
} catch (TimeoutException e) {
throw new AvroBaseException("Timed out", e);
}
}
@Override
public Iterable<Row<T, String>> scan(String startRow, String stopRow) throws AvroBaseException {
throw new NotImplementedException();
}
}