From 243664c4973952f459c85116e7d7c52df4fe1927 Mon Sep 17 00:00:00 2001
From: Alex Evanczuk <alexevanczuk@gmail.com>
Date: Fri, 1 Dec 2023 08:34:02 -0500
Subject: [PATCH] Add ability to configure error_handler in pg

---
 docs/GettingStarted.md                                |  1 +
 .../tracing/contrib/pg/configuration/settings.rb      |  5 +++++
 lib/datadog/tracing/contrib/pg/instrumentation.rb     |  2 ++
 spec/datadog/tracing/contrib/pg/patcher_spec.rb       | 11 +++++++++++
 4 files changed, 19 insertions(+)

diff --git a/docs/GettingStarted.md b/docs/GettingStarted.md
index 884e09e770c..e6f580c8f4e 100644
--- a/docs/GettingStarted.md
+++ b/docs/GettingStarted.md
@@ -1315,6 +1315,7 @@ end
 | `service_name`        | `DD_TRACE_PG_SERVICE_NAME` | Name of application running the `pg` instrumentation. May be overridden by `global_default_service_name`. [See *Additional Configuration* for more details](#additional-configuration)                                                                                                                                                                                  | `pg`         |
 | `peer_service`        | `DD_TRACE_PG_PEER_SERVICE` | Name of external service the application connects to                                                                                                                                                                                                                                                                                                                    | `nil`        |
 | `comment_propagation` | `DD_DBM_PROPAGATION_MODE`  | SQL comment propagation mode  for database monitoring. <br />(example: `disabled` \| `service`\| `full`). <br /><br />**Important**: *Note that enabling sql comment propagation results in potentially confidential data (service names) being stored in the databases which can then be accessed by other 3rd parties that have been granted access to the database.* | `'disabled'` |
+| `error_handler` | Custom error handler invoked when a job raises an error. Provided `span` and `error` as arguments. Sets error on the span by default. Useful for ignoring errors from Postgres that are handled at the application level. | `proc { |span, error| span.set_error(error) unless span.nil? }` |
 
 ### Presto
 
diff --git a/lib/datadog/tracing/contrib/pg/configuration/settings.rb b/lib/datadog/tracing/contrib/pg/configuration/settings.rb
index 272d37e155a..b6e16cb297f 100644
--- a/lib/datadog/tracing/contrib/pg/configuration/settings.rb
+++ b/lib/datadog/tracing/contrib/pg/configuration/settings.rb
@@ -25,6 +25,11 @@ class Settings < Contrib::Configuration::Settings
               o.default false
             end
 
+            option :error_handler do |o|
+              o.type :proc
+              o.default_proc(&Tracing::SpanOperation::Events::DEFAULT_ON_ERROR)
+            end
+
             option :analytics_sample_rate do |o|
               o.type :float
               o.env Ext::ENV_ANALYTICS_SAMPLE_RATE
diff --git a/lib/datadog/tracing/contrib/pg/instrumentation.rb b/lib/datadog/tracing/contrib/pg/instrumentation.rb
index 087ebc0184d..c03b215bfaa 100644
--- a/lib/datadog/tracing/contrib/pg/instrumentation.rb
+++ b/lib/datadog/tracing/contrib/pg/instrumentation.rb
@@ -81,10 +81,12 @@ def sync_exec_prepared(statement_name, params, *args, &block)
 
             def trace(name, sql: nil, statement_name: nil, block: nil)
               service = Datadog.configuration_for(self, :service_name) || datadog_configuration[:service_name]
+              error_handler = datadog_configuration[:error_handler]
               resource = statement_name || sql
 
               Tracing.trace(
                 name,
+                on_error: error_handler,
                 service: service,
                 resource: resource,
                 type: Tracing::Metadata::Ext::SQL::TYPE
diff --git a/spec/datadog/tracing/contrib/pg/patcher_spec.rb b/spec/datadog/tracing/contrib/pg/patcher_spec.rb
index cd03a6d4901..0a039a137d1 100644
--- a/spec/datadog/tracing/contrib/pg/patcher_spec.rb
+++ b/spec/datadog/tracing/contrib/pg/patcher_spec.rb
@@ -160,6 +160,17 @@
           it_behaves_like 'configured peer service span', 'DD_TRACE_PG_PEER_SERVICE', error: PG::Error do
             let(:configuration_options) { {} }
           end
+
+          context 'when there is custom error handling' do
+            let(:configuration_options) { { error_handler: error_handler } }
+            let(:error_handler) { ->(_span, _error) { false } }
+
+            it 'calls the error handler' do
+              expect { exec }.to raise_error(PG::Error)
+              expect(spans.count).to eq(1)
+              expect(span).to_not have_error
+            end
+          end
         end
       end