org.springframework.data.cassandra.core.cql.legacy.AsyncResultStream Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of spring-data-cassandra Show documentation
Cassandra support for Spring Data
There is a newer version: 4.3.2
Show newest version
/*
 * Copyright 2019-2023 the original author or 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
 *
 *      https://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.springframework.data.cassandra.core.cql.legacy;

import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.stream.Collector;

import org.springframework.data.cassandra.core.cql.RowMapper;
import org.springframework.util.Assert;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.util.concurrent.SettableListenableFuture;

import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
import com.datastax.oss.driver.api.core.cql.Row;

/**
 * Asynchronous supplied sequence of elements supporting sequential operations over a {@link AsyncResultSet a result
 * set}. An asynchronous stream represents a pipeline of operations to process a {@link AsyncResultSet}.
 *
 * @author Mark Paluch
 * @since 4.0
 */
@Deprecated(since = "4.0", forRemoval = true)
class AsyncResultStream {

	private final AsyncResultSet resultSet;

	private final RowMapper mapper;

	private AsyncResultStream(AsyncResultSet resultSet, RowMapper mapper) {
		this.resultSet = resultSet;
		this.mapper = mapper;
	}

	/**
	 * Creates a {@link AsyncResultStream} given {@link AsyncResultSet}.
	 *
	 * @param resultSet the result set to process.
	 * @return a new {@link AsyncResultStream} instance.
	 */
	static AsyncResultStream from(AsyncResultSet resultSet) {

		Assert.notNull(resultSet, "AsyncResultSet must not be null");

		return new AsyncResultStream<>(resultSet, (row, rowNum) -> row);
	}

	/**
	 * Returns a stream consisting of the results of applying the given function to the elements of this stream.
	 * 
	 * This is an intermediate operation.
	 *
	 * @param  The element type of the new stream
	 * @param mapper a non-interfering, stateless {@link RowMapper}.
	 */
	 AsyncResultStream map(RowMapper mapper) {

		Assert.notNull(mapper, "RowMapper must not be null");

		return new AsyncResultStream<>(resultSet, mapper);
	}

	/**
	 * Performs a mutable reduction operation on the elements of this stream using a {@link Collector} resulting in a
	 * {@link ListenableFuture}.
	 * 

	 * This is a terminal operation.
	 *
	 * @param  the type of the result
	 * @param  the intermediate accumulation type of the {@link Collector}
	 * @param collector the {@link Collector} describing the reduction
	 * @return the result of the reduction
	 */
	 ListenableFuture collect(Collector collector) {

		Assert.notNull(collector, "Collector must not be null");

		SettableListenableFuture future = new SettableListenableFuture<>();
		CollectState collectState = new CollectState<>(collector);

		collectState.collectAsync(future, this.resultSet);

		return future;
	}

	/**
	 * Performs an action for each element of this stream. This method returns a {@link ListenableFuture} that completes
	 * without a value ({@code null}) once all elements have been processed.
	 * 

	 * This is a terminal operation.
	 * 
	 * If the action accesses shared state, it is responsible for providing the required synchronization.
	 *
	 * @param action a non-interfering action to perform on the elements.
	 */
	ListenableFuture forEach(Consumer action) {

		Assert.notNull(action, "Action must not be null");

		SettableListenableFuture future = new SettableListenableFuture<>();
		ForwardLoopState loopState = new ForwardLoopState(action);

		loopState.forEachAsync(future, this.resultSet);

		return future;
	}

	/**
	 * State object for forward-looping using {@code forEach}.
	 */
	class ForwardLoopState {

		private final AtomicInteger rowNumber = new AtomicInteger();
		private final Consumer consumer;

		ForwardLoopState(Consumer consumer) {
			this.consumer = consumer;
		}

		void peekRow(Iterable rows) {
			rows.forEach(row -> consumer.accept(mapper.mapRow(row, rowNumber.incrementAndGet())));
		}

		/**
		 * Recursive async iteration.
		 *
		 * @param target
		 * @param resultSet
		 */
		void forEachAsync(SettableListenableFuture target, AsyncResultSet resultSet) {

			if (target.isCancelled()) {
				return;
			}

			try {
				peekRow(resultSet.currentPage());
			} catch (RuntimeException e) {
				target.setException(e);
				return;
			}

			if (!resultSet.hasMorePages()) {
				target.set(null);
			} else {

				CompletionStage nextPage = resultSet.fetchNextPage();

				nextPage.whenComplete((nextResultSet, throwable) -> {

					if (throwable != null) {
						target.setException(throwable);
					} else {
						forEachAsync(target, nextResultSet);
					}
				});
			}
		}
	}

	/**
	 * State object for collecting rows using {@code collect}.
	 */
	class CollectState {

		private final AtomicInteger rowNumber = new AtomicInteger();
		private volatile A intermediate;
		private final Collector collector;

		CollectState(Collector collector) {
			this.collector = collector;
			this.intermediate = collector.supplier().get();
		}

		void collectPage(Iterable rows) {

			for (Row row : rows) {
				collector.accumulator().accept(intermediate, mapper.mapRow(row, rowNumber.incrementAndGet()));
			}
		}

		R finish() {
			return collector.finisher().apply(intermediate);
		}

		/**
		 * Recursive collection.
		 *
		 * @param target
		 * @param resultSet
		 */
		void collectAsync(SettableListenableFuture target, AsyncResultSet resultSet) {

			if (target.isCancelled()) {
				return;
			}

			try {
				collectPage(resultSet.currentPage());
			} catch (RuntimeException e) {
				target.setException(e);
				return;
			}

			if (!resultSet.hasMorePages()) {
				target.set(finish());
			} else {

				CompletionStage nextPage = resultSet.fetchNextPage();

				nextPage.whenComplete((nextResultSet, throwable) -> {

					if (throwable != null) {
						target.setException(throwable);
					} else {
						collectAsync(target, nextResultSet);
					}
				});
			}
		}
	}
}