Examples with WritableByteChannel java.nio.channels.WritableByteChannel used on opensource projects

Example 51 with WritableByteChannel

use of java.nio.channels.WritableByteChannel in project beam by apache.

the class BigQueryIOTest method testWriteTables.

@Test
public void testWriteTables() throws Exception {
    p.enableAbandonedNodeEnforcement(false);
    FakeDatasetService datasetService = new FakeDatasetService();
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(new FakeJobService()).withDatasetService(datasetService);
    datasetService.createDataset("project-id", "dataset-id", "", "");
    long numTables = 3;
    long numPartitions = 3;
    long numFilesPerPartition = 10;
    String jobIdToken = "jobIdToken";
    String stepUuid = "stepUuid";
    Map<TableDestination, List<String>> expectedTempTables = Maps.newHashMap();
    Path baseDir = Files.createTempDirectory(tempFolder, "testWriteTables");
    List<KV<ShardedKey<String>, List<String>>> partitions = Lists.newArrayList();
    for (int i = 0; i < numTables; ++i) {
        String tableName = String.format("project-id:dataset-id.table%05d", i);
        TableDestination tableDestination = new TableDestination(tableName, tableName);
        for (int j = 0; j < numPartitions; ++j) {
            String tempTableId = BigQueryHelpers.createJobId(jobIdToken, tableDestination, j);
            List<String> filesPerPartition = Lists.newArrayList();
            for (int k = 0; k < numFilesPerPartition; ++k) {
                String filename = Paths.get(baseDir.toString(), String.format("files0x%08x_%05d", tempTableId.hashCode(), k)).toString();
                ResourceId fileResource = FileSystems.matchNewResource(filename, false);
                try (WritableByteChannel channel = FileSystems.create(fileResource, MimeTypes.TEXT)) {
                    try (OutputStream output = Channels.newOutputStream(channel)) {
                        TableRow tableRow = new TableRow().set("name", tableName);
                        TableRowJsonCoder.of().encode(tableRow, output, Context.OUTER);
                        output.write("\n".getBytes(StandardCharsets.UTF_8));
                    }
                }
                filesPerPartition.add(filename);
            }
            partitions.add(KV.of(ShardedKey.of(tableDestination.getTableSpec(), j), filesPerPartition));
            List<String> expectedTables = expectedTempTables.get(tableDestination);
            if (expectedTables == null) {
                expectedTables = Lists.newArrayList();
                expectedTempTables.put(tableDestination, expectedTables);
            }
            String json = String.format("{\"datasetId\":\"dataset-id\",\"projectId\":\"project-id\",\"tableId\":\"%s\"}", tempTableId);
            expectedTables.add(json);
        }
    }
    PCollectionView<String> jobIdTokenView = p.apply("CreateJobId", Create.of("jobId")).apply(View.<String>asSingleton());
    PCollectionView<Map<String, String>> schemaMapView = p.apply("CreateEmptySchema", Create.empty(new TypeDescriptor<KV<String, String>>() {
    })).apply(View.<String, String>asMap());
    WriteTables<String> writeTables = new WriteTables<>(false, fakeBqServices, jobIdTokenView, schemaMapView, WriteDisposition.WRITE_EMPTY, CreateDisposition.CREATE_IF_NEEDED, new IdentityDynamicTables());
    DoFnTester<KV<ShardedKey<String>, List<String>>, KV<TableDestination, String>> tester = DoFnTester.of(writeTables);
    tester.setSideInput(jobIdTokenView, GlobalWindow.INSTANCE, jobIdToken);
    tester.setSideInput(schemaMapView, GlobalWindow.INSTANCE, ImmutableMap.<String, String>of());
    tester.getPipelineOptions().setTempLocation("tempLocation");
    for (KV<ShardedKey<String>, List<String>> partition : partitions) {
        tester.processElement(partition);
    }
    Map<TableDestination, List<String>> tempTablesResult = Maps.newHashMap();
    for (KV<TableDestination, String> element : tester.takeOutputElements()) {
        List<String> tables = tempTablesResult.get(element.getKey());
        if (tables == null) {
            tables = Lists.newArrayList();
            tempTablesResult.put(element.getKey(), tables);
        }
        tables.add(element.getValue());
    }
    assertEquals(expectedTempTables, tempTablesResult);
}

Example 52 with WritableByteChannel

use of java.nio.channels.WritableByteChannel in project beam by apache.

the class FileBasedSinkTest method writeValuesWithWritableByteChannelFactory.

