embulk-input-pubsub 0.0.1

data/src/main/java/com/embulk/input/pubsub/checkpoint/CheckpointOrBuilder.java

@@ -0,0 +1,33 @@
+// Generated by the protocol buffer compiler.  DO NOT EDIT!
+// source: src/main/resources/checkpoint.proto
+
+package com.embulk.input.pubsub.checkpoint;
+
+public interface CheckpointOrBuilder extends
+    // @@protoc_insertion_point(interface_extends:Checkpoint)
+    com.google.protobuf.MessageOrBuilder {
+
+  /**
+   * <code>repeated .google.pubsub.v1.PubsubMessage messages = 1;</code>
+   */
+  java.util.List<com.google.pubsub.v1.PubsubMessage> 
+      getMessagesList();
+  /**
+   * <code>repeated .google.pubsub.v1.PubsubMessage messages = 1;</code>
+   */
+  com.google.pubsub.v1.PubsubMessage getMessages(int index);
+  /**
+   * <code>repeated .google.pubsub.v1.PubsubMessage messages = 1;</code>
+   */
+  int getMessagesCount();
+  /**
+   * <code>repeated .google.pubsub.v1.PubsubMessage messages = 1;</code>
+   */
+  java.util.List<? extends com.google.pubsub.v1.PubsubMessageOrBuilder> 
+      getMessagesOrBuilderList();
+  /**
+   * <code>repeated .google.pubsub.v1.PubsubMessage messages = 1;</code>
+   */
+  com.google.pubsub.v1.PubsubMessageOrBuilder getMessagesOrBuilder(
+      int index);
+}

data/src/main/java/com/embulk/input/pubsub/checkpoint/CheckpointProtos.java

@@ -0,0 +1,61 @@
+// Generated by the protocol buffer compiler.  DO NOT EDIT!
+// source: src/main/resources/checkpoint.proto
+
+package com.embulk.input.pubsub.checkpoint;
+
+public final class CheckpointProtos {
+  private CheckpointProtos() {}
+  public static void registerAllExtensions(
+      com.google.protobuf.ExtensionRegistryLite registry) {
+  }
+
+  public static void registerAllExtensions(
+      com.google.protobuf.ExtensionRegistry registry) {
+    registerAllExtensions(
+        (com.google.protobuf.ExtensionRegistryLite) registry);
+  }
+  static final com.google.protobuf.Descriptors.Descriptor
+    internal_static_Checkpoint_descriptor;
+  static final 
+    com.google.protobuf.GeneratedMessageV3.FieldAccessorTable
+      internal_static_Checkpoint_fieldAccessorTable;
+
+  public static com.google.protobuf.Descriptors.FileDescriptor
+      getDescriptor() {
+    return descriptor;
+  }
+  private static  com.google.protobuf.Descriptors.FileDescriptor
+      descriptor;
+  static {
+    java.lang.String[] descriptorData = {
+      "\n#src/main/resources/checkpoint.proto\032(g" +
+      "oogleapis/google/pubsub/v1/pubsub.proto\"" +
+      "?\n\nCheckpoint\0221\n\010messages\030\001 \003(\0132\037.google" +
+      ".pubsub.v1.PubsubMessageB8\n\"com.embulk.i" +
+      "nput.pubsub.checkpointB\020CheckpointProtos" +
+      "P\001b\006proto3"
+    };
+    com.google.protobuf.Descriptors.FileDescriptor.InternalDescriptorAssigner assigner =
+        new com.google.protobuf.Descriptors.FileDescriptor.    InternalDescriptorAssigner() {
+          public com.google.protobuf.ExtensionRegistry assignDescriptors(
+              com.google.protobuf.Descriptors.FileDescriptor root) {
+            descriptor = root;
+            return null;
+          }
+        };
+    com.google.protobuf.Descriptors.FileDescriptor
+      .internalBuildGeneratedFileFrom(descriptorData,
+        new com.google.protobuf.Descriptors.FileDescriptor[] {
+          com.google.pubsub.v1.PubsubProto.getDescriptor(),
+        }, assigner);
+    internal_static_Checkpoint_descriptor =
+      getDescriptor().getMessageTypes().get(0);
+    internal_static_Checkpoint_fieldAccessorTable = new
+      com.google.protobuf.GeneratedMessageV3.FieldAccessorTable(
+        internal_static_Checkpoint_descriptor,
+        new java.lang.String[] { "Messages", });
+    com.google.pubsub.v1.PubsubProto.getDescriptor();
+  }
+
+  // @@protoc_insertion_point(outer_class_scope)
+}

data/src/main/resources/checkpoint.proto

