Examples with CopyIn org.postgresql.copy.CopyIn used on opensource projects

Example 1 with CopyIn

use of org.postgresql.copy.CopyIn in project ReplicaDB by osalvador.

the class PostgresqlManager method insertDataToTable.

@Override
public int insertDataToTable(ResultSet resultSet, int taskId) throws SQLException, IOException {
    CopyIn copyIn = null;
    int totalRows = 0;
    try {
        ResultSetMetaData rsmd = resultSet.getMetaData();
        String tableName;
        // Get table name and columns
        if (options.getMode().equals(ReplicationMode.COMPLETE.getModeText())) {
            tableName = getSinkTableName();
        } else {
            tableName = getQualifiedStagingTableName();
        }
        String allColumns = getAllSinkColumns(rsmd);
        // Get Postgres COPY meta-command manager
        PgConnection copyOperationConnection = this.connection.unwrap(PgConnection.class);
        CopyManager copyManager = new CopyManager(copyOperationConnection);
        String copyCmd = getCopyCommand(tableName, allColumns);
        copyIn = copyManager.copyIn(copyCmd);
        char unitSeparator = 0x1F;
        char nullAscii = 0x00;
        int columnsNumber = rsmd.getColumnCount();
        StringBuilder row = new StringBuilder();
        StringBuilder cols = new StringBuilder();
        byte[] bytes;
        String colValue;
        if (resultSet.next()) {
            // Create Bandwidth Throttling
            bandwidthThrottlingCreate(resultSet, rsmd);
            do {
                bandwidthThrottlingAcquiere();
                // Get Columns values
                for (int i = 1; i <= columnsNumber; i++) {
                    if (i > 1)
                        cols.append(unitSeparator);
                    switch(rsmd.getColumnType(i)) {
                        case Types.CLOB:
                            colValue = clobToString(resultSet.getClob(i));
                            break;
                        // case Types.BINARY:
                        case Types.BLOB:
                            colValue = blobToPostgresHex(resultSet.getBlob(i));
                            break;
                        default:
                            colValue = resultSet.getString(i);
                            if (colValue == null)
                                colValue = String.valueOf(nullAscii);
                            break;
                    }
                    if (!resultSet.wasNull() || colValue != null)
                        cols.append(colValue);
                }
                // Escape special chars
                if (this.options.isSinkDisableEscape())
                    row.append(cols.toString().replace("\u0000", "\\N"));
                else
                    row.append(cols.toString().replace("\\", "\\\\").replace("\n", "\\n").replace("\r", "\\r").replace("\u0000", "\\N"));
                // Row ends with \n
                row.append("\n");
                // Copy data to postgres
                bytes = row.toString().getBytes(StandardCharsets.UTF_8);
                copyIn.writeToCopy(bytes, 0, bytes.length);
                // Clear StringBuilders
                // set length of buffer to 0
                row.setLength(0);
                row.trimToSize();
                // set length of buffer to 0
                cols.setLength(0);
                cols.trimToSize();
                totalRows++;
            } while (resultSet.next());
        }
        copyIn.endCopy();
    } catch (Exception e) {
        if (copyIn != null && copyIn.isActive()) {
            copyIn.cancelCopy();
        }
        this.connection.rollback();
        throw e;
    } finally {
        if (copyIn != null && copyIn.isActive()) {
            copyIn.cancelCopy();
        }
    }
    this.getConnection().commit();
    return totalRows;
}

Example 2 with CopyIn

use of org.postgresql.copy.CopyIn in project openGauss-connector-jdbc by opengauss-mirror.

the class CopyTest method testLockReleaseOnCancelFailure.

