Minimize AWS RDS Proxy connection pinning for PG

Author: i7an

activerecord/CHANGELOG.md

@@ -1,3 +1,19 @@
+*   Minimize AWS RDS Proxy connection pinning caused by SET for PostgreSQL.
+
+    When configuring a PostgreSQL connection, the adapter calls `SET` commands
+    to configure timezone, encoding, and other parameters. In AWS RDS Proxy,
+    these `SET` commands cause connection pinning.
+
+    This fix skips `SET` commands when the parameter value is already set to the
+    desired value, minimizing pinning. It's expected that you set these
+    variables in the RDS Proxy initialization query:
+
+    ```SQL
+    SET intervalstyle=iso_8601;SET client_min_messages=warning;
+    ```
+
+    *Ivan Yurchanka*
+
 *   Fix PostgreSQL schema dumping to handle schema-qualified table names in foreign_key references that span different schemas.
 
         # before

activerecord/lib/active_record/connection_adapters/postgresql/configuration_parameters.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module PostgreSQL
+      # Helper methods for working with configuration parameters.
+      module ConfigurationParameters
+        private
+          def ensure_parameter(name, value)
+            return if value.nil?
+            return if parameter_set_to?(name, value)
+
+            if block_given?
+              yield value
+            else
+              set_parameter(name, value)
+            end
+          end
+
+          def parameter_set_to?(name, value)
+            validate_parameter!(name)
+
+            normalized_value = case value
+            when TrueClass
+              "on"
+            when FalseClass
+              "off"
+            else
+              value.to_s
+            end
+
+            current_value = query_value("SHOW #{name}", "SCHEMA")
+
+            normalized_value == current_value
+          end
+
+          def set_parameter(name, value)
+            validate_parameter!(name)
+
+            if value == :default
+              query_command("SET SESSION #{name} TO DEFAULT", "SCHEMA")
+            else
+              query_command("SET SESSION #{name} TO #{quote(value)}", "SCHEMA")
+            end
+          end
+
+          def validate_parameter!(name)
+            raise ArgumentError, "Parameter name '#{name}' is invalid" unless name.match?(/\A[a-zA-Z0-9_.]+\z/)
+          end
+      end
+    end
+  end
+end

activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb

@@ -15,7 +15,8 @@ def query_rows(sql, name = "SCHEMA", allow_retry: true, materialize_transactions
           intent = internal_build_intent(sql, name, allow_retry:, materialize_transactions:)
           intent.execute!
           result = intent.raw_result
-          result.map_types!(@type_map_for_results).values
+          result.map_types!(@type_map_for_results) if @type_map_for_results
+          result.values
         end
 
         READ_QUERY = ActiveRecord::ConnectionAdapters::AbstractAdapter.build_read_query_regexp(

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -7,6 +7,7 @@
 require "active_record/connection_adapters/abstract_adapter"
 require "active_record/connection_adapters/statement_pool"
 require "active_record/connection_adapters/postgresql/column"
+require "active_record/connection_adapters/postgresql/configuration_parameters"
 require "active_record/connection_adapters/postgresql/database_statements"
 require "active_record/connection_adapters/postgresql/explain_pretty_printer"
 require "active_record/connection_adapters/postgresql/oid"
@@ -202,6 +203,7 @@ def dbconsole(config, options = {})
       include PostgreSQL::ReferentialIntegrity
       include PostgreSQL::SchemaStatements
       include PostgreSQL::DatabaseStatements
+      include PostgreSQL::ConfigurationParameters
 
       def supports_bulk_alter?
         true
@@ -443,7 +445,7 @@ def self.native_database_types # :nodoc:
       end
 
       def set_standard_conforming_strings
-        query_command("SET standard_conforming_strings = on", "SCHEMA")
+        ensure_parameter("standard_conforming_strings", "on")
       end
 
       def supports_ddl_transactions?
@@ -1029,11 +1031,13 @@ def reconnect
         def configure_connection
           super
 
-          if @config[:encoding]
-            @raw_connection.set_client_encoding(@config[:encoding])
+          set_client_encoding
+          ensure_parameter("client_min_messages", @config[:min_messages] || "warning") do |value|
+            self.client_min_messages = value
+          end
+          ensure_parameter("search_path", @config[:schema_search_path] || @config[:schema_order]) do |value|
+            self.schema_search_path = value
           end
-          self.client_min_messages = @config[:min_messages] || "warning"
-          self.schema_search_path = @config[:schema_search_path] || @config[:schema_order]
 
           unless ActiveRecord.db_warnings_action.nil?
             @raw_connection.set_notice_receiver do |result|
@@ -1047,19 +1051,18 @@ def configure_connection
           # Use standard-conforming strings so we don't have to do the E'...' dance.
           set_standard_conforming_strings
 
-          variables = @config.fetch(:variables, {}).stringify_keys
-
           # Set interval output format to ISO 8601 for ease of parsing by ActiveSupport::Duration.parse
-          query_command("SET intervalstyle = iso_8601", "SCHEMA")
+          ensure_parameter("intervalstyle", "iso_8601")
 
           # SET statements from :variables config hash
           # https://www.postgresql.org/docs/current/static/sql-set.html
+          variables = @config.fetch(:variables, {}).stringify_keys
           variables.map do |k, v|
             if v == ":default" || v == :default
               # Sets the value to the global or compile default
-              query_command("SET SESSION #{k} TO DEFAULT", "SCHEMA")
+              set_parameter(k, :default)
             elsif !v.nil?
-              query_command("SET SESSION #{k} TO #{quote(v)}", "SCHEMA")
+              ensure_parameter(k, v)
             end
           end
 
@@ -1082,13 +1085,29 @@ def reconfigure_connection_timezone
           # If using Active Record's time zone support configure the connection
           # to return TIMESTAMP WITH ZONE types in UTC.
           if default_timezone == :utc
-            intent = QueryIntent.new(adapter: self, processed_sql: "SET SESSION timezone TO 'UTC'", name: "SCHEMA")
-            intent.execute!
-            intent.finish
+            # For simplicity, ignore UTC aliases (e.g., UCT, Zulu, GMT, GMT0, ...).
+            is_utc = parameter_set_to?("timezone", "UTC") ||
+              parameter_set_to?("timezone", "Etc/UTC")
+            set_parameter("timezone", "UTC") unless is_utc
           else
-            intent = QueryIntent.new(adapter: self, processed_sql: "SET SESSION timezone TO DEFAULT", name: "SCHEMA")
-            intent.execute!
-            intent.finish
+            set_parameter("timezone", :default)
+          end
+        end
+
+        def set_client_encoding
+          return unless @config[:encoding]
+
+          normalized_encoding = @config[:encoding].upcase
+
+          # For simplicity, handle only the common 'unicode' alias.
+          # See https://www.postgresql.org/docs/current/multibyte.html#CHARSET-TABLE
+          normalized_encoding = "UTF8" if normalized_encoding == "UNICODE"
+
+          # PostgreSQL accepts loose input forms like 'utf@8'. We assume
+          # canonical form for simplicity.
+          # See https://github.com/postgres/postgres/blob/master/src/common/encnames.c
+          if @raw_connection.get_client_encoding != normalized_encoding
+            @raw_connection.set_client_encoding(normalized_encoding)
           end
         end
 

activerecord/test/cases/adapters/postgresql/connection_test.rb

@@ -200,6 +200,21 @@ def test_set_session_timezone
       end
     end
 
+    def test_set_invalid_session_variable
+      run_without_connection do |orig_connection|
+        assert_raises(ArgumentError) do
+          ActiveRecord::Base.establish_connection(orig_connection.deep_merge(variables: { 'sql injection': true }))
+        end
+      end
+    end
+
+    def test_set_session_encoding
+      run_without_connection do |orig_connection|
+        ActiveRecord::Base.establish_connection(orig_connection.deep_merge(encoding: "LATIN1"))
+        assert_equal "LATIN1", ActiveRecord::Base.lease_connection.select_value("SHOW client_encoding")
+      end
+    end
+
     def test_get_and_release_advisory_lock
       lock_id = 5295901941911233559
       list_advisory_locks = <<~SQL

activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb

@@ -26,6 +26,35 @@ def test_connection_error
         assert_kind_of ActiveRecord::ConnectionAdapters::NullPool, error.connection_pool
       end
 
+      def test_rds_proxy_compatible_connection_avoids_pinning
+        db_config = ActiveRecord::Base.configurations.configs_for(env_name: "arunit", name: "primary")
+        configuration = db_config.configuration_hash.merge(
+          options: [
+            "-c client_min_messages=warning",
+            "-c standard_conforming_strings=on",
+            "-c intervalstyle=iso_8601",
+            "-c timezone=UTC",
+            "-c debug_print_plan=on"
+          ].join(" "),
+          encoding: "unicode",
+          variables: {
+            debug_print_plan: true
+          }
+        )
+        connection = ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.new(configuration)
+
+        set_commands = []
+        subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |*, payload|
+          set_commands << payload[:sql] if payload[:sql].match?(/\bSET\b/i)
+        end
+
+        connection.connect!
+
+        assert_empty set_commands
+      ensure
+        ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+      end
+
       def test_reconnection_error
         fake_connection = Class.new do
           def async_exec(*)