@@ -0,0 +1,11 @@
+syntax = "proto3";
+
+option java_multiple_files = true;
+option java_outer_classname = "CheckpointProtos";
+option java_package = "com.embulk.input.pubsub.checkpoint";
+
+import "googleapis/google/pubsub/v1/pubsub.proto";
+
+message Checkpoint {
+  repeated google.pubsub.v1.PubsubMessage messages = 1;
+}

data/src/main/scala/org/embulk/input/pubsub/PluginTask.scala

@@ -0,0 +1,42 @@
+package org.embulk.input.pubsub
+
+import java.util.Optional
+
+import org.embulk.config.{Config, ConfigDefault, ConfigInject, Task}
+import org.embulk.spi.BufferAllocator
+
+trait PluginTask extends Task {
+
+  @Config("project_id")
+  def getProjectId: String
+
+  @Config("subscription_id")
+  def getSubscriptionId: String
+
+  @Config("json_keyfile")
+  def getJsonKeyfile: String
+
+  @Config("num_tasks")
+  @ConfigDefault("1")
+  def getNumTasks: Int
+
+  @Config("max_messages")
+  @ConfigDefault("10")
+  def getMaxMessages: Int
+
+  @Config("checkpoint_basedir")
+  @ConfigDefault("null")
+  def getCheckpointBasedir: Optional[String]
+
+  @Config("checkpoint")
+  @ConfigDefault("null")
+  def getCheckpoint: Optional[String]
+  def setCheckpoint(checkpoint: Optional[String]): Unit
+
+  @Config("payload_encoding")
+  @ConfigDefault("\"string\"")
+  def getPayloadEncoding: String
+
+  @ConfigInject
+  def getBufferAllocator: BufferAllocator
+}

data/src/main/scala/org/embulk/input/pubsub/PubsubBatchSubscriber.scala

@@ -0,0 +1,103 @@
+package org.embulk.input.pubsub
+
+import java.io.FileInputStream
+
+import com.google.api.gax.core.FixedCredentialsProvider
+import com.google.auth.oauth2.GoogleCredentials
+import com.google.cloud.pubsub.v1.stub.{
+  GrpcSubscriberStub,
+  SubscriberStubSettings
+}
+import com.google.protobuf.Empty
+import com.google.pubsub.v1.{
+  AcknowledgeRequest,
+  ProjectSubscriptionName,
+  PullRequest,
+  PullResponse
+}
+import org.embulk.input.pubsub.checkpoint.StoredCheckpoint
+
+import scala.jdk.CollectionConverters._
+import scala.util.{Success, Try}
+
+/**
+  * A subscriber for Cloud Pub/Sub calls batch based pulls with checkpoint.
+  *
+  * @param projectId
+  * @param subscriptionName
+  * @param pathToCredJson
+  */
+case class PubsubBatchSubscriber private (
+    projectId: String,
+    subscriptionName: String,
+    pathToCredJson: String
+) {
+  private val credentials =
+    GoogleCredentials.fromStream(new FileInputStream(pathToCredJson))
+  private val settings = SubscriberStubSettings
+    .newBuilder()
+    .setCredentialsProvider(FixedCredentialsProvider.create(credentials))
+    .setTransportChannelProvider(
+      SubscriberStubSettings.defaultGrpcTransportProviderBuilder().build()
+    )
+    .build()
+
+  def pull(count: Int, checkpointDir: Option[String]): Try[StoredCheckpoint] = {
+    val subscription =
+      ProjectSubscriptionName.of(projectId, subscriptionName).toString
+    val subscriber = GrpcSubscriberStub.create(settings)
+
+    for {
+      res <- pullImpl(subscriber, subscription, count)
+      messages = res.getReceivedMessagesList.asScala
+      checkpoint <- StoredCheckpoint.create(
+        messages.map(_.getMessage).toSeq,
+        checkpointDir
+      )
+      _ <- ackImpl(subscriber, subscription, messages.map(_.getAckId))
+    } yield checkpoint
+  }
+
+  private def pullImpl(
+      subscriber: GrpcSubscriberStub,
+      subscription: String,
+      count: Int
+  ): Try[PullResponse] = {
+    val req = PullRequest
+      .newBuilder()
+      .setSubscription(subscription)
+      .setReturnImmediately(true)
+      .setMaxMessages(count)
+      .build()
+
+    Try(subscriber.pullCallable().call(req))
+  }
+
+  private def ackImpl(
+      subscriber: GrpcSubscriberStub,
+      subscription: String,
+      ackIds: Iterable[String]
+  ): Try[Empty] = {
+    if (ackIds.nonEmpty) {
+      val ack = AcknowledgeRequest
+        .newBuilder()
+        .setSubscription(subscription)
+        .addAllAckIds(ackIds.asJava)
+        .build()
+
+      Try(subscriber.acknowledgeCallable().call(ack))
+    } else {
+      Success(Empty.getDefaultInstance)
+    }
+  }
+
+}
+
+object PubsubBatchSubscriber {
+  def of(task: PluginTask): PubsubBatchSubscriber =
+    PubsubBatchSubscriber(
+      task.getProjectId,
+      task.getSubscriptionId,
+      task.getJsonKeyfile
+    )
+}