@Test
public void testLockReleaseOnCancelFailure() throws SQLException, InterruptedException {
    if (!TestUtil.haveMinimumServerVersion(con, ServerVersion.v8_4)) {
        // pg_backend_pid() requires PostgreSQL 8.4+
        return;
    }
    // This is a fairly complex test because it is testing a
    // deadlock that only occurs when the connection to postgres
    // is broken during a copy operation. We'll start a copy
    // operation, use pg_terminate_backend to rudely break it,
    // and then cancel. The test passes if a subsequent operation
    // on the Connection object fails to deadlock.
    con.setAutoCommit(false);
    Statement stmt = con.createStatement();
    ResultSet rs = stmt.executeQuery("select pg_backend_pid()");
    rs.next();
    int pid = rs.getInt(1);
    rs.close();
    stmt.close();
    CopyManager manager = con.unwrap(PGConnection.class).getCopyAPI();
    CopyIn copyIn = manager.copyIn("COPY copytest FROM STDIN with " + copyParams);
    try {
        killConnection(pid);
        byte[] bunchOfNulls = ",,\n".getBytes();
        while (true) {
            copyIn.writeToCopy(bunchOfNulls, 0, bunchOfNulls.length);
        }
    } catch (SQLException e) {
        acceptIOCause(e);
    } finally {
        if (copyIn.isActive()) {
            try {
                copyIn.cancelCopy();
                fail("cancelCopy should have thrown an exception");
            } catch (SQLException e) {
                acceptIOCause(e);
            }
        }
    }
    // Now we'll execute rollback on another thread so that if the
    // deadlock _does_ occur the testcase doesn't just hange forever.
    Rollback rollback = new Rollback(con);
    rollback.start();
    rollback.join(1000);
    if (rollback.isAlive()) {
        fail("rollback did not terminate");
    }
    SQLException rollbackException = rollback.exception();
    if (rollbackException == null) {
        fail("rollback should have thrown an exception");
    }
    acceptIOCause(rollbackException);
}

Example 3 with CopyIn

use of org.postgresql.copy.CopyIn in project openGauss-connector-jdbc by opengauss-mirror.

the class CopyTest method testCopyInByRow.

@Test
public void testCopyInByRow() throws SQLException {
    String sql = "COPY copytest FROM STDIN";
    CopyIn cp = copyAPI.copyIn(sql);
    for (String anOrigData : origData) {
        byte[] buf = anOrigData.getBytes();
        cp.writeToCopy(buf, 0, buf.length);
    }
    long count1 = cp.endCopy();
    long count2 = cp.getHandledRowCount();
    assertEquals(dataRows, count1);
    assertEquals(dataRows, count2);
    try {
        cp.cancelCopy();
    } catch (SQLException se) {
        // should fail with obsolete operation
        if (!PSQLState.OBJECT_NOT_IN_STATE.getState().equals(se.getSQLState())) {
            fail("should have thrown object not in state exception.");
        }
    }
    int rowCount = getCount();
    assertEquals(dataRows, rowCount);
}

Example 4 with CopyIn

use of org.postgresql.copy.CopyIn in project stream-applications by spring-cloud.

the class PgcopySinkConfiguration method datasetSinkMessageHandler.

