Skip to content

Commit

Permalink
Support listing/getting/creating/updating tables/datasets/jobs in dif…
Browse files Browse the repository at this point in the history 
…ferent project
  • Loading branch information
@mziccard
mziccard committed Aug 31, 2016
1 parent 2a4e63d commit d6576ce
Show file tree
Hide file tree
Showing 6 changed files with 474 additions and 262 deletions.
11 changes: 11 additions & 0 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,17 @@ public static QueryResultsOption maxWaitTime(long maxWaitTime) {
*/
Page<Dataset> listDatasets(DatasetListOption... options);

/**
* Lists the datasets in the provided project. This method returns partial information on each
* dataset: ({@link Dataset#datasetId()}, {@link Dataset#friendlyName()} and
* {@link Dataset#generatedId()}). To get complete information use either
* {@link #getDataset(String, DatasetOption...)} or
* {@link #getDataset(DatasetId, DatasetOption...)}.
*
* @throws BigQueryException upon failure
*/
Page<Dataset> listDatasets(String projectId, DatasetListOption... options);

/**
* Deletes the requested dataset.
*
Expand Down
95 changes: 59 additions & 36 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,17 +48,19 @@ private static class DatasetPageFetcher implements NextPageFetcher<Dataset> {
private static final long serialVersionUID = -3057564042439021278L;
private final Map<BigQueryRpc.Option, ?> requestOptions;
private final BigQueryOptions serviceOptions;
private final String projectId;

DatasetPageFetcher(BigQueryOptions serviceOptions, String cursor,
DatasetPageFetcher(String projectId, BigQueryOptions serviceOptions, String cursor,
Map<BigQueryRpc.Option, ?> optionMap) {
this.projectId = projectId;
this.requestOptions =
PageImpl.nextRequestOptions(BigQueryRpc.Option.PAGE_TOKEN, cursor, optionMap);
this.serviceOptions = serviceOptions;
}

@Override
public Page<Dataset> nextPage() {
return listDatasets(serviceOptions, requestOptions);
return listDatasets(projectId, serviceOptions, requestOptions);
}
}

Expand All @@ -67,19 +69,19 @@ private static class TablePageFetcher implements NextPageFetcher<Table> {
private static final long serialVersionUID = 8611248840504201187L;
private final Map<BigQueryRpc.Option, ?> requestOptions;
private final BigQueryOptions serviceOptions;
private final String dataset;
private final DatasetId datasetId;

TablePageFetcher(String dataset, BigQueryOptions serviceOptions, String cursor,
TablePageFetcher(DatasetId datasetId, BigQueryOptions serviceOptions, String cursor,
Map<BigQueryRpc.Option, ?> optionMap) {
this.requestOptions =
PageImpl.nextRequestOptions(BigQueryRpc.Option.PAGE_TOKEN, cursor, optionMap);
this.serviceOptions = serviceOptions;
this.dataset = dataset;
this.datasetId = datasetId;
}

@Override
public Page<Table> nextPage() {
return listTables(dataset, serviceOptions, requestOptions);
return listTables(datasetId, serviceOptions, requestOptions);
}
}

Expand Down Expand Up @@ -213,13 +215,15 @@ public Dataset getDataset(String datasetId, DatasetOption... options) {

@Override
public Dataset getDataset(final DatasetId datasetId, DatasetOption... options) {
final DatasetId completeDatasetId = datasetId.setProjectId(options().projectId());
final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
try {
com.google.api.services.bigquery.model.Dataset answer =
runWithRetries(new Callable<com.google.api.services.bigquery.model.Dataset>() {
@Override
public com.google.api.services.bigquery.model.Dataset call() {
return bigQueryRpc.getDataset(datasetId.dataset(), optionsMap);
return bigQueryRpc.getDataset(
completeDatasetId.project(), completeDatasetId.dataset(), optionsMap);
}
}, options().retryParams(), EXCEPTION_HANDLER, options().clock());
return answer == null ? null : Dataset.fromPb(this, answer);
Expand All @@ -230,24 +234,29 @@ public com.google.api.services.bigquery.model.Dataset call() {

@Override
public Page<Dataset> listDatasets(DatasetListOption... options) {
return listDatasets(options(), optionMap(options));
return listDatasets(options().projectId(), options);
}

private static Page<Dataset> listDatasets(final BigQueryOptions serviceOptions,
final Map<BigQueryRpc.Option, ?> optionsMap) {
@Override
public Page<Dataset> listDatasets(String projectId, DatasetListOption... options) {
return listDatasets(projectId, options(), optionMap(options));
}

private static Page<Dataset> listDatasets(final String projectId,
final BigQueryOptions serviceOptions, final Map<BigQueryRpc.Option, ?> optionsMap) {
try {
BigQueryRpc.Tuple<String, Iterable<com.google.api.services.bigquery.model.Dataset>> result =
runWithRetries(new Callable<BigQueryRpc.Tuple<String,
Iterable<com.google.api.services.bigquery.model.Dataset>>>() {
@Override
public BigQueryRpc.Tuple<String,
Iterable<com.google.api.services.bigquery.model.Dataset>> call() {
return serviceOptions.rpc().listDatasets(optionsMap);
return serviceOptions.rpc().listDatasets(projectId, optionsMap);
}
}, serviceOptions.retryParams(), EXCEPTION_HANDLER, serviceOptions.clock());
String cursor = result.x();
return new PageImpl<>(new DatasetPageFetcher(serviceOptions, cursor, optionsMap), cursor,
Iterables.transform(result.y(),
return new PageImpl<>(new DatasetPageFetcher(projectId, serviceOptions, cursor, optionsMap),
cursor, Iterables.transform(result.y(),
new Function<com.google.api.services.bigquery.model.Dataset, Dataset>() {
@Override
public Dataset apply(com.google.api.services.bigquery.model.Dataset dataset) {
Expand All @@ -265,13 +274,15 @@ public boolean delete(String datasetId, DatasetDeleteOption... options) {
}

@Override
public boolean delete(final DatasetId datasetId, DatasetDeleteOption... options) {
public boolean delete(DatasetId datasetId, DatasetDeleteOption... options) {
final DatasetId completeDatasetId = datasetId.setProjectId(options().projectId());
final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
try {
return runWithRetries(new Callable<Boolean>() {
@Override
public Boolean call() {
return bigQueryRpc.deleteDataset(datasetId.dataset(), optionsMap);
return bigQueryRpc.deleteDataset(
completeDatasetId.project(), completeDatasetId.dataset(), optionsMap);
}
}, options().retryParams(), EXCEPTION_HANDLER, options().clock());
} catch (RetryHelper.RetryHelperException e) {
Expand All @@ -285,12 +296,14 @@ public boolean delete(String datasetId, String tableId) {
}

@Override
public boolean delete(final TableId tableId) {
public boolean delete(TableId tableId) {
final TableId completeTableId = tableId.setProjectId(options().projectId());
try {
return runWithRetries(new Callable<Boolean>() {
@Override
public Boolean call() {
return bigQueryRpc.deleteTable(tableId.dataset(), tableId.table());
return bigQueryRpc.deleteTable(
completeTableId.project(), completeTableId.dataset(), completeTableId.table());
}
}, options().retryParams(), EXCEPTION_HANDLER, options().clock());
} catch (RetryHelper.RetryHelperException e) {
Expand Down Expand Up @@ -340,14 +353,16 @@ public Table getTable(final String datasetId, final String tableId, TableOption.
}

@Override
public Table getTable(final TableId tableId, TableOption... options) {
public Table getTable(TableId tableId, TableOption... options) {
final TableId completeTableId = tableId.setProjectId(options().projectId());
final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
try {
com.google.api.services.bigquery.model.Table answer =
runWithRetries(new Callable<com.google.api.services.bigquery.model.Table>() {
@Override
public com.google.api.services.bigquery.model.Table call() {
return bigQueryRpc.getTable(tableId.dataset(), tableId.table(), optionsMap);
return bigQueryRpc.getTable(completeTableId.project(), completeTableId.dataset(),
completeTableId.table(), optionsMap);
}
}, options().retryParams(), EXCEPTION_HANDLER, options().clock());
return answer == null ? null : Table.fromPb(this, answer);
Expand All @@ -358,24 +373,27 @@ public com.google.api.services.bigquery.model.Table call() {

@Override
public Page<Table> listTables(String datasetId, TableListOption... options) {
return listTables(datasetId, options(), optionMap(options));
return listTables(
DatasetId.of(options().projectId(), datasetId), options(), optionMap(options));
}

@Override
public Page<Table> listTables(DatasetId datasetId, TableListOption... options) {
return listTables(datasetId.dataset(), options(), optionMap(options));
DatasetId completeDatasetId = datasetId.setProjectId(options().projectId());
return listTables(completeDatasetId, options(), optionMap(options));
}

private static Page<Table> listTables(final String datasetId, final BigQueryOptions
serviceOptions, final Map<BigQueryRpc.Option, ?> optionsMap) {
private static Page<Table> listTables(final DatasetId datasetId,
final BigQueryOptions serviceOptions, final Map<BigQueryRpc.Option, ?> optionsMap) {
try {
BigQueryRpc.Tuple<String, Iterable<com.google.api.services.bigquery.model.Table>> result =
runWithRetries(new Callable<BigQueryRpc.Tuple<String,
Iterable<com.google.api.services.bigquery.model.Table>>>() {
@Override
public BigQueryRpc.Tuple<String, Iterable<com.google.api.services.bigquery.model.Table>>
call() {
return serviceOptions.rpc().listTables(datasetId, optionsMap);
return serviceOptions.rpc().listTables(
datasetId.project(), datasetId.dataset(), optionsMap);
}
}, serviceOptions.retryParams(), EXCEPTION_HANDLER, serviceOptions.clock());
String cursor = result.x();
Expand All @@ -395,7 +413,7 @@ public Table apply(com.google.api.services.bigquery.model.Table table) {

@Override
public InsertAllResponse insertAll(InsertAllRequest request) {
final TableId tableId = request.table();
final TableId tableId = request.table().setProjectId(options().projectId());
final TableDataInsertAllRequest requestPb = new TableDataInsertAllRequest();
requestPb.setIgnoreUnknownValues(request.ignoreUnknownValues());
requestPb.setSkipInvalidRows(request.skipInvalidRows());
Expand All @@ -408,7 +426,7 @@ public Rows apply(RowToInsert rowToInsert) {
});
requestPb.setRows(rowsPb);
return InsertAllResponse.fromPb(
bigQueryRpc.insertAll(tableId.dataset(), tableId.table(), requestPb));
bigQueryRpc.insertAll(tableId.project(), tableId.dataset(), tableId.table(), requestPb));
}

@Override
Expand All @@ -425,12 +443,14 @@ public Page<List<FieldValue>> listTableData(TableId tableId, TableDataListOption
private static Page<List<FieldValue>> listTableData(final TableId tableId,
final BigQueryOptions serviceOptions, final Map<BigQueryRpc.Option, ?> optionsMap) {
try {
final TableId completeTableId = tableId.setProjectId(serviceOptions.projectId());
BigQueryRpc.Tuple<String, Iterable<TableRow>> result =
runWithRetries(new Callable<BigQueryRpc.Tuple<String, Iterable<TableRow>>>() {
@Override
public BigQueryRpc.Tuple<String, Iterable<TableRow>> call() {
return serviceOptions.rpc()
.listTableData(tableId.dataset(), tableId.table(), optionsMap);
.listTableData(completeTableId.project(), completeTableId.dataset(),
completeTableId.table(), optionsMap);
}
}, serviceOptions.retryParams(), EXCEPTION_HANDLER, serviceOptions.clock());
String cursor = result.x();
Expand Down Expand Up @@ -458,14 +478,15 @@ public Job getJob(String jobId, JobOption... options) {
}

@Override
public Job getJob(final JobId jobId, JobOption... options) {
public Job getJob(JobId jobId, JobOption... options) {
final Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);
final JobId completeJobId = jobId.setProjectId(options().projectId());
try {
com.google.api.services.bigquery.model.Job answer =
runWithRetries(new Callable<com.google.api.services.bigquery.model.Job>() {
@Override
public com.google.api.services.bigquery.model.Job call() {
return bigQueryRpc.getJob(jobId.job(), optionsMap);
return bigQueryRpc.getJob(completeJobId.project(), completeJobId.job(), optionsMap);
}
}, options().retryParams(), EXCEPTION_HANDLER, options().clock());
return answer == null ? null : Job.fromPb(this, answer);
Expand All @@ -487,7 +508,7 @@ private static Page<Job> listJobs(final BigQueryOptions serviceOptions,
@Override
public BigQueryRpc.Tuple<String, Iterable<com.google.api.services.bigquery.model.Job>>
call() {
return serviceOptions.rpc().listJobs(optionsMap);
return serviceOptions.rpc().listJobs(serviceOptions.projectId(), optionsMap);
}
}, serviceOptions.retryParams(), EXCEPTION_HANDLER, serviceOptions.clock());
String cursor = result.x();
Expand All @@ -507,12 +528,13 @@ public boolean cancel(String jobId) {
}

@Override
public boolean cancel(final JobId jobId) {
public boolean cancel(JobId jobId) {
final JobId completeJobId = jobId.setProjectId(options().projectId());
try {
return runWithRetries(new Callable<Boolean>() {
@Override
public Boolean call() {
return bigQueryRpc.cancel(jobId.job());
return bigQueryRpc.cancel(completeJobId.project(), completeJobId.job());
}
}, options().retryParams(), EXCEPTION_HANDLER, options().clock());
} catch (RetryHelper.RetryHelperException e) {
Expand Down Expand Up @@ -565,19 +587,20 @@ public QueryResponse getQueryResults(JobId job, QueryResultsOption... options) {
return getQueryResults(job, options(), optionsMap);
}

private static QueryResponse getQueryResults(final JobId jobId,
private static QueryResponse getQueryResults(JobId jobId,
final BigQueryOptions serviceOptions, final Map<BigQueryRpc.Option, ?> optionsMap) {
final JobId completeJobId = jobId.setProjectId(serviceOptions.projectId());
try {
GetQueryResultsResponse results =
runWithRetries(new Callable<GetQueryResultsResponse>() {
@Override
public GetQueryResultsResponse call() {
return serviceOptions.rpc().getQueryResults(jobId.job(), optionsMap);
return serviceOptions.rpc().getQueryResults(
completeJobId.project(), completeJobId.job(), optionsMap);
}
}, serviceOptions.retryParams(), EXCEPTION_HANDLER, serviceOptions.clock());
QueryResponse.Builder builder = QueryResponse.builder();
JobId completeJobId = JobId.fromPb(results.getJobReference());
builder.jobId(completeJobId);
builder.jobId(JobId.fromPb(results.getJobReference()));
builder.etag(results.getEtag());
builder.jobCompleted(results.getJobComplete());
List<TableRow> rowsPb = results.getRows();
Expand Down
Loading

0 comments on commit d6576ce

Please sign in to comment.