data/src/main/scala/org/embulk/input/pubsub/PubsubInputPlugin.scala

@@ -0,0 +1,142 @@
+package org.embulk.input.pubsub
+
+import java.nio.charset.StandardCharsets
+import java.util.{Base64, Optional, List => JList}
+
+import com.fasterxml.jackson.databind.ObjectMapper
+import org.embulk.config.{
+  ConfigDiff,
+  ConfigException,
+  ConfigSource,
+  TaskReport,
+  TaskSource
+}
+import org.embulk.input.pubsub.checkpoint.StoredCheckpoint
+import org.embulk.spi.`type`.Types
+import org.embulk.spi.{
+  DataException,
+  Exec,
+  InputPlugin,
+  PageBuilder,
+  PageOutput,
+  Schema
+}
+import org.embulk.spi.json.JsonParser
+import org.slf4j.LoggerFactory
+
+import scala.jdk.OptionConverters._
+import scala.jdk.CollectionConverters._
+import scala.util.{Failure, Success}
+
+case class PubsubInputPlugin() extends InputPlugin {
+  private val logger = LoggerFactory.getLogger(this.getClass)
+
+  private val jsonParser = new JsonParser()
+  private val objectMapper = new ObjectMapper()
+
+  private val schema = Schema
+    .builder()
+    .add("payload", Types.STRING) // string or base64 encoded bytes
+    .add("attribute", Types.JSON)
+    .build()
+
+  override def transaction(
+      config: ConfigSource,
+      control: InputPlugin.Control
+  ): ConfigDiff = {
+    val task = config.loadConfig(classOf[PluginTask])
+
+    if (!task.getCheckpoint.isPresent) {
+      val sub = PubsubBatchSubscriber.of(task)
+      val checkpoint =
+        sub.pull(task.getMaxMessages, task.getCheckpointBasedir.toScala).get
+      task.setCheckpoint(Optional.of(checkpoint.id))
+
+      logger.info(s"Created a new checkpoint! : ${checkpoint.id}")
+    }
+
+    resume(task.dump(), schema, task.getNumTasks, control)
+  }
+
+  override def resume(
+      taskSource: TaskSource,
+      schema: Schema,
+      taskCount: Int,
+      control: InputPlugin.Control
+  ): ConfigDiff = {
+    control.run(taskSource, schema, taskCount)
+    Exec.newConfigDiff()
+  }
+
+  override def cleanup(
+      taskSource: TaskSource,
+      schema: Schema,
+      taskCount: Int,
+      successTaskReports: JList[TaskReport]
+  ): Unit = {
+    val task = taskSource.loadTask(classOf[PluginTask])
+
+    val checkpointId = task.getCheckpoint.get()
+    val checkpoint =
+      StoredCheckpoint.from(checkpointId, task.getCheckpointBasedir.isPresent)
+    checkpoint match {
+      case Success(sc) =>
+        sc.cleanup match {
+          case Success(_) =>
+          case Failure(e) => logger.error(s"failed to cleanup: ${e.toString}")
+        }
+      case Failure(e) =>
+        logger.error(s"failed to fetch checkpoint: ${e.toString}")
+    }
+  }
+
+  override def run(
+      taskSource: TaskSource,
+      schema: Schema,
+      taskIndex: Int,
+      output: PageOutput
+  ): TaskReport = {
+    val task = taskSource.loadTask(classOf[PluginTask])
+    val allocator = task.getBufferAllocator
+    val pageBuilder = new PageBuilder(allocator, schema, output)
+
+    val encoder = task.getPayloadEncoding match {
+      case "string" =>
+        (data: Array[Byte]) => new String(data, StandardCharsets.UTF_8)
+      case "binary" =>
+        (data: Array[Byte]) => Base64.getEncoder.encodeToString(data)
+      case e => throw new ConfigException(s"unsupported encoding: ${e}")
+    }
+
+    val checkpointId = task.getCheckpoint.get()
+    val checkpoint =
+      StoredCheckpoint.from(checkpointId, task.getCheckpointBasedir.isPresent)
+    val messages = checkpoint match {
+      case Success(cp) => cp.content.getMessagesList.asScala
+      case _ =>
+        throw new DataException(s"unexpected checkpoint state: ${checkpoint}")
+    }
+
+    messages.foreach { msg =>
+      pageBuilder.setString(
+        pageBuilder.getSchema.getColumn(0),
+        encoder(msg.getData.toByteArray)
+      )
+
+      val json = objectMapper.writeValueAsString(msg.getAttributesMap)
+      pageBuilder.setJson(
+        pageBuilder.getSchema.getColumn(1),
+        jsonParser.parse(json)
+      )
+
+      pageBuilder.addRecord()
+    }
+    pageBuilder.finish()
+
+    Exec.newTaskReport()
+  }
+
+  override def guess(config: ConfigSource): ConfigDiff =
+    Exec.newConfigDiff()
+
+}

