Examples with Operation com.google.cloud.pso.bq_pii_classifier.entities.Operation used on opensource projects

Example 1 with Operation

use of com.google.cloud.pso.bq_pii_classifier.entities.Operation in project bq-pii-classifier by GoogleCloudPlatform.

the class DispatcherTest method testWithInput.

// @Test
// public void testDispatcher_withDatasets () throws IOException {
// 
// String jsonPayLoad = "{\"tablesInclude\":\"\""
// + ",\"tablesExclude\":\"p1.d1.t1\""
// + ",\"datasetsInclude\":\"p1.d1, p1.d2\""
// + ",\"datasetsExclude\":\"\""
// + ",\"projectsInclude\":\"p2\"" // should have no effect
// + "}";
// 
// List<String> expectedOutput = Lists.newArrayList("p1.d1.t2", "p1.d2.t1", "p1.d2.t2");
// List<String> actualOutput = testWithInput(jsonPayLoad);
// 
// assertEquals(expectedOutput, actualOutput);
// }
// 
// @Test
// public void testDispatcher_withProjects () throws IOException {
// 
// String jsonPayLoad = "{\"tablesInclude\":\"\""
// + ",\"tablesExclude\":\"p1.d2.t1\""
// + ",\"datasetsInclude\":\"\""
// + ",\"datasetsExclude\":\"p1.d1\""
// + ",\"projectsInclude\":\"p1, p2\"" // should have no effect
// + "}";
// 
// List<String> expectedOutput = Lists.newArrayList("p1.d2.t2", "p2.d1.t1", "p2.d1.t2");
// List<String> actualOutput = testWithInput(jsonPayLoad);
// 
// assertEquals(expectedOutput, actualOutput);
// }
private List<String> testWithInput(BigQueryScope bigQueryScope) throws IOException, NonRetryableApplicationException, InterruptedException {
    // Dispatcher function = new Dispatcher(envMock, bqServiceMock, cloudTasksServiceMock);
    PubSubPublishResults results = function.execute(bigQueryScope, "");
    PubSubServiceImpl pubSubServiceMock = mock(PubSubServiceImpl.class);
    lenient().when(pubSubServiceMock.publishTableOperationRequests(anyString(), anyString(), any())).thenReturn(new PubSubPublishResults(Arrays.asList(new SuccessPubSubMessage(new Operation("p1.d1.t1", "runId", "trackingId"), "publishedMessageId"), new SuccessPubSubMessage(new Operation("p1.d1.t2", "runId", "trackingId"), "publishedMessageId")), Arrays.asList(new FailedPubSubMessage(new Operation("", "", ""), new Exception("test fail message")))));
    return results.getSuccessMessages().stream().map(x -> ((Operation) x.getMsg()).getEntityKey()).collect(Collectors.toList());
}

Example 2 with Operation

use of com.google.cloud.pso.bq_pii_classifier.entities.Operation in project bq-pii-classifier by GoogleCloudPlatform.

the class InspectorController method receiveMessage.

@RequestMapping(value = "/", method = RequestMethod.POST)
public ResponseEntity receiveMessage(@RequestBody PubSubEvent requestBody) {
    String trackingId = "0000000000000-z";
    DlpService dlpService = null;
    BigQueryService bqService = null;
    try {
        if (requestBody == null || requestBody.getMessage() == null) {
            String msg = "Bad Request: invalid message format";
            logger.logSevereWithTracker(trackingId, msg);
            throw new NonRetryableApplicationException("Request body or message is Null.");
        }
        String requestJsonString = requestBody.getMessage().dataToUtf8String();
        // remove any escape characters (e.g. from Terraform
        requestJsonString = requestJsonString.replace("\\", "");
        logger.logInfoWithTracker(trackingId, String.format("Received payload: %s", requestJsonString));
        Operation operation = gson.fromJson(requestJsonString, Operation.class);
        trackingId = operation.getTrackingId();
        logger.logInfoWithTracker(trackingId, String.format("Parsed Request: %s", operation.toString()));
        dlpService = new DlpServiceImpl();
        bqService = new BigQueryServiceImpl();
        Inspector inspector = new Inspector(environment.toConfig(), dlpService, bqService, new GCSPersistentSetImpl(environment.getGcsFlagsBucket()), "inspector-flags");
        inspector.execute(operation, trackingId, requestBody.getMessage().getMessageId());
        return new ResponseEntity("Process completed successfully.", HttpStatus.OK);
    } catch (Exception e) {
        return ControllerExceptionHelper.handleException(e, logger, trackingId);
    } finally {
        if (dlpService != null) {
            dlpService.shutDown();
        }
    }
}

Example 3 with Operation

use of com.google.cloud.pso.bq_pii_classifier.entities.Operation in project bq-pii-classifier by GoogleCloudPlatform.

the class TaggerController method parseEvent.

