001// Licensed under the MIT license. See LICENSE file in the project root for full license information.
002
003package de.bytefish.pgbulkinsert.pgsql;
004
005import de.bytefish.pgbulkinsert.exceptions.BinaryWriteFailedException;
006import de.bytefish.pgbulkinsert.pgsql.handlers.IValueHandler;
007
008import java.io.BufferedOutputStream;
009import java.io.DataOutputStream;
010import java.io.IOException;
011import java.io.OutputStream;
012
013public class PgBinaryWriter implements AutoCloseable {
014
015    private final transient DataOutputStream buffer;
016
017    public PgBinaryWriter(final OutputStream out) {
018        this(out, 65536);
019    }
020
021    public PgBinaryWriter(final OutputStream out, final int bufferSize) {
022                buffer = new DataOutputStream(new BufferedOutputStream(out, bufferSize));
023                writeHeader();
024    }
025
026    public void startRow(int numColumns) {
027        try {
028            buffer.writeShort(numColumns);
029                } catch (IOException e) {
030                        Throwable t = e.getCause();
031                        if (null != t) {
032                                throw new BinaryWriteFailedException(t);
033                        } else {
034                                throw new BinaryWriteFailedException(e);
035                        }
036                }
037    }
038
039    public <TTargetType> void write(final IValueHandler<TTargetType> handler, final TTargetType value) {
040        handler.handle(buffer, value);
041    }
042
043    /**
044     * Writes primitive boolean to the output stream
045     *
046     * @param value value to write
047     *
048     */
049        public void writeBoolean(boolean value) {
050                try {
051                        buffer.writeInt(1);
052                        if (value) {
053                                buffer.writeByte(1);
054                        } else {
055                                buffer.writeByte(0);
056                        }
057                } catch (IOException e) {
058                        Throwable t = e.getCause();
059                        if (null != t) {
060                                throw new BinaryWriteFailedException(t);
061                        } else {
062                                throw new BinaryWriteFailedException(e);
063                        }
064                }
065        }
066
067
068    /**
069     * Writes primitive byte to the output stream
070     *
071     * @param value value to write
072     *
073     */
074        public void writeByte(int value) {
075                try {
076                        buffer.writeInt(1);
077                        buffer.writeByte(value);
078                } catch (IOException e) {
079                        Throwable t = e.getCause();
080                        if (null != t) {
081                                throw new BinaryWriteFailedException(t);
082                        } else {
083                                throw new BinaryWriteFailedException(e);
084                        }
085                }
086        }
087
088    /**
089     * Writes primitive short to the output stream
090     *
091     * @param value value to write
092     *
093     */
094        public void writeShort(int value) {
095                try {
096                        buffer.writeInt(2);
097                        buffer.writeShort(value);
098                } catch (IOException e) {
099                        Throwable t = e.getCause();
100                        if (null != t) {
101                                throw new BinaryWriteFailedException(t);
102                        } else {
103                                throw new BinaryWriteFailedException(e);
104                        }
105                }
106        }
107
108    /**
109     * Writes primitive integer to the output stream
110     *
111     * @param value value to write
112     *
113     */
114        public void writeInt(int value) {
115                try {
116                        buffer.writeInt(4);
117                        buffer.writeInt(value);
118                } catch (IOException e) {
119                        Throwable t = e.getCause();
120                        if (null != t) {
121                                throw new BinaryWriteFailedException(t);
122                        } else {
123                                throw new BinaryWriteFailedException(e);
124                        }
125                }
126        }
127
128    /**
129     * Writes primitive long to the output stream
130     *
131     * @param value value to write
132     *
133     */
134        public void writeLong(long value) {
135                try {
136                        buffer.writeInt(8);
137                        buffer.writeLong(value);
138                } catch (IOException e) {
139                        Throwable t = e.getCause();
140                        if (null != t) {
141                                throw new BinaryWriteFailedException(t);
142                        } else {
143                                throw new BinaryWriteFailedException(e);
144                        }
145                }
146        }
147
148    /**
149     * Writes primitive float to the output stream
150     *
151     * @param value value to write
152     *
153     */
154        public void writeFloat(float value) {
155                try {
156                buffer.writeInt(4);
157                buffer.writeFloat(value);
158                } catch (IOException e) {
159                        Throwable t = e.getCause();
160                        if (null != t) {
161                                throw new BinaryWriteFailedException(t);
162                        } else {
163                                throw new BinaryWriteFailedException(e);
164                        }
165                }
166        }
167
168    /**
169     * Writes primitive double to the output stream
170     *
171     * @param value value to write
172     *
173     */
174        public void writeDouble(double value) {
175                try {
176                        buffer.writeInt(8);
177                        buffer.writeDouble(value);
178                } catch (IOException e) {
179                        Throwable t = e.getCause();
180                        if (null != t) {
181                                throw new BinaryWriteFailedException(t);
182                        } else {
183                                throw new BinaryWriteFailedException(e);
184                        }
185                }
186        }
187
188        /**
189         * Writes a Null Value.
190         */
191        public void writeNull() {
192                try {
193                        buffer.writeInt(-1);
194                } catch (IOException e) {
195                        Throwable t = e.getCause();
196                        if (null != t) {
197                                throw new BinaryWriteFailedException(t);
198                        } else {
199                                throw new BinaryWriteFailedException(e);
200                        }
201                }
202        }
203
204    @Override
205    public void close() {
206        try {
207            buffer.writeShort(-1);
208
209            buffer.flush();
210            buffer.close();
211                } catch (Exception e) {
212                        Throwable t = e.getCause();
213                        if (null != t) {
214                                throw new BinaryWriteFailedException(t);
215                        } else {
216                                throw new BinaryWriteFailedException(e);
217                        }
218                }
219    }
220
221    private void writeHeader() {
222        try {
223
224            // 11 bytes required header
225            buffer.writeBytes("PGCOPY\n\377\r\n\0");
226            // 32 bit integer indicating no OID
227            buffer.writeInt(0);
228            // 32 bit header extension area length
229            buffer.writeInt(0);
230
231                } catch (IOException e) {
232                        Throwable t = e.getCause();
233                        if (null != t) {
234                                throw new BinaryWriteFailedException(t);
235                        } else {
236                                throw new BinaryWriteFailedException(e);
237                        }
238                }
239    }
240}