Add type partition for bad row partition persistence (closes #146)

aldemirenes · aldemirenes · commit 79e9b915a466 · 2019-09-10T15:53:30.000+03:00
diff --git a/build.sbt b/build.sbt
@@ -39,6 +39,7 @@ lazy val root = project.in(file("."))
       Dependencies.Libraries.json4sJackson,
       Dependencies.Libraries.snowplowTracker,
       Dependencies.Libraries.pureconfig,
+      Dependencies.Libraries.igluCoreJson4s,
       // Scala (test only)
       Dependencies.Libraries.specs2,
       // Thrift (test only)
diff --git a/project/Dependencies.scala b/project/Dependencies.scala
@@ -39,6 +39,7 @@ object Dependencies {
     val scalaz7         = "7.2.13"
     val snowplowTracker = "0.3.0"
     val pureconfig      = "0.8.0"
+    val igluCore        = "0.5.0"
     // Scala (test only)
     val specs2          = "3.9.1"
   }
@@ -74,6 +75,7 @@ object Dependencies {
     val scalaz7          = "org.scalaz"                %% "scalaz-core"               % V.scalaz7
     val snowplowTracker  = "com.snowplowanalytics"     %% "snowplow-scala-tracker"    % V.snowplowTracker
     val pureconfig       = "com.github.pureconfig"     %% "pureconfig"                % V.pureconfig
+    val igluCoreJson4s   = "com.snowplowanalytics"     %% "iglu-core-json4s"          % V.igluCore
     // Scala (test only)
     val specs2           = "org.specs2"                %% "specs2-core"               % V.specs2           % "test"
   }
diff --git a/src/main/scala/com.snowplowanalytics.s3/loader/KinesisS3Emitter.scala b/src/main/scala/com.snowplowanalytics.s3/loader/KinesisS3Emitter.scala
@@ -14,6 +14,8 @@ package com.snowplowanalytics.s3.loader
 
 // Scala
 import scala.collection.JavaConverters._
+import scala.util.Try
+import scala.util.{Success => TrySuccess}
 
 // Java libs
 import java.util.Calendar
@@ -32,10 +34,21 @@ import scala.collection.JavaConversions._
 // Tracker
 import com.snowplowanalytics.snowplow.scalatracker.Tracker
 
+// Json4s
+import org.json4s.jackson.JsonMethods.parse
+
+// Iglu Core
+import com.snowplowanalytics.iglu.core._
+import com.snowplowanalytics.iglu.core.json4s.implicits._
+
+// Scalaz
+import scalaz._
+
 // This project
 import sinks._
 import serializers._
 import model._
+import KinesisS3Emitter._
 
 /**
  * Emitter for flushing Kinesis event data to S3.
@@ -48,19 +61,12 @@ class KinesisS3Emitter(
   badSink: ISink, 
   serializer: ISerializer, 
   maxConnectionTime: Long, 
-  tracker: Option[Tracker]
+  tracker: Option[Tracker],
+  partition: Boolean,
+  partitionErrorDir: String
 ) extends IEmitter[EmitterInput]  {
 
   val s3Emitter = new S3Emitter(s3Config, provider, badSink, maxConnectionTime, tracker)
-  val dateFormat = new SimpleDateFormat("yyyy-MM-dd");
-
-  /**
-   * Determines the filename in S3, which is the corresponding
-   * Kinesis sequence range of records in the file.
-   */
-  protected def getBaseFilename(firstSeq: String, lastSeq: String): String =
-    dateFormat.format(Calendar.getInstance().getTime()) +
-      "-" + firstSeq + "-" + lastSeq
 
   /**
    * Reads items from a buffer and saves them to s3.
@@ -77,7 +83,19 @@ class KinesisS3Emitter(
     s3Emitter.log.info(s"Flushing buffer with ${buffer.getRecords.size} records.")
 
     val records = buffer.getRecords().asScala.toList
-    val baseFilename = getBaseFilename(buffer.getFirstSequenceNumber, buffer.getLastSequenceNumber)
+    if (partition) {
+      partitionWithSchemaKey(records, partitionErrorDir).foldLeft(List[EmitterInput]()) {
+        case (acc, (prefix, l)) =>
+          val baseFilename = getBaseFilename(buffer.getFirstSequenceNumber, buffer.getLastSequenceNumber, Some(prefix.getName))
+          acc ::: emitRecords(l, baseFilename)
+      }
+    } else {
+      val baseFilename = getBaseFilename(buffer.getFirstSequenceNumber, buffer.getLastSequenceNumber)
+      emitRecords(records, baseFilename)
+    }
+  }
+
+  private def emitRecords(records: List[EmitterInput], baseFilename: String) = {
     val serializationResults = serializer.serialize(records, baseFilename)
     val (successes, failures) = serializationResults.results.partition(_.isSuccess)
 
@@ -108,3 +126,65 @@ class KinesisS3Emitter(
   override def fail(records: java.util.List[EmitterInput]): Unit =
     s3Emitter.sendFailures(records)
 }
+
+object KinesisS3Emitter {
+
+  /** Type of row which determined according to schema of self describing data */
+  sealed trait RowType extends Product with Serializable {
+    def getName: String
+  }
+
+  object RowType {
+    /** Represents cases where row type could not be determined
+      * since either row is not valid json or it is not self
+      * describing json
+      */
+    case class PartitionError(errorDir: String) extends RowType {
+      override def getName: String = errorDir
+    }
+
+    /** Represents cases where type of row can be determined successfully
+      * e.g. does have proper schema key
+      */
+    case class SelfDescribing(rowType: String) extends RowType {
+      override def getName: String = rowType
+    }
+
+    case object UnexpectedError extends RowType {
+      override def getName: String = "unexpected_error"
+    }
+  }
+
+  val dateFormat = new SimpleDateFormat("yyyy-MM-dd");
+
+  /**
+    * Determines the filename in S3, which is the corresponding
+    * Kinesis sequence range of records in the file.
+    */
+  private def getBaseFilename(firstSeq: String, lastSeq: String, prefix: Option[String] = None): String =
+    prefix.map(p => if (p.isEmpty) "" else p + "/").getOrElse("") + dateFormat.format(Calendar.getInstance().getTime()) +
+      "-" + firstSeq + "-" + lastSeq
+
+  /**
+    * Assume records are self describing data and group them according
+    * to their schema key. Put records which are not self describing data
+    * to under "old bad row type".
+    */
+  private[loader] def partitionWithSchemaKey(records: List[EmitterInput], errorDir: String) = {
+    records.groupBy {
+      case Success(byteRecord) =>
+        val strRecord = new String(byteRecord, "UTF-8")
+        Try(parse(strRecord)) match {
+          case TrySuccess(e) =>
+            val json = parse(strRecord)
+            val schemaKey = SchemaKey.extract(json)
+            schemaKey.fold(
+              e => RowType.PartitionError(errorDir),
+              k => RowType.SelfDescribing(s"${k.vendor}.${k.name}")
+            )
+          case _ => RowType.PartitionError(errorDir)
+        }
+      case _ => RowType.UnexpectedError
+    }
+  }
+}
diff --git a/src/main/scala/com.snowplowanalytics.s3/loader/KinesisS3Pipeline.scala b/src/main/scala/com.snowplowanalytics.s3/loader/KinesisS3Pipeline.scala
@@ -28,9 +28,9 @@ import model._
 /**
  * S3Pipeline class sets up the Emitter/Buffer/Transformer/Filter
  */
