Use behaviour for KafkaUtils

lib/kafkaex_lag_exporter/kafka_utils.ex

@@ -1,14 +1,18 @@
 # source code taken from https://github.com/reachfh/brod_group_subscriber_example
 
 defmodule KafkaexLagExporter.KafkaUtils do
+  @behaviour KafkaexLagExporter.KafkaUtils.Behaviour
+
   @moduledoc "Utility functions for dealing with Kafka"
 
+  alias KafkaexLagExporter.KafkaWrapper.Behaviour, as: KafkaWrapper
+
+  require Logger
+
   @default_client :client1
 
-  @type endpoint() :: {host :: atom(), port :: non_neg_integer()}
-
   def connection, do: connection(@default_client)
-  @spec connection(atom) :: {list({charlist, non_neg_integer}), Keyword.t()}
+  @impl true
   def connection(client) do
     clients = Application.get_env(:brod, :clients)
     config = clients[client]
@@ -27,31 +31,31 @@ defmodule KafkaexLagExporter.KafkaUtils do
     {endpoints, sock_opts}
   end
 
-  @spec resolve_offsets(binary(), :earliest | :latest, atom()) ::
+  @impl true
-          list({non_neg_integer(), integer()})
   def resolve_offsets(topic, type, client) do
     {endpoints, sock_opts} = connection(client)
 
-    {:ok, partitions_count} = :brod.get_partitions_count(client, topic)
+    {:ok, partitions_count} = KafkaWrapper.get_partitions_count(client, topic)
 
     for i <- Range.new(0, partitions_count - 1),
-        {:ok, offset} = :brod.resolve_offset(endpoints, topic, i, type, sock_opts) do
+        {:ok, offset} =
+          KafkaWrapper.resolve_offset(endpoints, topic, i, type, sock_opts) do
       {i, offset}
     end
   end
 
-  @spec fetch_committed_offsets(binary(), binary(), atom()) ::
+  @impl true
-          {non_neg_integer(), non_neg_integer()}
   def fetch_committed_offsets(_topic, consumer_group, client) do
     {endpoints, sock_opts} = connection(client)
-    {:ok, response} = :brod.fetch_committed_offsets(endpoints, sock_opts, consumer_group)
+
+    {:ok, response} = KafkaWrapper.fetch_committed_offsets(endpoints, sock_opts, consumer_group)
 
     for r <- response,
         pr <- r[:partitions],
         do: {pr[:partition_index], pr[:committed_offset]}
   end
 
-  @spec lag(binary(), binary(), atom()) :: list({non_neg_integer(), integer()})
+  @impl true
   def lag(topic, consumer_group, client) do
     offsets =
       resolve_offsets(topic, :latest, client)
@@ -66,10 +70,37 @@ defmodule KafkaexLagExporter.KafkaUtils do
     end
   end
 
-  @spec lag_total(binary(), binary(), atom()) :: non_neg_integer()
+  @impl true
   def lag_total(topic, consumer_group, client) do
     for {_part, recs} <- lag(topic, consumer_group, client), reduce: 0 do
       acc -> acc + recs
     end
   end
+
+  @impl true
+  def get_consumer_group_names({host, port}) do
+    [{_, groups} | _] = KafkaWrapper.list_all_groups([{host, port}], [])
+
+    groups
+    |> Enum.filter(fn {_, _, protocol} -> protocol == "consumer" end)
+    |> Enum.map(fn {_, group_name, "consumer"} -> group_name end)
+  end
+
+  @impl true
+  def topic_names_for_consumer_groups(endpoint, list \\ [], consumer_group_names) do
+    KafkaWrapper.describe_groups(endpoint, list, consumer_group_names)
+    |> get_topic_names_for_consumer_groups
+  end
+
+  defp get_topic_names_for_consumer_groups({:ok, group_descriptions}) do
+    group_descriptions
+    |> Enum.map(fn %{group_id: consumer_group, members: members} -> [consumer_group, members] end)
+    |> Enum.map(fn [consumer_group, members] -> {consumer_group, get_topic_names(members)} end)
+  end
+
+  defp get_topic_names(members) do
+    Enum.flat_map(members, fn member ->
+      KafkaexLagExporter.TopicNameParser.parse_topic_names(member.member_assignment)
+    end)
+  end
 end