private Operation parseEvent(PubSubEvent event) throws NonRetryableApplicationException {
    String defaultTrackingId = "0000000000000-z";
    // check if the message is sent from a Standard DLP inspection job
    if (event.getMessage().getAttributes() != null) {
        String dlpJobName = event.getMessage().getAttributes().getOrDefault("DlpJobName", "");
        if (!dlpJobName.isBlank()) {
            logger.logInfoWithTracker(defaultTrackingId, String.format("Parsed DlpJobName '%s'", dlpJobName));
            String trackingId = TrackingHelper.extractTrackingIdFromJobName(dlpJobName);
            String runId = TrackingHelper.parseRunIdAsPrefix(trackingId);
            Operation taggerRequest = new Operation(dlpJobName, runId, trackingId);
            logger.logInfoWithTracker(taggerRequest.getTrackingId(), String.format("Parsed Request from Standard DLP: %s", taggerRequest.toString()));
            return taggerRequest;
        }
    }
    try {
        String requestJsonString = event.getMessage().dataToUtf8String();
        // remove any escape characters (e.g. from Terraform
        requestJsonString = requestJsonString.replace("\\", "");
        logger.logInfoWithTracker(defaultTrackingId, String.format("Received payload: %s", requestJsonString));
        Operation taggerRequest = gson.fromJson(requestJsonString, Operation.class);
        logger.logInfoWithTracker(taggerRequest.getTrackingId(), String.format("Parsed Request from Dispatcher: %s", taggerRequest.toString()));
        return taggerRequest;
    } catch (Exception ex) {
        try {
            byte[] data = event.getMessage().getData();
            DataProfilePubSubMessage dataProfilePubSubMessage = DataProfilePubSubMessage.parseFrom(data);
            logger.logInfoWithTracker(defaultTrackingId, String.format("Parsed DataProfilePubSubMessage= '%s'", dataProfilePubSubMessage));
            TableSpec targetTable = TableSpec.fromFullResource(dataProfilePubSubMessage.getProfile().getFullResource());
            String runId = TrackingHelper.generateOneTimeTaggingSuffix();
            String trackingId = TrackingHelper.generateTrackingId(runId, targetTable.toSqlString());
            Operation taggerRequest = new Operation(targetTable.toSqlString(), runId, trackingId);
            logger.logInfoWithTracker(taggerRequest.getTrackingId(), String.format("Parsed Request from Auto DLP notifications : %s", taggerRequest.toString()));
            return taggerRequest;
        } catch (InvalidProtocolBufferException invalidProtocolBufferException) {
            throw new NonRetryableApplicationException(String.format("Couldn't parse PubSub event as Proto: %s : %s", invalidProtocolBufferException.getClass().getSimpleName(), invalidProtocolBufferException.getMessage()));
        }
    }
}

Example 4 with Operation

use of com.google.cloud.pso.bq_pii_classifier.entities.Operation in project bq-pii-classifier by GoogleCloudPlatform.

the class TaggerController method receiveMessage.

@RequestMapping(value = "/", method = RequestMethod.POST)
public ResponseEntity receiveMessage(@RequestBody PubSubEvent requestBody) {
    String defaultTrackingId = "0000000000000-z";
    Operation taggerRequest = null;
    try {
        if (requestBody == null || requestBody.getMessage() == null) {
            String msg = "Bad Request: invalid message format";
            logger.logSevereWithTracker(defaultTrackingId, msg);
            throw new NonRetryableApplicationException("Request body or message is Null.");
        }
        // The pubsub message could come from different sources with different formats
        // 1. From the Dispatcher as JSON encoded "Operation" object (in standard and auto-dlp modes)
        // 2. From Standard DLP inspection job notifications (PubSub Actions) as a string message with the DLP job name
        // 3. From Auto DLP notifications (PubSub Actions) as Protobuf encoded "DataProfilePubSubMessage" (in AutoDLP mode only)
        // Here we try and parse it directly as an "Operation" object or indirectly via the "DataProfilePubSubMessage" proto
        taggerRequest = parseEvent(requestBody);
        BigQueryService bigQueryService = new BigQueryServiceImpl();
        // determine the Type of DLP FindingsReader based on env
        FindingsReader findingsReader = FindingsReaderFactory.getNewReader(FindingsReaderFactory.findReader(environment.getIsAutoDlpMode(), environment.getPromoteMixedTypes()), bigQueryService, environment.getProjectId(), environment.getDlpDataset(), environment.getIsAutoDlpMode() ? environment.getDlpTableAuto() : environment.getDlpTableStandard(), environment.getConfigViewDatasetDomainMap(), environment.getConfigViewProjectDomainMap(), environment.getConfigViewInfoTypePolicyTagsMap());
        Tagger tagger = new Tagger(environment.toConfig(), bigQueryService, findingsReader, new GCSPersistentSetImpl(environment.getGcsFlagsBucket()), "tagger-flags");
        tagger.execute(taggerRequest, requestBody.getMessage().getMessageId());
        return new ResponseEntity("Process completed successfully.", HttpStatus.OK);
    } catch (Exception e) {
        String trackingId = taggerRequest == null ? defaultTrackingId : taggerRequest.getTrackingId();
        return ControllerExceptionHelper.handleException(e, logger, trackingId);
    }
}