use of bio.terra.stairway.exception.RetryException in project terra-workspace-manager by DataBiosphere.
the class CreateTableCopyJobsStep method doStep.
/**
* Create one BigQuery copy job for each table in the source dataset. Keep a running map from
* table ID to job ID as new jobs are created, and only create jobs for tables that aren't in the
* map already. Rerun the step after every table is processed so that the map may be persisted
* incrementally.
*
* <p>On retry, create the jobs for any tables that don't have them. Use WRITE_TRUNCATE to avoid
* the possibility of duplicate data.
*/
@Override
public StepResult doStep(FlightContext flightContext) throws InterruptedException, RetryException {
final FlightMap workingMap = flightContext.getWorkingMap();
final CloningInstructions effectiveCloningInstructions = flightContext.getInputParameters().get(ControlledResourceKeys.CLONING_INSTRUCTIONS, CloningInstructions.class);
if (CloningInstructions.COPY_RESOURCE != effectiveCloningInstructions) {
return StepResult.getStepResultSuccess();
}
// Gather inputs
final DatasetCloneInputs sourceInputs = getSourceInputs();
workingMap.put(ControlledResourceKeys.SOURCE_CLONE_INPUTS, sourceInputs);
final DatasetCloneInputs destinationInputs = getDestinationInputs(flightContext);
workingMap.put(ControlledResourceKeys.DESTINATION_CLONE_INPUTS, destinationInputs);
final BigQueryCow bigQueryCow = crlService.createWsmSaBigQueryCow();
// TODO(jaycarlton): remove usage of this client when it's all in CRL PF-942
final Bigquery bigQueryClient = crlService.createWsmSaNakedBigQueryClient();
try {
// Get a list of all tables in the source dataset
final TableList sourceTables = bigQueryCow.tables().list(sourceInputs.getProjectId(), sourceInputs.getDatasetName()).execute();
// Start a copy job for each source table
final Map<String, String> tableToJobId = Optional.ofNullable(workingMap.get(ControlledResourceKeys.TABLE_TO_JOB_ID_MAP, new TypeReference<Map<String, String>>() {
})).orElseGet(HashMap::new);
final List<Tables> tables = Optional.ofNullable(sourceTables.getTables()).orElse(Collections.emptyList());
// Find the first table whose ID isn't a key in the map.
final Optional<Tables> tableMaybe = tables.stream().filter(t -> null != t.getId() && !tableToJobId.containsKey(t.getId())).findFirst();
if (tableMaybe.isPresent()) {
final Tables table = tableMaybe.get();
checkStreamingBuffer(sourceInputs, bigQueryCow, table);
final Job inputJob = buildTableCopyJob(sourceInputs, destinationInputs, table);
// bill the job to the destination project
final Job submittedJob = bigQueryClient.jobs().insert(destinationInputs.getProjectId(), inputJob).execute();
// Update the map, which will be persisted
tableToJobId.put(table.getId(), submittedJob.getId());
workingMap.put(ControlledResourceKeys.TABLE_TO_JOB_ID_MAP, tableToJobId);
return new StepResult(StepStatus.STEP_RESULT_RERUN);
} else {
// All tables have entries in the map, so all jobs are started.
workingMap.put(ControlledResourceKeys.TABLE_TO_JOB_ID_MAP, // in case it's empty
tableToJobId);
return StepResult.getStepResultSuccess();
}
} catch (IOException e) {
return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);
}
}
use of bio.terra.stairway.exception.RetryException in project terra-resource-buffer by DataBiosphere.
the class CreateConsumerDefinedQuotaForBigQueryDailyUsageStep method doStep.
/**
* Apply a Consumer Quota Override for the BigQuery Query Usage Quota.
*/
@Override
public StepResult doStep(FlightContext context) throws InterruptedException, RetryException {
Optional<Long> overrideValue = GoogleProjectConfigUtils.bigQueryDailyUsageOverrideValueMebibytes(gcpProjectConfig);
if (overrideValue.isEmpty()) {
// Do not apply any quota override
return StepResult.getStepResultSuccess();
}
long projectNumber = Optional.ofNullable(context.getWorkingMap().get(GOOGLE_PROJECT_NUMBER, Long.class)).orElseThrow();
QuotaOverride overridePerProjectPerDay = buildQuotaOverride(projectNumber, overrideValue.get());
// parent format and other details obtained by hitting the endpoint
// https://serviceusage.googleapis.com/v1beta1/projects/${PROJECT_NUMBER}/services/bigquery.googleapis.com/consumerQuotaMetrics
String parent = String.format("projects/%d/services/bigquery.googleapis.com/consumerQuotaMetrics/" + "bigquery.googleapis.com%%2Fquota%%2Fquery%%2Fusage/limits/%%2Fd%%2Fproject", projectNumber);
try {
// We are decreasing the quota by more than 10%, so we must tell Service Usage to bypass the
// check with the force flag.
Operation createOperation = serviceUsageCow.services().consumerQuotaMetrics().limits().consumerOverrides().create(parent, overridePerProjectPerDay).setForce(true).execute();
OperationCow<Operation> operationCow = serviceUsageCow.operations().operationCow(createOperation);
pollUntilSuccess(operationCow, Duration.ofSeconds(3), Duration.ofMinutes(5));
} catch (IOException e) {
throw new RetryException(e);
}
return StepResult.getStepResultSuccess();
}
use of bio.terra.stairway.exception.RetryException in project terra-resource-buffer by DataBiosphere.
the class CreateGkeDefaultSAStep method doStep.
@Override
public StepResult doStep(FlightContext flightContext) throws RetryException {
if (!createGkeDefaultSa(gcpProjectConfig)) {
return StepResult.getStepResultSuccess();
}
String projectId = flightContext.getWorkingMap().get(GOOGLE_PROJECT_ID, String.class);
CreateServiceAccountRequest createRequest = new CreateServiceAccountRequest().setAccountId(GKE_SA_NAME).setServiceAccount(new ServiceAccount().setDescription("Default service account can be used on GKE node. "));
try {
iamCow.projects().serviceAccounts().create("projects/" + projectId, createRequest).execute();
} catch (GoogleJsonResponseException e) {
// Otherwise throw a retry exception.
if (e.getStatusCode() != HttpStatus.CONFLICT.value()) {
throw new RetryException(e);
}
logger.warn("Service account {} already created for notebook instance.", GKE_SA_NAME);
} catch (IOException e) {
throw new RetryException(e);
}
// Grants permission that a GKE node runner needs
String serviceAccountEmail = ServiceAccountName.emailFromAccountId(GKE_SA_NAME, projectId);
try {
Policy policy = rmCow.projects().getIamPolicy(projectId, new GetIamPolicyRequest()).execute();
GKE_SA_ROLES.forEach(r -> policy.getBindings().add(new Binding().setRole(r).setMembers(Collections.singletonList("serviceAccount:" + serviceAccountEmail))));
// Duplicating bindings is harmless (e.g. on retry). GCP de-duplicates.
rmCow.projects().setIamPolicy(projectId, new SetIamPolicyRequest().setPolicy(policy)).execute();
} catch (IOException e) {
logger.info("Error when setting IAM policy for GKE default node SA", e);
return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);
}
return StepResult.getStepResultSuccess();
}
use of bio.terra.stairway.exception.RetryException in project terra-workspace-manager by DataBiosphere.
the class CreateCustomGcpRolesStep method createCustomRole.
/**
* Utility for creating custom roles in GCP from WSM's CustomGcpIamRole objects. These roles will
* be defined at the project level in the specified by projectId.
*/
private void createCustomRole(CustomGcpIamRole customRole, String projectId) throws RetryException {
try {
Role gcpRole = new Role().setIncludedPermissions(customRole.getIncludedPermissions()).setTitle(customRole.getRoleName());
CreateRoleRequest request = new CreateRoleRequest().setRole(gcpRole).setRoleId(customRole.getRoleName());
logger.debug("Creating role {} with permissions {} in project {}", customRole.getRoleName(), customRole.getIncludedPermissions(), projectId);
iamCow.projects().roles().create("projects/" + projectId, request).execute();
} catch (GoogleJsonResponseException googleEx) {
// of role names must be due to duplicate step execution.
if (googleEx.getStatusCode() != HttpStatus.CONFLICT.value()) {
throw new RetryException(googleEx);
}
} catch (IOException e) {
// Retry on IO exceptions thrown by CRL.
throw new RetryException(e);
}
}
use of bio.terra.stairway.exception.RetryException in project terra-workspace-manager by DataBiosphere.
the class CreateAiNotebookInstanceStep method undoStep.
@Override
public StepResult undoStep(FlightContext flightContext) throws InterruptedException {
final GcpCloudContext gcpCloudContext = flightContext.getWorkingMap().get(ControlledResourceKeys.GCP_CLOUD_CONTEXT, GcpCloudContext.class);
InstanceName instanceName = resource.toInstanceName(gcpCloudContext.getGcpProjectId());
AIPlatformNotebooksCow notebooks = crlService.getAIPlatformNotebooksCow();
try {
OperationCow<Operation> deletionOperation;
try {
deletionOperation = notebooks.operations().operationCow(notebooks.instances().delete(instanceName).execute());
} catch (GoogleJsonResponseException e) {
// The AI notebook instance may never have been created or have already been deleted.
if (e.getStatusCode() == HttpStatus.NOT_FOUND.value()) {
logger.debug("No notebook instance {} to delete.", instanceName.formatName());
return StepResult.getStepResultSuccess();
}
return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);
}
GcpUtils.pollUntilSuccess(deletionOperation, Duration.ofSeconds(20), Duration.ofMinutes(12));
} catch (IOException | RetryException e) {
return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);
}
return StepResult.getStepResultSuccess();
}
Aggregations