Adding parallelism to sink

epiphanous · Feb 24, 2023 · 0f15d71 · 0f15d71
1 parent 22966be
commit 0f15d71
Show file tree

Hide file tree

Showing 11 changed files with 21 additions and 8 deletions.
diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/CassandraSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/CassandraSinkConfig.scala
@@ -62,6 +62,7 @@ case class CassandraSinkConfig[ADT <: FlinkEvent](
       .addSink(new CassandraScalaProductSink[E](query, clusterBuilder))
       .uid(label)
       .name(label)
+      .setParallelism(parallelism)
   }
 
   override def getAvroSink[
@@ -82,5 +83,6 @@ case class CassandraSinkConfig[ADT <: FlinkEvent](
       )
       .uid(label)
       .name(label)
+      .setParallelism(parallelism)
 
 }
diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/ElasticsearchSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/ElasticsearchSinkConfig.scala
@@ -94,7 +94,7 @@ case class ElasticsearchSinkConfig[ADT <: FlinkEvent](
     bulkFlushMaxActions.foreach(esb.setBulkFlushMaxActions)
     bulkFlushMaxSizeMb.foreach(esb.setBulkFlushMaxSizeMb)
     bulkFlushIntervalMs.foreach(esb.setBulkFlushInterval)
-    dataStream.sinkTo(esb.build()).uid(label).name(label)
+    dataStream.sinkTo(esb.build()).uid(label).name(label).setParallelism(parallelism)
   }
 
   override def getSink[E <: ADT: TypeInformation](

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/FileSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/FileSinkConfig.scala
@@ -119,7 +119,7 @@ case class FileSinkConfig[ADT <: FlinkEvent](
         .withRollingPolicy(getCheckpointRollingPolicy)
         .withOutputFileConfig(getOutputFileConfig)
         .build()
-    )
+    ).setParallelism(parallelism)
 
   /** Create an bulk avro parquet file sink and send the data stream to it.
     * @param dataStream
@@ -165,7 +165,7 @@ case class FileSinkConfig[ADT <: FlinkEvent](
           s"Invalid format for getAvroSink: $format"
         )
     }
-    dataStream.sinkTo(sink)
+    dataStream.sinkTo(sink).setParallelism(parallelism)
   }
 
   def getBucketCheckInterval: Long =

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/FirehoseSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/FirehoseSinkConfig.scala
@@ -99,7 +99,7 @@ case class FirehoseSinkConfig[ADT <: FlinkEvent: TypeInformation](
         .map(k.setMaxBatchSizeInBytes)
         .getOrElse(k)
     }.build()
-    dataStream.sinkTo(kfs)
+    dataStream.sinkTo(kfs).setParallelism(parallelism)
   }
 
   def getSerializationSchema[E <: ADT: TypeInformation]

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/JdbcSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/JdbcSinkConfig.scala
@@ -686,6 +686,7 @@ case class JdbcSinkConfig[ADT <: FlinkEvent](
       )
       .uid(label)
       .name(label)
+      .setParallelism(parallelism)
   }
 
   override def getSink[E <: ADT: TypeInformation](

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/KafkaSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/KafkaSinkConfig.scala
@@ -139,11 +139,11 @@ case class KafkaSinkConfig[ADT <: FlinkEvent: TypeInformation](
       E <: ADT with EmbeddedAvroRecord[A]: TypeInformation,
       A <: GenericRecord: TypeInformation](
       dataStream: DataStream[E]): DataStreamSink[E] =
-    dataStream.sinkTo(_getSink[E](getAvroSerializationSchema[E, A]))
+    dataStream.sinkTo(_getSink[E](getAvroSerializationSchema[E, A])).setParallelism(parallelism)
 
   override def getSink[E <: ADT: TypeInformation](
       dataStream: DataStream[E]): DataStreamSink[E] =
-    dataStream.sinkTo(_getSink[E](getSerializationSchema[E]))
+    dataStream.sinkTo(_getSink[E](getSerializationSchema[E])).setParallelism(parallelism)
 
   def _getSink[E <: ADT: TypeInformation](
       serializer: KafkaRecordSerializationSchema[E]): KafkaSink[E] =

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/KinesisSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/KinesisSinkConfig.scala
@@ -79,6 +79,7 @@ case class KinesisSinkConfig[ADT <: FlinkEvent: TypeInformation](
       .sinkTo(ks)
       .uid(label)
       .name(label)
+      .setParallelism(parallelism)
   }
 
   override def getSink[E <: ADT: TypeInformation](

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/RabbitMQSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/RabbitMQSinkConfig.scala
@@ -58,7 +58,7 @@ case class RabbitMQSinkConfig[ADT <: FlinkEvent: TypeInformation](
             )
         }
     }
-    dataStream.addSink(sink).uid(label).name(label)
+    dataStream.addSink(sink).uid(label).name(label).setParallelism(parallelism)
   }
 
   override def getSink[E <: ADT: TypeInformation](

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/SinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/SinkConfig.scala
@@ -50,6 +50,10 @@ trait SinkConfig[ADT <: FlinkEvent] extends LazyLogging {
 
   lazy val label: String = s"${connector.entryName.toLowerCase}/$name"
 
+  lazy val parallelism: Int = config
+    .getIntOpt(pfx("parallelism"))
+    .getOrElse(config.globalParallelism)
+
   def getSink[E <: ADT: TypeInformation](
       dataStream: DataStream[E]): DataStreamSink[E]
 

diff --git a/src/main/scala/io/epiphanous/flinkrunner/model/sink/SocketSinkConfig.scala b/src/main/scala/io/epiphanous/flinkrunner/model/sink/SocketSinkConfig.scala
@@ -84,5 +84,5 @@ case class SocketSinkConfig[ADT <: FlinkEvent](
         maxRetries,
         autoFlush
       )
-    )
+    ).setParallelism(parallelism)
 }
diff --git a/src/test/scala/io/epiphanous/flinkrunner/model/source/KinesisSourceConfigSpec.scala b/src/test/scala/io/epiphanous/flinkrunner/model/source/KinesisSourceConfigSpec.scala
@@ -94,6 +94,11 @@ class KinesisSourceConfigSpec extends PropSpec {
     defaultConfigPlus("efo.enabled = false").useEfo shouldBe false
   }
 
+  property("parallelism property") {
+    defaultConfigPlus("parallelism = 10").parallelism shouldBe 10
+    defaultConfigPlus("parallelism = 10.5").parallelism shouldBe 10
+  }
+
   property("missing stream property") {
     the[Exception] thrownBy noProvidedConfig should have message "kinesis source kinesis-test is missing required 'stream' property"
   }