com.boozallen.aiops.mda.pattern.VoidInboundAndNativeOutboundAsyncBase Maven / Gradle / Ivy

Go to download
package com.boozallen.aiops.mda.pattern;

/*-
 * #%L
 * aiSSEMBLE::Test::MDA::Data Delivery Spark
 * %%
 * Copyright (C) 2021 Booz Allen
 * %%
 * This software package is licensed under the Booz Allen Public License. All Rights Reserved.
 * #L%
 */

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;

import java.time.Duration;

import com.boozallen.aiops.mda.pattern.pipeline.PipelineBase;
import org.aeonbits.owner.KrauseningConfigFactory;
import io.smallrye.mutiny.Multi;
import java.util.concurrent.CompletionStage;
import org.eclipse.microprofile.reactive.messaging.Outgoing;
import com.boozallen.aissemble.core.filestore.AbstractFileStore;
import com.boozallen.aissemble.core.metadata.MetadataAPI;
import com.boozallen.aissemble.core.metadata.MetadataModel;
import com.boozallen.aissemble.core.metadata.producer.MetadataProducer;

import com.boozallen.aissemble.alerting.core.Alert;
import com.boozallen.aissemble.alerting.core.AlertProducerApi;
import javax.enterprise.inject.spi.CDI;

import com.boozallen.aiops.data.delivery.spark.SparkConfig;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.api.java.UDF2;
import com.boozallen.aissemble.data.encryption.policy.config.EncryptAlgorithm;
import org.apache.spark.sql.types.DataTypes;

import javax.inject.Inject;
import java.time.Clock;
import java.time.ZonedDateTime;
import java.time.ZoneOffset;
import java.util.UUID;

import com.boozallen.aissemble.data.lineage.util.LineageUtil.LineageEventData;
import static com.boozallen.aissemble.data.lineage.util.LineageUtil.recordLineage;

import io.openlineage.client.OpenLineage.ParentRunFacet;

import com.boozallen.aissemble.data.encryption.AiopsEncrypt;
import com.boozallen.aissemble.data.encryption.SimpleAesEncrypt;
import com.boozallen.aissemble.data.encryption.VaultEncrypt;
import com.boozallen.aissemble.data.encryption.policy.EncryptionPolicy;
import com.boozallen.aissemble.data.encryption.policy.EncryptionPolicyManager;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.functions;
import org.apache.spark.sql.types.StructType;
import java.util.stream.Collectors;
import static org.apache.spark.sql.functions.col;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.HashMap;
import java.util.HashSet;
import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import static org.apache.spark.sql.functions.col;
import static org.apache.spark.sql.functions.lit;
import org.apache.commons.lang.NotImplementedException;
import org.aeonbits.owner.KrauseningConfigFactory;

/**
 * The smallrye connector consumes one message at a time - this class writes multiple source data records within one
 * message so that they can be ingested in batches.
 * 
 * GENERATED CODE - DO NOT MODIFY (add your customizations in VoidInboundAndNativeOutboundAsync).
 * 
 * Generated from: templates/data-delivery-spark/asynchronous.processor.base.java.vm
 */
public abstract class VoidInboundAndNativeOutboundAsyncBase extends AbstractPipelineStep {

    private static final Logger logger = LoggerFactory.getLogger(VoidInboundAndNativeOutboundAsyncBase.class);

    protected static final SparkConfig config = KrauseningConfigFactory.create(SparkConfig.class);

    protected static final String stepPhase = "VoidInboundAndNativeOutboundAsync";

    protected MetadataProducer metadataProducer;

    protected AlertProducerApi alertProducer;

    protected VoidInboundAndNativeOutboundAsyncBase(String subject, String action) {
        super(subject, action);
    }

