diff --git a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQuery.java b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQuery.java index 28fb33dcc58c..da8f23e9a0ba 100644 --- a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQuery.java +++ b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQuery.java @@ -16,7 +16,19 @@ package com.google.gcloud.bigquery; +import static com.google.common.base.Preconditions.checkArgument; + +import com.google.common.base.Function; +import com.google.common.base.Joiner; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.common.collect.Sets; +import com.google.gcloud.Page; import com.google.gcloud.Service; +import com.google.gcloud.spi.BigQueryRpc; + +import java.util.List; +import java.util.Set; /** * An interface for Google Cloud BigQuery. @@ -25,5 +37,628 @@ */ public interface BigQuery extends Service { - // TODO(mziccard) add missing methods + /** + * Fields of a BigQuery Dataset resource. + * + * @see Dataset + * Resource + */ + enum DatasetField { + ACCESS("access"), + CREATION_TIME("creationTime"), + DATASET_REFERENCE("datasetReference"), + DEFAULT_TABLE_EXPIRATION_MS("defaultTableExpirationMsS"), + DESCRIPTION("description"), + ETAG("etag"), + FRIENDLY_NAME("friendlyName"), + ID("id"), + LAST_MODIFIED_TIME("lastModifiedTime"), + LOCATION("location"), + SELF_LINK("selfLink"); + + private final String selector; + + DatasetField(String selector) { + this.selector = selector; + } + + public String selector() { + return selector; + } + + static String selector(DatasetField... fields) { + Set fieldStrings = Sets.newHashSetWithExpectedSize(fields.length + 1); + fieldStrings.add(DATASET_REFERENCE.selector()); + for (DatasetField field : fields) { + fieldStrings.add(field.selector()); + } + return Joiner.on(',').join(fieldStrings); + } + } + + /** + * Fields of a BigQuery Table resource. + * + * @see Table + * Resource + */ + enum TableField { + CREATION_TIME("creationTime"), + DESCRIPTION("description"), + ETAG("etag"), + EXPIRATION_TIME("expirationTime"), + EXTERNAL_DATA_CONFIGURATION("externalDataConfiguration"), + FRIENDLY_NAME("friendlyName"), + ID("id"), + LAST_MODIFIED_TIME("lastModifiedTime"), + LOCATION("location"), + NUM_BYTES("numBytes"), + NUM_ROWS("numRows"), + SCHEMA("schema"), + SELF_LINK("selfLink"), + STREAMING_BUFFER("streamingBuffer"), + TABLE_REFERENCE("tableReference"), + TYPE("type"), + VIEW("view"); + + private final String selector; + + TableField(String selector) { + this.selector = selector; + } + + public String selector() { + return selector; + } + + static String selector(TableField... fields) { + Set fieldStrings = Sets.newHashSetWithExpectedSize(fields.length + 2); + fieldStrings.add(TABLE_REFERENCE.selector()); + fieldStrings.add(TYPE.selector()); + for (TableField field : fields) { + fieldStrings.add(field.selector()); + } + return Joiner.on(',').join(fieldStrings); + } + } + + /** + * Fields of a BigQuery Job resource. + * + * @see Job Resource + * + */ + enum JobField { + CONFIGURATION("configuration"), + ETAG("etag"), + ID("id"), + JOB_REFERENCE("jobReference"), + SELF_LINK("selfLink"), + STATISTICS("statistics"), + STATUS("status"), + USER_EMAIL("user_email"); + + private final String selector; + + JobField(String selector) { + this.selector = selector; + } + + public String selector() { + return selector; + } + + static String selector(JobField... fields) { + Set fieldStrings = Sets.newHashSetWithExpectedSize(fields.length + 2); + fieldStrings.add(JOB_REFERENCE.selector()); + fieldStrings.add(CONFIGURATION.selector()); + for (JobField field : fields) { + fieldStrings.add(field.selector()); + } + return Joiner.on(',').join(fieldStrings); + } + } + + /** + * Class for specifying dataset list options. + */ + class DatasetListOption extends Option { + + private static final long serialVersionUID = 8660294969063340498L; + + private DatasetListOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the maximum number of datasets to be returned. + */ + public static DatasetListOption maxResults(long maxResults) { + return new DatasetListOption(BigQueryRpc.Option.MAX_RESULTS, maxResults); + } + + /** + * Returns an option to specify the page token from which to start listing datasets. + */ + public static DatasetListOption startPageToken(String pageToken) { + return new DatasetListOption(BigQueryRpc.Option.PAGE_TOKEN, pageToken); + } + + /** + * Returns an options to list all datasets, even hidden ones. + */ + public static DatasetListOption all() { + return new DatasetListOption(BigQueryRpc.Option.ALL_DATASETS, true); + } + } + + /** + * Class for specifying dataset get, create and update options. + */ + class DatasetOption extends Option { + + private static final long serialVersionUID = 1674133909259913250L; + + private DatasetOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the dataset's fields to be returned by the RPC call. If this + * option is not provided all dataset's fields are returned. {@code DatasetOption.fields} can + * be used to specify only the fields of interest. {@link DatasetInfo#datasetId()} is always + * returned, even if not specified. + */ + public static DatasetOption fields(DatasetField... fields) { + return new DatasetOption(BigQueryRpc.Option.FIELDS, DatasetField.selector(fields)); + } + } + + /** + * Class for specifying dataset delete options. + */ + class DatasetDeleteOption extends Option { + + private static final long serialVersionUID = -7166083569900951337L; + + private DatasetDeleteOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to delete a dataset even if non-empty. If not provided, attempting to + * delete a non-empty dataset will result in a {@link BigQueryException} being thrown. + */ + public static DatasetDeleteOption deleteContents() { + return new DatasetDeleteOption(BigQueryRpc.Option.DELETE_CONTENTS, true); + } + } + + /** + * Class for specifying table list options. + */ + class TableListOption extends Option { + + private static final long serialVersionUID = 8660294969063340498L; + + private TableListOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the maximum number of tables to be returned. + */ + public static TableListOption maxResults(long maxResults) { + checkArgument(maxResults >= 0); + return new TableListOption(BigQueryRpc.Option.MAX_RESULTS, maxResults); + } + + /** + * Returns an option to specify the page token from which to start listing tables. + */ + public static TableListOption startPageToken(String pageToken) { + return new TableListOption(BigQueryRpc.Option.PAGE_TOKEN, pageToken); + } + } + + /** + * Class for specifying table get, create and update options. + */ + class TableOption extends Option { + + private static final long serialVersionUID = -1723870134095936772L; + + private TableOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the table's fields to be returned by the RPC call. If this + * option is not provided all table's fields are returned. {@code TableOption.fields} can be + * used to specify only the fields of interest. {@link BaseTableInfo#tableId()} and + * {@link BaseTableInfo#type()} are always returned, even if not specified. + */ + public static TableOption fields(TableField... fields) { + return new TableOption(BigQueryRpc.Option.FIELDS, TableField.selector(fields)); + } + } + + /** + * Class for specifying table data list options. + */ + class TableDataListOption extends Option { + + private static final long serialVersionUID = 8488823381738864434L; + + private TableDataListOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the maximum number of rows to be returned. + */ + public static TableDataListOption maxResults(long maxResults) { + checkArgument(maxResults >= 0); + return new TableDataListOption(BigQueryRpc.Option.MAX_RESULTS, maxResults); + } + + /** + * Returns an option to specify the page token from which to start listing table data. + */ + public static TableDataListOption startPageToken(String pageToken) { + return new TableDataListOption(BigQueryRpc.Option.PAGE_TOKEN, pageToken); + } + + /** + * Returns an option that sets the zero-based index of the row from which to start listing table + * data. + */ + public static TableDataListOption startIndex(long index) { + checkArgument(index >= 0); + return new TableDataListOption(BigQueryRpc.Option.START_INDEX, index); + } + } + + /** + * Class for specifying job list options. + */ + class JobListOption extends Option { + + private static final long serialVersionUID = -8207122131226481423L; + + private JobListOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to list all jobs, even the ones issued by other users. + */ + public static JobListOption allUsers() { + return new JobListOption(BigQueryRpc.Option.ALL_USERS, true); + } + + /** + * Returns an option to list only jobs that match the provided state filters. + */ + public static JobListOption stateFilter(JobStatus.State... stateFilters) { + List stringFilters = Lists.transform(ImmutableList.copyOf(stateFilters), + new Function() { + @Override + public String apply(JobStatus.State state) { + return state.name().toLowerCase(); + } + }); + return new JobListOption(BigQueryRpc.Option.STATE_FILTER, stringFilters); + } + + /** + * Returns an option to specify the maximum number of jobs to be returned. + */ + public static JobListOption maxResults(long maxResults) { + checkArgument(maxResults >= 0); + return new JobListOption(BigQueryRpc.Option.MAX_RESULTS, maxResults); + } + + /** + * Returns an option to specify the page token from which to start listing jobs. + */ + public static JobListOption startPageToken(String pageToken) { + return new JobListOption(BigQueryRpc.Option.PAGE_TOKEN, pageToken); + } + + /** + * Returns an option to specify the job's fields to be returned by the RPC call. If this option + * is not provided all job's fields are returned. {@code JobOption.fields()} can be used to + * specify only the fields of interest. {@link JobInfo#jobId()}, {@link JobStatus#state()}, + * {@link JobStatus#error()} as well as type-specific configuration (e.g. + * {@link QueryJobInfo#query()} for Query Jobs) are always returned, even if not specified. + * {@link JobField#SELF_LINK} and {@link JobField#ETAG} can not be selected when listing jobs. + */ + public static JobListOption fields(JobField... fields) { + String selector = JobField.selector(fields); + StringBuilder builder = new StringBuilder(); + builder.append("etag,jobs(").append(selector).append(",state,errorResult),nextPageToken"); + return new JobListOption(BigQueryRpc.Option.FIELDS, builder.toString()); + } + } + + /** + * Class for specifying table get and create options. + */ + class JobOption extends Option { + + private static final long serialVersionUID = -3111736712316353665L; + + private JobOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the job's fields to be returned by the RPC call. If this option + * is not provided all job's fields are returned. {@code JobOption.fields()} can be used to + * specify only the fields of interest. {@link JobInfo#jobId()} as well as type-specific + * configuration (e.g. {@link QueryJobInfo#query()} for Query Jobs) are always returned, even if + * not specified. + */ + public static JobOption fields(JobField... fields) { + return new JobOption(BigQueryRpc.Option.FIELDS, JobField.selector(fields)); + } + } + + /** + * Class for specifying query results options. + */ + class QueryResultsOption extends Option { + + private static final long serialVersionUID = 3788898503226985525L; + + private QueryResultsOption(BigQueryRpc.Option option, Object value) { + super(option, value); + } + + /** + * Returns an option to specify the maximum number of rows to be returned. + */ + public static QueryResultsOption maxResults(long maxResults) { + checkArgument(maxResults >= 0); + return new QueryResultsOption(BigQueryRpc.Option.MAX_RESULTS, maxResults); + } + + /** + * Returns an option to specify the page token from which to start getting query results. + */ + public static QueryResultsOption startPageToken(String pageToken) { + return new QueryResultsOption(BigQueryRpc.Option.PAGE_TOKEN, pageToken); + } + + /** + * Returns an option that sets the zero-based index of the row from which to start getting query + * results. + */ + public static QueryResultsOption startIndex(long startIndex) { + checkArgument(startIndex >= 0); + return new QueryResultsOption(BigQueryRpc.Option.START_INDEX, startIndex); + } + + /** + * Returns an option that sets how long to wait for the query to complete, in milliseconds, + * before returning. Default is 10 seconds. If the timeout passes before the job completes, + * {@link QueryResponse#jobComplete()} will be {@code false}. + */ + public static QueryResultsOption maxWaitTime(long maxWaitTime) { + checkArgument(maxWaitTime >= 0); + return new QueryResultsOption(BigQueryRpc.Option.TIMEOUT, maxWaitTime); + } + } + + /** + * Creates a new dataset. + * + * @throws BigQueryException upon failure + */ + DatasetInfo create(DatasetInfo dataset, DatasetOption... options) throws BigQueryException; + + /** + * Creates a new table. + * + * @throws BigQueryException upon failure + */ + BaseTableInfo create(BaseTableInfo table, TableOption... options) throws BigQueryException; + + /** + * Creates a new job. + * + * @throws BigQueryException upon failure + */ + JobInfo create(JobInfo job, JobOption... options) throws BigQueryException; + + /** + * Returns the requested dataset or {@code null} if not found. + * + * @throws BigQueryException upon failure + */ + DatasetInfo getDataset(String datasetId, DatasetOption... options) throws BigQueryException; + + /** + * Returns the requested dataset or {@code null} if not found. + * + * @throws BigQueryException upon failure + */ + DatasetInfo getDataset(DatasetId datasetId, DatasetOption... options) throws BigQueryException; + + /** + * Lists the project's datasets. This method returns partial information on each dataset + * ({@link DatasetInfo#datasetId()}, {@link DatasetInfo#friendlyName()} and + * {@link DatasetInfo#id()}). To get complete information use either + * {@link #getDataset(String, DatasetOption...)} or + * {@link #getDataset(DatasetId, DatasetOption...)}. + * + * @throws BigQueryException upon failure + */ + Page listDatasets(DatasetListOption... options) throws BigQueryException; + + /** + * Deletes the requested dataset. + * + * @return {@code true} if dataset was deleted, {@code false} if it was not found. + * @throws BigQueryException upon failure + */ + boolean delete(String datasetId, DatasetDeleteOption... options) throws BigQueryException; + + /** + * Deletes the requested dataset. + * + * @return {@code true} if dataset was deleted, {@code false} if it was not found. + * @throws BigQueryException upon failure + */ + boolean delete(DatasetId datasetId, DatasetDeleteOption... options) throws BigQueryException; + + /** + * Deletes the requested table. + * + * @return {@code true} if table was deleted, {@code false} if it was not found. + * @throws BigQueryException upon failure + */ + boolean delete(String datasetId, String tableId) throws BigQueryException; + + /** + * Deletes the requested table. + * + * @return {@code true} if table was deleted, {@code false} if it was not found. + * @throws BigQueryException upon failure + */ + boolean delete(TableId tableId) throws BigQueryException; + + /** + * Updates dataset information. + * + * @throws BigQueryException upon failure + */ + DatasetInfo update(DatasetInfo dataset, DatasetOption... options) throws BigQueryException; + + /** + * Updates table information. + * + * @throws BigQueryException upon failure + */ + BaseTableInfo update(BaseTableInfo table, TableOption... options) throws BigQueryException; + + /** + * Returns the requested table or {@code null} if not found. + * + * @throws BigQueryException upon failure + */ + BaseTableInfo getTable(String datasetId, String tableId, TableOption... options) + throws BigQueryException; + + /** + * Returns the requested table or {@code null} if not found. + * + * @throws BigQueryException upon failure + */ + BaseTableInfo getTable(TableId tableId, TableOption... options) throws BigQueryException; + + /** + * Lists the tables in the dataset. This method returns partial information on each table + * ({@link BaseTableInfo#tableId()}, {@link BaseTableInfo#friendlyName()}, + * {@link BaseTableInfo#id()} and {@link BaseTableInfo#type()}). To get complete information use + * either {@link #getTable(TableId, TableOption...)} or + * {@link #getTable(String, String, TableOption...)}. + * + * @throws BigQueryException upon failure + */ + Page listTables(String datasetId, TableListOption... options) + throws BigQueryException; + + /** + * Lists the tables in the dataset. This method returns partial information on each table + * ({@link BaseTableInfo#tableId()}, {@link BaseTableInfo#friendlyName()}, + * {@link BaseTableInfo#id()} and {@link BaseTableInfo#type()}). To get complete information use + * either {@link #getTable(TableId, TableOption...)} or + * {@link #getTable(String, String, TableOption...)}. + * + * @throws BigQueryException upon failure + */ + Page listTables(DatasetId datasetId, TableListOption... options) + throws BigQueryException; + + /** + * Sends an insert all request. + * + * @throws BigQueryException upon failure + */ + InsertAllResponse insertAll(InsertAllRequest request) throws BigQueryException; + + /** + * Lists the table's rows. + * + * @throws BigQueryException upon failure + */ + Page> listTableData(String datasetId, String tableId, + TableDataListOption... options) throws BigQueryException; + + /** + * Lists the table's rows. + * + * @throws BigQueryException upon failure + */ + Page> listTableData(TableId tableId, TableDataListOption... options) + throws BigQueryException; + + /** + * Returns the requested job or {@code null} if not found. + * + * @throws BigQueryException upon failure + */ + JobInfo getJob(String jobId, JobOption... options) throws BigQueryException; + + /** + * Returns the requested job or {@code null} if not found. + * + * @throws BigQueryException upon failure + */ + JobInfo getJob(JobId jobId, JobOption... options) throws BigQueryException; + + /** + * Lists the jobs. + * + * @throws BigQueryException upon failure + */ + Page listJobs(JobListOption... options) throws BigQueryException; + + /** + * Sends a job cancel request. This call will return immediately. The job status can then be + * checked using either {@link #getJob(JobId, JobOption...)} or + * {@link #getJob(String, JobOption...)}). + * + * @return {@code true} if cancel was requested successfully, {@code false} if the job was not + * found + * @throws BigQueryException upon failure + */ + boolean cancel(String jobId) throws BigQueryException; + + /** + * Sends a job cancel request. This call will return immediately. The job status can then be + * checked using either {@link #getJob(JobId, JobOption...)} or + * {@link #getJob(String, JobOption...)}). + * + * @return {@code true} if cancel was requested successfully, {@code false} if the job was not + * found + * @throws BigQueryException upon failure + */ + boolean cancel(JobId tableId) throws BigQueryException; + + /** + * Runs the query associated with the request. + * + * @throws BigQueryException upon failure + */ + QueryResponse query(QueryRequest request) throws BigQueryException; + + /** + * Returns results of the query associated with the provided job. + * + * @throws BigQueryException upon failure + */ + QueryResponse getQueryResults(JobId job, QueryResultsOption... options) throws BigQueryException; } diff --git a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryFactory.java b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryFactory.java index 2fc98125f4be..90e7bbccd483 100644 --- a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryFactory.java +++ b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryFactory.java @@ -16,7 +16,6 @@ package com.google.gcloud.bigquery; - import com.google.gcloud.ServiceFactory; /** diff --git a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryImpl.java b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryImpl.java new file mode 100644 index 000000000000..e2f8f890c6a3 --- /dev/null +++ b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryImpl.java @@ -0,0 +1,714 @@ +/* + * Copyright 2015 Google Inc. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.gcloud.bigquery; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.gcloud.RetryHelper.runWithRetries; + +import com.google.api.services.bigquery.model.Dataset; +import com.google.api.services.bigquery.model.GetQueryResultsResponse; +import com.google.api.services.bigquery.model.Job; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableDataInsertAllRequest; +import com.google.api.services.bigquery.model.TableDataInsertAllRequest.Rows; +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; +import com.google.gcloud.BaseService; +import com.google.gcloud.ExceptionHandler; +import com.google.gcloud.ExceptionHandler.Interceptor; +import com.google.gcloud.Page; +import com.google.gcloud.PageImpl; +import com.google.gcloud.RetryHelper; +import com.google.gcloud.bigquery.InsertAllRequest.RowToInsert; +import com.google.gcloud.spi.BigQueryRpc; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.Callable; + +final class BigQueryImpl extends BaseService implements BigQuery { + + private static final Interceptor EXCEPTION_HANDLER_INTERCEPTOR = new Interceptor() { + + private static final long serialVersionUID = -7478333733015750774L; + + @Override + public RetryResult afterEval(Exception exception, RetryResult retryResult) { + return Interceptor.RetryResult.CONTINUE_EVALUATION; + } + + @Override + public RetryResult beforeEval(Exception exception) { + if (exception instanceof BigQueryException) { + boolean retriable = ((BigQueryException) exception).retryable(); + return retriable ? Interceptor.RetryResult.RETRY : Interceptor.RetryResult.NO_RETRY; + } + return Interceptor.RetryResult.CONTINUE_EVALUATION; + } + }; + static final ExceptionHandler EXCEPTION_HANDLER = ExceptionHandler.builder() + .abortOn(RuntimeException.class).interceptor(EXCEPTION_HANDLER_INTERCEPTOR).build(); + + private abstract static class BasePageFetcher implements PageImpl.NextPageFetcher { + + private static final long serialVersionUID = -338124488600215401L; + + protected final Map requestOptions; + protected final BigQueryOptions serviceOptions; + + BasePageFetcher(BigQueryOptions serviceOptions, String cursor, + Map optionMap) { + this.serviceOptions = serviceOptions; + ImmutableMap.Builder builder = ImmutableMap.builder(); + if (cursor != null) { + builder.put(BigQueryRpc.Option.PAGE_TOKEN, cursor); + } + for (Map.Entry option : optionMap.entrySet()) { + if (option.getKey() != BigQueryRpc.Option.PAGE_TOKEN) { + builder.put(option.getKey(), option.getValue()); + } + } + this.requestOptions = builder.build(); + } + } + + private static class DatasetPageFetcher extends BasePageFetcher { + + private static final long serialVersionUID = 3030824397616608646L; + + DatasetPageFetcher(BigQueryOptions serviceOptions, String cursor, + Map optionMap) { + super(serviceOptions, cursor, optionMap); + } + + @Override + public Page nextPage() { + return listDatasets(serviceOptions, requestOptions); + } + } + + private static class TablePageFetcher extends BasePageFetcher { + + private static final long serialVersionUID = 5908129355985236115L; + private final String dataset; + + TablePageFetcher(String dataset, BigQueryOptions serviceOptions, String cursor, + Map optionMap) { + super(serviceOptions, cursor, optionMap); + this.dataset = dataset; + } + + @Override + public Page nextPage() { + return listTables(dataset, serviceOptions, requestOptions); + } + } + + private static class JobPageFetcher extends BasePageFetcher { + + private static final long serialVersionUID = -4984845360519279880L; + + JobPageFetcher(BigQueryOptions serviceOptions, String cursor, + Map optionMap) { + super(serviceOptions, cursor, optionMap); + } + + @Override + public Page nextPage() { + return listJobs(serviceOptions, requestOptions); + } + } + + private static class TableDataPageFetcher extends BasePageFetcher> { + + private static final long serialVersionUID = 1281938239570262432L; + private final TableId table; + + TableDataPageFetcher(TableId table, BigQueryOptions serviceOptions, String cursor, + Map optionMap) { + super(serviceOptions, cursor, optionMap); + this.table = table; + } + + @Override + public Page> nextPage() { + return listTableData(table, serviceOptions, requestOptions); + } + } + + private static class QueryResultsPageFetcherImpl extends BasePageFetcher> + implements QueryResult.QueryResultsPageFetcher { + + private static final long serialVersionUID = 6713948754731557486L; + private final JobId job; + + QueryResultsPageFetcherImpl(JobId job, BigQueryOptions serviceOptions, String cursor, + Map optionMap) { + super(serviceOptions, cursor, optionMap); + this.job = job; + } + + @Override + public QueryResult nextPage() { + return getQueryResults(job, serviceOptions, requestOptions).result(); + } + } + + private final BigQueryRpc bigQueryRpc; + + BigQueryImpl(BigQueryOptions options) { + super(options); + bigQueryRpc = options.rpc(); + } + + @Override + public DatasetInfo create(DatasetInfo dataset, DatasetOption... options) + throws BigQueryException { + final Dataset datasetPb = setProjectId(dataset).toPb(); + final Map optionsMap = optionMap(options); + try { + return DatasetInfo.fromPb(runWithRetries(new Callable() { + @Override + public Dataset call() { + return bigQueryRpc.create(datasetPb, optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER)); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public BaseTableInfo create(BaseTableInfo table, TableOption... options) + throws BigQueryException { + final Table tablePb = setProjectId(table).toPb(); + final Map optionsMap = optionMap(options); + try { + return BaseTableInfo.fromPb(runWithRetries(new Callable() { + @Override + public Table call() { + return bigQueryRpc.create(tablePb, optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER)); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public JobInfo create(JobInfo job, JobOption... options) throws BigQueryException { + final Job jobPb = setProjectId(job).toPb(); + final Map optionsMap = optionMap(options); + try { + return JobInfo.fromPb(runWithRetries(new Callable() { + @Override + public Job call() { + return bigQueryRpc.create(jobPb, optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER)); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public DatasetInfo getDataset(String datasetId, DatasetOption... options) + throws BigQueryException { + return getDataset(DatasetId.of(datasetId), options); + } + + @Override + public DatasetInfo getDataset(final DatasetId datasetId, DatasetOption... options) + throws BigQueryException { + final Map optionsMap = optionMap(options); + try { + Dataset answer = runWithRetries(new Callable() { + @Override + public Dataset call() { + return bigQueryRpc.getDataset(datasetId.dataset(), optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER); + return answer == null ? null : DatasetInfo.fromPb(answer); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public Page listDatasets(DatasetListOption... options) throws BigQueryException { + return listDatasets(options(), optionMap(options)); + } + + private static Page listDatasets(final BigQueryOptions serviceOptions, + final Map optionsMap) { + try { + BigQueryRpc.Tuple> result = + runWithRetries(new Callable>>() { + @Override + public BigQueryRpc.Tuple> call() { + return serviceOptions.rpc().listDatasets(optionsMap); + } + }, serviceOptions.retryParams(), EXCEPTION_HANDLER); + String cursor = result.x(); + return new PageImpl<>(new DatasetPageFetcher(serviceOptions, cursor, optionsMap), cursor, + Iterables.transform(result.y(), DatasetInfo.FROM_PB_FUNCTION)); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public boolean delete(String datasetId, DatasetDeleteOption... options) throws BigQueryException { + return delete(DatasetId.of(datasetId), options); + } + + @Override + public boolean delete(final DatasetId datasetId, DatasetDeleteOption... options) + throws BigQueryException { + final Map optionsMap = optionMap(options); + try { + return runWithRetries(new Callable() { + @Override + public Boolean call() { + return bigQueryRpc.deleteDataset(datasetId.dataset(), optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public boolean delete(String datasetId, String tableId) throws BigQueryException { + return delete(TableId.of(datasetId, tableId)); + } + + @Override + public boolean delete(final TableId tableId) throws BigQueryException { + try { + return runWithRetries(new Callable() { + @Override + public Boolean call() { + return bigQueryRpc.deleteTable(tableId.dataset(), tableId.table()); + } + }, options().retryParams(), EXCEPTION_HANDLER); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public DatasetInfo update(DatasetInfo dataset, DatasetOption... options) + throws BigQueryException { + final Dataset datasetPb = setProjectId(dataset).toPb(); + final Map optionsMap = optionMap(options); + try { + return DatasetInfo.fromPb(runWithRetries(new Callable() { + @Override + public Dataset call() { + return bigQueryRpc.patch(datasetPb, optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER)); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public BaseTableInfo update(BaseTableInfo table, TableOption... options) + throws BigQueryException { + final Table tablePb = setProjectId(table).toPb(); + final Map optionsMap = optionMap(options); + try { + return BaseTableInfo.fromPb(runWithRetries(new Callable
() { + @Override + public Table call() { + return bigQueryRpc.patch(tablePb, optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER)); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public BaseTableInfo getTable(final String datasetId, final String tableId, + TableOption... options) throws BigQueryException { + return getTable(TableId.of(datasetId, tableId), options); + } + + @Override + public BaseTableInfo getTable(final TableId tableId, TableOption... options) + throws BigQueryException { + final Map optionsMap = optionMap(options); + try { + Table answer = runWithRetries(new Callable
() { + @Override + public Table call() { + return bigQueryRpc.getTable(tableId.dataset(), tableId.table(), optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER); + return answer == null ? null : BaseTableInfo.fromPb(answer); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public Page listTables(String datasetId, TableListOption... options) + throws BigQueryException { + return listTables(datasetId, options(), optionMap(options)); + } + + @Override + public Page listTables(DatasetId datasetId, TableListOption... options) + throws BigQueryException { + return listTables(datasetId.dataset(), options(), optionMap(options)); + } + + private static Page listTables(final String datasetId, final BigQueryOptions + serviceOptions, final Map optionsMap) { + try { + BigQueryRpc.Tuple> result = + runWithRetries(new Callable>>() { + @Override + public BigQueryRpc.Tuple> call() { + return serviceOptions.rpc().listTables(datasetId, optionsMap); + } + }, serviceOptions.retryParams(), EXCEPTION_HANDLER); + String cursor = result.x(); + Iterable tables = Iterables.transform(result.y(), + BaseTableInfo.FROM_PB_FUNCTION); + return new PageImpl<>(new TablePageFetcher(datasetId, serviceOptions, cursor, optionsMap), + cursor, tables); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public InsertAllResponse insertAll(InsertAllRequest request) throws BigQueryException { + final TableId tableId = request.table(); + final TableDataInsertAllRequest requestPb = new TableDataInsertAllRequest(); + requestPb.setIgnoreUnknownValues(request.ignoreUnknownValues()); + requestPb.setSkipInvalidRows(request.skipInvalidRows()); + List rowsPb = Lists.transform(request.rows(), new Function() { + @Override + public Rows apply(RowToInsert rowToInsert) { + return new Rows().setInsertId(rowToInsert.id()).setJson(rowToInsert.content()); + } + }); + requestPb.setRows(rowsPb); + return InsertAllResponse.fromPb( + bigQueryRpc.insertAll(tableId.dataset(), tableId.table(), requestPb)); + } + + @Override + public Page> listTableData(String datasetId, String tableId, + TableDataListOption... options) throws BigQueryException { + return listTableData(TableId.of(datasetId, tableId), options(), optionMap(options)); + } + + @Override + public Page> listTableData(TableId tableId, TableDataListOption... options) + throws BigQueryException { + return listTableData(tableId, options(), optionMap(options)); + } + + private static Page> listTableData(final TableId tableId, + final BigQueryOptions serviceOptions, final Map optionsMap) { + try { + BigQueryRpc.Tuple> result = + runWithRetries(new Callable>>() { + @Override + public BigQueryRpc.Tuple> call() { + return serviceOptions.rpc() + .listTableData(tableId.dataset(), tableId.table(), optionsMap); + } + }, serviceOptions.retryParams(), EXCEPTION_HANDLER); + String cursor = result.x(); + return new PageImpl<>(new TableDataPageFetcher(tableId, serviceOptions, cursor, optionsMap), + cursor, transformTableData(result.y())); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + private static List> transformTableData(Iterable tableDataPb) { + return ImmutableList.copyOf( + Iterables.transform(tableDataPb != null ? tableDataPb : ImmutableList.of(), + new Function>() { + @Override + public List apply(TableRow rowPb) { + return Lists.transform(rowPb.getF(), FieldValue.FROM_PB_FUNCTION); + } + })); + } + + @Override + public JobInfo getJob(String jobId, JobOption... options) throws BigQueryException { + return getJob(JobId.of(jobId), options); + } + + @Override + public JobInfo getJob(final JobId jobId, JobOption... options) throws BigQueryException { + final Map optionsMap = optionMap(options); + try { + Job answer = runWithRetries(new Callable() { + @Override + public Job call() { + return bigQueryRpc.getJob(jobId.job(), optionsMap); + } + }, options().retryParams(), EXCEPTION_HANDLER); + return answer == null ? null : JobInfo.fromPb(answer); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public Page listJobs(JobListOption... options) throws BigQueryException { + return listJobs(options(), optionMap(options)); + } + + private static Page listJobs(final BigQueryOptions serviceOptions, + final Map optionsMap) { + BigQueryRpc.Tuple> result = + runWithRetries(new Callable>>() { + @Override + public BigQueryRpc.Tuple> call() { + return serviceOptions.rpc().listJobs(optionsMap); + } + }, serviceOptions.retryParams(), EXCEPTION_HANDLER); + String cursor = result.x(); + Iterable jobs = Iterables.transform(result.y(), JobInfo.FROM_PB_FUNCTION); + return new PageImpl<>(new JobPageFetcher(serviceOptions, cursor, optionsMap), cursor, jobs); + } + + @Override + public boolean cancel(String jobId) throws BigQueryException { + return cancel(JobId.of(jobId)); + } + + @Override + public boolean cancel(final JobId jobId) throws BigQueryException { + try { + return runWithRetries(new Callable() { + @Override + public Boolean call() { + return bigQueryRpc.cancel(jobId.job()); + } + }, options().retryParams(), EXCEPTION_HANDLER); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public QueryResponse query(final QueryRequest request) throws BigQueryException { + try { + com.google.api.services.bigquery.model.QueryResponse results = + runWithRetries(new Callable() { + @Override + public com.google.api.services.bigquery.model.QueryResponse call() { + return bigQueryRpc.query(setProjectId(request).toPb()); + } + }, options().retryParams(), EXCEPTION_HANDLER); + QueryResponse.Builder builder = QueryResponse.builder(); + JobId completeJobId = JobId.fromPb(results.getJobReference()); + builder.jobId(completeJobId); + builder.jobComplete(results.getJobComplete()); + List rowsPb = results.getRows(); + if (results.getJobComplete()) { + builder.jobComplete(true); + QueryResult.Builder resultBuilder = transformQueryResults(completeJobId, rowsPb, + results.getPageToken(), options(), ImmutableMap.of()); + resultBuilder.totalBytesProcessed(results.getTotalBytesProcessed()); + resultBuilder.cacheHit(results.getCacheHit()); + if (results.getSchema() != null) { + resultBuilder.schema(Schema.fromPb(results.getSchema())); + } + if (results.getTotalRows() != null) { + resultBuilder.totalRows(results.getTotalRows().longValue()); + } + builder.result(resultBuilder.build()); + } + if (results.getErrors() != null) { + builder.executionErrors( + Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION)); + } + return builder.build(); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + @Override + public QueryResponse getQueryResults(JobId job, QueryResultsOption... options) + throws BigQueryException { + Map optionsMap = optionMap(options); + return getQueryResults(job, options(), optionsMap); + } + + private static QueryResponse getQueryResults(final JobId jobId, + final BigQueryOptions serviceOptions, final Map optionsMap) { + try { + GetQueryResultsResponse results = + runWithRetries(new Callable() { + @Override + public GetQueryResultsResponse call() { + return serviceOptions.rpc().getQueryResults(jobId.job(), optionsMap); + } + }, serviceOptions.retryParams(), EXCEPTION_HANDLER); + QueryResponse.Builder builder = QueryResponse.builder(); + JobId completeJobId = JobId.fromPb(results.getJobReference()); + builder.jobId(completeJobId); + builder.etag(results.getEtag()); + builder.jobComplete(results.getJobComplete()); + List rowsPb = results.getRows(); + if (results.getJobComplete()) { + QueryResult.Builder resultBuilder = transformQueryResults(completeJobId, rowsPb, + results.getPageToken(), serviceOptions, ImmutableMap.of()); + resultBuilder.totalBytesProcessed(results.getTotalBytesProcessed()); + resultBuilder.cacheHit(results.getCacheHit()); + if (results.getSchema() != null) { + resultBuilder.schema(Schema.fromPb(results.getSchema())); + } + if (results.getTotalRows() != null) { + resultBuilder.totalRows(results.getTotalRows().longValue()); + } + builder.result(resultBuilder.build()); + } + if (results.getErrors() != null) { + builder.executionErrors( + Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION)); + } + return builder.build(); + } catch (RetryHelper.RetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + } + + private static QueryResult.Builder transformQueryResults(JobId jobId, List rowsPb, + String cursor, BigQueryOptions serviceOptions, Map optionsMap) { + QueryResultsPageFetcherImpl nextPageFetcher = + new QueryResultsPageFetcherImpl(jobId, serviceOptions, cursor, optionsMap); + return QueryResult.builder() + .pageFetcher(nextPageFetcher) + .cursor(cursor) + .results(transformTableData(rowsPb)); + } + + private Map optionMap(Option... options) { + Map optionMap = Maps.newEnumMap(BigQueryRpc.Option.class); + for (Option option : options) { + Object prev = optionMap.put(option.rpcOption(), option.value()); + checkArgument(prev == null, "Duplicate option %s", option); + } + return optionMap; + } + + private DatasetInfo setProjectId(DatasetInfo dataset) { + DatasetInfo.Builder datasetBuilder = dataset.toBuilder(); + datasetBuilder.datasetId(setProjectId(dataset.datasetId())); + if (dataset.acl() != null) { + List acls = Lists.newArrayListWithCapacity(dataset.acl().size()); + for (Acl acl : dataset.acl()) { + if (acl.entity().type() == Acl.Entity.Type.VIEW) { + Dataset.Access accessPb = acl.toPb(); + TableReference viewReferencePb = accessPb.getView(); + if (viewReferencePb.getProjectId() == null) { + viewReferencePb.setProjectId(options().projectId()); + } + acls.add(new Acl(new Acl.View(TableId.fromPb(viewReferencePb)))); + } else { + acls.add(acl); + } + } + datasetBuilder.acl(acls); + } + return datasetBuilder.build(); + } + + private DatasetId setProjectId(DatasetId dataset) { + return dataset.project() != null ? dataset + : DatasetId.of(options().projectId(), dataset.dataset()); + } + + private BaseTableInfo setProjectId(BaseTableInfo table) { + return table.toBuilder().tableId(setProjectId(table.tableId())).build(); + } + + private TableId setProjectId(TableId table) { + return table.project() != null ? table + : TableId.of(options().projectId(), table.dataset(), table.table()); + } + + private JobInfo setProjectId(JobInfo job) { + if (job instanceof CopyJobInfo) { + CopyJobInfo copyJob = (CopyJobInfo) job; + CopyJobInfo.Builder copyBuilder = copyJob.toBuilder(); + copyBuilder.destinationTable(setProjectId(copyJob.destinationTable())); + copyBuilder.sourceTables( + Lists.transform(copyJob.sourceTables(), new Function() { + @Override + public TableId apply(TableId tableId) { + return setProjectId(tableId); + } + })); + return copyBuilder.build(); + } + if (job instanceof QueryJobInfo) { + QueryJobInfo queryJob = (QueryJobInfo) job; + QueryJobInfo.Builder queryBuilder = queryJob.toBuilder(); + if (queryJob.destinationTable() != null) { + queryBuilder.destinationTable(setProjectId(queryJob.destinationTable())); + } + if (queryJob.defaultDataset() != null) { + queryBuilder.defaultDataset(setProjectId(queryJob.defaultDataset())); + } + return queryBuilder.build(); + } + if (job instanceof ExtractJobInfo) { + ExtractJobInfo extractJob = (ExtractJobInfo) job; + ExtractJobInfo.Builder extractBuilder = extractJob.toBuilder(); + extractBuilder.sourceTable(setProjectId(extractJob.sourceTable())); + return extractBuilder.build(); + } + if (job instanceof LoadJobInfo) { + LoadJobInfo loadJob = (LoadJobInfo) job; + LoadJobInfo.Builder loadBuilder = loadJob.toBuilder(); + loadBuilder.destinationTable(setProjectId(loadJob.destinationTable())); + return loadBuilder.build(); + } + return job; + } + + private QueryRequest setProjectId(QueryRequest request) { + QueryRequest.Builder builder = request.toBuilder(); + if (request.defaultDataset() != null) { + builder.defaultDataset(setProjectId(request.defaultDataset())); + } + return builder.build(); + } +} diff --git a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryOptions.java b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryOptions.java index 59a4b3229f68..71d43cfbe565 100644 --- a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryOptions.java +++ b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/BigQueryOptions.java @@ -18,9 +18,9 @@ import com.google.common.collect.ImmutableSet; import com.google.gcloud.ServiceOptions; -import com.google.gcloud.spi.DefaultBigQueryRpc; import com.google.gcloud.spi.BigQueryRpc; import com.google.gcloud.spi.BigQueryRpcFactory; +import com.google.gcloud.spi.DefaultBigQueryRpc; import java.util.Set; @@ -36,8 +36,7 @@ public static class DefaultBigqueryFactory implements BigQueryFactory { @Override public BigQuery create(BigQueryOptions options) { - // TODO(mziccard) return new BigqueryImpl(options); - return null; + return new BigQueryImpl(options); } } @@ -47,8 +46,7 @@ public static class DefaultBigQueryRpcFactory implements BigQueryRpcFactory { @Override public BigQueryRpc create(BigQueryOptions options) { - // TODO(mziccard) return new DefaultBigqueryRpc(options); - return null; + return new DefaultBigQueryRpc(options); } } @@ -106,6 +104,10 @@ public boolean equals(Object obj) { return baseEquals(other); } + public static BigQueryOptions defaultInstance() { + return builder().build(); + } + public static Builder builder() { return new Builder(); } diff --git a/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/Option.java b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/Option.java new file mode 100644 index 000000000000..d88820fe5a29 --- /dev/null +++ b/gcloud-java-bigquery/src/main/java/com/google/gcloud/bigquery/Option.java @@ -0,0 +1,72 @@ +/* + * Copyright 2015 Google Inc. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.gcloud.bigquery; + +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.common.base.MoreObjects; +import com.google.gcloud.spi.BigQueryRpc; + +import java.io.Serializable; +import java.util.Objects; + +/** + * Base class for BigQuery operation option. + */ +class Option implements Serializable { + + private static final long serialVersionUID = -6647817677804099207L; + + private final BigQueryRpc.Option rpcOption; + private final Object value; + + Option(BigQueryRpc.Option rpcOption, Object value) { + this.rpcOption = checkNotNull(rpcOption); + this.value = value; + } + + BigQueryRpc.Option rpcOption() { + return rpcOption; + } + + Object value() { + return value; + } + + @Override + public boolean equals(Object obj) { + if (!(obj instanceof Option)) { + return false; + } + Option other = (Option) obj; + return Objects.equals(rpcOption, other.rpcOption) + && Objects.equals(value, other.value); + } + + @Override + public int hashCode() { + return Objects.hash(rpcOption, value); + } + + @Override + public String toString() { + return MoreObjects.toStringHelper(this) + .add("name", rpcOption.value()) + .add("value", value) + .toString(); + } +} diff --git a/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/BigQueryImplTest.java b/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/BigQueryImplTest.java new file mode 100644 index 000000000000..c9c46e0a632d --- /dev/null +++ b/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/BigQueryImplTest.java @@ -0,0 +1,1042 @@ +/* + * Copyright 2015 Google Inc. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.gcloud.bigquery; + +import static org.easymock.EasyMock.capture; +import static org.easymock.EasyMock.eq; +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +import com.google.api.services.bigquery.model.Dataset; +import com.google.api.services.bigquery.model.ErrorProto; +import com.google.api.services.bigquery.model.Job; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableCell; +import com.google.api.services.bigquery.model.TableDataInsertAllRequest; +import com.google.api.services.bigquery.model.TableDataInsertAllResponse; +import com.google.api.services.bigquery.model.TableRow; +import com.google.common.base.Function; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.common.collect.Lists; +import com.google.gcloud.AuthCredentials; +import com.google.gcloud.Page; +import com.google.gcloud.RetryParams; +import com.google.gcloud.bigquery.InsertAllRequest.RowToInsert; +import com.google.gcloud.spi.BigQueryRpc; +import com.google.gcloud.spi.BigQueryRpc.Tuple; +import com.google.gcloud.spi.BigQueryRpcFactory; + +import org.easymock.Capture; +import org.easymock.EasyMock; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; + +import java.math.BigInteger; +import java.util.List; +import java.util.Map; + +public class BigQueryImplTest { + + private static final String PROJECT = "project"; + private static final String DATASET = "dataset"; + private static final String TABLE = "table"; + private static final String JOB = "job"; + private static final String OTHER_TABLE = "otherTable"; + private static final String OTHER_DATASET = "otherDataset"; + private static final List ACCESS_RULES = ImmutableList.of( + new Acl(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER), + new Acl(new Acl.View(TableId.of("dataset", "table")), Acl.Role.WRITER)); + private static final List ACCESS_RULES_WITH_PROJECT = ImmutableList.of( + new Acl(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER), + new Acl(new Acl.View(TableId.of(PROJECT, "dataset", "table")))); + private static final DatasetInfo DATASET_INFO = DatasetInfo.builder(DATASET) + .acl(ACCESS_RULES) + .description("description") + .build(); + private static final DatasetInfo DATASET_INFO_WITH_PROJECT = DatasetInfo.builder(PROJECT, DATASET) + .acl(ACCESS_RULES_WITH_PROJECT) + .description("description") + .build(); + private static final DatasetInfo OTHER_DATASET_INFO = DatasetInfo.builder(PROJECT, OTHER_DATASET) + .acl(ACCESS_RULES) + .description("other description") + .build(); + private static final TableId TABLE_ID = TableId.of(DATASET, TABLE); + private static final TableId OTHER_TABLE_ID = TableId.of(PROJECT, DATASET, OTHER_TABLE); + private static final TableId TABLE_ID_WITH_PROJECT = TableId.of(PROJECT, DATASET, TABLE); + private static final Field FIELD_SCHEMA1 = + Field.builder("StringField", Field.Type.string()) + .mode(Field.Mode.NULLABLE) + .description("FieldDescription1") + .build(); + private static final Field FIELD_SCHEMA2 = + Field.builder("IntegerField", Field.Type.integer()) + .mode(Field.Mode.REPEATED) + .description("FieldDescription2") + .build(); + private static final Field FIELD_SCHEMA3 = + Field.builder("RecordField", Field.Type.record(FIELD_SCHEMA1, FIELD_SCHEMA2)) + .mode(Field.Mode.REQUIRED) + .description("FieldDescription3") + .build(); + private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA1, FIELD_SCHEMA2, FIELD_SCHEMA3); + private static final BaseTableInfo TABLE_INFO = TableInfo.of(TABLE_ID, TABLE_SCHEMA); + private static final BaseTableInfo OTHER_TABLE_INFO = TableInfo.of(OTHER_TABLE_ID, TABLE_SCHEMA); + private static final BaseTableInfo TABLE_INFO_WITH_PROJECT = + TableInfo.of(TABLE_ID_WITH_PROJECT, TABLE_SCHEMA); + private static final JobInfo LOAD_JOB = LoadJobInfo.of(TABLE_ID, "URI"); + private static final JobInfo LOAD_JOB_WITH_PROJECT = LoadJobInfo.of(TABLE_ID_WITH_PROJECT, "URI"); + private static final JobInfo COMPLETE_LOAD_JOB = LoadJobInfo.builder(TABLE_ID_WITH_PROJECT, "URI") + .jobId(JobId.of(PROJECT, JOB)) + .build(); + private static final JobInfo COPY_JOB = + CopyJobInfo.of(TABLE_ID, ImmutableList.of(TABLE_ID, TABLE_ID)); + private static final JobInfo COPY_JOB_WITH_PROJECT = + CopyJobInfo.of(TABLE_ID_WITH_PROJECT, ImmutableList.of(TABLE_ID_WITH_PROJECT, + TABLE_ID_WITH_PROJECT)); + private static final JobInfo COMPLETE_COPY_JOB = + CopyJobInfo.builder(TABLE_ID_WITH_PROJECT, ImmutableList.of(TABLE_ID_WITH_PROJECT, + TABLE_ID_WITH_PROJECT)) + .jobId(JobId.of(PROJECT, JOB)) + .build(); + private static final JobInfo QUERY_JOB = QueryJobInfo.builder("SQL") + .defaultDataset(DatasetId.of(DATASET)) + .destinationTable(TABLE_ID) + .build(); + private static final JobInfo QUERY_JOB_WITH_PROJECT = QueryJobInfo.builder("SQL") + .defaultDataset(DatasetId.of(PROJECT, DATASET)) + .destinationTable(TABLE_ID_WITH_PROJECT) + .build(); + private static final JobInfo COMPLETE_QUERY_JOB = QueryJobInfo.builder("SQL") + .defaultDataset(DatasetId.of(PROJECT, DATASET)).destinationTable(TABLE_ID_WITH_PROJECT) + .jobId(JobId.of(PROJECT, JOB)) + .build(); + private static final JobInfo EXTRACT_JOB = ExtractJobInfo.of(TABLE_ID, "URI"); + private static final JobInfo EXTRACT_JOB_WITH_PROJECT = + ExtractJobInfo.of(TABLE_ID_WITH_PROJECT, "URI"); + private static final JobInfo COMPLETE_EXTRACT_JOB = + ExtractJobInfo.builder(TABLE_ID_WITH_PROJECT, "URI") + .jobId(JobId.of(PROJECT, JOB)) + .build(); + private static final TableCell BOOLEAN_FIELD = new TableCell().setV("false"); + private static final TableCell INTEGER_FIELD = new TableCell().setV("1"); + private static final TableRow TABLE_ROW = + new TableRow().setF(ImmutableList.of(BOOLEAN_FIELD, INTEGER_FIELD)); + private static final QueryRequest QUERY_REQUEST = QueryRequest.builder("SQL") + .maxResults(42L) + .useQueryCache(false) + .defaultDataset(DatasetId.of(DATASET)) + .build(); + private static final QueryRequest QUERY_REQUEST_WITH_PROJECT = QueryRequest.builder("SQL") + .maxResults(42L) + .useQueryCache(false) + .defaultDataset(DatasetId.of(PROJECT, DATASET)) + .build(); + + // Empty BigQueryRpc options + private static final Map EMPTY_RPC_OPTIONS = ImmutableMap.of(); + + // Dataset options + private static final BigQuery.DatasetOption DATASET_OPTION_FIELDS = + BigQuery.DatasetOption.fields(BigQuery.DatasetField.ACCESS, BigQuery.DatasetField.ETAG); + + // Dataset list options + private static final BigQuery.DatasetListOption DATASET_LIST_ALL = + BigQuery.DatasetListOption.all(); + private static final BigQuery.DatasetListOption DATASET_LIST_PAGE_TOKEN = + BigQuery.DatasetListOption.startPageToken("cursor"); + private static final BigQuery.DatasetListOption DATASET_LIST_MAX_RESULTS = + BigQuery.DatasetListOption.maxResults(42L); + private static final Map DATASET_LIST_OPTIONS = ImmutableMap.of( + BigQueryRpc.Option.ALL_DATASETS, true, + BigQueryRpc.Option.PAGE_TOKEN, "cursor", + BigQueryRpc.Option.MAX_RESULTS, 42L); + + // Dataset delete options + private static final BigQuery.DatasetDeleteOption DATASET_DELETE_CONTENTS = + BigQuery.DatasetDeleteOption.deleteContents(); + private static final Map DATASET_DELETE_OPTIONS = ImmutableMap.of( + BigQueryRpc.Option.DELETE_CONTENTS, true); + + // Table options + private static final BigQuery.TableOption TABLE_OPTION_FIELDS = + BigQuery.TableOption.fields(BigQuery.TableField.SCHEMA, BigQuery.TableField.ETAG); + + // Table list options + private static final BigQuery.TableListOption TABLE_LIST_MAX_RESULTS = + BigQuery.TableListOption.maxResults(42L); + private static final BigQuery.TableListOption TABLE_LIST_PAGE_TOKEN = + BigQuery.TableListOption.startPageToken("cursor"); + private static final Map TABLE_LIST_OPTIONS = ImmutableMap.of( + BigQueryRpc.Option.MAX_RESULTS, 42L, + BigQueryRpc.Option.PAGE_TOKEN, "cursor"); + + // TableData list options + private static final BigQuery.TableDataListOption TABLE_DATA_LIST_MAX_RESULTS = + BigQuery.TableDataListOption.maxResults(42L); + private static final BigQuery.TableDataListOption TABLE_DATA_LIST_PAGE_TOKEN = + BigQuery.TableDataListOption.startPageToken("cursor"); + private static final BigQuery.TableDataListOption TABLE_DATA_LIST_START_INDEX = + BigQuery.TableDataListOption.startIndex(0L); + private static final Map TABLE_DATA_LIST_OPTIONS = ImmutableMap.of( + BigQueryRpc.Option.MAX_RESULTS, 42L, + BigQueryRpc.Option.PAGE_TOKEN, "cursor", + BigQueryRpc.Option.START_INDEX, 0L); + + // Job options + private static final BigQuery.JobOption JOB_OPTION_FIELDS = + BigQuery.JobOption.fields(BigQuery.JobField.USER_EMAIL); + + // Job list options + private static final BigQuery.JobListOption JOB_LIST_OPTION_FIELD = + BigQuery.JobListOption.fields(BigQuery.JobField.STATISTICS); + private static final BigQuery.JobListOption JOB_LIST_ALL_USERS = + BigQuery.JobListOption.allUsers(); + private static final BigQuery.JobListOption JOB_LIST_STATE_FILTER = + BigQuery.JobListOption.stateFilter(JobStatus.State.DONE, JobStatus.State.PENDING); + private static final BigQuery.JobListOption JOB_LIST_PAGE_TOKEN = + BigQuery.JobListOption.startPageToken("cursor"); + private static final BigQuery.JobListOption JOB_LIST_MAX_RESULTS = + BigQuery.JobListOption.maxResults(42L); + private static final Map JOB_LIST_OPTIONS = ImmutableMap.of( + BigQueryRpc.Option.ALL_USERS, true, + BigQueryRpc.Option.STATE_FILTER, ImmutableList.of("done", "pending"), + BigQueryRpc.Option.PAGE_TOKEN, "cursor", + BigQueryRpc.Option.MAX_RESULTS, 42L); + + // Query Results options + private static final BigQuery.QueryResultsOption QUERY_RESULTS_OPTION_TIME = + BigQuery.QueryResultsOption.maxWaitTime(42L); + private static final BigQuery.QueryResultsOption QUERY_RESULTS_OPTION_INDEX = + BigQuery.QueryResultsOption.startIndex(1024L); + private static final BigQuery.QueryResultsOption QUERY_RESULTS_OPTION_PAGE_TOKEN = + BigQuery.QueryResultsOption.startPageToken("cursor"); + private static final BigQuery.QueryResultsOption QUERY_RESULTS_OPTION_MAX_RESULTS = + BigQuery.QueryResultsOption.maxResults(0L); + private static final Map QUERY_RESULTS_OPTIONS = ImmutableMap.of( + BigQueryRpc.Option.TIMEOUT, 42L, + BigQueryRpc.Option.START_INDEX, 1024L, + BigQueryRpc.Option.PAGE_TOKEN, "cursor", + BigQueryRpc.Option.MAX_RESULTS, 0L); + + private BigQueryOptions options; + private BigQueryRpcFactory rpcFactoryMock; + private BigQueryRpc bigqueryRpcMock; + private BigQuery bigquery; + + @Rule + public ExpectedException thrown = ExpectedException.none(); + + @Before + public void setUp() { + rpcFactoryMock = EasyMock.createMock(BigQueryRpcFactory.class); + bigqueryRpcMock = EasyMock.createMock(BigQueryRpc.class); + EasyMock.expect(rpcFactoryMock.create(EasyMock.anyObject(BigQueryOptions.class))) + .andReturn(bigqueryRpcMock); + EasyMock.replay(rpcFactoryMock); + options = BigQueryOptions.builder() + .projectId(PROJECT) + .authCredentials(AuthCredentials.noCredentials()) + .serviceRpcFactory(rpcFactoryMock) + .build(); + } + + @After + public void tearDown() { + EasyMock.verify(rpcFactoryMock, bigqueryRpcMock); + } + + @Test + public void testGetOptions() { + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertSame(options, bigquery.options()); + } + + @Test + public void testCreateDataset() { + EasyMock.expect(bigqueryRpcMock.create(DATASET_INFO_WITH_PROJECT.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.create(DATASET_INFO); + assertEquals(DATASET_INFO_WITH_PROJECT, dataset); + } + + @Test + public void testCreateDatasetWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + EasyMock.expect( + bigqueryRpcMock.create(eq(DATASET_INFO_WITH_PROJECT.toPb()), capture(capturedOptions))) + .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.create(DATASET_INFO, DATASET_OPTION_FIELDS); + String selector = (String) capturedOptions.getValue().get(DATASET_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("datasetReference")); + assertTrue(selector.contains("access")); + assertTrue(selector.contains("etag")); + assertEquals(28, selector.length()); + assertEquals(DATASET_INFO_WITH_PROJECT, dataset); + } + + @Test + public void testGetDataset() { + EasyMock.expect(bigqueryRpcMock.getDataset(DATASET, EMPTY_RPC_OPTIONS)) + .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.getDataset(DATASET); + assertEquals(DATASET_INFO_WITH_PROJECT, dataset); + } + + @Test + public void testGetDatasetFromDatasetId() { + EasyMock.expect(bigqueryRpcMock.getDataset(DATASET, EMPTY_RPC_OPTIONS)) + .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.getDataset(DatasetId.of(PROJECT, DATASET)); + assertEquals(DATASET_INFO_WITH_PROJECT, dataset); + } + + @Test + public void testGetDatasetWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + EasyMock.expect(bigqueryRpcMock.getDataset(eq(DATASET), capture(capturedOptions))) + .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.getDataset(DATASET, DATASET_OPTION_FIELDS); + String selector = (String) capturedOptions.getValue().get(DATASET_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("datasetReference")); + assertTrue(selector.contains("access")); + assertTrue(selector.contains("etag")); + assertEquals(28, selector.length()); + assertEquals(DATASET_INFO_WITH_PROJECT, dataset); + } + + @Test + public void testListDatasets() { + String cursor = "cursor"; + ImmutableList datasetList = ImmutableList.of(DATASET_INFO_WITH_PROJECT, + OTHER_DATASET_INFO); + Tuple> result = + Tuple.of(cursor, Iterables.transform(datasetList, DatasetInfo.TO_PB_FUNCTION)); + EasyMock.expect(bigqueryRpcMock.listDatasets(EMPTY_RPC_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listDatasets(); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(datasetList.toArray(), Iterables.toArray(page.values(), DatasetInfo.class)); + } + + @Test + public void testListEmptyDatasets() { + ImmutableList datasets = ImmutableList.of(); + Tuple> result = Tuple.>of(null, datasets); + EasyMock.expect(bigqueryRpcMock.listDatasets(EMPTY_RPC_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listDatasets(); + assertNull(page.nextPageCursor()); + assertArrayEquals(ImmutableList.of().toArray(), + Iterables.toArray(page.values(), DatasetInfo.class)); + } + + @Test + public void testListDatasetsWithOptions() { + String cursor = "cursor"; + ImmutableList datasetList = ImmutableList.of(DATASET_INFO_WITH_PROJECT, + OTHER_DATASET_INFO); + Tuple> result = + Tuple.of(cursor, Iterables.transform(datasetList, DatasetInfo.TO_PB_FUNCTION)); + EasyMock.expect(bigqueryRpcMock.listDatasets(DATASET_LIST_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listDatasets(DATASET_LIST_ALL, DATASET_LIST_PAGE_TOKEN, + DATASET_LIST_MAX_RESULTS); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(datasetList.toArray(), Iterables.toArray(page.values(), DatasetInfo.class)); + } + + @Test + public void testDeleteDataset() { + EasyMock.expect(bigqueryRpcMock.deleteDataset(DATASET, EMPTY_RPC_OPTIONS)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.delete(DATASET)); + } + + @Test + public void testDeleteDatasetFromDatasetId() { + EasyMock.expect(bigqueryRpcMock.deleteDataset(DATASET, EMPTY_RPC_OPTIONS)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.delete(DatasetId.of(PROJECT, DATASET))); + } + + @Test + public void testDeleteDatasetWithOptions() { + EasyMock.expect(bigqueryRpcMock.deleteDataset(DATASET, DATASET_DELETE_OPTIONS)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.delete(DATASET, DATASET_DELETE_CONTENTS)); + } + + @Test + public void testUpdateDataset() { + DatasetInfo updatedDatasetInfo = DATASET_INFO.toBuilder().description("newDescription").build(); + DatasetInfo updatedDatasetInfoWithProject = DATASET_INFO_WITH_PROJECT.toBuilder() + .description("newDescription") + .build(); + EasyMock.expect(bigqueryRpcMock.patch(updatedDatasetInfoWithProject.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(updatedDatasetInfoWithProject.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.update(updatedDatasetInfo); + assertEquals(updatedDatasetInfoWithProject, dataset); + } + + @Test + public void testUpdateDatasetWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + DatasetInfo updatedDatasetInfo = DATASET_INFO.toBuilder().description("newDescription").build(); + DatasetInfo updatedDatasetInfoWithProject = DATASET_INFO_WITH_PROJECT.toBuilder() + .description("newDescription") + .build(); + EasyMock.expect( + bigqueryRpcMock.patch(eq(updatedDatasetInfoWithProject.toPb()), capture(capturedOptions))) + .andReturn(updatedDatasetInfoWithProject.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + DatasetInfo dataset = bigquery.update(updatedDatasetInfo, DATASET_OPTION_FIELDS); + String selector = (String) capturedOptions.getValue().get(DATASET_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("datasetReference")); + assertTrue(selector.contains("access")); + assertTrue(selector.contains("etag")); + assertEquals(28, selector.length()); + assertEquals(updatedDatasetInfoWithProject, dataset); + } + + @Test + public void testCreateTable() { + EasyMock.expect(bigqueryRpcMock.create(TABLE_INFO_WITH_PROJECT.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.create(TABLE_INFO); + assertEquals(TABLE_INFO_WITH_PROJECT, table); + } + + @Test + public void testCreateTableWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + EasyMock.expect( + bigqueryRpcMock.create(eq(TABLE_INFO_WITH_PROJECT.toPb()), capture(capturedOptions))) + .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.create(TABLE_INFO, TABLE_OPTION_FIELDS); + String selector = (String) capturedOptions.getValue().get(TABLE_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("tableReference")); + assertTrue(selector.contains("schema")); + assertTrue(selector.contains("etag")); + assertEquals(31, selector.length()); + assertEquals(TABLE_INFO_WITH_PROJECT, table); + } + + @Test + public void testGetTable() { + EasyMock.expect(bigqueryRpcMock.getTable(DATASET, TABLE, EMPTY_RPC_OPTIONS)) + .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.getTable(DATASET, TABLE); + assertEquals(TABLE_INFO_WITH_PROJECT, table); + } + + @Test + public void testGetTableFromTableId() { + EasyMock.expect(bigqueryRpcMock.getTable(DATASET, TABLE, EMPTY_RPC_OPTIONS)) + .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.getTable(TABLE_ID); + assertEquals(TABLE_INFO_WITH_PROJECT, table); + } + + @Test + public void testGetTableWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + EasyMock.expect(bigqueryRpcMock.getTable(eq(DATASET), eq(TABLE), capture(capturedOptions))) + .andReturn(TABLE_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.getTable(TABLE_ID, TABLE_OPTION_FIELDS); + String selector = (String) capturedOptions.getValue().get(TABLE_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("tableReference")); + assertTrue(selector.contains("schema")); + assertTrue(selector.contains("etag")); + assertEquals(31, selector.length()); + assertEquals(TABLE_INFO_WITH_PROJECT, table); + } + + @Test + public void testListTables() { + String cursor = "cursor"; + ImmutableList tableList = ImmutableList.of(TABLE_INFO_WITH_PROJECT, + OTHER_TABLE_INFO); + Tuple> result = + Tuple.of(cursor, Iterables.transform(tableList, BaseTableInfo.TO_PB_FUNCTION)); + EasyMock.expect(bigqueryRpcMock.listTables(DATASET, EMPTY_RPC_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listTables(DATASET); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(tableList.toArray(), Iterables.toArray(page.values(), BaseTableInfo.class)); + } + + @Test + public void testListTablesFromDatasetId() { + String cursor = "cursor"; + ImmutableList tableList = ImmutableList.of(TABLE_INFO_WITH_PROJECT, + OTHER_TABLE_INFO); + Tuple> result = + Tuple.of(cursor, Iterables.transform(tableList, BaseTableInfo.TO_PB_FUNCTION)); + EasyMock.expect(bigqueryRpcMock.listTables(DATASET, EMPTY_RPC_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listTables(DatasetId.of(PROJECT, DATASET)); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(tableList.toArray(), Iterables.toArray(page.values(), BaseTableInfo.class)); + } + + @Test + public void testListTablesWithOptions() { + String cursor = "cursor"; + ImmutableList tableList = ImmutableList.of(TABLE_INFO_WITH_PROJECT, + OTHER_TABLE_INFO); + Tuple> result = + Tuple.of(cursor, Iterables.transform(tableList, BaseTableInfo.TO_PB_FUNCTION)); + EasyMock.expect(bigqueryRpcMock.listTables(DATASET, TABLE_LIST_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listTables(DATASET, TABLE_LIST_MAX_RESULTS, + TABLE_LIST_PAGE_TOKEN); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(tableList.toArray(), Iterables.toArray(page.values(), BaseTableInfo.class)); + } + + @Test + public void testDeleteTable() { + EasyMock.expect(bigqueryRpcMock.deleteTable(DATASET, TABLE)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.delete(DATASET, TABLE)); + } + + @Test + public void testDeleteTableFromTableId() { + EasyMock.expect(bigqueryRpcMock.deleteTable(DATASET, TABLE)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.delete(TABLE_ID)); + } + + @Test + public void testUpdateTable() { + BaseTableInfo updatedTableInfo = TABLE_INFO.toBuilder().description("newDescription").build(); + BaseTableInfo updatedTableInfoWithProject = TABLE_INFO_WITH_PROJECT.toBuilder() + .description("newDescription") + .build(); + EasyMock.expect(bigqueryRpcMock.patch(updatedTableInfoWithProject.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(updatedTableInfoWithProject.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.update(updatedTableInfo); + assertEquals(updatedTableInfoWithProject, table); + } + + @Test + public void testUpdateTableWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + BaseTableInfo updatedTableInfo = TABLE_INFO.toBuilder().description("newDescription").build(); + BaseTableInfo updatedTableInfoWithProject = TABLE_INFO_WITH_PROJECT.toBuilder() + .description("newDescription") + .build(); + EasyMock.expect(bigqueryRpcMock.patch(eq(updatedTableInfoWithProject.toPb()), + capture(capturedOptions))).andReturn(updatedTableInfoWithProject.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + BaseTableInfo table = bigquery.update(updatedTableInfo, TABLE_OPTION_FIELDS); + String selector = (String) capturedOptions.getValue().get(TABLE_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("tableReference")); + assertTrue(selector.contains("schema")); + assertTrue(selector.contains("etag")); + assertEquals(31, selector.length()); + assertEquals(updatedTableInfoWithProject, table); + } + + @Test + public void testInsertAll() { + Map row1 = ImmutableMap.of("field", "value1"); + Map row2 = ImmutableMap.of("field", "value2"); + List rows = ImmutableList.of( + new RowToInsert("row1", row1), + new RowToInsert("row2", row2) + ); + InsertAllRequest request = InsertAllRequest.builder(TABLE_ID) + .rows(rows) + .skipInvalidRows(false) + .ignoreUnknownValues(true) + .build(); + TableDataInsertAllRequest requestPb = new TableDataInsertAllRequest().setRows( + Lists.transform(rows, new Function() { + @Override + public TableDataInsertAllRequest.Rows apply(RowToInsert rowToInsert) { + return new TableDataInsertAllRequest.Rows().setInsertId(rowToInsert.id()) + .setJson(rowToInsert.content()); + } + }) + ).setSkipInvalidRows(false).setIgnoreUnknownValues(true); + TableDataInsertAllResponse responsePb = new TableDataInsertAllResponse().setInsertErrors( + ImmutableList.of(new TableDataInsertAllResponse.InsertErrors().setIndex(0L).setErrors( + ImmutableList.of(new ErrorProto().setMessage("ErrorMessage"))))); + EasyMock.expect(bigqueryRpcMock.insertAll(DATASET, TABLE, requestPb)) + .andReturn(responsePb); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + InsertAllResponse response = bigquery.insertAll(request); + assertNotNull(response.errorsFor(0L)); + assertNull(response.errorsFor(1L)); + assertEquals(1, response.errorsFor(0L).size()); + assertEquals("ErrorMessage", response.errorsFor(0L).get(0).message()); + } + + @Test + public void testListTableData() { + String cursor = "cursor"; + com.google.api.services.bigquery.model.TableCell cell1 = + new com.google.api.services.bigquery.model.TableCell().setV("Value1"); + com.google.api.services.bigquery.model.TableCell cell2 = + new com.google.api.services.bigquery.model.TableCell().setV("Value2"); + ImmutableList> tableData = ImmutableList.of( + (List) ImmutableList.of(FieldValue.fromPb(cell1)), + ImmutableList.of(FieldValue.fromPb(cell2))); + Tuple> result = + Tuple.>of(cursor, + ImmutableList.of( + new com.google.api.services.bigquery.model.TableRow().setF( + ImmutableList.of(new com.google.api.services.bigquery.model.TableCell() + .setV("Value1"))), + new com.google.api.services.bigquery.model.TableRow().setF( + ImmutableList.of(new com.google.api.services.bigquery.model.TableCell() + .setV("Value2"))))); + EasyMock.expect(bigqueryRpcMock.listTableData(DATASET, TABLE, EMPTY_RPC_OPTIONS)) + .andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page> page = bigquery.listTableData(DATASET, TABLE); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(tableData.toArray(), Iterables.toArray(page.values(), List.class)); + } + + @Test + public void testListTableDataFromTableId() { + String cursor = "cursor"; + com.google.api.services.bigquery.model.TableCell cell1 = + new com.google.api.services.bigquery.model.TableCell().setV("Value1"); + com.google.api.services.bigquery.model.TableCell cell2 = + new com.google.api.services.bigquery.model.TableCell().setV("Value2"); + ImmutableList> tableData = ImmutableList.of( + (List) ImmutableList.of(FieldValue.fromPb(cell1)), + ImmutableList.of(FieldValue.fromPb(cell2))); + Tuple> result = + Tuple.>of(cursor, + ImmutableList.of( + new com.google.api.services.bigquery.model.TableRow().setF( + ImmutableList.of(new com.google.api.services.bigquery.model.TableCell() + .setV("Value1"))), + new com.google.api.services.bigquery.model.TableRow().setF( + ImmutableList.of(new com.google.api.services.bigquery.model.TableCell() + .setV("Value2"))))); + EasyMock.expect(bigqueryRpcMock.listTableData(DATASET, TABLE, EMPTY_RPC_OPTIONS)) + .andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page> page = bigquery.listTableData(TableId.of(DATASET, TABLE)); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(tableData.toArray(), Iterables.toArray(page.values(), List.class)); + } + + @Test + public void testListTableDataWithOptions() { + String cursor = "cursor"; + com.google.api.services.bigquery.model.TableCell cell1 = + new com.google.api.services.bigquery.model.TableCell().setV("Value1"); + com.google.api.services.bigquery.model.TableCell cell2 = + new com.google.api.services.bigquery.model.TableCell().setV("Value2"); + ImmutableList> tableData = ImmutableList.of( + (List) ImmutableList.of(FieldValue.fromPb(cell1)), + ImmutableList.of(FieldValue.fromPb(cell2))); + Tuple> result = + Tuple.>of(cursor, + ImmutableList.of( + new com.google.api.services.bigquery.model.TableRow().setF( + ImmutableList.of(new com.google.api.services.bigquery.model.TableCell() + .setV("Value1"))), + new com.google.api.services.bigquery.model.TableRow().setF( + ImmutableList.of(new com.google.api.services.bigquery.model.TableCell() + .setV("Value2"))))); + EasyMock.expect(bigqueryRpcMock.listTableData(DATASET, TABLE, TABLE_DATA_LIST_OPTIONS)) + .andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page> page = bigquery.listTableData(DATASET, TABLE, + TABLE_DATA_LIST_MAX_RESULTS, TABLE_DATA_LIST_PAGE_TOKEN, TABLE_DATA_LIST_START_INDEX); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(tableData.toArray(), Iterables.toArray(page.values(), List.class)); + } + + @Test + public void testCreateQueryJob() { + EasyMock.expect(bigqueryRpcMock.create(QUERY_JOB_WITH_PROJECT.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_QUERY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.create(QUERY_JOB); + assertEquals(COMPLETE_QUERY_JOB, job); + } + + @Test + public void testCreateLoadJob() { + EasyMock.expect(bigqueryRpcMock.create(LOAD_JOB_WITH_PROJECT.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_LOAD_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.create(LOAD_JOB); + assertEquals(COMPLETE_LOAD_JOB, job); + } + + @Test + public void testCreateCopyJob() { + EasyMock.expect(bigqueryRpcMock.create(COPY_JOB_WITH_PROJECT.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_COPY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.create(COPY_JOB); + assertEquals(COMPLETE_COPY_JOB, job); + } + + @Test + public void testCreateExtractJob() { + EasyMock.expect(bigqueryRpcMock.create(EXTRACT_JOB_WITH_PROJECT.toPb(), EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_EXTRACT_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.create(EXTRACT_JOB); + assertEquals(COMPLETE_EXTRACT_JOB, job); + } + + @Test + public void testCreateJobWithSelectedFields() { + Capture> capturedOptions = Capture.newInstance(); + EasyMock.expect( + bigqueryRpcMock.create(eq(QUERY_JOB_WITH_PROJECT.toPb()), capture(capturedOptions))) + .andReturn(COMPLETE_QUERY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.create(QUERY_JOB, JOB_OPTION_FIELDS); + assertEquals(COMPLETE_QUERY_JOB, job); + String selector = (String) capturedOptions.getValue().get(JOB_OPTION_FIELDS.rpcOption()); + assertTrue(selector.contains("jobReference")); + assertTrue(selector.contains("configuration")); + assertTrue(selector.contains("user_email")); + assertEquals(37, selector.length()); + } + + @Test + public void testGetJob() { + EasyMock.expect(bigqueryRpcMock.getJob(JOB, EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_COPY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.getJob(JOB); + assertEquals(COMPLETE_COPY_JOB, job); + } + + @Test + public void testGetJobFromJobId() { + EasyMock.expect(bigqueryRpcMock.getJob(JOB, EMPTY_RPC_OPTIONS)) + .andReturn(COMPLETE_COPY_JOB.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + JobInfo job = bigquery.getJob(JobId.of(PROJECT, JOB)); + assertEquals(COMPLETE_COPY_JOB, job); + } + + @Test + public void testListJobs() { + String cursor = "cursor"; + ImmutableList jobList = ImmutableList.of(QUERY_JOB_WITH_PROJECT, + LOAD_JOB_WITH_PROJECT); + Tuple> result = + Tuple.of(cursor, Iterables.transform(jobList, new Function() { + @Override + public Job apply(JobInfo jobInfo) { + return jobInfo.toPb(); + } + })); + EasyMock.expect(bigqueryRpcMock.listJobs(EMPTY_RPC_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listJobs(); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(jobList.toArray(), Iterables.toArray(page.values(), JobInfo.class)); + } + + @Test + public void testListJobsWithOptions() { + String cursor = "cursor"; + ImmutableList jobList = ImmutableList.of(QUERY_JOB_WITH_PROJECT, + LOAD_JOB_WITH_PROJECT); + Tuple> result = + Tuple.of(cursor, Iterables.transform(jobList, new Function() { + @Override + public Job apply(JobInfo jobInfo) { + return jobInfo.toPb(); + } + })); + EasyMock.expect(bigqueryRpcMock.listJobs(JOB_LIST_OPTIONS)).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listJobs(JOB_LIST_ALL_USERS, JOB_LIST_STATE_FILTER, + JOB_LIST_PAGE_TOKEN, JOB_LIST_MAX_RESULTS); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(jobList.toArray(), Iterables.toArray(page.values(), JobInfo.class)); + } + + @Test + public void testListJobsWithSelectedFields() { + String cursor = "cursor"; + Capture> capturedOptions = Capture.newInstance(); + ImmutableList jobList = ImmutableList.of(QUERY_JOB_WITH_PROJECT, + LOAD_JOB_WITH_PROJECT); + Tuple> result = + Tuple.of(cursor, Iterables.transform(jobList, new Function() { + @Override + public Job apply(JobInfo jobInfo) { + return jobInfo.toPb(); + } + })); + EasyMock.expect(bigqueryRpcMock.listJobs(capture(capturedOptions))).andReturn(result); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + Page page = bigquery.listJobs(JOB_LIST_OPTION_FIELD); + assertEquals(cursor, page.nextPageCursor()); + assertArrayEquals(jobList.toArray(), Iterables.toArray(page.values(), JobInfo.class)); + String selector = (String) capturedOptions.getValue().get(JOB_OPTION_FIELDS.rpcOption()); + System.out.println(selector); + assertTrue(selector.contains("etag,jobs(")); + assertTrue(selector.contains("configuration")); + assertTrue(selector.contains("jobReference")); + assertTrue(selector.contains("statistics")); + assertTrue(selector.contains("state,errorResult),nextPageToken")); + assertEquals(80, selector.length()); + } + + @Test + public void testCancelJob() { + EasyMock.expect(bigqueryRpcMock.cancel(JOB)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.cancel(JOB)); + } + + @Test + public void testCancelJobFromJobId() { + EasyMock.expect(bigqueryRpcMock.cancel(JOB)).andReturn(true); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + assertTrue(bigquery.cancel(JobId.of(PROJECT, JOB))); + } + + @Test + public void testQueryRequest() { + JobId queryJob = JobId.of(PROJECT, JOB); + com.google.api.services.bigquery.model.QueryResponse responsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setJobReference(queryJob.toPb()) + .setJobComplete(false); + EasyMock.expect(bigqueryRpcMock.query(QUERY_REQUEST_WITH_PROJECT.toPb())).andReturn(responsePb); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + QueryResponse response = bigquery.query(QUERY_REQUEST); + assertNull(response.etag()); + assertNull(response.result()); + assertEquals(queryJob, response.jobId()); + assertEquals(false, response.jobComplete()); + assertEquals(ImmutableList.of(), response.executionErrors()); + assertFalse(response.hasErrors()); + assertEquals(null, response.result()); + } + + @Test + public void testQueryRequestCompleted() { + JobId queryJob = JobId.of(PROJECT, JOB); + com.google.api.services.bigquery.model.QueryResponse responsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setJobReference(queryJob.toPb()) + .setRows(ImmutableList.of(TABLE_ROW)) + .setJobComplete(true) + .setCacheHit(false) + .setPageToken("cursor") + .setTotalBytesProcessed(42L) + .setTotalRows(BigInteger.valueOf(1L)); + EasyMock.expect(bigqueryRpcMock.query(QUERY_REQUEST_WITH_PROJECT.toPb())).andReturn(responsePb); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + QueryResponse response = bigquery.query(QUERY_REQUEST); + assertNull(response.etag()); + assertEquals(queryJob, response.jobId()); + assertEquals(true, response.jobComplete()); + assertEquals(false, response.result().cacheHit()); + assertEquals(ImmutableList.of(), response.executionErrors()); + assertFalse(response.hasErrors()); + assertEquals(null, response.result().schema()); + assertEquals(42L, response.result().totalBytesProcessed()); + assertEquals(1L, response.result().totalRows()); + for (List row : response.result().values()) { + assertEquals(false, row.get(0).booleanValue()); + assertEquals(1L, row.get(1).longValue()); + } + assertEquals("cursor", response.result().nextPageCursor()); + } + + @Test + public void testGetQueryResults() { + JobId queryJob = JobId.of(PROJECT, JOB); + com.google.api.services.bigquery.model.GetQueryResultsResponse responsePb = + new com.google.api.services.bigquery.model.GetQueryResultsResponse() + .setEtag("etag") + .setJobReference(queryJob.toPb()) + .setRows(ImmutableList.of(TABLE_ROW)) + .setJobComplete(true) + .setCacheHit(false) + .setPageToken("cursor") + .setTotalBytesProcessed(42L) + .setTotalRows(BigInteger.valueOf(1L)); + EasyMock.expect(bigqueryRpcMock.getQueryResults(JOB, EMPTY_RPC_OPTIONS)).andReturn(responsePb); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + QueryResponse response = bigquery.getQueryResults(queryJob); + assertEquals("etag", response.etag()); + assertEquals(queryJob, response.jobId()); + assertEquals(true, response.jobComplete()); + assertEquals(false, response.result().cacheHit()); + assertEquals(ImmutableList.of(), response.executionErrors()); + assertFalse(response.hasErrors()); + assertEquals(null, response.result().schema()); + assertEquals(42L, response.result().totalBytesProcessed()); + assertEquals(1L, response.result().totalRows()); + for (List row : response.result().values()) { + assertEquals(false, row.get(0).booleanValue()); + assertEquals(1L, row.get(1).longValue()); + } + assertEquals("cursor", response.result().nextPageCursor()); + } + + @Test + public void testGetQueryResultsWithOptions() { + JobId queryJob = JobId.of(PROJECT, JOB); + com.google.api.services.bigquery.model.GetQueryResultsResponse responsePb = + new com.google.api.services.bigquery.model.GetQueryResultsResponse() + .setJobReference(queryJob.toPb()) + .setRows(ImmutableList.of(TABLE_ROW)) + .setJobComplete(true) + .setCacheHit(false) + .setPageToken("cursor") + .setTotalBytesProcessed(42L) + .setTotalRows(BigInteger.valueOf(1L)); + EasyMock.expect(bigqueryRpcMock.getQueryResults(JOB, QUERY_RESULTS_OPTIONS)) + .andReturn(responsePb); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.service(); + QueryResponse response = bigquery.getQueryResults(queryJob, QUERY_RESULTS_OPTION_TIME, + QUERY_RESULTS_OPTION_INDEX, QUERY_RESULTS_OPTION_MAX_RESULTS, + QUERY_RESULTS_OPTION_PAGE_TOKEN); + assertEquals(queryJob, response.jobId()); + assertEquals(true, response.jobComplete()); + assertEquals(false, response.result().cacheHit()); + assertEquals(ImmutableList.of(), response.executionErrors()); + assertFalse(response.hasErrors()); + assertEquals(null, response.result().schema()); + assertEquals(42L, response.result().totalBytesProcessed()); + assertEquals(1L, response.result().totalRows()); + for (List row : response.result().values()) { + assertEquals(false, row.get(0).booleanValue()); + assertEquals(1L, row.get(1).longValue()); + } + assertEquals("cursor", response.result().nextPageCursor()); + } + + @Test + public void testRetryableException() { + EasyMock.expect(bigqueryRpcMock.getDataset(DATASET, EMPTY_RPC_OPTIONS)) + .andThrow(new BigQueryException(500, "InternalError", true)) + .andReturn(DATASET_INFO_WITH_PROJECT.toPb()); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.toBuilder().retryParams(RetryParams.getDefaultInstance()).build().service(); + DatasetInfo dataset = bigquery.getDataset(DATASET); + assertEquals(DATASET_INFO_WITH_PROJECT, dataset); + } + + @Test + public void testNonRetryableException() { + String exceptionMessage = "Not Implemented"; + EasyMock.expect(bigqueryRpcMock.getDataset(DATASET, EMPTY_RPC_OPTIONS)) + .andThrow(new BigQueryException(501, exceptionMessage, false)); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.toBuilder().retryParams(RetryParams.getDefaultInstance()).build().service(); + thrown.expect(BigQueryException.class); + thrown.expectMessage(exceptionMessage); + bigquery.getDataset(DatasetId.of(DATASET)); + } + + @Test + public void testRuntimeException() { + String exceptionMessage = "Artificial runtime exception"; + EasyMock.expect(bigqueryRpcMock.getDataset(DATASET, EMPTY_RPC_OPTIONS)) + .andThrow(new RuntimeException(exceptionMessage)); + EasyMock.replay(bigqueryRpcMock); + bigquery = options.toBuilder().retryParams(RetryParams.getDefaultInstance()).build().service(); + thrown.expect(BigQueryException.class); + thrown.expectMessage(exceptionMessage); + bigquery.getDataset(DATASET); + } +} diff --git a/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/OptionTest.java b/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/OptionTest.java new file mode 100644 index 000000000000..d337f5a67849 --- /dev/null +++ b/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/OptionTest.java @@ -0,0 +1,38 @@ +/* + * Copyright 2015 Google Inc. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.gcloud.bigquery; + +import static org.junit.Assert.assertEquals; + +import com.google.gcloud.spi.BigQueryRpc; + +import org.junit.Test; + +public class OptionTest { + + @Test + public void testOption() { + Option option = new Option(BigQueryRpc.Option.PAGE_TOKEN, "token"); + assertEquals(BigQueryRpc.Option.PAGE_TOKEN, option.rpcOption()); + assertEquals("token", option.value()); + } + + @Test(expected=NullPointerException.class) + public void testNullRpcOption() { + new Option(null, "token"); + } +} diff --git a/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/SerializationTest.java b/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/SerializationTest.java index 3576769f4007..3766ef493064 100644 --- a/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/SerializationTest.java +++ b/gcloud-java-bigquery/src/test/java/com/google/gcloud/bigquery/SerializationTest.java @@ -233,7 +233,11 @@ public void testModelAndRequests() throws Exception { TABLE_SCHEMA, TABLE_INFO, VIEW_INFO, EXTERNAL_TABLE_INFO, INLINE_FUNCTION, URI_FUNCTION, JOB_STATISTICS, EXTRACT_STATISTICS, LOAD_STATISTICS, QUERY_STATISTICS, BIGQUERY_ERROR, JOB_STATUS, JOB_ID, COPY_JOB, EXTRACT_JOB, LOAD_JOB, QUERY_JOB, INSERT_ALL_REQUEST, - INSERT_ALL_RESPONSE, FIELD_VALUE, QUERY_REQUEST, QUERY_RESPONSE}; + INSERT_ALL_RESPONSE, FIELD_VALUE, QUERY_REQUEST, QUERY_RESPONSE, + BigQuery.DatasetOption.fields(), BigQuery.DatasetDeleteOption.deleteContents(), + BigQuery.DatasetListOption.all(), BigQuery.TableOption.fields(), + BigQuery.TableListOption.maxResults(42L), BigQuery.JobOption.fields(), + BigQuery.JobListOption.allUsers()}; for (Serializable obj : objects) { Object copy = serializeAndDeserialize(obj); assertEquals(obj, obj);