@Bean
@ServiceActivator(inputChannel = "toSink")
public MessageHandler datasetSinkMessageHandler(final JdbcTemplate jdbcTemplate, final PlatformTransactionManager platformTransactionManager) {
    final TransactionTemplate txTemplate = new TransactionTemplate(platformTransactionManager);
    if (StringUtils.hasText(properties.getErrorTable())) {
        verifyErrorTable(jdbcTemplate, txTemplate);
    }
    StringBuilder columns = new StringBuilder();
    for (String col : properties.getColumns()) {
        if (columns.length() > 0) {
            columns.append(",");
        }
        columns.append(col);
    }
    // the copy command
    final StringBuilder sql = new StringBuilder("COPY " + properties.getTableName());
    if (columns.length() > 0) {
        sql.append(" (" + columns + ")");
    }
    sql.append(" FROM STDIN");
    StringBuilder options = new StringBuilder();
    if (properties.getFormat() == PgcopySinkProperties.Format.CSV) {
        options.append("CSV");
    }
    if (properties.getDelimiter() != null) {
        options.append(escapedOptionCharacterValue(options.length(), "DELIMITER", properties.getDelimiter()));
    }
    if (properties.getNullString() != null) {
        options.append((options.length() > 0 ? " " : "") + "NULL '" + properties.getNullString() + "'");
    }
    if (properties.getQuote() != null) {
        options.append(quotedOptionCharacterValue(options.length(), "QUOTE", properties.getQuote()));
    }
    if (properties.getEscape() != null) {
        options.append(quotedOptionCharacterValue(options.length(), "ESCAPE", properties.getEscape()));
    }
    if (options.length() > 0) {
        sql.append(" WITH " + options.toString());
    }
    return new MessageHandler() {

        @Override
        public void handleMessage(Message<?> message) throws MessagingException {
            Object payload = message.getPayload();
            if (payload instanceof Collection<?>) {
                final Collection<?> payloads = (Collection<?>) payload;
                if (logger.isDebugEnabled()) {
                    logger.debug("Executing batch of size " + payloads.size() + " for " + sql);
                }
                try {
                    long rows = doCopy(payloads, txTemplate);
                    if (logger.isDebugEnabled()) {
                        logger.debug("Wrote " + rows + " rows");
                    }
                } catch (DataAccessException e) {
                    logger.error("Error while copying batch of data: " + e.getMessage());
                    logger.error("Switching to single row copy for current batch");
                    long rows = 0;
                    for (Object singlePayload : payloads) {
                        try {
                            rows = rows + doCopy(Collections.singletonList(singlePayload), txTemplate);
                        } catch (DataAccessException e2) {
                            logger.error("Copy for single row caused error: " + e2.getMessage());
                            logger.error("Bad Data: \n" + singlePayload);
                            if (StringUtils.hasText(properties.getErrorTable())) {
                                writeError(e2, singlePayload);
                            }
                        }
                    }
                    if (logger.isDebugEnabled()) {
                        logger.debug("Re-tried batch and wrote " + rows + " rows");
                    }
                }
            } else {
                throw new IllegalStateException("Expected a collection of strings but received " + message.getPayload().getClass().getName());
            }
        }

        private void writeError(final DataAccessException exception, final Object payload) {
            final String message;
            if (exception.getCause() != null) {
                message = exception.getCause().getMessage();
            } else {
                message = exception.getMessage();
            }
            try {
                txTemplate.execute(new TransactionCallback<Long>() {

                    @Override
                    public Long doInTransaction(TransactionStatus transactionStatus) {
                        jdbcTemplate.update("insert into " + properties.getErrorTable() + " (table_name, error_message, payload) values (?, ?, ?)", new Object[] { properties.getTableName(), message, payload });
                        return null;
                    }
                });
            } catch (DataAccessException e) {
                logger.error("Writing to error table failed: " + e.getMessage());
            }
        }

        private long doCopy(final Collection<?> payloads, TransactionTemplate txTemplate) {
            Long rows = txTemplate.execute(transactionStatus -> jdbcTemplate.execute(new ConnectionCallback<Long>() {

                @Override
                public Long doInConnection(Connection connection) throws SQLException, DataAccessException {
                    CopyManager cm = connection.unwrap(BaseConnection.class).getCopyAPI();
                    CopyIn ci = cm.copyIn(sql.toString());
                    for (Object payloadData : payloads) {
                        String textPayload = (payloadData instanceof byte[]) ? new String((byte[]) payloadData) : (String) payloadData;
                        byte[] data = (textPayload + "\n").getBytes();
                        ci.writeToCopy(data, 0, data.length);
                    }
                    return Long.valueOf(ci.endCopy());
                }
            }));
            return rows;
        }
    };
}

Example 5 with CopyIn

use of org.postgresql.copy.CopyIn in project questdb by bluestreak01.

the class PGJobContextTest method testCopyIn.

@Test
@Ignore
public void testCopyIn() throws SQLException {
    try (final PGWireServer ignored = createPGServer(2);
        final Connection connection = getConnection(false, true)) {
        PreparedStatement stmt = connection.prepareStatement("create table tab (a int, b int)");
        stmt.execute();
        CopyManager copyManager = new CopyManager((BaseConnection) connection);
        CopyIn copyIn = copyManager.copyIn("copy tab from STDIN");
        String text = "a,b\r\n" + "10,20";
        byte[] bytes = text.getBytes();
        copyIn.writeToCopy(bytes, 0, bytes.length);
        copyIn.endCopy();
    }
}