    /**
     * Receives messages and processes them asynchronously.
     */
    public CompletionStage> executeStep() {
        UUID runId = UUID.randomUUID();
        ParentRunFacet parentRunFacet = PipelineBase.getInstance().getPipelineRunAsParentRunFacet();
        String jobName = getJobName();
        String defaultNamespace = getDefaultNamespace();
        LineageEventData eventData = createBaseEventData();
        ZonedDateTime eventStartTime = ZonedDateTime.now(ZoneOffset.UTC);
        ZonedDateTime eventEndTime = null;
        Map eventParams = new HashMap();
        eventParams.put("startTime", eventStartTime);
        recordLineage(createLineageStartEvent(runId, jobName, defaultNamespace, parentRunFacet, eventData, eventParams));
        try {
            recordProvenance();

            CompletionStage> outboundPayload = executeStepImpl();
            if (eventEndTime == null) {
                eventEndTime = ZonedDateTime.now(ZoneOffset.UTC);
                eventParams.put("endTime", eventEndTime);
                recordLineage(createLineageCompleteEvent(runId, jobName, defaultNamespace, parentRunFacet, eventData,
                        eventParams));
            }
            sendAlert(Alert.Status.SUCCESS, getSuccessMessage());
            return outboundPayload;
        } catch (Exception e) {
            logger.error("Step failed to complete", e);
            sendAlert(Alert.Status.FAILURE, getErrorMessage(e));
            eventEndTime = ZonedDateTime.now(ZoneOffset.UTC);
            eventParams.put("error", e);
            eventParams.put("endTime", eventEndTime);
            recordLineage(createLineageFailEvent(runId, jobName, defaultNamespace, parentRunFacet, eventData,
                    eventParams));
            PipelineBase.getInstance().recordPipelineLineageFailEvent();
            throw e;
        }
    }

    /**
     * Performs the processing of inbound data. This method is called when a message is received.
     */
    protected abstract CompletionStage> executeStepImpl();

    /**
     * Controls creating the metadata that will be recorded for provenance purposes.
     * 
     * @param resource
     *            the identifier of the data
     * @param subject
     *            the thing acting on the resource
     * @param action
     *            the name of the activity being performed on the resource
     * @return the provenance metadata to create
     */
    protected abstract MetadataModel createProvenanceMetadata(String resource, String subject, String action);

    /**
     * Returns the resource identifier that will be used to track this record set via provenance.
     * 
     * @return resource identifier
     */
    protected String getProvenanceResourceIdentifier() {
        return "Unspecified VoidInboundAndNativeOutboundAsync resource";
    }

    /**
     * Records provenance for this step.
     */
    protected void recordProvenance() {
        logger.info("Sending provenance...");

        String resource = getProvenanceResourceIdentifier();
        MetadataModel model = createProvenanceMetadata(resource, subject, action);
        metadataProducer.send(model);

        logger.info("Provenance sent");
    }

    @Inject
    public void setMetadataProducer(MetadataProducer metadataProducer) {
        this.metadataProducer = metadataProducer;
    }

    @Inject
    public void setAlertProducer(AlertProducerApi alertProducer) {
        this.alertProducer = alertProducer;
    }

    /**
     * Send an alert with a given status and message.
     * 
     * @param status
     *            the status of the alert
     * @param message
     *            the message
     */
    protected void sendAlert(Alert.Status status, String message) {
        if (alertProducer != null) {
            alertProducer.sendAlert(status, message);
        } else {
            logger.error("Alert cannot be sent without a valid Alert Producer! Please add an Alert Producer to the PipelinesCdiContext!");
        }
    }

    /**
     * Invoked in VoidInboundAndNativeOutboundAsyncBase. For static messages update this method; however, if more
     * complex messaging is desired turn off alerting and implement as desired in VoidInboundAndNativeOutboundAsync.
     */
    protected String getSuccessMessage() {
        return action + " completed asynchronously at " + Clock.systemUTC().instant();
    }

    /**
     * Invoked in VoidInboundAndNativeOutboundAsyncBase. For static messages update this method; however, if more
     * complex messaging is desired turn off alerting and implement as desired in VoidInboundAndNativeOutboundAsync.
     */
    protected String getErrorMessage(Exception e) {
        return action + " failed due to error: " + e.getMessage();
    }

    /**
     * Spark User Defined Function for running encryption on columns. Note: must be registered with the spark session.
     * 
     * @return The cipher text
     */
    protected UDF2 encryptUDF() {
        return (plainText, encryptAlgorithm) -> {
            if (plainText != null) {
                // Default algorithm is AES
                AiopsEncrypt aiopsEncrypt = new SimpleAesEncrypt();

                if(encryptAlgorithm.equals("VAULT_ENCRYPT")) {
                    aiopsEncrypt = new VaultEncrypt();
                }

                return aiopsEncrypt.encryptValue(plainText);
            } else {
                return "";
            }
        };
    }

}