001package de.bytefish.pgbulkinsert.row;
002
003import de.bytefish.pgbulkinsert.exceptions.BinaryWriteFailedException;
004import de.bytefish.pgbulkinsert.pgsql.PgBinaryWriter;
005import de.bytefish.pgbulkinsert.pgsql.handlers.ValueHandlerProvider;
006import de.bytefish.pgbulkinsert.util.PostgreSqlUtils;
007import de.bytefish.pgbulkinsert.util.StringUtils;
008import org.postgresql.PGConnection;
009import org.postgresql.copy.PGCopyOutputStream;
010
011import java.sql.SQLException;
012import java.util.Arrays;
013import java.util.HashMap;
014import java.util.Map;
015import java.util.function.Consumer;
016import java.util.function.Function;
017import java.util.stream.Collectors;
018
019public class SimpleRowWriter implements  AutoCloseable {
020
021    public static class Table {
022
023        private final String schema;
024        private final String table;
025        private final String[] columns;
026
027        public Table(String table, String... columns) {
028            this(null, table, columns);
029        }
030
031        public Table(String schema, String table, String... columns) {
032            this.schema = schema;
033            this.table = table;
034            this.columns = columns;
035        }
036
037        public String getSchema() {
038            return schema;
039        }
040
041        public String getTable() {
042            return table;
043        }
044
045        public String[] getColumns() {
046            return columns;
047        }
048
049        public String getFullyQualifiedTableName(boolean usePostgresQuoting) {
050            return PostgreSqlUtils.getFullyQualifiedTableName(schema, table, usePostgresQuoting);
051        }
052
053        public String getCopyCommand(boolean usePostgresQuoting) {
054
055            String commaSeparatedColumns = Arrays.stream(columns)
056                    .map(x -> usePostgresQuoting ? PostgreSqlUtils.quoteIdentifier(x) : x)
057                    .collect(Collectors.joining(", "));
058
059            return String.format("COPY %1$s(%2$s) FROM STDIN BINARY",
060                    getFullyQualifiedTableName(usePostgresQuoting),
061                    commaSeparatedColumns);
062        }
063    }
064
065    private final Table table;
066    private final PgBinaryWriter writer;
067    private final ValueHandlerProvider provider;
068    private final Map<String, Integer> lookup;
069
070    private Function<String, String> nullCharacterHandler;
071    private boolean isOpened;
072    private boolean isClosed;
073
074    public SimpleRowWriter(final Table table, final PGConnection connection) throws SQLException {
075        this(table, connection, false);
076    }
077
078    public SimpleRowWriter(final Table table, final PGConnection connection, final boolean usePostgresQuoting) throws SQLException {
079        this.table = table;
080        this.isClosed = false;
081        this.isOpened = false;
082        this.nullCharacterHandler = (val) -> val;
083
084        this.provider = new ValueHandlerProvider();
085        this.lookup = new HashMap<>();
086
087        for (int ordinal = 0; ordinal < table.columns.length; ordinal++) {
088            lookup.put(table.columns[ordinal], ordinal);
089        }
090
091        this.writer = new PgBinaryWriter(new PGCopyOutputStream(connection, table.getCopyCommand(usePostgresQuoting), 1));
092
093        isClosed = false;
094        isOpened = true;
095    }
096
097    public synchronized void startRow(Consumer<SimpleRow> consumer) {
098
099        // We try to write a Row, but the underlying Stream to PostgreSQL has not
100        // been opened yet. We should not proceed and throw an Exception:
101        if(!isOpened) {
102            throw new BinaryWriteFailedException("The SimpleRowWriter has not been opened");
103        }
104
105        // We try to write a Row, but the underlying Stream to PostgreSQL has already
106        // been closed. We should not proceed and throw an Exception:
107        if(isClosed) {
108            throw new BinaryWriteFailedException("The PGCopyOutputStream has already been closed");
109        }
110
111        try {
112
113            writer.startRow(table.columns.length);
114
115            SimpleRow row = new SimpleRow(provider, lookup, nullCharacterHandler);
116
117            consumer.accept(row);
118
119            row.writeRow(writer);
120
121        } catch(Exception e) {
122
123            try {
124                close();
125            } catch(Exception ex) {
126                // There is nothing more we can do ...
127            }
128
129            throw e;
130        }
131    }
132
133    @Override
134    public void close()  {
135
136        // This stream shouldn't be reused, so let's store a flag here:
137        isOpened = false;
138        isClosed = true;
139
140        writer.close();
141    }
142
143    public void enableNullCharacterHandler() {
144        this.nullCharacterHandler = (val) -> StringUtils.removeNullCharacter(val);
145    }
146
147    public void setNullCharacterHandler(Function<String, String> nullCharacterHandler) {
148        this.nullCharacterHandler = nullCharacterHandler;
149    }
150}