COPY Operations

package org.postgresql.copy;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Reader;
import java.io.Writer;
import java.sql.SQLException;

/**
 * Manager for PostgreSQL COPY bulk data operations.
 * Access via PGConnection.getCopyAPI().
 */
public class CopyManager {
    /**
     * Constructs a CopyManager for the given connection.
     *
     * @param connection PostgreSQL connection
     * @throws SQLException if CopyManager cannot be created
     */
    public CopyManager(org.postgresql.core.BaseConnection connection)
            throws SQLException;

    /**
     * Starts a COPY FROM STDIN operation using low-level API.
     * Provides fine control over data transfer.
     *
     * @param sql COPY FROM STDIN command
     * @return CopyIn interface for writing data
     * @throws SQLException if operation cannot be started
     */
    public CopyIn copyIn(String sql) throws SQLException;

    /**
     * Starts a COPY TO STDOUT operation using low-level API.
     * Provides fine control over data transfer.
     *
     * @param sql COPY TO STDOUT command
     * @return CopyOut interface for reading data
     * @throws SQLException if operation cannot be started
     */
    public CopyOut copyOut(String sql) throws SQLException;

    /**
     * Starts a bidirectional COPY operation (for replication).
     *
     * @param sql COPY command
     * @return CopyDual interface
     * @throws SQLException if operation cannot be started
     */
    public CopyDual copyDual(String sql) throws SQLException;

    /**
     * Copies data from database to Writer (text format).
     * High-level convenience method.
     *
     * @param sql COPY TO STDOUT command
     * @param to Writer to receive data
     * @return Number of rows (for server 8.2+; -1 for older)
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyOut(String sql, Writer to) throws SQLException, IOException;

    /**
     * Copies data from database to OutputStream (binary or text format).
     * High-level convenience method.
     *
     * @param sql COPY TO STDOUT command
     * @param to OutputStream to receive data
     * @return Number of rows (for server 8.2+; -1 for older)
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyOut(String sql, OutputStream to) throws SQLException, IOException;

    /**
     * Copies data from Reader to database (text format).
     * High-level convenience method.
     *
     * @param sql COPY FROM STDIN command
     * @param from Reader providing data
     * @return Number of rows loaded
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyIn(String sql, Reader from) throws SQLException, IOException;

    /**
     * Copies data from Reader to database with buffer size.
     *
     * @param sql COPY FROM STDIN command
     * @param from Reader providing data
     * @param bufferSize Buffer size in bytes
     * @return Number of rows loaded
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyIn(String sql, Reader from, int bufferSize)
            throws SQLException, IOException;

    /**
     * Copies data from InputStream to database (binary or text format).
     * High-level convenience method.
     *
     * @param sql COPY FROM STDIN command
     * @param from InputStream providing data
     * @return Number of rows loaded
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyIn(String sql, InputStream from) throws SQLException, IOException;

    /**
     * Copies data from InputStream to database with buffer size.
     *
     * @param sql COPY FROM STDIN command
     * @param from InputStream providing data
     * @param bufferSize Buffer size in bytes
     * @return Number of rows loaded
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyIn(String sql, InputStream from, int bufferSize)
            throws SQLException, IOException;

    /**
     * Copies data from ByteStreamWriter to database.
     * Efficient for programmatic data generation.
     *
     * @param sql COPY FROM STDIN command
     * @param from ByteStreamWriter providing data
     * @return Number of rows loaded
     * @throws SQLException if operation fails
     * @throws IOException if I/O error occurs
     */
    public long copyIn(String sql, org.postgresql.util.ByteStreamWriter from)
            throws SQLException, IOException;
}

import org.postgresql.PGConnection;
import org.postgresql.copy.CopyManager;
import java.sql.Connection;
import java.sql.SQLException;
import java.io.*;

// Example 1: Copy data FROM file to database
public class CopyFromFileExample {
    public static long loadFromCSV(Connection conn, String tableName,
                                   String filename) throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        CopyManager copyManager = pgConn.getCopyAPI();

        // COPY command
        String sql = String.format("COPY %s FROM STDIN WITH (FORMAT CSV, HEADER)", tableName);

        // Load from file
        try (FileReader reader = new FileReader(filename)) {
            return copyManager.copyIn(sql, reader);
        }
    }
}

// Example 2: Copy data TO file from database
public class CopyToFileExample {
    public static long exportToCSV(Connection conn, String tableName,
                                   String filename) throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        CopyManager copyManager = pgConn.getCopyAPI();

