com.acxiom.gcp.steps.PubSubSteps.scala Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of metalus-gcp_2.12-spark_3.0 Show documentation
Metalus GCP Step Library
There is a newer version: 1.9.8
Show newest version
package com.acxiom.gcp.steps

import com.acxiom.gcp.utils.GCPUtilities
import com.acxiom.pipeline.PipelineContext
import com.acxiom.pipeline.annotations.{StepFunction, StepObject, StepParameter, StepParameters}
import com.google.auth.oauth2.GoogleCredentials
import org.apache.spark.sql.DataFrame

@StepObject
object PubSubSteps {
  @StepFunction("451d4dc8-9bce-4cb4-a91d-1a09e0efd9b8",
    "Write DataFrame to a PubSub Topic",
    "This step will write a DataFrame to a PubSub Topic",
    "Pipeline",
    "GCP")
  @StepParameters(Map("dataFrame" -> StepParameter(None, Some(true), None, None, None, None, Some("The DataFrame to post to the Pub/Sub topic")),
    "topicName" -> StepParameter(None, Some(true), None, None, None, None, Some("The topic within the Pub/Sub")),
    "separator" -> StepParameter(None, Some(false), None, None, None, None, Some("The separator character to use when combining the column data")),
    "credentials" -> StepParameter(None, Some(false), None, None, None, None, Some("The optional credentials to use for Pub/Sub access"))))
  def writeToStreamWithCredentials(dataFrame: DataFrame,
                    topicName: String,
                    separator: String = ",",
                    credentials: Option[Map[String, String]] = None): Unit = {
    val creds = GCPUtilities.generateCredentials(credentials)
    publishDataFrame(dataFrame, separator, topicName, creds)
  }

  @StepFunction("aaa880e1-4190-4ffe-9fda-4150680f17c9",
    "Write DataFrame to a PubSub Topic Using Global Credentials",
    "This step will write a DataFrame to a PubSub Topic using the CredentialProvider to get Credentials",
    "Pipeline",
    "GCP")
  @StepParameters(Map("dataFrame" -> StepParameter(None, Some(true), None, None, None, None, Some("The DataFrame to post to the Pub/Sub topic")),
    "topicName" -> StepParameter(None, Some(true), None, None, None, None, Some("The topic within the Pub/Sub")),
    "separator" -> StepParameter(None, Some(false), None, None, None, None, Some("The separator character to use when combining the column data"))))
  def writeToStream(dataFrame: DataFrame,
                    topicName: String,
                    separator: String = ",",
                    pipelineContext: PipelineContext): Unit = {
    val creds = GCPUtilities.getCredentialsFromPipelineContext(pipelineContext)
    publishDataFrame(dataFrame, separator, topicName, creds)
  }

  @StepFunction("2c937e74-8735-46d6-abfe-0c040ae8f435",
    "Write a single message to a PubSub Topic Using Provided Credentials",
    "This step will write a DataFrame to a PubSub Topic using the providedt Credentials",
    "Pipeline",
    "GCP")
  @StepParameters(Map("message" -> StepParameter(None, Some(true), None, None, None, None, Some("The message to post to the Pub/Sub topic")),
    "topicName" -> StepParameter(None, Some(true), None, None, None, None, Some("The topic within the Pub/Sub")),
    "credentials" -> StepParameter(None, Some(false), None, None, None, None, Some("The optional credentials to use when posting"))))
  def postMessage(message: String, topicName: String, credentials: Option[Map[String, String]] = None): Unit = {
    val creds: Option[GoogleCredentials] = GCPUtilities.generateCredentials(credentials)
    GCPUtilities.postMessage(topicName, creds, message)
  }

  @StepFunction("b359130d-8e11-44e4-b552-9cef6150bc2b",
    "Write a single message to a PubSub Topic Using Global Credentials",
    "This step will write a message to a PubSub Topic using the CredentialProvider to get Credentials",
    "Pipeline",
    "GCP")
  @StepParameters(Map("message" -> StepParameter(None, Some(true), None, None, None, None, Some("The message to post to the Pub/Sub topic")),
    "topicName" -> StepParameter(None, Some(true), None, None, None, None, Some("The topic within the Pub/Sub"))))
  def postMessage(message: String, topicName: String, pipelineContext: PipelineContext): Unit = {
    GCPUtilities.postMessage(topicName, GCPUtilities.getCredentialsFromPipelineContext(pipelineContext), message)
  }

  /**
    * Publish the DataFrame to the given Pub/Sub topic using the optional credentials
    *
    * @param dataFrame The DataFrame to publish
    * @param topicName The Pub/Sub topic name
    * @param separator The separator to use when combining the column data
    * @param creds     The optional GoogleCredentials
    */
  private def publishDataFrame(dataFrame: DataFrame, topicName: String, separator: String, creds: Option[GoogleCredentials]): Unit = {
    dataFrame.rdd.foreach(row => {
      val rowData = row.mkString(separator)
      GCPUtilities.postMessage(topicName, creds, rowData)
    })
  }
}