static <T> BigQueryQuerySource<T> create( String stepUuid, ValueProvider<String> query, Boolean flattenResults, Boolean useLegacySql, BigQueryServices bqServices, Coder<T> coder, SerializableFunction<SchemaAndRecord, T> parseFn, QueryPriority priority, String location) { return new BigQueryQuerySource<>( stepUuid, query, flattenResults, useLegacySql, bqServices, coder, parseFn, priority, location); }
private synchronized JobStatistics dryRunQueryIfNeeded(BigQueryOptions bqOptions) throws InterruptedException, IOException { if (dryRunJobStats.get() == null) { JobStatistics jobStats = bqServices .getJobService(bqOptions) .dryRunQuery(bqOptions.getProject(), createBasicQueryConfig(), this.location); dryRunJobStats.compareAndSet(null, jobStats); } return dryRunJobStats.get(); }
@Override public long getEstimatedSizeBytes(PipelineOptions options) throws Exception { BigQueryOptions bqOptions = options.as(BigQueryOptions.class); return dryRunQueryIfNeeded(bqOptions).getQuery().getTotalBytesProcessed(); }
private BigQuerySourceBase<T> createSource(String jobUuid, Coder<T> coder) { BigQuerySourceBase<T> source; if (getQuery() == null) { source = BigQueryTableSource.create( jobUuid, getTableProvider(), getBigQueryServices(), coder, getParseFn()); } else { source = BigQueryQuerySource.create( jobUuid, getQuery(), getFlattenResults(), getUseLegacySql(), getBigQueryServices(), coder, getParseFn(), MoreObjects.firstNonNull(getQueryPriority(), QueryPriority.BATCH), getQueryLocation()); } return source; }
@Test public void testBigQueryQuerySourceEstimatedSize() throws Exception { List<TableRow> data = ImmutableList.of( new TableRow().set("name", "A").set("number", 10L), new TableRow().set("name", "B").set("number", 11L), new TableRow().set("name", "C").set("number", 12L)); PipelineOptions options = PipelineOptionsFactory.create(); BigQueryOptions bqOptions = options.as(BigQueryOptions.class); bqOptions.setProject("project"); String stepUuid = "testStepUuid"; String query = FakeBigQueryServices.encodeQuery(data); BigQueryQuerySource<TableRow> bqSource = BigQueryQuerySource.create( stepUuid, ValueProvider.StaticValueProvider.of(query), true /* flattenResults */, true /* useLegacySql */, fakeBqServices, TableRowJsonCoder.of(), BigQueryIO.TableRowParser.INSTANCE, QueryPriority.BATCH, null); fakeJobService.expectDryRunQuery( bqOptions.getProject(), query, new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(100L))); assertEquals(100, bqSource.getEstimatedSizeBytes(bqOptions)); }
dryRunQueryIfNeeded(bqOptions).getQuery().getReferencedTables(); if (referencedTables != null && !referencedTables.isEmpty()) { TableReference queryTable = referencedTables.get(0); executeQuery( jobIdToken, bqOptions.getProject(),
BigQueryQuerySource.create( stepUuid, ValueProvider.StaticValueProvider.of(query),
BigQueryQuerySource.create( stepUuid, ValueProvider.StaticValueProvider.of(query),
createBasicQueryConfig() .setAllowLargeResults(true) .setCreateDisposition("CREATE_IF_NEEDED")