handle partition_by option

Author: tkusumo

README.md

@@ -60,6 +60,11 @@ KinesisClient.Stream.start_link(opts)
 `MyShardConsumer` needs to implement the `Broadway` behaviour. You will want to
 start the `KinesisClient.Stream` in your application's supervision tree.
 
+## partition_by option
+
+If you want to include the partition_by option to the Broadway pipeline 
+then you need to implement a partition_by/1 function in the consumer `MyShardConsumer`.
+
 
 ## Things to keep in mind...
 

lib/kinesis_client/stream/shard/pipeline.ex

@@ -46,18 +46,18 @@ defmodule KinesisClient.Stream.Shard.Pipeline do
       |> Keyword.get(:pipeline_context, %{})
       |> Map.put(:shard_consumer, opts[:shard_consumer])
 
-    pipeline_opts = [
-      name: register_name(__MODULE__, opts[:app_name], opts[:stream_name], [opts[:shard_id]]),
-      producer: [
-        module: {Producer, producer_opts},
-        concurrency: 1
-      ],
-      context: pipeline_context,
-      processors: processor_opts,
-      batchers: batcher_opts
-    ]
-
-    pipeline_opts = optional_kw(pipeline_opts, :partition_by, Keyword.get(opts, :partition_by))
+    pipeline_opts =
+      [
+        name: register_name(__MODULE__, opts[:app_name], opts[:stream_name], [opts[:shard_id]]),
+        producer: [
+          module: {Producer, producer_opts},
+          concurrency: 1
+        ],
+        context: pipeline_context,
+        processors: processor_opts,
+        batchers: batcher_opts
+      ]
+      |> partition_by(opts)
 
     Broadway.start_link(__MODULE__, pipeline_opts)
   end
@@ -108,6 +108,17 @@ defmodule KinesisClient.Stream.Shard.Pipeline do
     end
   end
 
+  @impl Broadway
+  def prepare_messages(messages, ctx) do
+    module = Map.get(ctx, :shard_consumer)
+
+    if function_exported?(module, :prepare_messages, 2) do
+      module.prepare_messages(messages, ctx)
+    else
+      messages
+    end
+  end
+
   @impl Broadway
   def handle_message(processor, msg, ctx) do
     module = Map.get(ctx, :shard_consumer)
@@ -127,4 +138,16 @@ defmodule KinesisClient.Stream.Shard.Pipeline do
 
     module.handle_failed(messages, context)
   end
+
+  defp partition_by(pipeline_opts, opts) do
+    shard_consumer =
+      opts[:shard_consumer]
+      |> Code.ensure_loaded!()
+
+    if function_exported?(shard_consumer, :partition_by, 1) do
+      Keyword.put(pipeline_opts, :partition_by, &shard_consumer.partition_by/1)
+    else
+      pipeline_opts
+    end
+  end
 end

test/kinesis_client/stream/coordinator_test.exs

@@ -249,6 +249,7 @@ defmodule KinesisClient.Stream.CoordinatorTest do
       notify_pid: self(),
       kinesis_opts: [adapter: KinesisClient.KinesisMock],
       shard_args: [
+        shard_consumer: __MODULE__,
         app_name: app_name,
         coordinator_name: coordinator_name,
         lease_owner: worker_ref(),

test/kinesis_client/stream/shard/lease_test.exs

@@ -27,7 +27,6 @@ defmodule KinesisClient.Stream.Shard.LeaseTest do
 
     assert_receive {:initialized, lease_state}, 1_000
 
-    inspect(lease_state, label: "lease_state")
     assert lease_state.lease_holder == true
     assert lease_state.lease_count == 1
 

test/kinesis_client/stream/shard/pipeline_test.exs

@@ -71,7 +71,12 @@ defmodule KinesisClient.Stream.Shard.PipelineTest do
     stream_name = "pipeline-test-stream"
     shard_id = "shard-1"
 
-    opts = [app_name: app_name, stream_name: stream_name, shard_id: shard_id]
+    opts = [
+      app_name: app_name,
+      stream_name: stream_name,
+      shard_id: shard_id,
+      shard_consumer: __MODULE__
+    ]
 
     {:ok, pid} = start_supervised({Pipeline, opts})