Merge pull request #3303 from alexevanczuk/ae-add-ability-to-configur…

…e-pg-error-handler
DataDog · Dec 6, 2023 · f3358c0 · f3358c0
2 parents ff53df3 + 243664c
commit f3358c0
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git 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
@@ -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
@@ -99,10 +99,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
@@ -171,6 +171,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