package com.pardot.rhombus.functional;
import com.datastax.driver.core.utils.UUIDs;
import com.google.common.collect.Maps;
import com.pardot.rhombus.ConnectionManager;
import com.pardot.rhombus.Criteria;
import com.pardot.rhombus.ObjectMapper;
import com.pardot.rhombus.cobject.CKeyspaceDefinition;
import com.pardot.rhombus.util.JsonUtil;
import org.joda.time.DateTime;
import org.junit.BeforeClass;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.*;
import static org.junit.Assert.*;
public class ObjectMapperShardingITCase extends RhombusFunctionalTest {
private static Logger logger = LoggerFactory.getLogger(ObjectMapperShardingITCase.class);
private static CKeyspaceDefinition keyspaceDefinition;
@BeforeClass
public static void setupKeyspace()
{
try {
ConnectionManager cm = getConnectionManagerStatic();
keyspaceDefinition = JsonUtil.objectFromJsonResource(CKeyspaceDefinition.class, ObjectMapperShardingITCase.class.getClassLoader(), "ShardedKeyspace.js");
cm.buildKeyspace(keyspaceDefinition, true);
} catch(Exception e) {
throw new RuntimeException(e);
}
}
@Test
public void testShardedQueries() throws Exception {
logger.debug("Starting testShardedQueries");
// Get an object mapper for the keyspace and truncate the data
ConnectionManager cm = getConnectionManager();
ObjectMapper om = cm.getObjectMapper(keyspaceDefinition);
om.truncateTables();
om.setLogCql(true);
//Set up test data
List<Map<String, Object>> values = JsonUtil.rhombusMapFromResource(this.getClass().getClassLoader(), "ShardedTestData.js");
for(Map<String, Object> object : values) {
Map<String, Object> updatedObject = JsonUtil.rhombusMapFromJsonMap(object, keyspaceDefinition.getDefinitions().get("object1"));
Long createdAt = ((Date)(updatedObject.get("created_at"))).getTime();
logger.debug("Inserting object with created_at: {}", createdAt);
UUID id = (UUID)om.insert("object1", updatedObject, createdAt);
logger.debug("Inserted object with uuid unix time: {}", UUIDs.unixTimestamp(id));
}
//Query it back out
//Make sure that we have the proper number of results
SortedMap<String, Object> indexValues = Maps.newTreeMap();
indexValues.put("account_id", UUID.fromString("00000003-0000-0030-0040-000000030000"));
indexValues.put("user_id", UUID.fromString("00000003-0000-0030-0040-000000030000"));
Criteria criteria = new Criteria();
criteria.setIndexKeys(indexValues);
criteria.setLimit(50L);
List<Map<String, Object>> results = om.list("object1", criteria);
assertEquals(3, results.size());
cm.teardown();
}
@Test
public void testQueriesUseShardIndex() throws Exception {
logger.debug("Starting testQueriesUseShardIndex");
// Get an object mapper for the keyspace and truncate the data
ConnectionManager cm = getConnectionManager();
ObjectMapper om = cm.getObjectMapper(keyspaceDefinition);
om.truncateTables();
om.setLogCql(true);
// Insert a record that is more than ObjectMapper.reasonableStatementLimit months old
DateTime dateTime = DateTime.now().minusMonths(om.getReasonableStatementLimit() * 2);
UUID id = UUIDs.startOf(dateTime.getMillis());
UUID accountId = UUID.fromString("00000003-0000-0030-0040-000000030000");
UUID userId = UUID.fromString("00000003-0000-0030-0040-000000030000");
Map<String, Object> record = Maps.newHashMap();
record.put("id", id);
record.put("account_id", accountId);
record.put("user_id", userId);
om.insert("object1", record);
//Query it back out, make sure that we are using the shard index, because if we are not, we will surpass reasonable statement limit
SortedMap<String, Object> indexValues = Maps.newTreeMap();
indexValues.put("account_id", accountId);
indexValues.put("user_id", userId);
Criteria criteria = new Criteria();
criteria.setIndexKeys(indexValues);
criteria.setEndTimestamp(DateTime.now().getMillis());
criteria.setLimit(50L);
List<Map<String, Object>> results = om.list("object1", criteria);
assertEquals(1, results.size());
cm.teardown();
}
}