Skip to content

Commit

Permalink
Feat: Add universe domain (#3090)
Browse files Browse the repository at this point in the history 
* Universe domain feat

* Feat: Add universe domain feat

* Feat: Add universe domain feat

* Feat: Add universe domain feat
  • Loading branch information
@PhongChuong
PhongChuong committed Jan 23, 2024
1 parent eb773b9 commit b2814a2
Show file tree
Hide file tree
Showing 2 changed files with 177 additions and 2 deletions.
45 changes: 43 additions & 2 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/spi/v2/HttpBigQueryRpc.java
View file
Expand Up @@ -105,7 +105,7 @@ public HttpBigQueryRpc(BigQueryOptions options) {
this.options = options;
bigquery =
new Bigquery.Builder(transport, new GsonFactory(), initializer)
.setRootUrl(options.getHost())
.setRootUrl(options.getResolvedApiaryHost("bigquery"))
.setApplicationName(options.getApplicationName())
.build();
}
Expand All @@ -114,9 +114,16 @@ private static BigQueryException translate(IOException exception) {
return new BigQueryException(exception);
}

private void validateRPC() throws BigQueryException, IOException {
if (!this.options.hasValidUniverseDomain()) {
throw new BigQueryException(BigQueryException.UNKNOWN_CODE, "Invalid universe domain");
}
}

@Override
public Dataset getDataset(String projectId, String datasetId, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.datasets()
.get(projectId, datasetId)
Expand All @@ -135,6 +142,7 @@ public Dataset getDataset(String projectId, String datasetId, Map<Option, ?> opt
@Override
public Tuple<String, Iterable<Dataset>> listDatasets(String projectId, Map<Option, ?> options) {
try {
validateRPC();
DatasetList datasetsList =
bigquery
.datasets()
Expand All @@ -159,6 +167,7 @@ public Tuple<String, Iterable<Dataset>> listDatasets(String projectId, Map<Optio
@Override
public Dataset create(Dataset dataset, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.datasets()
.insert(dataset.getDatasetReference().getProjectId(), dataset)
Expand All @@ -173,6 +182,7 @@ public Dataset create(Dataset dataset, Map<Option, ?> options) {
@Override
public Table create(Table table, Map<Option, ?> options) {
try {
validateRPC();
// unset the type, as it is output only
table.setType(null);
TableReference reference = table.getTableReference();
Expand All @@ -190,6 +200,7 @@ public Table create(Table table, Map<Option, ?> options) {
@Override
public Routine create(Routine routine, Map<Option, ?> options) {
try {
validateRPC();
RoutineReference reference = routine.getRoutineReference();
return bigquery
.routines()
Expand All @@ -205,6 +216,7 @@ public Routine create(Routine routine, Map<Option, ?> options) {
@Override
public Job create(Job job, Map<Option, ?> options) {
try {
validateRPC();
String projectId =
job.getJobReference() != null
? job.getJobReference().getProjectId()
Expand All @@ -223,6 +235,7 @@ public Job create(Job job, Map<Option, ?> options) {
@Override
public Job createJobForQuery(Job job) {
try {
validateRPC();
String projectId =
job.getJobReference() != null
? job.getJobReference().getProjectId()
Expand All @@ -236,6 +249,7 @@ public Job createJobForQuery(Job job) {
@Override
public boolean deleteDataset(String projectId, String datasetId, Map<Option, ?> options) {
try {
validateRPC();
bigquery
.datasets()
.delete(projectId, datasetId)
Expand All @@ -255,6 +269,7 @@ public boolean deleteDataset(String projectId, String datasetId, Map<Option, ?>
@Override
public Dataset patch(Dataset dataset, Map<Option, ?> options) {
try {
validateRPC();
DatasetReference reference = dataset.getDatasetReference();
return bigquery
.datasets()
Expand All @@ -270,6 +285,7 @@ public Dataset patch(Dataset dataset, Map<Option, ?> options) {
@Override
public Table patch(Table table, Map<Option, ?> options) {
try {
validateRPC();
// unset the type, as it is output only
table.setType(null);
TableReference reference = table.getTableReference();
Expand All @@ -289,6 +305,7 @@ public Table patch(Table table, Map<Option, ?> options) {
public Table getTable(
String projectId, String datasetId, String tableId, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.tables()
.get(projectId, datasetId, tableId)
Expand Down Expand Up @@ -316,6 +333,7 @@ private String getTableMetadataOption(Map<Option, ?> options) {
public Tuple<String, Iterable<Table>> listTables(
String projectId, String datasetId, Map<Option, ?> options) {
try {
validateRPC();
TableList tableList =
bigquery
.tables()
Expand Down Expand Up @@ -351,6 +369,7 @@ public Table apply(TableList.Tables tablePb) {
@Override
public boolean deleteTable(String projectId, String datasetId, String tableId) {
try {
validateRPC();
bigquery.tables().delete(projectId, datasetId, tableId).execute();
return true;
} catch (IOException ex) {
Expand All @@ -365,6 +384,7 @@ public boolean deleteTable(String projectId, String datasetId, String tableId) {
@Override
public Model patch(Model model, Map<Option, ?> options) {
try {
validateRPC();
// unset the type, as it is output only
ModelReference reference = model.getModelReference();
return bigquery
Expand All @@ -382,6 +402,7 @@ public Model patch(Model model, Map<Option, ?> options) {
public Model getModel(
String projectId, String datasetId, String modelId, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.models()
.get(projectId, datasetId, modelId)
Expand All @@ -401,6 +422,7 @@ public Model getModel(
public Tuple<String, Iterable<Model>> listModels(
String projectId, String datasetId, Map<Option, ?> options) {
try {
validateRPC();
ListModelsResponse modelList =
bigquery
.models()
Expand All @@ -420,6 +442,7 @@ public Tuple<String, Iterable<Model>> listModels(
@Override
public boolean deleteModel(String projectId, String datasetId, String modelId) {
try {
validateRPC();
bigquery.models().delete(projectId, datasetId, modelId).execute();
return true;
} catch (IOException ex) {
Expand All @@ -434,6 +457,7 @@ public boolean deleteModel(String projectId, String datasetId, String modelId) {
@Override
public Routine update(Routine routine, Map<Option, ?> options) {
try {
validateRPC();
RoutineReference reference = routine.getRoutineReference();
return bigquery
.routines()
Expand All @@ -451,6 +475,7 @@ public Routine update(Routine routine, Map<Option, ?> options) {
public Routine getRoutine(
String projectId, String datasetId, String routineId, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.routines()
.get(projectId, datasetId, routineId)
Expand All @@ -470,6 +495,7 @@ public Routine getRoutine(
public Tuple<String, Iterable<Routine>> listRoutines(
String projectId, String datasetId, Map<Option, ?> options) {
try {
validateRPC();
ListRoutinesResponse routineList =
bigquery
.routines()
Expand All @@ -491,6 +517,7 @@ public Tuple<String, Iterable<Routine>> listRoutines(
@Override
public boolean deleteRoutine(String projectId, String datasetId, String routineId) {
try {
validateRPC();
bigquery.routines().delete(projectId, datasetId, routineId).execute();
return true;
} catch (IOException ex) {
Expand All @@ -506,6 +533,7 @@ public boolean deleteRoutine(String projectId, String datasetId, String routineI
public TableDataInsertAllResponse insertAll(
String projectId, String datasetId, String tableId, TableDataInsertAllRequest request) {
try {
validateRPC();
return bigquery
.tabledata()
.insertAll(projectId, datasetId, tableId, request)
Expand All @@ -520,6 +548,7 @@ public TableDataInsertAllResponse insertAll(
public TableDataList listTableData(
String projectId, String datasetId, String tableId, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.tabledata()
.list(projectId, datasetId, tableId)
Expand All @@ -544,6 +573,7 @@ public TableDataList listTableDataWithRowLimit(
Integer maxResultPerPage,
String pageToken) {
try {
validateRPC();
return bigquery
.tabledata()
.list(projectId, datasetId, tableId)
Expand All @@ -559,6 +589,7 @@ public TableDataList listTableDataWithRowLimit(
@Override
public Job getJob(String projectId, String jobId, String location, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.jobs()
.get(projectId, jobId)
Expand All @@ -578,6 +609,7 @@ public Job getJob(String projectId, String jobId, String location, Map<Option, ?
@Override
public Job getQueryJob(String projectId, String jobId, String location) {
try {
validateRPC();
return bigquery
.jobs()
.get(projectId, jobId)
Expand All @@ -596,6 +628,7 @@ public Job getQueryJob(String projectId, String jobId, String location) {
@Override
public Tuple<String, Iterable<Job>> listJobs(String projectId, Map<Option, ?> options) {
try {
validateRPC();
Bigquery.Jobs.List request =
bigquery
.jobs()
Expand Down Expand Up @@ -650,6 +683,7 @@ public Job apply(JobList.Jobs jobPb) {
@Override
public boolean cancel(String projectId, String jobId, String location) {
try {
validateRPC();
bigquery
.jobs()
.cancel(projectId, jobId)
Expand All @@ -669,6 +703,7 @@ public boolean cancel(String projectId, String jobId, String location) {
@Override
public boolean deleteJob(String projectId, String jobName, String location) {
try {
validateRPC();
bigquery
.jobs()
.delete(projectId, jobName)
Expand All @@ -685,6 +720,7 @@ public boolean deleteJob(String projectId, String jobName, String location) {
public GetQueryResultsResponse getQueryResults(
String projectId, String jobId, String location, Map<Option, ?> options) {
try {
validateRPC();
return bigquery
.jobs()
.getQueryResults(projectId, jobId)
Expand All @@ -707,6 +743,7 @@ public GetQueryResultsResponse getQueryResults(
public GetQueryResultsResponse getQueryResultsWithRowLimit(
String projectId, String jobId, String location, Integer maxResultPerPage, Long timeoutMs) {
try {
validateRPC();
return bigquery
.jobs()
.getQueryResults(projectId, jobId)
Expand All @@ -723,6 +760,7 @@ public GetQueryResultsResponse getQueryResultsWithRowLimit(
@Override
public QueryResponse queryRpc(String projectId, QueryRequest content) {
try {
validateRPC();
return bigquery.jobs().query(projectId, content).execute();
} catch (IOException ex) {
throw translate(ex);
Expand All @@ -732,7 +770,7 @@ public QueryResponse queryRpc(String projectId, QueryRequest content) {
@Override
public String open(Job loadJob) {
try {
String builder = options.getHost();
String builder = options.getResolvedApiaryHost("bigquery");
if (!builder.endsWith("/")) {
builder += "/";
}
Expand Down Expand Up @@ -807,6 +845,7 @@ public Job write(
@Override
public Policy getIamPolicy(String resourceId, Map<Option, ?> options) {
try {
validateRPC();
GetIamPolicyRequest policyRequest = new GetIamPolicyRequest();
if (null != Option.REQUESTED_POLICY_VERSION.getLong(options)) {
policyRequest =
Expand All @@ -828,6 +867,7 @@ public Policy getIamPolicy(String resourceId, Map<Option, ?> options) {
@Override
public Policy setIamPolicy(String resourceId, Policy policy, Map<Option, ?> options) {
try {
validateRPC();
SetIamPolicyRequest policyRequest = new SetIamPolicyRequest().setPolicy(policy);
return bigquery
.tables()
Expand All @@ -843,6 +883,7 @@ public Policy setIamPolicy(String resourceId, Policy policy, Map<Option, ?> opti
public TestIamPermissionsResponse testIamPermissions(
String resourceId, List<String> permissions, Map<Option, ?> options) {
try {
validateRPC();
TestIamPermissionsRequest permissionsRequest =
new TestIamPermissionsRequest().setPermissions(permissions);
return bigquery
Expand Down

0 comments on commit b2814a2

Please sign in to comment.