Running a COPY IN within a Transactional completely halts the process until force quit

[favna]

Hi there,

I'll preface by stating that I'm in the same team that @ArjanSchouten used to be in when he created #500 (he has since departed us, lest I'd first ask him directly)

I am updating one of our applications that uses the in #500 added COPY IN support from Java 11 to Java 17. As part of that I am now using the by this library provided methods in an update of our utility library that uses those functions. So far I've had success with this in other applications, however, the current one has the unique challenge of having to run within an @Transactional and no matter what I try the process always completely halts by the time it gets to the part where it's supposed to do the COPY IN query.

Our utility/wrapper class is as follows:

package reverse.dns.postgresqlsupport;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Joiner;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.r2dbc.postgresql.api.PostgresqlConnection;
import io.r2dbc.spi.Connection;
import io.r2dbc.spi.ConnectionFactory;
import io.r2dbc.spi.Wrapped;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.buffer.DataBuffer;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.ByteBuffer;
import java.time.Instant;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.time.temporal.Temporal;
import java.util.Date;
import java.util.List;
import java.util.Objects;
import java.util.function.Consumer;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.stream.Collectors.joining;

/**
 * Helper component for inserting data in batch using COPY IN.
 */
@Slf4j
@RequiredArgsConstructor
public class CopyInTemplate {

	private static final String DELIMITER = "|";

	private static final String CSV_DELIMITER = ";";

	private static final String ROW_DELIMITER = "\n";

	private static final String NULL_VALUE = "<<null>>";

	private static final String ESCAPE = "\\";

	private static final Joiner COLUMN_JOINER = Joiner.on(",");

	private final ConnectionFactory connectionFactory;

	/**
	 * Copy data into the columns of the given table.
	 *
	 * <p>
	 * For timestamps only {@link ZonedDateTime} and {@link Instant} are supported, representing the time in UTC.
	 * Before insert these timestamps are converted to ISO-8601 format using
	 * {@link DateTimeFormatter#ISO_ZONED_DATE_TIME} and {@link DateTimeFormatter#ISO_INSTANT} respectively
	 * </p>
	 *
	 * @param tableName Name of the table to copy data into.
	 * @param columnNames Names of the columns to copy data into.
	 * @param rows Individual rows that can be written/staged, representing the values for the supplied columns.
	 * @return The number of rows copied wrapped in a {@link Mono}.
	 */
	public Mono<Long> copyIn(final String tableName, final List<String> columnNames, final Flux<List<?>> rows) {
		final Flux<ByteBuf> frames = rows.map(row -> byteArrayToByteBuf(rowToString(row).getBytes(UTF_8)));

		return Mono.usingWhen(
						connectionFactory.create(),
						c -> ((Wrapped<PostgresqlConnection>) c)
								.unwrap()
								.copyIn(copyInSql(tableName, columnNames))
								.fromMany(frames.map(Mono::just))
								.build(),
						Connection::close
				)
				.switchIfEmpty(Mono.error(new RuntimeException("No connection found")));
	}

	/**
	 * Copy a CSV into the greenplum database.
	 *
	 * <p>
	 * For timestamps only {@link ZonedDateTime} and {@link Instant} are supported, representing the time in UTC.
	 * Before insert these timestamps are converted to ISO-8601 format using
	 * {@link DateTimeFormatter#ISO_ZONED_DATE_TIME} and {@link DateTimeFormatter#ISO_INSTANT} respectively
	 * </p>
	 *
	 * @param tableName Name of the table to copy data into.
	 * @param columnNames Names of the columns to copy data into.
	 * @param dataBuffers The CSV data buffer that contains the data to be copied.
	 * @return The number of rows copied wrapped in a {@link Mono}.
	 */
	public Mono<Long> copyInCsv(final String tableName, final List<String> columnNames, final Flux<DataBuffer> dataBuffers) {
		final Flux<ByteBuf> byteBuffers = dataBuffers.map(this::dataBufferToByteBuf);

		return Mono.usingWhen(
						connectionFactory.create(),
						c -> ((Wrapped<PostgresqlConnection>) c)
								.unwrap()
								.copyIn(copyInCsvSql(tableName, columnNames))
								.fromMany(byteBuffers.map(Mono::just))
								.build(),
						Connection::close
				)
				.switchIfEmpty(Mono.error(new RuntimeException("No connection found")));
	}

