Java Code Examples for org.apache.beam.sdk.transforms.SerializableFunction#apply()

private static void allLeafFields(
    Schema schema,
    List<String> nameComponents,
    SerializableFunction<List<String>, String> nameFn,
    Map<String, String> fieldsSelected) {
  for (Field field : schema.getFields()) {
    nameComponents.add(field.getName());
    if (field.getType().getTypeName().isCompositeType()) {
      allLeafFields(field.getType().getRowSchema(), nameComponents, nameFn, fieldsSelected);
    } else {
      String newName = nameFn.apply(nameComponents);
      fieldsSelected.put(String.join(".", nameComponents), newName);
    }
    nameComponents.remove(nameComponents.size() - 1);
  }
}
 

@Test
public void onlyPaneMultiplePanesFails() {
  SerializableFunction<Iterable<ValueInSingleWindow<Integer>>, Iterable<Integer>> extractor =
      PaneExtractors.onlyPane(PAssert.PAssertionSite.capture(""));
  Iterable<ValueInSingleWindow<Integer>> multipleFiring =
      ImmutableList.of(
          ValueInSingleWindow.of(
              4,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(true, false, Timing.EARLY)),
          ValueInSingleWindow.of(
              2,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.ON_TIME, 1L, 0L)),
          ValueInSingleWindow.of(
              1,
              new Instant(0L),
              GlobalWindow.INSTANCE,
              PaneInfo.createPane(false, false, Timing.LATE, 2L, 1L)));

  thrown.expectMessage("trigger that fires at most once");
  extractor.apply(multipleFiring);
}
 

@VisibleForTesting
static <T> T executeWithRetries(
    AbstractGoogleClientRequest<T> request,
    String errorMessage,
    Sleeper sleeper,
    BackOff backoff,
    SerializableFunction<IOException, Boolean> shouldRetry)
    throws IOException, InterruptedException {
  Exception lastException = null;
  do {
    try {
      return request.execute();
    } catch (IOException e) {
      lastException = e;
      if (!shouldRetry.apply(e)) {
        break;
      }
      LOG.info("Ignore the error and retry the request.", e);
    }
  } while (nextBackOff(sleeper, backoff));
  throw new IOException(errorMessage, lastException);
}
 

protected static <ActualT> SuccessOrFailure doChecks(
    PAssertionSite site, ActualT actualContents, SerializableFunction<ActualT, Void> checkerFn) {
  try {
    checkerFn.apply(actualContents);
    return SuccessOrFailure.success();
  } catch (Throwable t) {
    return SuccessOrFailure.failure(site, t);
  }
}
 

static SerializableFunction<BigtableOptions.Builder, BigtableOptions.Builder>
    enableBulkApiConfigurator(
        final @Nullable SerializableFunction<BigtableOptions.Builder, BigtableOptions.Builder>
                userConfigurator) {
  return optionsBuilder -> {
    if (userConfigurator != null) {
      optionsBuilder = userConfigurator.apply(optionsBuilder);
    }

    return optionsBuilder.setBulkOptions(
        optionsBuilder.build().getBulkOptions().toBuilder().setUseBulkApi(true).build());
  };
}
 

protected Instant computeTargetTimestamp(Instant time) {
  Instant result = time;
  for (SerializableFunction<Instant, Instant> timestampMapper : timestampMappers) {
    result = timestampMapper.apply(result);
  }
  return result;
}
 

/**
 * Since we are not able to track the instances of the iterators created here and close them
 * properly, we need to load the content into memory.
 */
private <OutputT> Iterable<OutputT> createIterable(
    SerializableFunction<org.apache.samza.storage.kv.Entry<ByteArray, byte[]>, OutputT> fn) {
  final ByteArray maxKey = createMaxKey();
  final KeyValueIterator<ByteArray, byte[]> kvIter = store.range(getEncodedStoreKey(), maxKey);
  final List<Entry<ByteArray, byte[]>> iterable = ImmutableList.copyOf(kvIter);
  kvIter.close();

  return new Iterable<OutputT>() {
    @Override
    public Iterator<OutputT> iterator() {
      final Iterator<Entry<ByteArray, byte[]>> iter = iterable.iterator();

      return new Iterator<OutputT>() {
        @Override
        public boolean hasNext() {
          return iter.hasNext();
        }

        @Override
        public OutputT next() {
          return fn.apply(iter.next());
        }
      };
    }
  };
}
 

private static <KeyT, ValueT, OutT>
    SerializableFunction<KafkaRecord<KeyT, ValueT>, OutT> unwrapKafkaAndThen(
        final SerializableFunction<KV<KeyT, ValueT>, OutT> fn) {
  return record -> fn.apply(record.getKV());
}
 

@Override
public Object schemaElement(int index) {
  SerializableFunction converter = doFnSchemaInformation.getElementConverters().get(index);
  return converter.apply(element());
}
 