private File writeValuesWithWritableByteChannelFactory(final WritableByteChannelFactory factory, String... values) throws IOException {
    final File file = tmpFolder.newFile("test.gz");
    final WritableByteChannel channel = factory.create(Channels.newChannel(new FileOutputStream(file)));
    for (String value : values) {
        channel.write(ByteBuffer.wrap((value + "\n").getBytes(StandardCharsets.UTF_8)));
    }
    channel.close();
    return file;
}

Example 53 with WritableByteChannel

use of java.nio.channels.WritableByteChannel in project beam by apache.

the class PackageUtil method stageOnePackage.

/**
   * Utility to verify whether a package has already been staged and, if not, copy it to the
   * staging location.
   */
private static void stageOnePackage(PackageAttributes attributes, AtomicInteger numUploaded, AtomicInteger numCached, Sleeper retrySleeper, CreateOptions createOptions) {
    String source = attributes.getSourcePath();
    String target = attributes.getDataflowPackage().getLocation();
    // always using MimeTypes.BINARY?
    try {
        try {
            long remoteLength = FileSystems.matchSingleFileSpec(target).sizeBytes();
            if (remoteLength == attributes.getSize()) {
                LOG.debug("Skipping classpath element already staged: {} at {}", attributes.getSourcePath(), target);
                numCached.incrementAndGet();
                return;
            }
        } catch (FileNotFoundException expected) {
        // If the file doesn't exist, it means we need to upload it.
        }
        // Upload file, retrying on failure.
        BackOff backoff = BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff());
        while (true) {
            try {
                LOG.debug("Uploading classpath element {} to {}", source, target);
                try (WritableByteChannel writer = makeWriter(target, createOptions)) {
                    copyContent(source, writer);
                }
                numUploaded.incrementAndGet();
                break;
            } catch (IOException e) {
                if (ERROR_EXTRACTOR.accessDenied(e)) {
                    String errorMessage = String.format("Uploaded failed due to permissions error, will NOT retry staging " + "of classpath %s. Please verify credentials are valid and that you have " + "write access to %s. Stale credentials can be resolved by executing " + "'gcloud auth application-default login'.", source, target);
                    LOG.error(errorMessage);
                    throw new IOException(errorMessage, e);
                }
                long sleep = backoff.nextBackOffMillis();
                if (sleep == BackOff.STOP) {
                    // Rethrow last error, to be included as a cause in the catch below.
                    LOG.error("Upload failed, will NOT retry staging of classpath: {}", source, e);
                    throw e;
                } else {
                    LOG.warn("Upload attempt failed, sleeping before retrying staging of classpath: {}", source, e);
                    retrySleeper.sleep(sleep);
                }
            }
        }
    } catch (Exception e) {
        throw new RuntimeException("Could not stage classpath element: " + source, e);
    }
}

Example 54 with WritableByteChannel

use of java.nio.channels.WritableByteChannel in project robovm by robovm.

the class ChannelsTest method testNewWriterWritableByteChannelString_InputNull.

/*
     * this test cannot be passed when buffer set to 0!
     */
public void testNewWriterWritableByteChannelString_InputNull() throws IOException {
    this.fouts = new FileOutputStream(tmpFile);
    WritableByteChannel wbChannel = Channels.newChannel(this.fouts);
    Writer testWriter = Channels.newWriter(wbChannel, Charset.forName(CODE_SET).newEncoder(), 1);
    //$NON-NLS-1$
    String writebuf = "";
    for (int val = 0; val < this.writebufSize / 2; val++) {
        writebuf = writebuf + ((char) (val + 64));
    }
    // can write to buffer
    testWriter.write(writebuf);
    testWriter.flush();
    testWriter.close();
}

Example 55 with WritableByteChannel

use of java.nio.channels.WritableByteChannel in project robovm by robovm.

the class ChannelsTest method testNewOutputStreamWritableByteChannel.

public void testNewOutputStreamWritableByteChannel() throws Exception {
    byte[] writebuf = new byte[this.testNum];
    ByteBuffer writebcbuf = ByteBuffer.allocateDirect(this.testNum);
    this.fouts = new FileOutputStream(tmpFile);
    WritableByteChannel writebc = this.fouts.getChannel();
    assertTrue(writebc.isOpen());
    OutputStream testouts = Channels.newOutputStream(writebc);
    // read in testins and fins use the same pointer
    testouts.write(writebuf);
    this.assertFileSizeSame(tmpFile, this.testNum);
    writebc.write(writebcbuf);
    this.assertFileSizeSame(tmpFile, this.testNum * 2);
    testouts.write(writebuf);
    this.assertFileSizeSame(tmpFile, this.testNum * 3);
    // readbc.close() affect testins
    writebc.close();
    assertFalse(writebc.isOpen());
    try {
        testouts.write(writebuf);
        fail();
    } catch (ClosedChannelException e) {
    // correct
    }
}