-class KinesisS3Pipeline(s3Config: S3Config, badSink: ISink, serializer: ISerializer, maxConnectionTime: Long, tracker: Option[Tracker]) extends IKinesisConnectorPipeline[ValidatedRecord, EmitterInput] {
+class KinesisS3Pipeline(s3Config: S3Config, badSink: ISink, serializer: ISerializer, maxConnectionTime: Long, tracker: Option[Tracker], partition: Boolean, partitionErrorDir: String) extends IKinesisConnectorPipeline[ValidatedRecord, EmitterInput] {
 
-  override def getEmitter(configuration: KinesisConnectorConfiguration) = new KinesisS3Emitter(s3Config, configuration.AWS_CREDENTIALS_PROVIDER, badSink, serializer, maxConnectionTime, tracker)
+  override def getEmitter(configuration: KinesisConnectorConfiguration) = new KinesisS3Emitter(s3Config, configuration.AWS_CREDENTIALS_PROVIDER, badSink, serializer, maxConnectionTime, tracker, partition, partitionErrorDir)
 
   override def getBuffer(configuration: KinesisConnectorConfiguration) = new BasicMemoryBuffer[ValidatedRecord](configuration)
 
diff --git a/src/main/scala/com.snowplowanalytics.s3/loader/KinesisSourceExecutor.scala b/src/main/scala/com.snowplowanalytics.s3/loader/KinesisSourceExecutor.scala
@@ -49,7 +49,9 @@ class KinesisSourceExecutor(
   badSink: ISink,
   serializer: ISerializer,
   maxConnectionTime: Long,
-  tracker: Option[Tracker]
+  tracker: Option[Tracker],
+  partition: Boolean,
+  partitionErrorDir: String
 ) extends KinesisConnectorExecutorBase[ValidatedRecord, EmitterInput] {
 
   val LOG = LoggerFactory.getLogger(getClass)
@@ -113,5 +115,5 @@ class KinesisSourceExecutor(
   initialize(config, null)
 
   override def getKinesisConnectorRecordProcessorFactory =
-    new KinesisConnectorRecordProcessorFactory[ValidatedRecord, EmitterInput](new KinesisS3Pipeline(s3Config, badSink, serializer, maxConnectionTime, tracker), config)
+    new KinesisConnectorRecordProcessorFactory[ValidatedRecord, EmitterInput](new KinesisS3Pipeline(s3Config, badSink, serializer, maxConnectionTime, tracker, partition, partitionErrorDir), config)
 }
diff --git a/src/main/scala/com.snowplowanalytics.s3/loader/SinkApp.scala b/src/main/scala/com.snowplowanalytics.s3/loader/SinkApp.scala
@@ -149,7 +149,9 @@ object SinkApp {
                                     badSink,
                                     serializer,
                                     maxConnectionTime,
-                                    tracker
+                                    tracker,
+                                    config.partition.getOrElse(false),
+                                    config.partitionErrorDir.getOrElse("")
                                    ).success
         // Read records from NSQ
         case "nsq" =>
diff --git a/src/main/scala/com.snowplowanalytics.s3/loader/model.scala b/src/main/scala/com.snowplowanalytics.s3/loader/model.scala
@@ -91,6 +91,8 @@ package model {
     kinesis: KinesisConfig,
     streams: StreamsConfig,
     s3: S3Config,
-    monitoring: Option[MonitoringConfig]
+    monitoring: Option[MonitoringConfig],
+    partition: Option[Boolean],
+    partitionErrorDir: Option[String]
   )
 }
diff --git a/src/test/scala/com/snowplowanalytics/s3/loader/KinesisS3EmitterSpec.scala b/src/test/scala/com/snowplowanalytics/s3/loader/KinesisS3EmitterSpec.scala
@@ -0,0 +1,112 @@
+/*
+ * Copyright (c) 2014-2019 Snowplow Analytics Ltd. All rights reserved.
+ *
+ * This program is licensed to you under the Apache License Version 2.0,
+ * and you may not use this file except in compliance with the Apache License Version 2.0.
+ * You may obtain a copy of the Apache License Version 2.0 at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the Apache License Version 2.0 is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the Apache License Version 2.0 for the specific language governing permissions and limitations there under.
+ */
+package com.snowplowanalytics.s3.loader
+
+import scalaz._
+
+import org.specs2.mutable.Specification
+
+class KinesisS3EmitterSpec extends Specification {
+
+  "KinesisS3Emitter" should {
+
+    "partition records correctly according to schema key" in {
+      val dataType11 =
+        """
+          | {
+          |   "schema": "iglu:com.acme1/example1/jsonschema/2-0-1",
+          |   "data": "data1"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val dataType21 =
+        """
+          | {
+          |   "schema": "iglu:com.acme1/example2/jsonschema/2-0-0",
+          |   "data": "data1"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val dataType22 =
+        """
+          | {
+          |   "schema": "iglu:com.acme1/example2/jsonschema/2-0-1",
+          |   "data": "data2"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val dataType31 =
+        """
+          | {
+          |   "schema": "iglu:com.acme2/example1/jsonschema/2-0-0",
+          |   "data": "data1"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val dataType32 =
+        """
+          | {
+          |   "schema": "iglu:com.acme2/example1/jsonschema/2-0-1",
+          |   "data": "data2"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val dataType33 =
+        """
+          | {
+          |   "schema": "iglu:com.acme2/example1/jsonschema/2-0-1",
+          |   "data": "data3"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val nonSelfDescribingJson =
+        """
+          | {
+          |   "data": "data",
+          |   "key": "value"
+          | }
+        """.stripMargin.getBytes("UTF-8")
+
+      val nonJsonData = "nonJsonData".getBytes("UTF-8")
+
+      val failure = FailedRecord(List("error1", "error2"), "line")
+
+      val records = List(
+        Success(dataType11),
+        Success(dataType21),
+        Success(dataType22),
+        Success(dataType31),
+        Success(dataType32),
+        Success(dataType33),
+        Failure(failure),
+        Failure(failure),
+        Success(nonSelfDescribingJson),
+        Success(nonJsonData)
+      )
+
+      val res = KinesisS3Emitter.partitionWithSchemaKey(records, "partition_error_dir")
+        .map { case (rowType, l) => (rowType.getName, l)}
+
+      val expected = Map(
+        "com.acme1.example1" -> List(Success(dataType11)),
+        "com.acme1.example2" -> List(Success(dataType21), Success(dataType22)),
+        "com.acme2.example1" -> List(Success(dataType31), Success(dataType32), Success(dataType33)),
+        "partition_error_dir" -> List(Success(nonSelfDescribingJson), Success(nonJsonData)),
+        "unexpected_error" -> List(Failure(failure), Failure(failure))
+      )
+
+      res must beEqualTo(expected)
+    }
+  }
+
+}

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,8 @@ package model {`
`91`	`91`	`kinesis: KinesisConfig,`
`92`	`92`	`streams: StreamsConfig,`
`93`	`93`	`s3: S3Config,`
`94`		`- monitoring: Option[MonitoringConfig]`
	`94`	`+ monitoring: Option[MonitoringConfig],`
	`95`	`+ partition: Option[Boolean],`
	`96`	`+ partitionErrorDir: Option[String]`
`95`	`97`	`)`
`96`	`98`	`}`