Skip to content

Commit 0b874dd

Browse files
pratickchokhanipratickchokhani
committed
chore: Fixed constructor argument.
1 parent 3bc834f commit 0b874dd

File tree

7 files changed

+250
-236
lines changed

7 files changed

+250
-236
lines changed

‎v1/src/main/java/com/google/cloud/teleport/spanner/ExportPipeline.java‎

Lines changed: 35 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -26,20 +26,21 @@
2626
import com.google.cloud.teleport.spanner.iam.IAMCheckResult;
2727
import com.google.cloud.teleport.spanner.iam.IAMPermissionsChecker;
2828
import com.google.cloud.teleport.spanner.iam.IAMRequirementsCreator;
29-
import com.google.cloud.teleport.spanner.iam.IAMResourceRequirements;
3029
import com.google.cloud.teleport.spanner.spannerio.SpannerConfig;
31-
import java.util.Collections;
30+
import java.io.IOException;
31+
import java.security.GeneralSecurityException;
3232
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
3333
import org.apache.beam.sdk.Pipeline;
3434
import org.apache.beam.sdk.PipelineResult;
35-
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
3635
import org.apache.beam.sdk.options.Default;
3736
import org.apache.beam.sdk.options.Description;
3837
import org.apache.beam.sdk.options.PipelineOptions;
3938
import org.apache.beam.sdk.options.PipelineOptionsFactory;
4039
import org.apache.beam.sdk.options.ValueProvider;
4140
import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
4241
import org.apache.beam.sdk.transforms.SerializableFunction;
42+
import org.slf4j.Logger;
43+
import org.slf4j.LoggerFactory;
4344

4445
/**
4546
* Dataflow template that exports a Cloud Spanner database to Avro files in GCS.
@@ -78,6 +79,7 @@
7879
"In addition to the Identity and Access Management (IAM) roles necessary to run Dataflow jobs, you must also have the <a href=\"https://cloud.google.com/spanner/docs/export#iam\">appropriate IAM roles</a> for reading your Cloud Spanner data and writing to your Cloud Storage bucket."
7980
})
8081
public class ExportPipeline {
82+
private static final Logger LOG = LoggerFactory.getLogger(ExportPipeline.class);
8183

8284
/** Options for Export pipeline. */
8385
public interface ExportPipelineOptions extends PipelineOptions {
@@ -282,7 +284,7 @@ public static void main(String[] args) {
282284
options.getShouldExportRelatedTables(),
283285
options.getShouldExportTimestampAsLogicalType(),
284286
options.getAvroTempDirectory()));
285-
validateRequiredPermissions(options);
287+
validateRequiredPermissions(spannerConfig);
286288
PipelineResult result = p.run();
287289
if (options.getWaitUntilFinish()
288290
&&
@@ -294,25 +296,36 @@ public static void main(String[] args) {
294296
}
295297
}
296298