        // COPY command
        String sql = String.format("COPY %s TO STDOUT WITH (FORMAT CSV, HEADER)", tableName);

        // Export to file
        try (FileWriter writer = new FileWriter(filename)) {
            return copyManager.copyOut(sql, writer);
        }
    }
}

// Example 3: Copy data with query
public class CopyQueryExample {
    public static long exportQueryResults(Connection conn, String filename)
            throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        CopyManager copyManager = pgConn.getCopyAPI();

        // COPY a query result
        String sql = "COPY (SELECT id, name, email FROM users WHERE active = true) " +
                    "TO STDOUT WITH (FORMAT CSV, HEADER)";

        try (FileWriter writer = new FileWriter(filename)) {
            return copyManager.copyOut(sql, writer);
        }
    }
}

// Example 4: Binary format copy
public class BinaryCopyExample {
    public static long loadBinary(Connection conn) throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        CopyManager copyManager = pgConn.getCopyAPI();

        String sql = "COPY data_table FROM STDIN WITH (FORMAT BINARY)";

        try (FileInputStream in = new FileInputStream("data.bin")) {
            return copyManager.copyIn(sql, in);
        }
    }

    public static long exportBinary(Connection conn) throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        CopyManager copyManager = pgConn.getCopyAPI();

        String sql = "COPY data_table TO STDOUT WITH (FORMAT BINARY)";

        try (FileOutputStream out = new FileOutputStream("data.bin")) {
            return copyManager.copyOut(sql, out);
        }
    }
}

package org.postgresql.copy;

import java.sql.SQLException;

/**
 * Base interface for COPY operations.
 */
public interface CopyOperation {
    /**
     * Cancels the COPY operation.
     * Must be called if operation is not completed normally.
     *
     * @throws SQLException if cancellation fails
     */
    void cancelCopy() throws SQLException;

    /**
     * Returns the overall format: 0 for text, 1 for binary.
     *
     * @return Format code
     * @throws SQLException if format cannot be determined
     */
    int getFormat() throws SQLException;

    /**
     * Returns the format for a specific field.
     *
     * @param field Field number (0-based)
     * @return Format code for field
     * @throws SQLException if format cannot be determined
     */
    int getFieldFormat(int field) throws SQLException;

    /**
     * Returns whether the COPY operation is still active.
     *
     * @return true if operation is active
     */
    boolean isActive();

    /**
     * Returns the number of fields in each row.
     *
     * @return Field count
     * @throws SQLException if count cannot be determined
     */
    int getFieldCount() throws SQLException;

    /**
     * Returns the number of rows handled so far.
     *
     * @return Row count
     */
    long getHandledRowCount();
}

package org.postgresql.copy;

import org.postgresql.util.ByteStreamWriter;
import java.sql.SQLException;

/**
 * Interface for copying data FROM client TO database.
 */
public interface CopyIn extends CopyOperation {
    /**
     * Writes data to the COPY operation.
     *
     * @param buf Byte array containing data
     * @param off Offset in array
     * @param siz Number of bytes to write
     * @throws SQLException if write fails
     */
    void writeToCopy(byte[] buf, int off, int siz) throws SQLException;

    /**
     * Writes data from ByteStreamWriter.
     *
     * @param from Data source
     * @throws SQLException if write fails
     */
    void writeToCopy(ByteStreamWriter from) throws SQLException;

    /**
     * Flushes buffered data to server.
     *
     * @throws SQLException if flush fails
     */
    void flushCopy() throws SQLException;

    /**
     * Completes the COPY operation successfully.
     * Must be called to finalize the operation.
     *
     * @return Number of rows loaded
     * @throws SQLException if completion fails
     */
    long endCopy() throws SQLException;
}

import org.postgresql.PGConnection;
import org.postgresql.copy.CopyIn;
import java.sql.Connection;
import java.sql.SQLException;
import java.nio.charset.StandardCharsets;

// Example 1: Low-level COPY IN
public class LowLevelCopyInExample {
    public static void copyInManually(Connection conn) throws SQLException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        String sql = "COPY users (name, email) FROM STDIN WITH (FORMAT CSV)";

        CopyIn copyIn = pgConn.getCopyAPI().copyIn(sql);

