Examples with FlightContext bio.terra.stairway.FlightContext used on opensource projects

Example 1 with FlightContext

use of bio.terra.stairway.FlightContext in project jade-data-repo by DataBiosphere.

the class LoadUnitTest method getLoadTagFailTest.

@Test(expected = LoadLockFailureException.class)
public void getLoadTagFailTest() throws Exception {
    FlightMap inputParams = new FlightMap();
    FlightContext flightContext = new FlightContext(inputParams, null);
    loadService.getLoadTag(flightContext);
}

Example 2 with FlightContext

use of bio.terra.stairway.FlightContext 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);
    }
}

Example 3 with FlightContext

use of bio.terra.stairway.FlightContext in project jade-data-repo by DataBiosphere.

the class LoadUnitTest method getLoadTagTest.

@Test
public void getLoadTagTest() throws Exception {
    // Should get tag from working map
    FlightMap inputParams = new FlightMap();
    FlightContext flightContext = new FlightContext(inputParams, null);
    FlightMap workingMap = flightContext.getWorkingMap();
    workingMap.put(LoadMapKeys.LOAD_TAG, LoadTagsUsedByTest.LOADTAG_1.getTag());
    String loadTag = loadService.getLoadTag(flightContext);
    assertThat("working map load tag", loadTag, equalTo(LoadTagsUsedByTest.LOADTAG_1.getTag()));
    // Should get from input Params
    FlightMap inputParams1 = new FlightMap();
    inputParams1.put(LoadMapKeys.LOAD_TAG, LoadTagsUsedByTest.LOADTAG_1.getTag());
    flightContext = new FlightContext(inputParams1, null);
    workingMap = flightContext.getWorkingMap();
    workingMap.put(LoadMapKeys.LOAD_TAG, LoadTagsUsedByTest.LOADTAG_2.getTag());
    loadTag = loadService.getLoadTag(flightContext);
    assertThat("input params load tag", loadTag, equalTo(LoadTagsUsedByTest.LOADTAG_1.getTag()));
}

Example 4 with FlightContext

use of bio.terra.stairway.FlightContext in project jade-data-repo by DataBiosphere.

the class CreateSnapshotValidateAssetStep method doStep.

@Override
public StepResult doStep(FlightContext context) {
    /*
         * get dataset
         * get dataset asset list
         * get snapshot asset name
         * check that snapshot asset exists
         */
    Snapshot snapshot = snapshotService.retrieveByName(snapshotReq.getName());
    for (SnapshotSource snapshotSource : snapshot.getSnapshotSources()) {
        Dataset dataset = datasetService.retrieve(snapshotSource.getDataset().getId());
        List<String> datasetAssetNamesList = dataset.getAssetSpecifications().stream().map(assetSpec -> assetSpec.getName()).collect(Collectors.toList());
        String snapshotSourceAssetName = snapshotSource.getAssetSpecification().getName();
        if (!datasetAssetNamesList.contains(snapshotSourceAssetName)) {
            String datasetAssetNames = String.join("', '", datasetAssetNamesList);
            String message = String.format("Mismatched asset name: '%s' is not an asset in the asset list for dataset '%s'." + "Asset list is '%s'", snapshotSourceAssetName, dataset.getName(), datasetAssetNames);
            FlightUtils.setErrorResponse(context, message, HttpStatus.BAD_REQUEST);
            return new StepResult(StepStatus.STEP_RESULT_FAILURE_FATAL, new MismatchedValueException(message));
        }
    }
    return StepResult.getStepResultSuccess();
}

Example 5 with FlightContext

use of bio.terra.stairway.FlightContext in project terra-resource-buffer by DataBiosphere.

the class SetIamPolicyStep method doStep.

@Override
public StepResult doStep(FlightContext flightContext) throws RetryException {
    // Skip if IAM binding is not set.
    if (gcpProjectConfig.getIamBindings() == null || gcpProjectConfig.getIamBindings().isEmpty()) {
        return StepResult.getStepResultSuccess();
    }
    String projectId = flightContext.getWorkingMap().get(GOOGLE_PROJECT_ID, String.class);
    try {
        Policy policy = rmCow.projects().getIamPolicy(projectId, new GetIamPolicyRequest()).execute();
        gcpProjectConfig.getIamBindings().stream().map(iamBinding -> new Binding().setRole(iamBinding.getRole()).setMembers(iamBinding.getMembers())).forEach(policy.getBindings()::add);
        rmCow.projects().setIamPolicy(projectId, new SetIamPolicyRequest().setPolicy(policy)).execute();
    } catch (IOException e) {
        logger.info("Error when setting IAM policy", e);
        return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);
    }
    return StepResult.getStepResultSuccess();
}