Igor Kulikov
6 years ago
10 changed files with 184 additions and 106 deletions
@ -0,0 +1,131 @@ |
|||
/** |
|||
* Copyright © 2016-2020 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.dao.nosql; |
|||
|
|||
import com.datastax.oss.driver.api.core.cql.AsyncResultSet; |
|||
import com.datastax.oss.driver.api.core.cql.ColumnDefinitions; |
|||
import com.datastax.oss.driver.api.core.cql.ExecutionInfo; |
|||
import com.datastax.oss.driver.api.core.cql.Row; |
|||
import com.datastax.oss.driver.api.core.cql.Statement; |
|||
import com.google.common.collect.Lists; |
|||
import com.google.common.util.concurrent.FutureCallback; |
|||
import com.google.common.util.concurrent.Futures; |
|||
import com.google.common.util.concurrent.ListenableFuture; |
|||
import com.google.common.util.concurrent.MoreExecutors; |
|||
import com.google.common.util.concurrent.SettableFuture; |
|||
import edu.umd.cs.findbugs.annotations.NonNull; |
|||
import org.checkerframework.checker.nullness.qual.Nullable; |
|||
|
|||
import java.nio.ByteBuffer; |
|||
import java.util.ArrayList; |
|||
import java.util.List; |
|||
import java.util.concurrent.CompletionStage; |
|||
import java.util.concurrent.Executor; |
|||
import java.util.function.Function; |
|||
|
|||
public class TbResultSet implements AsyncResultSet { |
|||
|
|||
private final Statement originalStatement; |
|||
private final AsyncResultSet delegate; |
|||
private final Function<Statement, TbResultSetFuture> executeAsyncFunction; |
|||
|
|||
public TbResultSet(Statement originalStatement, AsyncResultSet delegate, |
|||
Function<Statement, TbResultSetFuture> executeAsyncFunction) { |
|||
this.originalStatement = originalStatement; |
|||
this.delegate = delegate; |
|||
this.executeAsyncFunction = executeAsyncFunction; |
|||
} |
|||
|
|||
@NonNull |
|||
@Override |
|||
public ColumnDefinitions getColumnDefinitions() { |
|||
return delegate.getColumnDefinitions(); |
|||
} |
|||
|
|||
@NonNull |
|||
@Override |
|||
public ExecutionInfo getExecutionInfo() { |
|||
return delegate.getExecutionInfo(); |
|||
} |
|||
|
|||
@Override |
|||
public int remaining() { |
|||
return delegate.remaining(); |
|||
} |
|||
|
|||
@NonNull |
|||
@Override |
|||
public Iterable<Row> currentPage() { |
|||
return delegate.currentPage(); |
|||
} |
|||
|
|||
@Override |
|||
public boolean hasMorePages() { |
|||
return delegate.hasMorePages(); |
|||
} |
|||
|
|||
@NonNull |
|||
@Override |
|||
public CompletionStage<AsyncResultSet> fetchNextPage() throws IllegalStateException { |
|||
return delegate.fetchNextPage(); |
|||
} |
|||
|
|||
@Override |
|||
public boolean wasApplied() { |
|||
return delegate.wasApplied(); |
|||
} |
|||
|
|||
public ListenableFuture<List<Row>> allRows(Executor executor) { |
|||
List<Row> allRows = new ArrayList<>(); |
|||
SettableFuture<List<Row>> resultFuture = SettableFuture.create(); |
|||
this.processRows(originalStatement, delegate, allRows, resultFuture, executor); |
|||
return resultFuture; |
|||
} |
|||
|
|||
private void processRows(Statement statement, |
|||
AsyncResultSet resultSet, |
|||
List<Row> allRows, |
|||
SettableFuture<List<Row>> resultFuture, |
|||
Executor executor) { |
|||
allRows.addAll(loadRows(resultSet)); |
|||
if (resultSet.hasMorePages()) { |
|||
ByteBuffer nextPagingState = resultSet.getExecutionInfo().getPagingState(); |
|||
Statement<?> nextStatement = statement.setPagingState(nextPagingState); |
|||
TbResultSetFuture resultSetFuture = executeAsyncFunction.apply(nextStatement); |
|||
Futures.addCallback(resultSetFuture, |
|||
new FutureCallback<TbResultSet>() { |
|||
@Override |
|||
public void onSuccess(@Nullable TbResultSet result) { |
|||
processRows(nextStatement, result, |
|||
allRows, resultFuture, executor); |
|||
} |
|||
|
|||
@Override |
|||
public void onFailure(Throwable t) { |
|||
resultFuture.setException(t); |
|||
} |
|||
}, executor != null ? executor : MoreExecutors.directExecutor() |
|||
); |
|||
} else { |
|||
resultFuture.set(allRows); |
|||
} |
|||
} |
|||
|
|||
List<Row> loadRows(AsyncResultSet resultSet) { |
|||
return Lists.newArrayList(resultSet.currentPage()); |
|||
} |
|||
|
|||
} |
|||
@ -1,65 +0,0 @@ |
|||
/** |
|||
* Copyright © 2016-2020 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.dao.nosql; |
|||
|
|||
import com.datastax.oss.driver.api.core.cql.AsyncResultSet; |
|||
import com.datastax.oss.driver.api.core.cql.Row; |
|||
import com.google.common.collect.Lists; |
|||
import com.google.common.util.concurrent.Futures; |
|||
import com.google.common.util.concurrent.ListenableFuture; |
|||
import com.google.common.util.concurrent.SettableFuture; |
|||
|
|||
import java.util.ArrayList; |
|||
import java.util.List; |
|||
import java.util.concurrent.CompletionStage; |
|||
import java.util.concurrent.Executor; |
|||
import java.util.stream.Collectors; |
|||
|
|||
public class ResultSetUtils { |
|||
|
|||
public static ListenableFuture<List<Row>> allRows(AsyncResultSet resultSet, Executor executor) { |
|||
List<ListenableFuture<AsyncResultSet>> futures = new ArrayList<>(); |
|||
futures.add(Futures.immediateFuture(resultSet)); |
|||
while (resultSet.hasMorePages()) { |
|||
futures.add(toListenable(resultSet.fetchNextPage())); |
|||
} |
|||
return Futures.transform( Futures.allAsList(futures), |
|||
resultSets -> resultSets.stream() |
|||
.map(rs -> loadRows(rs)) |
|||
.flatMap(rows -> rows.stream()) |
|||
.collect(Collectors.toList()), |
|||
executor |
|||
); |
|||
} |
|||
|
|||
private static <T> ListenableFuture<T> toListenable(CompletionStage<T> completable) { |
|||
SettableFuture<T> future = SettableFuture.create(); |
|||
completable.whenComplete( |
|||
(r, ex) -> { |
|||
if (ex != null) { |
|||
future.setException(ex); |
|||
} else { |
|||
future.set(r); |
|||
} |
|||
} |
|||
); |
|||
return future; |
|||
} |
|||
|
|||
private static List<Row> loadRows(AsyncResultSet resultSet) { |
|||
return Lists.newArrayList(resultSet.currentPage()); |
|||
} |
|||
} |
|||
Loading…
Reference in new issue