        try {
            // Write CSV rows
            String row1 = "John Doe,john@example.com\n";
            String row2 = "Jane Smith,jane@example.com\n";

            copyIn.writeToCopy(row1.getBytes(StandardCharsets.UTF_8), 0,
                    row1.getBytes(StandardCharsets.UTF_8).length);
            copyIn.writeToCopy(row2.getBytes(StandardCharsets.UTF_8), 0,
                    row2.getBytes(StandardCharsets.UTF_8).length);

            // Complete the operation
            long rows = copyIn.endCopy();
            System.out.println("Loaded " + rows + " rows");

        } catch (SQLException e) {
            // Cancel on error
            if (copyIn.isActive()) {
                copyIn.cancelCopy();
            }
            throw e;
        }
    }
}

// Example 2: Streaming data generation
public class StreamingCopyExample {
    public static void generateAndCopy(Connection conn, int numRows)
            throws SQLException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        String sql = "COPY large_table (id, value) FROM STDIN WITH (FORMAT CSV)";

        CopyIn copyIn = pgConn.getCopyAPI().copyIn(sql);

        try {
            StringBuilder batch = new StringBuilder();
            int batchSize = 1000;

            for (int i = 0; i < numRows; i++) {
                // Generate row
                batch.append(i).append(",").append("Value ").append(i).append("\n");

                // Write in batches
                if ((i + 1) % batchSize == 0 || i == numRows - 1) {
                    byte[] bytes = batch.toString().getBytes(StandardCharsets.UTF_8);
                    copyIn.writeToCopy(bytes, 0, bytes.length);
                    batch.setLength(0); // Clear buffer
                }
            }

            long rows = copyIn.endCopy();
            System.out.println("Generated and loaded " + rows + " rows");

        } catch (SQLException e) {
            if (copyIn.isActive()) {
                copyIn.cancelCopy();
            }
            throw e;
        }
    }
}

package org.postgresql.copy;

import java.sql.SQLException;

/**
 * Interface for copying data FROM database TO client.
 */
public interface CopyOut extends CopyOperation {
    /**
     * Reads a row of data from the COPY operation.
     * Blocks until data is available or operation completes.
     *
     * @return Byte array containing one row, or null if no more data
     * @throws SQLException if read fails
     */
    byte[] readFromCopy() throws SQLException;

    /**
     * Reads a row of data with optional blocking.
     *
     * @param block If true, blocks until data available; if false, returns immediately
     * @return Byte array containing one row, or null if no data available
     * @throws SQLException if read fails
     */
    byte[] readFromCopy(boolean block) throws SQLException;
}

import org.postgresql.PGConnection;
import org.postgresql.copy.CopyOut;
import java.sql.Connection;
import java.sql.SQLException;
import java.io.FileOutputStream;
import java.io.IOException;

// Example: Low-level COPY OUT
public class LowLevelCopyOutExample {
    public static void copyOutManually(Connection conn, String filename)
            throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        String sql = "COPY users TO STDOUT WITH (FORMAT CSV, HEADER)";

        CopyOut copyOut = pgConn.getCopyAPI().copyOut(sql);

        try (FileOutputStream out = new FileOutputStream(filename)) {
            byte[] row;
            while ((row = copyOut.readFromCopy()) != null) {
                out.write(row);
            }

            long rows = copyOut.getHandledRowCount();
            System.out.println("Exported " + rows + " rows");

        } catch (SQLException | IOException e) {
            if (copyOut.isActive()) {
                copyOut.cancelCopy();
            }
            throw e;
        }
    }
}

package org.postgresql.copy;

/**
 * Bidirectional COPY interface combining CopyIn and CopyOut.
 * Used primarily for PostgreSQL replication via COPY protocol.
 * Extends both CopyIn and CopyOut, providing all methods from both interfaces.
 *
 * @see CopyIn
 * @see CopyOut
 */
public interface CopyDual extends CopyIn, CopyOut {
    // Inherits all methods from CopyIn:
    // - writeToCopy(byte[], int, int)
    // - writeToCopy(ByteStreamWriter)
    // - flushCopy()
    // - endCopy()
    //
    // Inherits all methods from CopyOut:
    // - readFromCopy()
    // - readFromCopy(boolean)
    //
    // Inherits all methods from CopyOperation:
    // - cancelCopy()
    // - getFormat()
    // - getFieldFormat(int)
    // - isActive()
    // - getFieldCount()
    // - getHandledRowCount()
}

import org.postgresql.PGConnection;
import org.postgresql.copy.CopyDual;
import java.sql.Connection;
import java.sql.SQLException;

