Skip to content

Commit 4e945b9

Browse files
scmacdonscmacdon
committed
udpaed code in the scenario
1 parent 4876f77 commit 4e945b9

3 files changed

Lines changed: 206 additions & 105 deletions

File tree

‎javav2/example_code/controltower/src/main/java/com/example/controltower/scenario/ControlTowerActions.java‎

Lines changed: 131 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -338,7 +338,7 @@ public CompletableFuture<List<EnabledBaselineSummary>> listEnabledBaselinesAsync
338338
&& !response.enabledBaselines().isEmpty()) {
339339

340340
response.enabledBaselines().forEach(baseline -> {
341-
System.out.format("Enabled baseline: {}", baseline.baselineIdentifier());
341+
System.out.format("Enabled baseline: {}", baseline.arn());
342342
enabledBaselines.add(baseline);
343343
});
344344
} else {
@@ -401,26 +401,24 @@ public CompletableFuture<String> enableBaselineAsync(
401401
String baselineIdentifier,
402402
String baselineVersion
403403
) {
404-
// Build the enable request
405404
EnableBaselineRequest request = EnableBaselineRequest.builder()
406405
.baselineIdentifier(baselineIdentifier)
407406
.baselineVersion(baselineVersion)
408407
.targetIdentifier(targetIdentifier)
409408
.build();
410409

411410
return getAsyncClient().enableBaseline(request)
412-
.handle((response, exception) -> {
411+
.handle((resp, exception) -> {
413412
if (exception != null) {
414413
Throwable cause = exception.getCause() != null ? exception.getCause() : exception;
415-
416414
if (cause instanceof ControlTowerException e) {
417415
String code = e.awsErrorDetails() != null ? e.awsErrorDetails().errorCode() : "UNKNOWN";
418416
String msg = e.awsErrorDetails() != null ? e.awsErrorDetails().errorMessage() : e.getMessage();
419417

420-
// Already enabled → fetch ARN instead of failing
421418
if ("ValidationException".equals(code) && msg.contains("already enabled")) {
422419
System.out.println("Baseline is already enabled for this target → fetching ARN...");
423-
return fetchEnabledBaselineArn(targetIdentifier, baselineIdentifier);
420+
return fetchEnabledBaselineArn(targetIdentifier, baselineIdentifier)
421+
.join(); // fetch existing ARN synchronously
424422
}
425423

426424
throw new RuntimeException("Error enabling baseline: " + code + " - " + msg, e);
@@ -429,38 +427,54 @@ public CompletableFuture<String> enableBaselineAsync(
429427
throw new RuntimeException("Unexpected error enabling baseline: " + cause.getMessage(), cause);
430428
}
431429

432-
// Success → return ARN
433-
return response.arn();
430+
return resp;
434431
})
435-
// Unwrap: handle CompletableFuture returned by fetchEnabledBaselineArn
436-
.thenCompose(futOrArn -> {
437-
if (futOrArn instanceof CompletableFuture) {
438-
// Already-enabled branch
439-
return ((CompletableFuture<String>) futOrArn)
440-
.thenCompose(existingArn -> {
441-
if (existingArn == null) {
442-
// Retry fetching asynchronously after short delay
443-
System.out.println("Baseline already enabled but ARN not yet available, retrying...");
444-
return CompletableFuture.supplyAsync(() -> null,
445-
CompletableFuture.delayedExecutor(15, TimeUnit.SECONDS))
446-
.thenCompose(v -> fetchEnabledBaselineArn(targetIdentifier, baselineIdentifier));
447-
}
448-
return CompletableFuture.completedFuture(existingArn);
449-
});
450-
}
432+
.thenCompose(result -> {
433+
if (result instanceof EnableBaselineResponse resp) {
434+
String operationId = resp.operationIdentifier();
435+
String enabledBaselineArn = resp.arn();
436+
System.out.println("Baseline enable started. ARN: " + enabledBaselineArn
437+
+ ", operation ID: " + operationId);
438+
439+
// Inline polling
440+
return CompletableFuture.supplyAsync(() -> {
441+
while (true) {
442+
GetBaselineOperationRequest opReq = GetBaselineOperationRequest.builder()
443+
.operationIdentifier(operationId)
444+
.build();
451445

452-
String enabledBaselineArn = (String) futOrArn;
453-
if (enabledBaselineArn == null) {
454-
// Should not happen, return null
455-
return CompletableFuture.completedFuture(null);
446+
GetBaselineOperationResponse opResp = getAsyncClient().getBaselineOperation(opReq).join();
447+
BaselineOperation op = opResp.baselineOperation();
448+
BaselineOperationStatus status = op.status();
449+
System.out.println("Operation " + operationId + " status: " + status);
450+
451+
if (status == BaselineOperationStatus.SUCCEEDED) {
452+
return enabledBaselineArn;
453+
} else if (status == BaselineOperationStatus.FAILED) {
454+
String opId = op.operationIdentifier();
455+
String reason = op.statusMessage() != null ? op.statusMessage() : "No failure reason provided";
456+
throw new RuntimeException("Baseline operation failed (ID: " + opId + "), status: "
457+
+ status + ", reason: " + reason);
458+
}
459+
460+
try {
461+
Thread.sleep(Duration.ofSeconds(15).toMillis());
462+
} catch (InterruptedException e) {
463+
Thread.currentThread().interrupt();
464+
throw new RuntimeException(e);
465+
}
466+
}
467+
});
468+
} else if (result instanceof String existingArn) {
469+
// Already enabled branch
470+
return CompletableFuture.completedFuture(existingArn);
456471
}
457472

458-
// Poll the operation until it completes
459-
return pollBaselineOperationAsync(enabledBaselineArn)
460-
.thenApply(status -> enabledBaselineArn);
473+
return CompletableFuture.completedFuture(null);
461474
});
462475
}
463476

477+
464478
/**
465479
* Fetches the ARN of an already-enabled baseline for the target asynchronously.
466480
*/
@@ -505,60 +519,70 @@ private CompletableFuture<BaselineOperationStatus> pollBaselineOperationAsync(St
505519
*/
506520
public CompletableFuture<String> disableBaselineAsync(String enabledBaselineIdentifier) {
507521

508-
// Build the request
522+
System.out.println("Starting disable of enabled baseline…");
523+
System.out.println("This operation will check the status every 15 seconds until it completes (SUCCEEDED or FAILED).");
524+
509525
DisableBaselineRequest request = DisableBaselineRequest.builder()
510526
.enabledBaselineIdentifier(enabledBaselineIdentifier)
511527
.build();
512528

513-
// Call disableBaseline asynchronously
514529
return getAsyncClient().disableBaseline(request)
515-
.handle((response, exception) -> {
516-
if (exception != null) {
517-
// Determine the actual cause
518-
Throwable cause = exception.getCause() != null ? exception.getCause() : exception;
519-
520-
// AWS ControlTower-specific exceptions
521-
if (cause instanceof ControlTowerException e) {
522-
String errorCode = e.awsErrorDetails() != null
523-
? e.awsErrorDetails().errorCode()
524-
: "UNKNOWN";
525-
526-
switch (errorCode) {
527-
case "ConflictException":
528-
System.out.println("Baseline could not be disabled (conflict): " + e.getMessage());
529-
break;
530-
531-
case "ResourceNotFoundException":
532-
System.out.println("Baseline not found for disabling: " + e.getMessage());
533-
break;
530+
.thenCompose(response -> {
531+
String operationId = response.operationIdentifier();
532+
System.out.println("Disable baseline operation ID: " + operationId);
533+
534+
// CompletableFuture that will be completed when operation finishes
535+
CompletableFuture<String> resultFuture = new CompletableFuture<>();
536+
537+
// Polling loop
538+
Runnable poller = new Runnable() {
539+
@Override
540+
public void run() {
541+
getBaselineOperationAsync(operationId)
542+
.thenAccept(statusObj -> {
543+
String status = statusObj.toString(); // Convert enum/status to string for printing
544+
System.out.println("Current disable operation status: " + status + " → waiting for SUCCEEDED or FAILED...");
545+
546+
if ("SUCCEEDED".equalsIgnoreCase(status) || "FAILED".equalsIgnoreCase(status)) {
547+
System.out.println("Disable operation finished with status: " + status);
548+
resultFuture.complete(operationId);
549+
} else {
550+
// Schedule next poll in 15 seconds
551+
CompletableFuture.delayedExecutor(15, TimeUnit.SECONDS)
552+
.execute(this);
553+
}
554+
})
555+
.exceptionally(ex -> {
556+
System.out.println("Error checking baseline operation status: " + ex.getMessage());
557+
resultFuture.completeExceptionally(ex);
558+
return null;
559+
});
560+
}
561+
};
534562

535-
case "UnauthorizedException":
536-
case "AccessDeniedException":
537-
System.out.println("Unauthorized to disable baseline: " + e.getMessage());
538-
break;
563+
// Start first poll immediately
564+
poller.run();
539565

540-
default:
541-
System.out.println("ControlTower error disabling baseline (" + errorCode + "): " + e.getMessage());
542-
}
566+
return resultFuture;
567+
})
568+
.exceptionally(ex -> {
569+
Throwable cause = ex.getCause() != null ? ex.getCause() : ex;
543570

544-
// Always return null on exception
545-
return null;
546-
}
571+
if (cause instanceof ControlTowerException e) {
572+
String errorCode = e.awsErrorDetails() != null ? e.awsErrorDetails().errorCode() : "UNKNOWN";
573+
String errorMessage = e.awsErrorDetails() != null ? e.awsErrorDetails().errorMessage() : e.getMessage();
547574

548-
// AWS SDK exceptions (network, timeout, etc.)
549-
if (cause instanceof SdkException sdkEx) {
550-
System.out.println("SDK error disabling baseline: " + sdkEx.getMessage());
551-
return null;
552-
}
575+
System.out.println("ControlTowerException caught while disabling baseline: Code=" + errorCode + ", Message=" + errorMessage);
576+
return null;
577+
}
553578

554-
// Catch-all for any other exceptions
555-
System.out.println("Unexpected error disabling baseline: " + cause.getMessage());
579+
if (cause instanceof SdkException sdkEx) {
580+
System.out.println("SDK exception caught while disabling baseline: " + sdkEx.getMessage());
556581
return null;
557582
}
558583

559-
// Success → return operation ID
560-
System.out.println("Successfully initiated disable of baseline: " + response.operationIdentifier());
561-
return response.operationIdentifier();
584+
System.out.println("Unexpected exception while disabling baseline: " + cause.getMessage());
585+
return null;
562586
});
563587
}
564588

@@ -957,16 +981,49 @@ public CompletableFuture<List<ControlSummary>> listControlsAsync() {
957981
public CompletableFuture<String> resetEnabledBaselineAsync(String enabledBaselineIdentifier) {
958982

959983
System.out.println("Starting reset of enabled baseline…");
984+
System.out.println("This operation will check the status every 15 seconds until it completes (SUCCEEDED or FAILED).");
960985

961986
ResetEnabledBaselineRequest request = ResetEnabledBaselineRequest.builder()
962987
.enabledBaselineIdentifier(enabledBaselineIdentifier)
963988
.build();
964989

965990
return getAsyncClient().resetEnabledBaseline(request)
966-
.thenApply(response -> {
991+
.thenCompose(response -> {
967992
String operationId = response.operationIdentifier();
968993
System.out.println("Reset enabled baseline operation ID: " + operationId);
969-
return operationId;
994+
995+
// Polling loop
996+
CompletableFuture<String> resultFuture = new CompletableFuture<>();
997+
998+
Runnable poller = new Runnable() {
999+
@Override
1000+
public void run() {
1001+
getBaselineOperationAsync(operationId)
1002+
.thenAccept(statusObj -> {
1003+
String status = statusObj.toString(); // Convert enum/status to string for printing
1004+
System.out.println("Current baseline operation status: " + status + " → waiting for SUCCEEDED or FAILED...");
1005+
1006+
if ("SUCCEEDED".equalsIgnoreCase(status) || "FAILED".equalsIgnoreCase(status)) {
1007+
System.out.println("Baseline operation finished with status: " + status);
1008+
resultFuture.complete(operationId);
1009+
} else {
1010+
// Schedule next poll in 15 seconds
1011+
CompletableFuture.delayedExecutor(15, TimeUnit.SECONDS)
1012+
.execute(this);
1013+
}
1014+
})
1015+
.exceptionally(ex -> {
1016+
System.out.println("Error checking baseline operation status: " + ex.getMessage());
1017+
resultFuture.completeExceptionally(ex);
1018+
return null;
1019+
});
1020+
}
1021+
};
1022+
1023+
// Start first poll immediately
1024+
poller.run();
1025+
1026+
return resultFuture;
9701027
})
9711028
.exceptionally(ex -> {
9721029
Throwable cause = ex.getCause() != null ? ex.getCause() : ex;
@@ -976,8 +1033,6 @@ public CompletableFuture<String> resetEnabledBaselineAsync(String enabledBaselin
9761033
String errorMessage = e.awsErrorDetails() != null ? e.awsErrorDetails().errorMessage() : e.getMessage();
9771034

9781035
System.out.println("ControlTowerException caught: Code=" + errorCode + ", Message=" + errorMessage);
979-
980-
// Don't fail the pipeline — return null so downstream can continue
9811036
return null;
9821037
}
9831038

@@ -990,7 +1045,6 @@ public CompletableFuture<String> resetEnabledBaselineAsync(String enabledBaselin
9901045
return null;
9911046
});
9921047
}
993-
9941048
// snippet-end:[controltower.java2.reset_enabled_baseline.main]
9951049
}
9961050
// snippet-end:[controltower.java2.controltower_actions.main]

0 commit comments

Comments
 (0)