	@VisibleForTesting
	String valueToString(Object value) {
		if (value instanceof Temporal || value instanceof Date) {
			return temporalToCopyValue(value);
		}

		if (value instanceof final String stringValue) {
			value = stringValueToCopyValue(stringValue);
		}

		return Objects.toString(value, NULL_VALUE);
	}

	/**
	 * Converts a Spring {@link DataBuffer} to a netty {@link ByteBuf}
	 *
	 * @param dataBuffer The {@link DataBuffer} to convert
	 * @return The converted {@link ByteBuf}
	 */
	private ByteBuf dataBufferToByteBuf(DataBuffer dataBuffer) {
		// First convert the DataBuffer to a ByteBuffer
		ByteBuffer destinationByteBuffer = ByteBuffer.allocate(dataBuffer.readableByteCount());
		dataBuffer.toByteBuffer(destinationByteBuffer);

		// Then convert the ByteBuffer to a ByteBuf
		ByteBuf byteBuf = ByteBufAllocator.DEFAULT.buffer(dataBuffer.readableByteCount());
		byteBuf.writeBytes(destinationByteBuffer);
		return byteBuf;
	}

	/**
	 * Converts a native {@link byte[]} to a {@link ByteBuf}.
	 *
	 * @param byteArray The {@link byte[]} to convert
	 * @return The converted {@link ByteBuf}
	 */
	private ByteBuf byteArrayToByteBuf(byte[] byteArray) {
		ByteBufAllocator allocator = ByteBufAllocator.DEFAULT;
		ByteBuf byteBuf = allocator.buffer(byteArray.length);
		byteBuf.writeBytes(byteArray);
		return byteBuf;
	}

	private String rowToString(final List<?> rowValues) {
		final var row = rowValues.stream()
				.map(this::valueToString)
				.collect(joining(DELIMITER));

		log.trace("appendRow :: row={}", row);

		return row + ROW_DELIMITER;
	}

	private String temporalToCopyValue(final Object value) {
		if (value instanceof final ZonedDateTime zonedDateTime) {
			// We cannot use DateTimeFormatter.ISO_ZONED_DATE_TIME because it adds the timezone between
			// [square brackets] which is not supported by PostgreSQL
			final var formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSSSXXX");
			return zonedDateTime.format(formatter);
		} else if (value instanceof final Instant instant) {
			return DateTimeFormatter.ISO_INSTANT.format(instant);
		} else {
			throw new IllegalArgumentException(
					"Only ZonedDateTime and Instant supported for dates/timestamps, was: " +
					value.getClass().getName());
		}
	}

	private String stringValueToCopyValue(final String value) {
		return value
				.replace(ESCAPE, ESCAPE + ESCAPE)// Escape escape chars (\'s)
				.replace("|", ESCAPE + "|") // Escape |'s
				.replace("\n", ESCAPE + "n"); // Escape enters
	}

	private String copyInSql(final String tableName, final List<String> columnNames) {
		return copyInCsvSql(tableName, columnNames, bld -> bld
				.append("FROM STDIN WITH")
				.append("  DELIMITER '" + DELIMITER + "'")
				.append("  NULL '" + NULL_VALUE + "'")
				.append("  ESCAPE '" + ESCAPE + "'")
		);
	}

	private String copyInCsvSql(final String tableName, final List<String> columnNames) {
		return copyInCsvSql(tableName, columnNames, bld -> bld
				.append("FROM STDIN WITH CSV HEADER")
				.append("  DELIMITER '" + CSV_DELIMITER + "'")
				.append("  ESCAPE '" + ESCAPE + "'")
		);
	}

	private String copyInCsvSql(final String tableName, final List<String> columnNames, final Consumer<StringBuilder> consumer) {
		final var bld = new StringBuilder();

		bld.append("COPY ")
				.append(tableName)
				.append(" (")
				.append(COLUMN_JOINER.join(columnNames))
				.append(") ");

		consumer.accept(bld);

		final var statement = bld.toString();
		log.debug("copyInSql :: statement={}", statement);
		return statement;
	}

}

We call this function a few levels deep in our application so here's the relevant code bits:

public interface CustomExtractedDataRepository {
	/**
	 * Store the given query results into the configured table.
	 */
	Mono<Long> store(long extractionId, String loadCycleId, String queryId, Flux<BindingSet> queryResults);
}