static RetryJobIdResult getRetryJobId(
    RetryJobId currentJobId, SerializableFunction<RetryJobId, Job> lookupJob) {
  for (int retryIndex = currentJobId.getRetryIndex(); ; retryIndex++) {
    RetryJobId jobId = new RetryJobId(currentJobId.getJobIdPrefix(), retryIndex);
    try {
      Job loadJob = lookupJob.apply(jobId);
      if (loadJob == null) {
        LOG.info("job id {} not found, so retrying with that id", jobId);
        // This either means that the original job was never properly issued (on the first
        // iteration of the loop) or that we've found a retry id that has not been used yet. Try
        // again with this job id.
        return new RetryJobIdResult(jobId, true);
      }
      JobStatus jobStatus = loadJob.getStatus();
      if (jobStatus == null) {
        LOG.info("job status for {} not found, so retrying with that job id", jobId);
        return new RetryJobIdResult(jobId, true);
      }
      if ("PENDING".equals(jobStatus.getState()) || "RUNNING".equals(jobStatus.getState())) {
        // The job id has been issued and is currently pending. This can happen after receiving
        // an error from the load or copy job creation (e.g. that error might come because the
        // job already exists). Return to the caller which job id is pending (it might not be the
        // one passed in) so the caller can then wait for this job to finish.
        LOG.info("job {} in pending or running state, so continuing with that job id", jobId);
        return new RetryJobIdResult(jobId, false);
      }
      if (jobStatus.getErrorResult() == null
          && (jobStatus.getErrors() == null || jobStatus.getErrors().isEmpty())) {
        // Import succeeded. No retry needed.
        LOG.info("job {} succeeded, so not retrying ", jobId);
        return new RetryJobIdResult(jobId, false);
      }
      // This job has failed, so we assume the data cannot enter BigQuery. We will check the next
      // job in the sequence (with the same unique prefix) to see if is either pending/succeeded
      // or can be used to generate a retry job.
      LOG.info("job {} is failed. Checking the next job id", jobId);
    } catch (RuntimeException e) {
      LOG.info("caught exception while querying job {}", jobId);
      return new RetryJobIdResult(jobId, true);
    }
  }
}
 

public void copyToTable(SnowflakeServiceConfig config) throws SQLException {

    SerializableFunction<Void, DataSource> dataSourceProviderFn = config.getDataSourceProviderFn();
    List<String> filesList = config.getFilesList();
    String table = config.getTable();
    String query = config.getQuery();
    WriteDisposition writeDisposition = config.getWriteDisposition();
    String storageIntegrationName = config.getstorageIntegrationName();
    String stagingBucketDir = config.getStagingBucketDir();

    String source;
    if (query != null) {
      // Query must be surrounded with brackets
      source = String.format("(%s)", query);
    } else {
      source = String.format("'%s'", stagingBucketDir);
    }

    filesList = filesList.stream().map(e -> String.format("'%s'", e)).collect(Collectors.toList());
    String files = String.join(", ", filesList);
    files = files.replaceAll(stagingBucketDir, "");
    DataSource dataSource = dataSourceProviderFn.apply(null);

    prepareTableAccordingWriteDisposition(dataSource, table, writeDisposition);

    if (!storageIntegrationName.isEmpty()) {
      query =
          String.format(
              "COPY INTO %s FROM %s FILES=(%s) FILE_FORMAT=(TYPE=CSV FIELD_OPTIONALLY_ENCLOSED_BY='%s' COMPRESSION=GZIP) STORAGE_INTEGRATION=%s;",
              table,
              getProperBucketDir(source),
              files,
              CSV_QUOTE_CHAR_FOR_COPY,
              storageIntegrationName);
    } else {
      query =
          String.format(
              "COPY INTO %s FROM %s FILES=(%s) FILE_FORMAT=(TYPE=CSV FIELD_OPTIONALLY_ENCLOSED_BY='%s' COMPRESSION=GZIP);",
              table, source, files, CSV_QUOTE_CHAR_FOR_COPY);
    }

    runStatement(query, dataSource.getConnection(), null);
  }
 

private Connection getConnection(SerializableFunction<Void, DataSource> dataSourceProviderFn)
    throws SQLException {
  DataSource dataSource = dataSourceProviderFn.apply(null);
  return dataSource.getConnection();
}
 

private static <KeyT, ValueT, OutT>
    SerializableFunction<KafkaRecord<KeyT, ValueT>, OutT> unwrapKafkaAndThen(
        final SerializableFunction<KV<KeyT, ValueT>, OutT> fn) {
  return record -> fn.apply(record.getKV());
}
 

@Override
public Object schemaElement(int index) {
  SerializableFunction converter = doFnSchemaInformation.getElementConverters().get(index);
  return converter.apply(element());
}