From b85420afed5995c78304846e7ed56bb2b98df545 Mon Sep 17 00:00:00 2001
From: Josh Lee <joshleecreates@users.noreply.github.com>
Date: Tue, 13 Jun 2023 07:13:01 -0400
Subject: [PATCH] Don't record DB statements without sanitizaiton (#166)

* Don't record DB statements without sanitizaiton

This adds an option to OpentelemetryEcto.setup/1 that allows a query
sanitization function to be provided. If it is not provided, queries
are not captured (this is the default).

* test that db.statement isnt present unless query sanitizer is configured

* rename option to `:db_statement`

* run mix format
---
 .../lib/opentelemetry_ecto.ex                 | 23 ++++++++++++--
 .../test/opentelemetry_ecto_test.exs          | 30 +++++++++++++++----
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex b/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex
index 0635a4f..5a58ee0 100644
--- a/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex
+++ b/instrumentation/opentelemetry_ecto/lib/opentelemetry_ecto.ex
@@ -31,7 +31,6 @@ defmodule OpentelemetryEcto do
 
     * `:time_unit` - a time unit used to convert the values of query phase
       timings, defaults to `:microsecond`. See `System.convert_time_unit/3`
-
     * `:span_prefix` - the first part of the span name, as a `String.t`,
       defaults to the concatenation of the event name with periods, e.g.
       `"blog.repo.query"`. This will always be followed with a colon and the
@@ -39,6 +38,12 @@ defmodule OpentelemetryEcto do
     * `:additional_attributes` - additional attributes to include in the span. If there
       are conflits with default provided attributes, the ones provided with
       this config will have precedence.
+    * `:db_statement` - :disabled (default) | :enabled | fun
+      Whether or not to include db statements.
+      Optionally provide a function that takes a query string and returns a
+      sanitized version of it. This is useful for removing sensitive information from the
+      query string. Unless this option is `:enabled` or a function,
+      query statements will not be recorded on spans.
   """
   def setup(event_prefix, config \\ []) do
     event = event_prefix ++ [:query]
@@ -89,7 +94,6 @@ defmodule OpentelemetryEcto do
     # net.peer.name or net.peer.ip and net.peer.port
     base_attributes = %{
       "db.type": db_type,
-      "db.statement": query,
       source: source,
       "db.instance": database,
       "db.name": database,
@@ -97,6 +101,21 @@ defmodule OpentelemetryEcto do
       "total_time_#{time_unit}s": System.convert_time_unit(total_time, :native, time_unit)
     }
 
+    base_attributes =
+      case Keyword.fetch(config, :db_statement) do
+        {:ok, :enabled} ->
+          Map.put(base_attributes, :"db.statement", query)
+
+        {:ok, :disabled} ->
+          base_attributes
+
+        {:ok, sanitizer} ->
+          Map.put(base_attributes, :"db.statement", sanitizer.(query))
+
+        :error ->
+          base_attributes
+      end
+
     attributes =
       measurements
       |> Enum.reduce(%{}, fn
diff --git a/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs b/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs
index 082b522..77b17fe 100644
--- a/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs
+++ b/instrumentation/opentelemetry_ecto/test/opentelemetry_ecto_test.exs
@@ -50,8 +50,6 @@ defmodule OpentelemetryEctoTest do
 
     assert %{
              "db.instance": "opentelemetry_ecto_test",
-             "db.name": "opentelemetry_ecto_test",
-             "db.statement": "SELECT u0.\"id\", u0.\"email\" FROM \"users\" AS u0",
              "db.type": :sql,
              "db.url": "ecto://localhost",
              decode_time_microseconds: _,
@@ -62,7 +60,31 @@ defmodule OpentelemetryEctoTest do
            } = :otel_attributes.map(attributes)
   end
 
-  test "include additionaL_attributes" do
+  test "exclude unsantized query" do
+    attach_handler()
+    Repo.all(User)
+
+    assert_receive {:span, span(attributes: attributes)}
+    assert !Map.has_key?(:otel_attributes.map(attributes), :"db.statement")
+  end
+
+  test "include unsanitized query when enabled" do
+    attach_handler(db_statement: :enabled)
+    Repo.all(User)
+
+    assert_receive {:span, span(attributes: attributes)}
+    assert %{"db.statement": "SELECT u0.\"id\", u0.\"email\" FROM \"users\" AS u0"} = :otel_attributes.map(attributes)
+  end
+
+  test "include santized query with sanitizer function" do
+    attach_handler(db_statement: fn str -> String.replace(str, "SELECT", "") end)
+    Repo.all(User)
+
+    assert_receive {:span, span(attributes: attributes)}
+    assert %{"db.statement": " u0.\"id\", u0.\"email\" FROM \"users\" AS u0"} = :otel_attributes.map(attributes)
+  end
+
+  test "include additional_attributes" do
     attach_handler(additional_attributes: %{"config.attribute": "special value", "db.instance": "my_instance"})
     Repo.all(User)
 
@@ -83,8 +105,6 @@ defmodule OpentelemetryEctoTest do
 
     assert %{
              "db.instance": "opentelemetry_ecto_test",
-             "db.name": "opentelemetry_ecto_test",
-             "db.statement": "SELECT p0.\"id\", p0.\"body\", p0.\"user_id\" FROM \"posts\" AS p0",
              "db.type": :sql,
              "db.url": "ecto://localhost",
              decode_time_milliseconds: _,