@Slf4j
@RequiredArgsConstructor
class CustomExtractedDataRepositoryImpl implements CustomExtractedDataRepository {
	public static final String XSD_DECIMAL = XSD.DECIMAL.stringValue();
	public static final String XSD_DOUBLE = XSD.DOUBLE.stringValue();
	public static final String XSD_FLOAT = XSD.FLOAT.stringValue();
	private final CopyInTemplate copyInTemplate;
	private final QueryProperties queryProperties;

	@Override
	public Mono<Long> store(long extractionId, String loadCycleId, String queryId, Flux<BindingSet> queryResults) {
		log.debug("store :: extractionId={}, loadCycleId={}, queryId={}", extractionId, loadCycleId, queryId);

		var query = queryProperties.query(queryId);
		var tableName = SCHEMA_NAME + "." + query.getTargetTable();

		return copyInTemplate.copyIn(
				tableName,
				columns(query),
				queryResults.map(bindingSet -> toRow(bindingSet, query.getColumns(), extractionId, loadCycleId))
		);
	}
}

import reverse.dns.ExtractionResult;
import org.springframework.data.r2dbc.repository.R2dbcRepository;
import org.springframework.stereotype.Repository;

@Repository
public interface ExtractedDataRepository extends R2dbcRepository<ExtractionResult, Long>, CustomExtractedDataRepository {
}

public interface Transactions {
	String TRANSACTION_MANAGER = "connectionFactoryTransactionManager";
}

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import reverse.dns.ExtractionResult;
import reverse.dns.SparqlResult;
import reverse.dns.ExtractionEntity;
import reverse.dns.ExtractionRepository;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.Assert;
import reactor.core.publisher.Mono;

import java.util.List;

import static reverse.dns.Transactions.TRANSACTION_MANAGER;
import static reactor.core.publisher.Mono.defer;

@Slf4j
@Service
@RequiredArgsConstructor
class ExtractionServiceImpl implements ExtractionService {
	private final SparqlRunner sparqlRunner;
	private final ExtractionResultManager extractionResultManager;
	private final ExtractionRepository extractionRepository;

	@Override
	@Transactional(TRANSACTION_MANAGER)
	public Mono<List<ExtractionResult>> extract(String loadCycleId, String sparqlEndpoint, String querySetId) {
		log.info("extract :: loadCycleId={}, sparqlEndpoint={}, querySetId={}",
				loadCycleId, sparqlEndpoint, querySetId);

		Assert.notNull(loadCycleId, "Missing loadCycleId");
		Assert.notNull(sparqlEndpoint, "Missing sparqlEndpoint");
		Assert.notNull(querySetId, "Missing querySetId");

		return extractionRepository.findBy(loadCycleId, querySetId, sparqlEndpoint)
				.flatMap(existing -> Mono.error(new IllegalArgumentException(
						"Extractions already exists: " + existing)))
				.then(defer(() -> create(loadCycleId, sparqlEndpoint, querySetId)))
				.flatMap(created -> extract(created, querySetId, sparqlEndpoint));
	}

	private Mono<ExtractionEntity> create(String loadCycleId, String sparqlEndpoint, String querySetId) {
		var extraction = ExtractionEntity.builder()
				.loadCycleId(loadCycleId)
				.querySetId(querySetId)
				.sparqlEndpoint(sparqlEndpoint)
				.build();

		return extractionRepository.save(extraction);
	}

	private Mono<List<ExtractionResult>> extract(ExtractionEntity extraction, String querySetId, String endpoint) {
		return sparqlRunner.runQuerySet(querySetId, endpoint)
				.concatMap(result -> storeResult(result, extraction))
				.collectList();
	}

	private Mono<ExtractionResult> storeResult(SparqlResult result, ExtractionEntity extraction) {
		return extractionResultManager.storeResult(extraction.getId(), extraction.getLoadCycleId(), result);
	}
}

Lastly, this is where the extract function from the block above gets called and as you can see I do block the reactive stream here, so it isn't just as simple as not subscribing/blocking.

.then(extractionService.extract(loadCycleId, sparqlEndpoint, querySetLd))
				.doOnNext(results -> log.info("extract :: completed, results={}", results))
				.block();