// Example: Bidirectional COPY for replication
public class ReplicationCopyExample {
    public static void replicationCopy(Connection conn) throws SQLException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);

        // Start bidirectional COPY for replication
        String sql = "COPY pg_logical_slot_get_binary_changes('slot_name', NULL, NULL) TO STDOUT";

        CopyDual copyDual = pgConn.getCopyAPI().copyDual(sql);

        try {
            // Can both read and write through the same interface
            byte[] data;
            while ((data = copyDual.readFromCopy()) != null) {
                // Process replication data
                processReplicationData(data);

                // Can also write feedback/acknowledgments
                // copyDual.writeToCopy(feedback, 0, feedback.length);
                // copyDual.flushCopy();
            }

        } catch (SQLException e) {
            if (copyDual.isActive()) {
                copyDual.cancelCopy();
            }
            throw e;
        }
    }

    private static void processReplicationData(byte[] data) {
        // Process replication data
    }
}

package org.postgresql.util;

import java.io.IOException;
import java.io.OutputStream;

/**
 * Interface for efficiently writing byte streams to COPY operations.
 * Useful when data is stored off-heap or in non-contiguous buffers.
 */
public interface ByteStreamWriter {
    /**
     * Returns the total length of the stream in bytes.
     * Must be known ahead of time before calling writeTo().
     *
     * @return Number of bytes in the stream
     */
    int getLength();

    /**
     * Writes the data to the provided target.
     * Must not write more than getLength() bytes.
     *
     * @param target Target to write data to
     * @throws IOException if write fails or if more than getLength() bytes are written
     */
    void writeTo(ByteStreamTarget target) throws IOException;

    /**
     * Creates a ByteStreamWriter from one or more ByteBuffers.
     *
     * @param buf ByteBuffer(s) to write
     * @return ByteStreamWriter for the buffer(s)
     */
    static ByteStreamWriter of(java.nio.ByteBuffer... buf);

    /**
     * Target interface for writing byte streams.
     */
    interface ByteStreamTarget {
        /**
         * Provides an OutputStream to write bytes to.
         *
         * @return OutputStream for writing
         */
        OutputStream getOutputStream();
    }
}

import org.postgresql.util.ByteStreamWriter;
import java.nio.ByteBuffer;
import java.sql.PreparedStatement;

// Example: Using ByteStreamWriter with direct ByteBuffer
public class ByteStreamWriterExample {
    public static void useByteStreamWithBuffer(Connection conn, ByteBuffer buffer)
            throws SQLException {
        String sql = "INSERT INTO data_table (bytes) VALUES (?)";

        try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
            // Create ByteStreamWriter from ByteBuffer
            ByteStreamWriter writer = ByteStreamWriter.of(buffer);

            // Use setObject to pass the ByteStreamWriter
            pstmt.setObject(1, writer);
            pstmt.executeUpdate();
        }
    }

    // Example: Custom ByteStreamWriter implementation
    public static class CustomByteStreamWriter implements ByteStreamWriter {
        private final byte[] data;

        public CustomByteStreamWriter(byte[] data) {
            this.data = data;
        }

        @Override
        public int getLength() {
            return data.length;
        }

        @Override
        public void writeTo(ByteStreamTarget target) throws IOException {
            target.getOutputStream().write(data);
        }
    }
}

-- CSV format with header
COPY table FROM STDIN WITH (FORMAT CSV, HEADER)

-- CSV with custom delimiter
COPY table FROM STDIN WITH (FORMAT CSV, DELIMITER '|')

-- CSV with quote character
COPY table FROM STDIN WITH (FORMAT CSV, QUOTE '"')

-- CSV with NULL representation
COPY table FROM STDIN WITH (FORMAT CSV, NULL 'NULL')

-- Binary format (fastest)
COPY table FROM STDIN WITH (FORMAT BINARY)

-- Text format (default)
COPY table FROM STDIN WITH (FORMAT TEXT)

-- Specific columns only
COPY table (col1, col2, col3) FROM STDIN WITH (FORMAT CSV)

-- Copy from query result
COPY (SELECT * FROM table WHERE condition) TO STDOUT WITH (FORMAT CSV)

// Custom delimited format
public class CustomFormatExample {
    public static long loadPipeDelimited(Connection conn, Reader reader)
            throws SQLException, IOException {
        PGConnection pgConn = conn.unwrap(PGConnection.class);
        CopyManager copyManager = pgConn.getCopyAPI();

        String sql = "COPY users FROM STDIN WITH (FORMAT CSV, DELIMITER '|', " +
                    "HEADER false, NULL 'NULL')";

        return copyManager.copyIn(sql, reader);
    }
}