@Before
public void setUp() throws Exception {
ElasticsearchUtils.setMapper(mapper);
DataStore dataStore = TestUtils.getDataStore();
// Initializing Cache Factory
hazelcastInstance = new TestHazelcastInstanceFactory(1).newHazelcastInstance();
HazelcastConnection hazelcastConnection = Mockito.mock(HazelcastConnection.class);
when(hazelcastConnection.getHazelcast()).thenReturn(hazelcastInstance);
CacheUtils.setCacheFactory(new DistributedCacheFactory(hazelcastConnection, mapper));
elasticsearchServer = new MockElasticsearchServer(UUID.randomUUID().toString());
ElasticsearchConnection elasticsearchConnection = Mockito.mock(ElasticsearchConnection.class);
when(elasticsearchConnection.getClient()).thenReturn(elasticsearchServer.getClient());
ElasticsearchUtils.initializeMappings(elasticsearchServer.getClient());
// Ensure that table exists before saving/reading data from it
TableMetadataManager tableMetadataManager = Mockito.mock(TableMetadataManager.class);
when(tableMetadataManager.exists(TestUtils.TEST_TABLE_NAME)).thenReturn(true);
when(tableMetadataManager.get(anyString())).thenReturn(TestUtils.TEST_TABLE);
QueryStore queryStore =
new ElasticsearchQueryStore(tableMetadataManager, elasticsearchConnection, dataStore);
AnalyticsLoader analyticsLoader =
new AnalyticsLoader(tableMetadataManager, dataStore, queryStore, elasticsearchConnection);
TestUtils.registerActions(analyticsLoader, mapper);
ExecutorService executorService = Executors.newFixedThreadPool(1);
queryExecutor = new QueryExecutor(analyticsLoader, executorService);
queryStore.save(TestUtils.TEST_TABLE_NAME, TestUtils.getQueryDocuments(mapper));
}
@Override
public Set<String> loadAllKeys() {
logger.info("Load all keys called");
SearchResponse response =
elasticsearchConnection
.getClient()
.prepareSearch(TABLE_META_INDEX)
.setTypes(TABLE_META_TYPE)
.setQuery(QueryBuilders.matchAllQuery())
.setSearchType(SearchType.SCAN)
.setScroll(new TimeValue(30, TimeUnit.SECONDS))
.setNoFields()
.execute()
.actionGet();
Set<String> ids = Sets.newHashSet();
while (true) {
response =
elasticsearchConnection
.getClient()
.prepareSearchScroll(response.getScrollId())
.setScroll(new TimeValue(60000))
.execute()
.actionGet();
SearchHits hits = response.getHits();
for (SearchHit hit : hits) {
ids.add(hit.getId());
}
if (0 == response.getHits().hits().length) {
break;
}
}
logger.info("Loaded value count: " + ids.size());
return ids;
}
@Override
public void storeAll(Map<String, Table> map) {
if (map == null) {
throw new RuntimeException("Illegal Store Request - Null Map");
}
if (map.containsKey(null)) {
throw new RuntimeException("Illegal Store Request - Null Key is Present");
}
logger.info("Store all called for multiple values");
BulkRequestBuilder bulkRequestBuilder =
elasticsearchConnection
.getClient()
.prepareBulk()
.setConsistencyLevel(WriteConsistencyLevel.ALL)
.setRefresh(true);
for (Map.Entry<String, Table> mapEntry : map.entrySet()) {
try {
if (mapEntry.getValue() == null) {
throw new RuntimeException(
String.format(
"Illegal Store Request - Object is Null for Table - %s", mapEntry.getKey()));
}
bulkRequestBuilder.add(
elasticsearchConnection
.getClient()
.prepareIndex(TABLE_META_INDEX, TABLE_META_TYPE, mapEntry.getKey())
.setSource(objectMapper.writeValueAsString(mapEntry.getValue())));
} catch (JsonProcessingException e) {
throw new RuntimeException("Error bulk saving meta: ", e);
}
}
bulkRequestBuilder.execute().actionGet();
}
@Override
public void deleteAll(Collection<String> keys) {
logger.info(String.format("Delete all called for multiple values: %s", keys));
BulkRequestBuilder bulRequestBuilder =
elasticsearchConnection
.getClient()
.prepareBulk()
.setConsistencyLevel(WriteConsistencyLevel.ALL)
.setRefresh(true);
for (String key : keys) {
bulRequestBuilder.add(
elasticsearchConnection
.getClient()
.prepareDelete(TABLE_META_INDEX, TABLE_META_TYPE, key));
}
bulRequestBuilder.execute().actionGet();
logger.info(String.format("Deleted multiple values: %s", keys));
}
@Override
public void delete(String key) {
logger.info("Delete called for value: " + key);
elasticsearchConnection
.getClient()
.prepareDelete()
.setConsistencyLevel(WriteConsistencyLevel.ALL)
.setRefresh(true)
.setIndex(TABLE_META_INDEX)
.setType(TABLE_META_TYPE)
.setId(key)
.execute()
.actionGet();
logger.info("Deleted value: " + key);
}
@Override
public Table load(String key) {
logger.info("Load called for: " + key);
GetResponse response =
elasticsearchConnection
.getClient()
.prepareGet()
.setIndex(TABLE_META_INDEX)
.setType(TABLE_META_TYPE)
.setId(key)
.execute()
.actionGet();
if (!response.isExists()) {
return null;
}
try {
return objectMapper.readValue(response.getSourceAsBytes(), Table.class);
} catch (Exception e) {
throw new RuntimeException("Error getting data for table: " + key);
}
}
@Override
public void store(String key, Table value) {
if (key == null || value == null || value.getName() == null) {
throw new RuntimeException(String.format("Illegal Store Request - %s - %s", key, value));
}
logger.info("Storing key: " + key);
try {
elasticsearchConnection
.getClient()
.prepareIndex()
.setIndex(TABLE_META_INDEX)
.setType(TABLE_META_TYPE)
.setConsistencyLevel(WriteConsistencyLevel.ALL)
.setSource(objectMapper.writeValueAsString(value))
.setId(key)
.setRefresh(true)
.execute()
.actionGet();
} catch (JsonProcessingException e) {
throw new RuntimeException("Error saving meta: ", e);
}
}
@Override
public Map<String, Table> loadAll(Collection<String> keys) {
logger.info("Load all called for multiple keys");
MultiGetResponse response =
elasticsearchConnection
.getClient()
.prepareMultiGet()
.add(TABLE_META_INDEX, TABLE_META_TYPE, keys)
.execute()
.actionGet();
Map<String, Table> tables = Maps.newHashMap();
for (MultiGetItemResponse multiGetItemResponse : response) {
try {
Table table =
objectMapper.readValue(
multiGetItemResponse.getResponse().getSourceAsString(), Table.class);
tables.put(table.getName(), table);
} catch (Exception e) {
throw new RuntimeException("Error getting data for table: " + multiGetItemResponse.getId());
}
}
logger.info("Loaded value count: " + tables.size());
return tables;
}
