001// Licensed under the MIT license. See LICENSE file in the project root for full license information.
002
003package de.bytefish.pgbulkinsert;
004
005import de.bytefish.pgbulkinsert.configuration.Configuration;
006import de.bytefish.pgbulkinsert.configuration.IConfiguration;
007import de.bytefish.pgbulkinsert.exceptions.SaveEntityFailedException;
008import de.bytefish.pgbulkinsert.mapping.AbstractMapping;
009import de.bytefish.pgbulkinsert.pgsql.PgBinaryWriter;
010import org.postgresql.PGConnection;
011import org.postgresql.copy.PGCopyOutputStream;
012
013import java.sql.SQLException;
014import java.util.Collection;
015import java.util.Objects;
016import java.util.stream.Stream;
017
018public class PgBulkInsert<TEntity> implements IPgBulkInsert<TEntity> {
019
020    private final IConfiguration configuration;
021    private final AbstractMapping<TEntity> mapping;
022
023    public PgBulkInsert(AbstractMapping<TEntity> mapping) {
024        this(new Configuration(), mapping);
025    }
026
027    public PgBulkInsert(IConfiguration configuration, AbstractMapping<TEntity> mapping)
028    {
029        Objects.requireNonNull(configuration, "'configuration' has to be set");
030        Objects.requireNonNull(mapping, "'mapping' has to be set");
031
032        this.configuration = configuration;
033        this.mapping = mapping;
034    }
035
036    @Override
037    public void saveAll(PGConnection connection, Stream<TEntity> entities) throws SQLException {
038        // Wrap the CopyOutputStream in our own Writer:
039        try (PgBinaryWriter bw = new PgBinaryWriter(new PGCopyOutputStream(connection, mapping.getCopyCommand(), 1), configuration.getBufferSize())) {
040            // Insert Each Column:
041            entities.forEach(entity -> saveEntitySynchonized(bw, entity));
042        }
043    }
044
045    public void saveAll(PGConnection connection, Collection<TEntity> entities) throws SQLException {
046        saveAll(connection, entities.stream());
047    }
048
049    private void saveEntity(PgBinaryWriter bw, TEntity entity) throws SaveEntityFailedException {
050        // Start a new Row in PostgreSQL:
051        bw.startRow(mapping.getColumns().size());
052
053        try {
054                // Iterate over each column mapping:
055                mapping.getColumns().forEach(column -> {
056                        column.getWrite().accept(bw, entity);
057                });
058        } catch (Exception e) {
059            throw new SaveEntityFailedException(e);
060        }
061    }
062
063    private void saveEntitySynchonized(PgBinaryWriter bw, TEntity entity) throws SaveEntityFailedException {
064        synchronized (bw) {
065                saveEntity(bw, entity);
066        }
067    }
068}