/** * Helper function that either returns the mock Bigtable service supplied by {@link * #withBigtableService} or creates and returns an implementation that talks to {@code Cloud * Bigtable}. * * <p>Also populate the credentials option from {@link GcpOptions#getGcpCredential()} if the * default credentials are being used on {@link BigtableOptions}. */ @VisibleForTesting BigtableService getBigtableService(PipelineOptions pipelineOptions) { if (getBigtableService() != null) { return getBigtableService(); } BigtableOptions.Builder bigtableOptions = effectiveUserProvidedBigtableOptions(); bigtableOptions.setUserAgent(pipelineOptions.getUserAgent()); if (bigtableOptions.build().getCredentialOptions().getCredentialType() == CredentialOptions.CredentialType.DefaultCredentials) { bigtableOptions.setCredentialOptions( CredentialOptions.credential(pipelineOptions.as(GcpOptions.class).getGcpCredential())); } // Default option that should be forced bigtableOptions.setUseCachedDataPool(true); return new BigtableServiceImpl(bigtableOptions.build()); }
/** * Builds a Cloud Datastore client for the given pipeline options, project and an optional * locahost. */ public Datastore getDatastore( PipelineOptions pipelineOptions, String projectId, @Nullable String localhost) { Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential(); HttpRequestInitializer initializer; if (credential != null) { initializer = new ChainingHttpRequestInitializer( new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer()); } else { initializer = new RetryHttpRequestInitializer(); } DatastoreOptions.Builder builder = new DatastoreOptions.Builder().projectId(projectId).initializer(initializer); if (localhost != null) { builder.localHost(localhost); } else { builder.host("batch-datastore.googleapis.com"); } return DatastoreFactory.get().create(builder.build()); }
/** Build a new datastore client. */ static Datastore getDatastore(PipelineOptions pipelineOptions, String projectId) { Credentials credential = pipelineOptions.as(GcpOptions.class).getGcpCredential(); HttpRequestInitializer initializer; if (credential != null) { initializer = new ChainingHttpRequestInitializer( new HttpCredentialsAdapter(credential), new RetryHttpRequestInitializer()); } else { initializer = new RetryHttpRequestInitializer(); } DatastoreOptions.Builder builder = new DatastoreOptions.Builder().projectId(projectId).initializer(initializer); return DatastoreFactory.get().create(builder.build()); }
@Before public void setup() throws Exception { PipelineOptionsFactory.register(BigtableTestOptions.class); options = TestPipeline.testingPipelineOptions().as(BigtableTestOptions.class); project = options.as(GcpOptions.class).getProject(); bigtableOptions = new Builder() .setProjectId(project) .setInstanceId(options.getInstanceId()) .setUserAgent("apache-beam-test") .build(); session = new BigtableSession( bigtableOptions .toBuilder() .setCredentialOptions( CredentialOptions.credential(options.as(GcpOptions.class).getGcpCredential())) .build()); tableAdminClient = session.getTableAdminClient(); }