The last log line that I see before the process hangs is from log.debug("store :: extractionId={}, loadCycleId={}, queryId={}", extractionId, loadCycleId, queryId); and if I write that particular function like so:

	@Override
	public Mono<Long> store(long extractionId, String loadCycleId, String queryId, Flux<BindingSet> queryResults) {
		log.debug("store :: extractionId={}, loadCycleId={}, queryId={}", extractionId, loadCycleId, queryId);

		var query = queryProperties.query(queryId);
		var tableName = SCHEMA_NAME + "." + query.getTargetTable();
		Flux<List<?>> rows = queryResults.map(bindingSet -> toRow(bindingSet, query.getColumns(), extractionId, loadCycleId));
		
		rows.collectList().doOnNext(loggableRows -> log.debug("store :: loggableRows={}", loggableRows)).subscribe();

		return copyInTemplate.copyIn(
				tableName,
				columns(query),
				rows
		);
	}

Then I do see the rows being logged and they all look correct, which leads me to believe that the hanging process occurs somewhere in the COPY IN.

I have also tried to change the syntax of the copyIn method in CopyInTemplate to:

		return Mono.from(connectionFactory.create())
				.flatMap(connection -> Mono.from(((PostgresqlConnection) connection)
								.copyIn(copyInSql(tableName, columnNames))
								.fromMany(frames.map(Mono::just))
								.build()
						)
				)
				.switchIfEmpty(Mono.error(new RuntimeException("No connection found")));

so it doesn't use Mono.usingWhen() and doesn't close the connection in the asyncCleanup for that method but this still has the process hang.

I'm absolutely entirely out of ideas on how to solve this issue and COPY IN really is a big must for us because we're talking about inserting upwards of 150_000 rows at the minimum or 700_000 at a more common occurrence and even that is with a vastly limited dataset so it's only going to grow and grow.

All that said I could really use some help here to get the COPY IN functional within an @Transactional.

Some version info:

• org.postgresql:r2dbc-postgresql:1.0.2.RELEASE
• Java 17.0.9

[mp911de]

Thanks for reaching out. I'm not a user of COPY so I have no experience with the command. I don't know whether COPY is working in the context of a transaction (have you tried to spike on transactions with PGJDBC?). I suggest crafting a small sample with manual transaction control and a small data set, then capture the debug log to see which frames arrive at the client and when the client stops responding.

Looking at the code, there seems to be some interaction with entities while you work with COPY. I could well imagine that while COPY is active, you cannot run other SQL queries on the same connection (note that the connection is being reused across multiple calls within a transaction).

[favna]

Hmm, I'll have to look into these suggestions. I have not yet tried to spike on transactions and I'll see if I can get a small repro worked out. It'll be a few weeks at least because my vacation is starting soon and afterwards I got a busy agenda but I'll be sure to get back to this.

[ArjanSchouten]

Hey @favna , this is a batch processing job and was set up such that it could work without transactions for the COPY if I remember correctly. Is it later switched back to using transactions?

Essentially, COPY is a transaction in and of itself because only when the COPY ends, the data is written to the table.

And also important, it is a Greenplum database.
I assume the behavior for Greenplum is different compared to Postgresql. Can you test with plain Postgresql to verify this isn't caused by Greenplum?

And I don't understand why a Java 11 to Java 17 update would cause this problem.

[favna]

And I don't understand why a Java 11 to Java 17 update would cause this problem.

Well not necessarily caused by I suppose but I incorporated your work for R2DBC to support Copy In and we (finally) updated the postgresql-support library to use it instead of JDBC for Copy In. We then fully dropped JDBC from the library. This change was made in the same process as setting the target source to Java 17 and so included in the same major version.

Other than that I'll keep these remarks in mind for when I'm back from vacation. Thanks Arjan!

p.s. I was joking with the others "it would be funny if Arjan chimes in, I intentionally tagged him but I suspect he's too busy". You can betcha I'm sending them a message that you did reply. Hope you're doing well 😁!

[favna]

Updating this thread with a link to a WIP repro. Sadly I have to actually properly reproduce the issue, as I'm currently running into a completely different problem, namely:

io.r2dbc.postgresql.ExceptionFactory$PostgresqlNonTransientResourceException: COPY from stdin failed: Copy operation failed: For input string: "JohnDoe"
	at io.r2dbc.postgresql.ExceptionFactory.createException(ExceptionFactory.java:109) ~[r2dbc-postgresql-1.0.4.RELEASE.jar:1.0.4.RELEASE]
	Suppressed: reactor.core.publisher.FluxOnAssembly$OnAssemblyException: 

I suspect the dummy data I am generating is invalid somehow.

Link: https://github.com/favna/copy-in-repro