297-
private static void validateRequiredPermissions(ExportPipelineOptions options) {
298-
IAMResourceRequirements spannerRequirements =
299-
IAMRequirementsCreator.createSpannerResourceRequirement();
300-
301-
GcpOptions gcpOptions = options.as(GcpOptions.class);
302-
303-
IAMPermissionsChecker iamPermissionsChecker =
304-
new IAMPermissionsChecker(gcpOptions.getProject(), gcpOptions);
305-
IAMCheckResult missingPermission =
306-
iamPermissionsChecker.check(Collections.singletonList(spannerRequirements));
307-
if (missingPermission.isSuccess()) {
299+
private static void validateRequiredPermissions(SpannerConfig spannerConfig) {
300+
if (!spannerConfig.getInstanceId().isAccessible()) {
301+
// If instance id is not accessible, then the current context is not runtime. Hence,
302+
// validation is not applicable.
308303
return;
309304
}
310-
String errorString =
311-
"For resource: "
312-
+ missingPermission.getResourceName()
313-
+ ", missing permissions: "
314-
+ missingPermission.getMissingPermissions()
315-
+ ";";
316-
throw new RuntimeException(errorString);
305+
String instanceId = spannerConfig.getInstanceId().get();
306+
String databaseId = spannerConfig.getDatabaseId().get();
307+
IAMCheckResult iamCheckResult;
308+
try {
309+
IAMPermissionsChecker iamPermissionsChecker =
310+
new IAMPermissionsChecker(spannerConfig.getProjectId().get());
311+
iamCheckResult =
312+
iamPermissionsChecker.checkSpannerDatabaseRequirements(
313+
IAMRequirementsCreator.createSpannerWriteResourceRequirement(),
314+
instanceId,
315+
databaseId);
316+
if (iamCheckResult.isPermissionsAvailable()) {
317+
return;
318+
}
319+
String errorString =
320+
"For resource: "
321+
+ iamCheckResult.getResourceName()
322+
+ ", missing permissions: "
323+
+ iamCheckResult.getMissingPermissions()
324+
+ ";";
325+
throw new RuntimeException(errorString);
326+
} catch (GeneralSecurityException | IOException e) {
327+
LOG.error("Error while validating permissions for spanner", e);
328+
throw new RuntimeException(e);
329+
}
317330
}
318331
}

‎v1/src/main/java/com/google/cloud/teleport/spanner/ImportPipeline.java‎

Lines changed: 46 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package com.google.cloud.teleport.spanner;
1717

18+
import com.google.cloud.spanner.DatabaseNotFoundException;
1819
import com.google.cloud.spanner.Options.RpcPriority;
1920
import com.google.cloud.spanner.SpannerOptions;
2021
import com.google.cloud.teleport.metadata.Template;
@@ -26,20 +27,21 @@
2627
import com.google.cloud.teleport.spanner.iam.IAMCheckResult;
2728
import com.google.cloud.teleport.spanner.iam.IAMPermissionsChecker;
2829
import com.google.cloud.teleport.spanner.iam.IAMRequirementsCreator;
29-
import com.google.cloud.teleport.spanner.iam.IAMResourceRequirements;
3030
import com.google.cloud.teleport.spanner.spannerio.SpannerConfig;
31-
import java.util.Collections;
31+
import java.io.IOException;
32+
import java.security.GeneralSecurityException;
3233
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
3334
import org.apache.beam.sdk.Pipeline;
3435
import org.apache.beam.sdk.PipelineResult;
35-
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
3636
import org.apache.beam.sdk.options.Default;
3737
import org.apache.beam.sdk.options.Description;
3838
import org.apache.beam.sdk.options.PipelineOptions;
3939
import org.apache.beam.sdk.options.PipelineOptionsFactory;
4040
import org.apache.beam.sdk.options.ValueProvider;
4141
import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
4242
import org.apache.beam.sdk.transforms.SerializableFunction;
43+
import org.slf4j.Logger;
44+
import org.slf4j.LoggerFactory;
4345

4446
/**
4547
* Avro to Cloud Spanner Import pipeline.
@@ -65,6 +67,7 @@
6567
"The input Cloud Storage path must exist, and it must include a <a href=\"https://cloud.google.com/spanner/docs/import-non-spanner#create-export-json\">spanner-export.json</a> file that contains a JSON description of files to import."
6668
})
6769
public class ImportPipeline {
70+
private static final Logger LOG = LoggerFactory.getLogger(ImportPipeline.class);
6871

6972
/** Options for {@link ImportPipeline}. */
7073
public interface Options extends PipelineOptions {
@@ -258,7 +261,8 @@ public static void main(String[] args) {
258261
options.getDdlCreationTimeoutInMinutes(),
259262
options.getEarlyIndexCreateThreshold()));
260263

261-
validateRequiredPermissions(options);
264+
validateRequiredPermissions(spannerConfig);
265+
262266
PipelineResult result = p.run();
263267

264268
if (options.getWaitUntilFinish()
@@ -271,24 +275,45 @@ public static void main(String[] args) {
271275
}
272276
}
273277

274-
private static void validateRequiredPermissions(Options options) {
275-
IAMResourceRequirements spannerRequirements =
276-
IAMRequirementsCreator.createSpannerResourceRequirement();
277-
GcpOptions gcpOptions = options.as(GcpOptions.class);
278-
279-
IAMPermissionsChecker iamPermissionsChecker =
280-
new IAMPermissionsChecker(gcpOptions.getProject(), gcpOptions);
281-
IAMCheckResult missingPermission =
282-
iamPermissionsChecker.check(Collections.singletonList(spannerRequirements));
283-
if (missingPermission.isSuccess()) {
278+
private static void validateRequiredPermissions(SpannerConfig spannerConfig) {
279+
if (!spannerConfig.getInstanceId().isAccessible()) {
280+
// If instance id is not accessible, then the current context is not runtime. Hence,
281+
// validation is not applicable.
284282
return;
285283
}
286-
String errorString =
287-
"For resource: "
288-
+ missingPermission.getResourceName()
289-
+ ", missing permissions: "
290-
+ missingPermission.getMissingPermissions()
291-
+ ";";
292-
throw new RuntimeException(errorString);
284+
String instanceId = spannerConfig.getInstanceId().get();
285+
String databaseId = spannerConfig.getDatabaseId().get();
286+
IAMCheckResult iamCheckResult;
287+
try {
288+
IAMPermissionsChecker iamPermissionsChecker =
289+
new IAMPermissionsChecker(spannerConfig.getProjectId().get());
290+
try {
291+
292+
iamCheckResult =
293+
iamPermissionsChecker.checkSpannerDatabaseRequirements(
294+
IAMRequirementsCreator.createSpannerWriteResourceRequirement(),
295+
instanceId,
296+
databaseId);
297+
} catch (DatabaseNotFoundException e) {
298+
// If DatabaseNotFoundException exception is thrown, then validation at instance level need
299+
// to be performed.
300+
iamCheckResult =
301+
iamPermissionsChecker.checkSpannerInstanceRequirements(
302+
IAMRequirementsCreator.createSpannerWriteResourceRequirement(), instanceId);
303+
}
304+
if (iamCheckResult.isPermissionsAvailable()) {
305+
return;
306+
}
307+
String errorString =
308+
"For resource: "
309+
+ iamCheckResult.getResourceName()
310+
+ ", missing permissions: "
311+
+ iamCheckResult.getMissingPermissions()
312+
+ ";";
313+
throw new RuntimeException(errorString);
314+
} catch (GeneralSecurityException | IOException e) {
315+
LOG.error("Error while validating permissions for spanner", e);
316+
throw new RuntimeException(e);
317+
}
293318
}
294319
}

‎v1/src/main/java/com/google/cloud/teleport/spanner/iam/IAMCheckResult.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ public List<String> getMissingPermissions() {
3636
return new ArrayList<>(missingPermissions);
3737
}
3838

39-
public boolean isSuccess() {
39+
public boolean isPermissionsAvailable() {
4040
return missingPermissions.isEmpty();
4141
}
4242

@@ -49,7 +49,7 @@ public String toString() {
4949
+ ", missingPermissions="
5050
+ missingPermissions
5151
+ ", success="
52-
+ isSuccess()
52+
+ isPermissionsAvailable()
5353
+ '}';
5454
}
5555
}

0 commit comments

Comments
 (0)