Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support listing/getting/creating/updating tables/datasets/jobs in other projects #1217

Merged
merged 2 commits into from
Sep 1, 2016
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
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());

This comment was marked as spam.

This comment was marked as spam.

This comment was marked as spam.

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
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,10 @@ public String toString() {
return toPb().toString();
}

JobId setProjectId(String projectId) {
return project() != null ? this : JobId.of(projectId, job());
}

JobReference toPb() {
return new JobReference().setProjectId(project).setJobId(job);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -326,7 +326,11 @@ public boolean equals(Object obj) {
}

JobInfo setProjectId(String projectId) {
return toBuilder().configuration(configuration.setProjectId(projectId)).build();
Builder builder = toBuilder();
if (jobId != null) {
builder.jobId(jobId.setProjectId(projectId));
}
return builder.configuration(configuration.setProjectId(projectId)).build();
}

Job toPb() {
Expand Down
Loading