data/src/main/scala/org/embulk/input/pubsub/checkpoint/StoredCheckpoint.scala

@@ -0,0 +1,123 @@
+package org.embulk.input.pubsub.checkpoint
+
+import java.io.{File, FileInputStream, FileOutputStream}
+
+import com.embulk.input.pubsub.checkpoint.Checkpoint
+import com.google.pubsub.v1.PubsubMessage
+import org.embulk.config.ConfigException
+
+import scala.collection.mutable
+import scala.util.{Failure, Success, Try}
+import scala.jdk.CollectionConverters._
+
+/**
+  * A checkpoint stored in a (maybe)persistent storage.
+  */
+sealed trait StoredCheckpoint {
+  def id: String
+  def content: Checkpoint
+  def cleanup: Try[Unit]
+}
+
+object StoredCheckpoint {
+  def create(
+      messages: Seq[PubsubMessage],
+      dir: Option[String]
+  ): Try[StoredCheckpoint] = {
+    dir match {
+      case Some(d) =>
+        LocalFileStoredCheckpoint.withPersistency(d, messages)
+      case _ =>
+        val content = Checkpoint
+          .newBuilder()
+          .addAllMessages(messages.asJava)
+          .build()
+        Success(MemoryStoredStoredCheckpoint.withoutPersistency(content))
+    }
+  }
+
+  def from(id: String, persistent: Boolean): Try[StoredCheckpoint] = {
+    if (persistent) {
+      LocalFileStoredCheckpoint.from(id)
+    } else {
+      MemoryStoredStoredCheckpoint.from(id)
+    }
+  }
+}
+
+/**
+  * A checkpoint stored in only memory which doesn't have persistence.
+  *
+  * @param id
+  * @param content
+  */
+case class MemoryStoredStoredCheckpoint private (
+    id: String,
+    content: Checkpoint
+) extends StoredCheckpoint {
+  import MemoryStoredStoredCheckpoint._
+
+  override def cleanup: Try[Unit] = {
+    storage.remove(id) match {
+      case Some(_) => Success(())
+      case _ =>
+        Failure(new ConfigException(s"A checkpoint ${id} is not deletable"))
+    }
+  }
+}
+
+object MemoryStoredStoredCheckpoint {
+  private val storage = mutable.Map[String, MemoryStoredStoredCheckpoint]()
+
+  def from(key: String): Try[MemoryStoredStoredCheckpoint] = Try(storage(key))
+
+  def withoutPersistency(content: Checkpoint): StoredCheckpoint = {
+    val id = content.hashCode().toString
+    val checkpoint = MemoryStoredStoredCheckpoint(id, content)
+    storage.put(id, checkpoint)
+    checkpoint
+  }
+
+}
+
+/**
+  * A checkpoint stored in local filesystem.
+  *
+  * @param id
+  * @param content
+  */
+case class LocalFileStoredCheckpoint private (id: String, content: Checkpoint)
+    extends StoredCheckpoint {
+  override def cleanup: Try[Unit] = {
+    for {
+      f <- Try(new File(id))
+      _ <- Try(f.delete)
+    } yield ()
+  }
+}
+
+object LocalFileStoredCheckpoint {
+  def from(path: String): Try[LocalFileStoredCheckpoint] = {
+    for {
+      in <- Try(new FileInputStream(path))
+      c <- Try(Checkpoint.parseFrom(in))
+    } yield LocalFileStoredCheckpoint(path, c)
+  }
+
+  def withPersistency(
+      prefix: String,
+      messages: Seq[PubsubMessage]
+  ): Try[LocalFileStoredCheckpoint] = {
+    val path = s"${prefix}checkpoint-${messages.hashCode().toString}"
+    val content = Checkpoint
+      .newBuilder()
+      .addAllMessages(messages.asJava)
+      .build()
+
+    for {
+      out <- Try(new FileOutputStream(path))
+      _ <- Try(content.writeTo(out))
+      _ <- Try(out.close())
+    } yield LocalFileStoredCheckpoint(path, content)
+  }
+}