org.postgresql.copy.CopyManager Java Examples

@Override
public int[] executeBatch() throws SQLException {
    long res = 0;
    try {
        CopyManager cpManager = ((PGConnection) connection).getCopyAPI();
        PushbackReader reader = new PushbackReader(new StringReader(""),
                batchBuilder.length());
        reader.unread(batchBuilder.toString().toCharArray());
        res = cpManager.copyIn("COPY " + tableName +  " FROM STDIN WITH CSV", reader);
        batchBuilder.setLength(0);
        reader.close();
    } catch (IOException e) {
        throw new SQLException(e);
    }
    return new int[] { (int) res };
}
 

/**
 * Metadata is supposed to be presented in a four-tuple CSV format (docid, key,
 * value, type). @see uhh_lt.newsleak.reader.NewsleakReader should write a
 * temporary metadata file in that format (or assume it was produced by an
 * external process)
 * 
 * The CSV file is imported via postgres directly.
 * 
 * See <i>data/metadata_example.csv</i> for an example.
 */
private void metadataToPostgres() {

	try {
		// we need a mapping of document ids since ElasticsearchDocumentWriter generates
		// new Ids from an autoincrement-value
		String mappedMetadataFilepath = this.dataDirectory + File.separator + this.metadataFile + ".mapped";
		mappingIdsInMetadata(mappedMetadataFilepath);

		// import csv into postgres db
		CopyManager cpManager = new CopyManager((BaseConnection) conn);
		st.executeUpdate("TRUNCATE TABLE metadata;");
		this.logger.log(Level.INFO, "Importing metadata from " + mappedMetadataFilepath);
		Long n = cpManager.copyIn("COPY metadata FROM STDIN WITH CSV", new FileReader(mappedMetadataFilepath));
		this.logger.log(Level.INFO, n + " metadata imported");
	} catch (Exception e) {
		e.printStackTrace();
		System.exit(1);
	}
}
 

@Override
public int insertDataToTable(ResultSet resultSet, int taskId) throws SQLException, IOException {

    CopyIn copyIn = null;

    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;
        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);
                            break;
                    }

                    if (!resultSet.wasNull() || colValue != null) cols.append(colValue);
                }

                // Escape special chars
                if (this.options.isSinkDisableEscape())
                    row.append(cols.toString());
                else
                    row.append(cols.toString().replace("\\", "\\\\").replace("\n", "\\n").replace("\r", "\\r").replace("\u0000", ""));

                // 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
                row.setLength(0); // set length of buffer to 0
                row.trimToSize();
                cols.setLength(0); // set length of buffer to 0
                cols.trimToSize();
            } 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 0;
}
 

public synchronized long copyFromReader(TableSchema tableSchema, Reader reader) {
	try {
		String query = tableSchema.copyFromStdin();
		CopyManager copyManager = new CopyManager((BaseConnection) currentConnection);
		long numRows = copyManager.copyIn(query, reader);

		return numRows;
	} catch (Exception e) {
		throw new EmissionException(e);
	}
}
 

@Override
public Long call() throws SQLException, IOException {
  try {
    CopyManager mgr = new CopyManager((BaseConnection) conn);
    return mgr.copyIn(sql, pipeIn);
  } finally {
    try {
      pipeIn.close();
    } catch (IOException ignore) {
    }
  }
}
 

/**
 * This is a blocking call - run in an execBlocking statement
 * import data in a tab delimited file into columns of an existing table
 * Using only default values of the COPY FROM STDIN Postgres command
 * @param path - path to the file
 * @param tableName - name of the table to import the content into
 */
public void importFile(String path, String tableName) {

  long recordsImported[] = new long[]{-1};
  vertx.<String>executeBlocking(dothis -> {
    try {
      String host = postgreSQLClientConfig.getString(HOST);
      int port = postgreSQLClientConfig.getInteger(PORT);
      String user = postgreSQLClientConfig.getString(_USERNAME);
      String pass = postgreSQLClientConfig.getString(_PASSWORD);
      String db = postgreSQLClientConfig.getString(DATABASE);

      log.info("Connecting to " + db);

      Connection con = DriverManager.getConnection(
        "jdbc:postgresql://"+host+":"+port+"/"+db, user , pass);

      log.info("Copying text data rows from stdin");

      CopyManager copyManager = new CopyManager((BaseConnection) con);

      FileReader fileReader = new FileReader(path);
      recordsImported[0] = copyManager.copyIn("COPY "+tableName+" FROM STDIN", fileReader );

    } catch (Exception e) {
      log.error(messages.getMessage("en", MessageConsts.ImportFailed), e);
      dothis.fail(e);
    }
    dothis.complete("Done.");

  }, whendone -> {

    if(whendone.succeeded()){
      log.info("Done importing file: " + path + ". Number of records imported: " + recordsImported[0]);
    }
    else{
      log.info("Failed importing file: " + path);
    }

  });

}
 

@Test
@Ignore("not used right now in read code")
public void testCopy100RecordsIntoFile() throws Exception {
    System.out.println("Starting");

    int a = 0;

    long start = System.currentTimeMillis();
    try (Connection connection = dbManager.getConnection();
         PreparedStatement ps = connection.prepareStatement(ReportingDBDao.insertMinute)) {

        String userName = "[email protected]";
        long minute = (System.currentTimeMillis() / AverageAggregatorProcessor.MINUTE) * AverageAggregatorProcessor.MINUTE;

        for (int i = 0; i < 100; i++) {
            ReportingDBDao.prepareReportingInsert(ps, userName, 1, 0, (short) 0, PinType.VIRTUAL, minute, (double) i);
            ps.addBatch();
            minute += AverageAggregatorProcessor.MINUTE;
            a++;
        }

        ps.executeBatch();
        connection.commit();
    }

    System.out.println("Finished : " + (System.currentTimeMillis() - start)  + " millis. Executed : " + a);


    try (Connection connection = dbManager.getConnection();
         Writer gzipWriter = new OutputStreamWriter(new GZIPOutputStream(new FileOutputStream(new File("/home/doom369/output.csv.gz"))), "UTF-8")) {

        CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));


        String selectQuery = "select pintype || pin, ts, value from reporting_average_minute where project_id = 1 and email = '[email protected]'";
        long res = copyManager.copyOut("COPY (" + selectQuery + " ) TO STDOUT WITH (FORMAT CSV)", gzipWriter);
        System.out.println(res);
    }


}
 

/**
 * Uses the given {@link Connection connection} to COPY the table information from the table given
 * as table name to the directory given as testDataDirectory.
 *
 * @param connection The connection over which to execute the COPY from the database.
 * @param testDataDirectory The test data directory where the result of the COPY is stored.
 * @param tableName The name of the table which is currently dumped from the database.
 *
 * @throws SQLException if the COPY command on the SQL server fails.
 * @throws IOException if the dump result could not be written successfully.
 */
private static void dumpTableInformation(
    final Connection connection, final File testDataDirectory, final String tableName)
        throws SQLException, IOException {
  final CopyManager manager = new CopyManager((BaseConnection) connection);
  try (FileWriter fw = new FileWriter(new File(testDataDirectory, tableName + ".sql"))) {
    manager.copyOut("COPY " + tableName + " TO STDOUT", fw);
  } catch (final IOException exception) {
    CUtilityFunctions.logException(exception);
  } 
}