io.vertx.cassandra.impl.CassandraClientImpl Maven / Gradle / Ivy
/*
* Copyright 2018 The Vert.x Community.
*
* 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 io.vertx.cassandra.impl;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.CqlSessionBuilder;
import com.datastax.oss.driver.api.core.cql.*;
import com.datastax.oss.driver.api.core.metadata.Metadata;
import com.datastax.oss.driver.api.core.session.Session;
import io.vertx.cassandra.CassandraClient;
import io.vertx.cassandra.CassandraClientOptions;
import io.vertx.cassandra.CassandraRowStream;
import io.vertx.cassandra.ResultSet;
import io.vertx.cassandra.impl.tracing.QueryRequest;
import io.vertx.core.*;
import io.vertx.core.internal.ContextInternal;
import io.vertx.core.internal.VertxInternal;
import io.vertx.core.spi.tracing.SpanKind;
import io.vertx.core.spi.tracing.TagExtractor;
import io.vertx.core.spi.tracing.VertxTracer;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collector;
import static io.vertx.cassandra.impl.tracing.RequestTags.REQUEST_TAG_EXTRACTOR;
/**
* @author Pavel Drankou
* @author Thomas Segismont
*/
public class CassandraClientImpl implements CassandraClient {
public static final String HOLDERS_LOCAL_MAP_NAME = "__vertx.cassandraClient.sessionHolders";
final VertxInternal vertx;
private final VertxTracer tracer;
private final String clientName;
private final CassandraClientOptions options;
private final Map holders;
private final ContextInternal creatingContext;
private boolean closed;
public CassandraClientImpl(Vertx vertx, String clientName, CassandraClientOptions options) {
Objects.requireNonNull(vertx, "vertx");
Objects.requireNonNull(clientName, "clientName");
Objects.requireNonNull(options, "options");
this.vertx = (VertxInternal) vertx;
this.tracer = ((VertxInternal) vertx).tracer();
this.clientName = clientName;
this.options = options;
this.creatingContext = ((VertxInternal) vertx).getOrCreateContext();
holders = vertx.sharedData().getLocalMap(HOLDERS_LOCAL_MAP_NAME);
SessionHolder current = holders.compute(clientName, (k, h) -> h == null ? new SessionHolder() : h.increment());
creatingContext.addCloseHook(new Closeable() {
@Override
public void close(Promise completion) {
CassandraClientImpl.this.close().onComplete(completion);
}
});
}
@Override
public synchronized boolean isConnected() {
if (closed) {
return false;
}
Session s = holders.get(clientName).session;
return s != null && !s.isClosed();
}
@Override
public Future> executeWithFullFetch(String query) {
return executeWithFullFetch(SimpleStatement.newInstance(query));
}
@Override
public Future> executeWithFullFetch(Statement statement) {
return execute(statement)
.flatMap(ResultSet::all);
}
@Override
public Future execute(String query) {
return execute(SimpleStatement.newInstance(query));
}
@Override
public Future execute(String query, Collector collector) {
return execute(SimpleStatement.newInstance(query), collector);
}
@Override
public Future execute(Statement statement) {
return executeInternal(statement)
.map(rs -> new ResultSetImpl(rs, vertx));
}
private Future executeInternal(Statement statement) {
return getSession(vertx.getOrCreateContext())
.flatMap(session -> {
Object payload;
if (tracer != null) {
payload = sendRequest(session, statement);
} else {
payload = null;
}
Future future = Future.fromCompletionStage(session.executeAsync(statement), vertx.getContext());
if (tracer != null) {
future = future.onComplete(ar -> receiveResponse(payload, ar));
}
return future;
});
}
private Object sendRequest(CqlSession session, Statement statement) {
QueryRequest request = new QueryRequest(session, statement);
return tracer.sendRequest(vertx.getContext(), SpanKind.RPC, options.getTracingPolicy(), request, "Query", (k, v) -> {
}, REQUEST_TAG_EXTRACTOR);
}
private void receiveResponse(Object payload, AsyncResult asyncResult) {
tracer.receiveResponse(vertx.getContext(), null, payload, asyncResult.cause(), TagExtractor.empty());
}
@Override
public Future execute(Statement statement, Collector collector) {
return executeAndCollect(statement, collector);
}
private Future executeAndCollect(Statement statement, Collector collector) {
C container = collector.supplier().get();
BiConsumer accumulator = collector.accumulator();
Function finisher = collector.finisher();
return queryStream(statement)
.flatMap(cassandraRowStream -> {
Promise resultPromise = Promise.promise();
cassandraRowStream.endHandler(end -> {
R result = finisher.apply(container);
resultPromise.complete(result);
});
cassandraRowStream.handler(row -> {
accumulator.accept(container, row);
});
cassandraRowStream.exceptionHandler(resultPromise::fail);
return resultPromise.future();
});
}
@Override
public Future prepare(String query) {
return getSession(vertx.getOrCreateContext())
.flatMap(session -> Future.fromCompletionStage(session.prepareAsync(query), vertx.getContext()));
}
@Override
public Future prepare(SimpleStatement statement) {
return getSession(vertx.getOrCreateContext())
.flatMap(session -> Future.fromCompletionStage(session.prepareAsync(statement), vertx.getContext()));
}
@Override
public Future queryStream(String sql) {
return queryStream(SimpleStatement.newInstance(sql));
}
@Override
public Future queryStream(Statement statement) {
return executeInternal(statement)
.map(rs -> {
ResultSet resultSet = new ResultSetImpl(rs, vertx);
CassandraRowStreamImpl stream = new CassandraRowStreamImpl(vertx.getContext());
stream.init(resultSet);
return stream;
});
}
@Override
public Future close() {
ContextInternal context = vertx.getOrCreateContext();
if (raiseCloseFlag()) {
do {
SessionHolder current = holders.get(clientName);
SessionHolder next = current.decrement();
if (next.refCount == 0) {
if (holders.remove(clientName, current)) {
if (current.session != null) {
return Future.fromCompletionStage(current.session.closeAsync(), context);
}
break;
}
} else if (holders.replace(clientName, current, next)) {
break;
}
} while (true);
}
return context.succeededFuture();
}
@Override
public Future metadata() {
return getSession(vertx.getOrCreateContext()).map(Session::getMetadata);
}
private synchronized boolean raiseCloseFlag() {
if (!closed) {
closed = true;
return true;
}
return false;
}
synchronized Future getSession(ContextInternal context) {
if (closed) {
return context.failedFuture("Client is closed");
}
SessionHolder holder = holders.get(clientName);
if (holder.session != null) {
return context.succeededFuture(holder.session);
}
return context.executeBlocking(this::connect);
}
private CqlSession connect() {
SessionHolder current = holders.get(clientName);
if (current == null) {
throw new VertxException("Client closed while connecting", true);
}
if (current.session != null) {
return current.session;
}
CqlSessionBuilder builder = options.dataStaxClusterBuilder();
CqlSession session = builder.build();
current = holders.compute(clientName, (k, h) -> h == null ? null : h.connected(session));
if (current != null) {
return current.session;
} else {
try {
session.close();
} catch (Exception ignored) {
}
throw new VertxException("Client closed while connecting", true);
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy