Add Appsignal.set_headers helper method

I want to deprecate and remove the `Transaction#set_session_data` method. To be able to remove it without breaking things for applications calling it directly, add helpers for all types of sample data. This commit adds request headers helpers. This is similar to our Node.js sample data helpers.
appsignal · Jul 9, 2024 · 7d82dff · 7d82dff
1 parent 48c7663
commit 7d82dff
Show file tree

Hide file tree

Showing 5 changed files with 290 additions and 25 deletions.
diff --git a/.changesets/add-appsignal-set_headers-helper.md b/.changesets/add-appsignal-set_headers-helper.md
@@ -0,0 +1,10 @@
+---
+bump: patch
+type: add
+---
+
+Add `Appsignal.set_headers` helper. Set custom request headers on the current transaction with the `Appsignal.set_headers` helper. Note that this will overwrite any request headers that would be set automatically on the transaction. When this method is called multiple times, it will overwrite the previously set value.
+
+```ruby
+Appsignal.set_headers("PATH_INFO" => "/some-path", "HTTP_USER_AGENT" => "Firefox")
+```
diff --git a/lib/appsignal/helpers/instrumentation.rb b/lib/appsignal/helpers/instrumentation.rb
@@ -666,6 +666,67 @@ def set_session_data(session_data = nil, &block)
         transaction.set_session_data(session_data, &block)
       end
 
+      # Set request headers on the current transaction.
+      #
+      # Request headers are automatically set by most of our integrations. It
+      # should not be necessary to call this method unless you want to report
+      # different request headers.
+      #
+      # To filter request headers, see our session data filtering guide.
+      #
+      # When this method is called multiple times, it will overwrite the
+      # previously set value.
+      #
+      # A block can be given to this method to defer the fetching and parsing
+      # of the request headers until and only when the transaction is sampled.
+      #
+      # When both the `request_headers` argument and a block is given to this
+      # method, the `request_headers` argument is leading and the block will
+      # _not_ be called.
+      #
+      # @example Set request headers
+      #   Appsignal.set_headers(
+      #     "PATH_INFO" => "/some-path",
+      #     "HTTP_USER_AGENT" => "Firefox"
+      #   )
+      #
+      # @example Calling `set_headers` multiple times will only keep the last call
+      #   Appsignal.set_headers("PATH_INFO" => "/some-path")
+      #   Appsignal.set_headers("HTTP_USER_AGENT" => "Firefox")
+      #   # The request headers are: { "HTTP_USER_AGENT" => "Firefox" }
+      #
+      # @example Calling `set_headers` with a block
+      #   Appsignal.set_headers do
+      #     # Some slow code to parse request headers
+      #     JSON.parse('{"PATH_INFO": "/some-path"}')
+      #   end
+      #   # The session data is: { "PATH_INFO" => "/some-path" }
+      #
+      # @example Calling `set_headers` with a headers argument and a block
+      #   Appsignal.set_headers("PATH_INFO" => "/some-path") do
+      #     # Some slow code to parse session data
+      #     JSON.parse('{"PATH_INFO": "/block-path"}')
+      #   end
+      #   # The session data is: { "PATH_INFO" => "/some-path" }
+      #
+      # @since 3.11.0
+      # @param headers [Hash] The request headers to set on the transaction.
+      # @yield This block is called when the transaction is sampled. The block's
+      #   return value will become the new request headers.
+      # @see https://docs.appsignal.com/guides/custom-data/sample-data.html
+      #   Sample data guide
+      # @see https://docs.appsignal.com/guides/filter-data/filter-headers.html
+      #   Request headers filtering guide
+      # @see Transaction#set_headers
+      # @return [void]
+      def set_headers(headers = nil, &block)
+        return unless active?
+        return unless Appsignal::Transaction.current?
+
+        transaction = Appsignal::Transaction.current
+        transaction.set_headers(headers, &block)
+      end
+
       # Add breadcrumbs to the transaction.
       #
       # Breadcrumbs can be used to trace what path a user has taken

diff --git a/lib/appsignal/transaction.rb b/lib/appsignal/transaction.rb
@@ -103,6 +103,7 @@ def initialize(transaction_id, namespace, request, options = {})
       @options[:params_method] ||= :params
       @params = nil
       @session_data = nil
+      @headers = nil
 
       @ext = Appsignal::Extension.start_transaction(
         @transaction_id,
@@ -272,6 +273,45 @@ def set_session_data_if_nil(given_session_data = nil, &block)
       set_session_data(given_session_data, &block) unless @session_data
     end
 
+    # Set headers on the transaction.
+    #
+    # When both the `given_headers` and a block is given to this method,
+    # the `given_headers` argument is leading and the block will _not_ be
+    # called.
+    #
+    # @param given_headers [Hash] A hash containing headers.
+    # @yield This block is called when the transaction is sampled. The block's
+    #   return value will become the new headers.
+    # @return [void]
+    #
+    # @since 3.10.1
+    # @see Helpers::Instrumentation#set_headers
+    # @see https://docs.appsignal.com/guides/custom-data/sample-data.html
+    #   Sample data guide
+    def set_headers(given_headers = nil, &block)
+      @headers = block if block
+      @headers = given_headers if given_headers
+    end
+
+    # Set headers on the transaction if not already set.
+    #
+    # When both the `given_headers` and a block is given to this method,
+    # the `given_headers` argument is leading and the block will _not_ be
+    # called.
+    #
+    # @param given_headers [Hash] A hash containing headers.
+    # @yield This block is called when the transaction is sampled. The block's
+    #   return value will become the new headers.
+    # @return [void]
+    #
+    # @since 3.10.1
+    # @see #set_headers
+    # @see https://docs.appsignal.com/guides/custom-data/sample-data.html
+    #   Sample data guide
+    def set_headers_if_nil(given_headers = nil, &block)
+      set_headers(given_headers, &block) unless @headers
+    end
+
     # Set custom data on the transaction.
     #
     # When this method is called multiple times, it will overwrite the
@@ -661,24 +701,6 @@ def request_params
       end
     end
 
-    # Returns sanitized environment for a transaction.
-    #
-    # The environment of a transaction can contain a lot of information, not
-    # all of it useful for debugging.
-    #
-    # @return [nil] if no environment is present.
-    # @return [Hash<String, Object>]
-    def sanitized_environment
-      env = environment
-      return if env.empty?
-
-      {}.tap do |out|
-        Appsignal.config[:request_headers].each do |key|
-          out[key] = env[key] if env[key]
-        end
-      end
-    end
-
     def session_data
       if @session_data
         if @session_data.respond_to? :call
@@ -720,17 +742,36 @@ def sanitized_metadata
         .reject { |key, _value| Appsignal.config[:filter_metadata].include?(key) }
     end
 
-    # Returns the environment for a transaction.
+    def environment
+      if @headers
+        if @headers.respond_to? :call
+          @headers.call
+        else
+          @headers
+        end
+      elsif request.respond_to?(:env)
+        request.env || {}
+      else
+        {}
+      end
+    end
+
+    # Returns sanitized environment for a transaction.
     #
-    # Returns an empty Hash when the {#request} object doesn't listen to the
-    # `#env` method or the `#env` is nil.
+    # The environment of a transaction can contain a lot of information, not
+    # all of it useful for debugging.
     #
+    # @return [nil] if no environment is present.
     # @return [Hash<String, Object>]
-    def environment
-      return {} unless request.respond_to?(:env)
-      return {} unless request.env
+    def sanitized_environment
+      env = environment
+      return if env.empty?
 
-      request.env
+      {}.tap do |out|
+        Appsignal.config[:request_headers].each do |key|
+          out[key] = env[key] if env[key]
+        end
+      end
     end
 
     # Only keep tags if they meet the following criteria:

diff --git a/spec/lib/appsignal/transaction_spec.rb b/spec/lib/appsignal/transaction_spec.rb
@@ -599,6 +599,120 @@ def create_transaction(id = transaction_id)
       end
     end
 
+    describe "#set_headers" do
+      around { |example| keep_transactions { example.run } }
+
+      context "when the headers are set" do
+        it "updates the headers on the transaction" do
+          headers = { "PATH_INFO" => "value" }
+          transaction.set_headers(headers)
+
+          transaction._sample
+          expect(transaction).to include_environment(headers)
+        end
+
+        it "updates the session data on the transaction with a block" do
+          headers = { "PATH_INFO" => "value" }
+          transaction.set_headers { headers }
+
+          transaction._sample
+          expect(transaction).to include_environment(headers)
+        end
+
+        it "updates with the session data argument when both an argument and block are given" do
+          arg_data = { "PATH_INFO" => "/arg-path" }
+          block_data = { "PATH_INFO" => "/block-path" }
+          transaction.set_headers(arg_data) { block_data }
+
+          transaction._sample
+          expect(transaction).to include_environment(arg_data)
+        end
+
+        it "does not include filtered out session data" do
+          Appsignal.config[:request_headers] = ["MY_HEADER"]
+          transaction.set_headers("MY_HEADER" => "value1", "filtered_key" => "filtered_value")
+
+          transaction._sample
+          expect(transaction).to include_environment("MY_HEADER" => "value1")
+        end
+      end
+
+      context "when the given session data is nil" do
+        it "does not update the session data on the transaction" do
+          headers = { "PATH_INFO" => "value" }
+          transaction.set_headers(headers)
+          transaction.set_headers(nil)
+
+          transaction._sample
+          expect(transaction).to include_environment(headers)
+        end
+      end
+    end
+
+    describe "#set_headers_if_nil" do
+      around { |example| keep_transactions { example.run } }
+
+      context "when the params are not set" do
+        it "sets the params on the transaction" do
+          headers = { "PATH_INFO" => "value" }
+          transaction.set_headers_if_nil(headers)
+
+          transaction._sample
+          expect(transaction).to include_environment(headers)
+        end
+
+        it "updates the params on the transaction with a block" do
+          headers = { "PATH_INFO" => "value" }
+          transaction.set_headers_if_nil { headers }
+
+          transaction._sample
+          expect(transaction).to include_environment(headers)
+        end
+
+        it "updates with the params argument when both an argument and block are given" do
+          arg_data = { "PATH_INFO" => "/arg-path" }
+          block_data = { "PATH_INFO" => "/block-path" }
+          transaction.set_headers_if_nil(arg_data) { block_data }
+
+          transaction._sample
+          expect(transaction).to include_environment(arg_data)
+        end
+
+        context "when the given params is nil" do
+          it "does not update the params on the transaction" do
+            headers = { "PATH_INFO" => "value" }
+            transaction.set_headers(headers)
+            transaction.set_headers_if_nil(nil)
+
+            transaction._sample
+            expect(transaction).to include_environment(headers)
+          end
+        end
+      end
+
+      context "when the params are set" do
+        it "does not update the params on the transaction" do
+          preset_headers = { "PATH_INFO" => "/first-path" }
+          headers = { "PATH_INFO" => "/other-path" }
+          transaction.set_headers(preset_headers)
+          transaction.set_headers_if_nil(headers)
+
+          transaction._sample
+          expect(transaction).to include_environment(preset_headers)
+        end
+
+        it "does not update the params with a block on the transaction" do
+          preset_headers = { "PATH_INFO" => "/first-path" }
+          headers = { "PATH_INFO" => "/other-path" }
+          transaction.set_headers(preset_headers)
+          transaction.set_headers_if_nil { headers }
+
+          transaction._sample
+          expect(transaction).to include_environment(preset_headers)
+        end
+      end
+    end
+
     describe "#set_tags" do
       let(:long_string) { "a" * 10_001 }
 

diff --git a/spec/lib/appsignal_spec.rb b/spec/lib/appsignal_spec.rb
@@ -598,6 +598,45 @@
       end
     end
 
+    describe ".set_headers" do
+      before { start_agent }
+
+      context "with transaction" do
+        let(:transaction) { http_request_transaction }
+        before { set_current_transaction(transaction) }
+
+        it "sets request headers on the transaction" do
+          Appsignal.set_headers("PATH_INFO" => "/some-path")
+
+          transaction._sample
+          expect(transaction).to include_environment("PATH_INFO" => "/some-path")
+        end
+
+        it "overwrites the request headers if called multiple times" do
+          Appsignal.set_headers("PATH_INFO" => "/some-path1")
+          Appsignal.set_headers("PATH_INFO" => "/some-path2")
+
+          transaction._sample
+          expect(transaction).to include_environment("PATH_INFO" => "/some-path2")
+        end
+
+        it "sets request headers with a block on the transaction" do
+          Appsignal.set_headers { { "PATH_INFO" => "/some-path" } }
+
+          transaction._sample
+          expect(transaction).to include_environment("PATH_INFO" => "/some-path")
+        end
+      end
+
+      context "without transaction" do
+        it "does not set request headers on the transaction" do
+          Appsignal.set_headers("PATH_INFO" => "/some-path")
+
+          expect_any_instance_of(Appsignal::Transaction).to_not receive(:set_headers)
+        end
+      end
+    end
+
     describe ".set_custom_data" do
       before { start_agent }