fix: add support for dataloader kv sources

instrumentation/opentelemetry_dataloader/lib/opentelemetry_dataloader.ex

@@ -74,11 +74,30 @@ defmodule OpentelemetryDataloader do
       OpenTelemetry.Ctx.attach(parent_ctx)
     end
 
-    {batch_name, _batch_args} = metadata.batch
-
-    attributes = %{
-      "dataloader.batch_key" => inspect(batch_name)
-    }
+    attributes =
+      case metadata.batch do
+        # Ecto source
+        {batch_name, _batch_args} ->
+          %{
+            "dataloader.batch_key" => inspect(batch_name)
+          }
+
+        # KV source with atom key
+        batch_name when is_atom(batch_name) ->
+          %{
+            "dataloader.batch_key" => to_string(batch_name)
+          }
+
+        # KV source with binary key
+        batch_name when is_binary(batch_name) ->
+          %{
+            "dataloader.batch_key" => batch_name
+          }
+
+        # Unknown source type
+        _ ->
+          %{}
+      end
 
     OpentelemetryTelemetry.start_telemetry_span(
       config.tracer_id,

instrumentation/opentelemetry_dataloader/test/opentelemetry_dataloader_test.exs

@@ -66,4 +66,98 @@ defmodule OpentelemetryDataloaderTest do
     assert %{"dataloader.batch_key" => key} = :otel_attributes.map(attributes)
     assert key =~ ~r/OpentelemetryDataloader.TestModels.Post/
   end
+
+  test "captures dataloader kv source events with atom keys" do
+    OpentelemetryDataloader.setup()
+
+    loader =
+      Dataloader.new()
+      |> Dataloader.add_source(:kv, Dataloader.KV.new(&kv_query/2))
+      |> Dataloader.load(:kv, :test_maps, 1)
+      |> Dataloader.load_many(:kv, :test_maps, [4, 9])
+
+    Dataloader.run(loader)
+
+    assert_receive {:span,
+                    span(
+                      name: "dataloader.run",
+                      attributes: attributes,
+                      kind: :client
+                    )}
+
+    assert %{} = :otel_attributes.map(attributes)
+
+    assert_receive {:span,
+                    span(
+                      name: "dataloader.batch",
+                      attributes: attributes,
+                      kind: :client
+                    )}
+
+    assert %{"dataloader.batch_key" => "test_maps"} = :otel_attributes.map(attributes)
+  end
+
+  test "captures dataloader kv source events with modules" do
+    OpentelemetryDataloader.setup()
+
+    loader =
+      Dataloader.new()
+      |> Dataloader.add_source(:kv, Dataloader.KV.new(&kv_query/2))
+      |> Dataloader.load(:kv, TestModule, 1)
+      |> Dataloader.load_many(:kv, TestModule, [1, 3])
+
+    Dataloader.run(loader)
+
+    assert_receive {:span,
+                    span(
+                      name: "dataloader.run",
+                      attributes: attributes,
+                      kind: :client
+                    )}
+
+    assert %{} = :otel_attributes.map(attributes)
+
+    assert_receive {:span,
+                    span(
+                      name: "dataloader.batch",
+                      attributes: attributes,
+                      kind: :client
+                    )}
+
+    assert %{"dataloader.batch_key" => "Elixir.TestModule"} = :otel_attributes.map(attributes)
+  end
+
+  test "captures dataloader kv source events with binary names" do
+    OpentelemetryDataloader.setup()
+
+    loader =
+      Dataloader.new()
+      |> Dataloader.add_source(:kv, Dataloader.KV.new(&kv_query/2))
+      |> Dataloader.load(:kv, "test_keys", 1)
+      |> Dataloader.load_many(:kv, "test_keys", [1, 3])
+
+    Dataloader.run(loader)
+
+    assert_receive {:span,
+                    span(
+                      name: "dataloader.run",
+                      attributes: attributes,
+                      kind: :client
+                    )}
+
+    assert %{} = :otel_attributes.map(attributes)
+
+    assert_receive {:span,
+                    span(
+                      name: "dataloader.batch",
+                      attributes: attributes,
+                      kind: :client
+                    )}
+
+    assert %{"dataloader.batch_key" => "test_keys"} = :otel_attributes.map(attributes)
+  end
+
+  defp kv_query(:test_maps, ids), do: Enum.map(ids, &%{id: &1})
+  defp kv_query(TestModule, ids), do: Enum.map(ids, &%{id: &1})
+  defp kv_query("test_keys", ids), do: Enum.map(ids, & &1)
 end