/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License 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.facebook.presto.cassandra;
import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.google.common.collect.ImmutableList;
import com.google.common.io.Resources;
import io.airlift.json.JsonCodec;
import io.airlift.log.Logger;
import io.airlift.units.Duration;
import org.apache.cassandra.service.CassandraDaemon;
import javax.management.ObjectName;
import java.io.File;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.net.InetSocketAddress;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import static com.datastax.driver.core.ProtocolVersion.V3;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.io.Files.createTempDir;
import static com.google.common.io.Files.write;
import static com.google.common.io.Resources.getResource;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MINUTES;
import static org.testng.Assert.assertEquals;
public final class EmbeddedCassandra
{
private static Logger log = Logger.get(EmbeddedCassandra.class);
private static final String HOST = "127.0.0.1";
private static final int PORT = 9142;
private static CassandraSession session;
private static boolean initialized;
private EmbeddedCassandra() {}
public static synchronized void start()
throws Exception
{
if (initialized) {
return;
}
log.info("Starting cassandra...");
System.setProperty("cassandra.config", "file:" + prepareCassandraYaml());
System.setProperty("cassandra-foreground", "true");
System.setProperty("cassandra.native.epoll.enabled", "false");
CassandraDaemon cassandraDaemon = new CassandraDaemon();
cassandraDaemon.activate();
Cluster cluster = Cluster.builder()
.withProtocolVersion(V3)
.withClusterName("TestCluster")
.addContactPointsWithPorts(ImmutableList.of(
new InetSocketAddress(HOST, PORT)))
.build();
CassandraSession session = new NativeCassandraSession(
"EmbeddedCassandra",
JsonCodec.listJsonCodec(ExtraColumnMetadata.class),
cluster,
new Duration(1, MINUTES));
try {
checkConnectivity(session);
}
catch (RuntimeException e) {
cluster.close();
cassandraDaemon.deactivate();
throw e;
}
EmbeddedCassandra.session = session;
initialized = true;
}
private static String prepareCassandraYaml()
throws IOException
{
String original = Resources.toString(getResource("cu-cassandra.yaml"), UTF_8);
File tempDirFile = createTempDir();
tempDirFile.deleteOnExit();
Path tmpDirPath = tempDirFile.toPath();
Path dataDir = tmpDirPath.resolve("data");
Files.createDirectory(dataDir);
String modified = original.replaceAll("\\$\\{data_directory\\}", dataDir.toAbsolutePath().toString());
Path yamlLocation = tmpDirPath.resolve("cu-cassandra.yaml");
write(modified, yamlLocation.toFile(), UTF_8);
return yamlLocation.toAbsolutePath().toString();
}
public static synchronized CassandraSession getSession()
{
checkIsInitialized();
return requireNonNull(session, "cluster is null");
}
public static synchronized String getHost()
{
checkIsInitialized();
return HOST;
}
public static synchronized int getPort()
{
checkIsInitialized();
return PORT;
}
private static void checkIsInitialized()
{
checkState(initialized, "EmbeddedCassandra must be started with #start() method before retrieving the cluster retrieval");
}
private static void checkConnectivity(CassandraSession session)
{
ResultSet result = session.execute("SELECT release_version FROM system.local");
List<Row> rows = result.all();
assertEquals(rows.size(), 1);
String version = rows.get(0).getString(0);
log.info("Cassandra version: %s", version);
}
public static void flush(String keyspace, String table)
throws Exception
{
ManagementFactory
.getPlatformMBeanServer()
.invoke(
new ObjectName("org.apache.cassandra.db:type=StorageService"),
"forceKeyspaceFlush",
new Object[] {keyspace, new String[] {table}},
new String[] {"java.lang.String", "[Ljava.lang.String;"});
}
public static void refreshSizeEstimates()
throws Exception
{
ManagementFactory
.getPlatformMBeanServer()
.invoke(
new ObjectName("org.apache.cassandra.db:type=StorageService"),
"refreshSizeEstimates",
new Object[] {},
new String[] {});
}
}