@Override public String create(PipelineOptions options) { return options.as(GcpOptions.class).getProject(); } }
@Test public void testEmptyGcpTempLocation() throws Exception { GcpOptions options = PipelineOptionsFactory.as(GcpOptions.class); options.setGcpCredential(new TestCredential()); options.setProject(""); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("--project is a required option"); options.getGcpTempLocation(); }
@Test public void testDefaultGcpTempLocation() throws Exception { GcpOptions options = PipelineOptionsFactory.as(GcpOptions.class); String tempLocation = "gs://bucket"; options.setTempLocation(tempLocation); options.as(GcsOptions.class).setPathValidatorClass(NoopPathValidator.class); assertEquals(tempLocation, options.getGcpTempLocation()); }
@Test public void testDefaultGcpTempLocationInvalid() throws Exception { GcpOptions options = PipelineOptionsFactory.as(GcpOptions.class); options.setTempLocation("file://"); thrown.expect(IllegalArgumentException.class); thrown.expectMessage( "Error constructing default value for gcpTempLocation: tempLocation is not" + " a valid GCS path"); options.getGcpTempLocation(); }
@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(); }
/** * 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()); }
/** Tests that credentials are not used from PipelineOptions if supplied by BigtableOptions. */ @Test public void testDontUsePipelineOptionsCredentialsIfSpecifiedInBigtableOptions() throws Exception { BigtableOptions options = BIGTABLE_OPTIONS .toBuilder() .setCredentialOptions(CredentialOptions.nullCredential()) .build(); GcpOptions pipelineOptions = PipelineOptionsFactory.as(GcpOptions.class); pipelineOptions.setGcpCredential(new TestCredential()); BigtableService readService = BigtableIO.read() .withBigtableOptions(options) .withTableId("TEST-TABLE") .getBigtableConfig() .getBigtableService(pipelineOptions); BigtableService writeService = BigtableIO.write() .withBigtableOptions(options) .withTableId("TEST-TABLE") .getBigtableConfig() .getBigtableService(pipelineOptions); assertEquals( CredentialType.None, readService.getBigtableOptions().getCredentialOptions().getCredentialType()); assertEquals( CredentialType.None, writeService.getBigtableOptions().getCredentialOptions().getCredentialType()); }
@Override public Credentials create(PipelineOptions options) { GcpOptions gcpOptions = options.as(GcpOptions.class); try { CredentialFactory factory = InstanceBuilder.ofType(CredentialFactory.class) .fromClass(gcpOptions.getCredentialFactoryClass()) .fromFactoryMethod("fromOptions") .withArg(PipelineOptions.class, options) .build(); return factory.getCredential(); } catch (IOException | GeneralSecurityException e) { throw new RuntimeException("Unable to obtain credential", e); } } }
/** * 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()); }
@Test public void testDefaultGcpTempLocationDoesNotExist() { GcpOptions options = PipelineOptionsFactory.as(GcpOptions.class); String tempLocation = "gs://does/not/exist"; options.setTempLocation(tempLocation); thrown.expect(IllegalArgumentException.class); thrown.expectMessage( "Error constructing default value for gcpTempLocation: tempLocation is not" + " a valid GCS path"); thrown.expectCause( hasMessage(containsString("Output path does not exist or is not writeable"))); options.getGcpTempLocation(); }
/** Tests that credentials are used from PipelineOptions if not supplied by BigtableOptions. */ @Test public void testUsePipelineOptionsCredentialsIfNotSpecifiedInBigtableOptions() throws Exception { BigtableOptions options = BIGTABLE_OPTIONS .toBuilder() .setCredentialOptions(CredentialOptions.defaultCredentials()) .build(); GcpOptions pipelineOptions = PipelineOptionsFactory.as(GcpOptions.class); pipelineOptions.setGcpCredential(new TestCredential()); BigtableService readService = BigtableIO.read() .withBigtableOptions(options) .withTableId("TEST-TABLE") .getBigtableConfig() .getBigtableService(pipelineOptions); BigtableService writeService = BigtableIO.write() .withBigtableOptions(options) .withTableId("TEST-TABLE") .getBigtableConfig() .getBigtableService(pipelineOptions); assertEquals( CredentialType.SuppliedCredentials, readService.getBigtableOptions().getCredentialOptions().getCredentialType()); assertEquals( CredentialType.SuppliedCredentials, writeService.getBigtableOptions().getCredentialOptions().getCredentialType()); }
@Override public String create(PipelineOptions options) { return "projects/" + options.as(GcpOptions.class).getProject() + "/topics/" + options.getJobName(); } }
/** 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()); }
@Override public String create(PipelineOptions options) { return "projects/" + options.as(GcpOptions.class).getProject() + "/subscriptions/" + options.getJobName(); } }
projectPath = project.get(); } else { String projectId = options.as(GcpOptions.class).getProject(); checkState( projectId != null, "Cannot create subscription to topic %s because pipeline option 'project' not specified", topicPath); projectPath = PubsubClient.projectPathFromId(options.as(GcpOptions.class).getProject());
@Before public void setup() { PipelineOptionsFactory.register(V1TestOptions.class); options = TestPipeline.testingPipelineOptions().as(V1TestOptions.class); project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); ancestor = UUID.randomUUID().toString(); }
@Before public void setup() throws Exception { PipelineOptionsFactory.register(V1TestOptions.class); options = TestPipeline.testingPipelineOptions().as(V1TestOptions.class); project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); ancestor = UUID.randomUUID().toString(); // Create entities and write them to datastore writeEntitiesToDatastore(options, project, ancestor, numEntities); }
@Before public void setupBqEnvironment() { Long timeSeed = System.currentTimeMillis(); Integer random = new Random(timeSeed).nextInt(900) + 100; this.bigQueryDatasetId = "bq_query_to_table_" + timeSeed.toString() + "_" + random.toString(); PipelineOptionsFactory.register(BigQueryToTableOptions.class); options = TestPipeline.testingPipelineOptions().as(BigQueryToTableOptions.class); options.setTempLocation(options.getTempRoot() + "/bq_it_temp"); project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); bqOption = options.as(BigQueryOptions.class); bqClient = new BigqueryClient(bqOption.getAppName()); bqClient.createNewDataset(project, this.bigQueryDatasetId); outputTable = project + ":" + this.bigQueryDatasetId + "." + BigQueryToTableIT.OUTPUT_TABLE_NAME; }
private void setupTestEnvironment(String recordSize, boolean enableCustomBigquery) { PipelineOptionsFactory.register(BigQueryIOReadOptions.class); options = TestPipeline.testingPipelineOptions().as(BigQueryIOReadOptions.class); options.setNumRecords(numOfRecords.get(recordSize)); options.setTempLocation(options.getTempRoot() + "/temp-it/"); project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); options.setInputTable(project + ":" + datasetId + "." + tablePrefix + recordSize); if (enableCustomBigquery) { options.setExperiments( ImmutableList.of("enable_custom_bigquery_sink", "enable_custom_bigquery_source")); } }
@Before public void setupTestEnvironment() { PipelineOptionsFactory.register(TrafficRoutesOptions.class); this.options = TestPipeline.testingPipelineOptions().as(TrafficRoutesOptions.class); this.projectId = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); this.bqClient = new BigqueryClient("TrafficRoutesIT"); this.bqClient.createNewDataset(this.projectId, this.outputDatasetId); }