aws_advanced_ruby_driver_wrapper 1.0.1 → 1.1.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/blue_green/status_provider.rb

@@ -18,6 +18,7 @@ require 'concurrent'
 require_relative '../../logging'
 require_relative '../../property_definition'
 require_relative '../../utils/rds_utils'
+require_relative 'drop_blue_connections_marker'
 require_relative 'iam_host_tracker'
 require_relative 'interval_rate'
 require_relative 'host_mapper'
@@ -61,6 +62,10 @@ module AwsAdvancedRubyDriverWrapper
           @process_status_lock   = Mutex.new
           @monitor_generation    = 0
 
+          @drop_blue_connections = PropertyDefinition::BG_DROP_BLUE_CONNECTIONS.get_bool(wrapper_props)
+          # Identity of the drop-blue-connections marker created for the current switchover, if any.
+          @drop_marker_id        = nil
+
           @status_check_interval_map = {
             IntervalRate::BASELINE => PropertyDefinition::BG_INTERVAL_BASELINE_SEC.get_float(wrapper_props) * 1000,
             IntervalRate::INCREASED => PropertyDefinition::BG_INTERVAL_INCREASED_SEC.get_float(wrapper_props) * 1000,
@@ -176,6 +181,18 @@ module AwsAdvancedRubyDriverWrapper
 
           logger.warn { 'Blue/Green switchover has timed out.' }
           @summary_status = @rollback ? @builder.created : @builder.completed(@rollback)
+          # A timeout means the wrapper stopped waiting, not that the switchover is confirmed. Only
+          # activate the marker if the switchover was last observed at POST (past the point of no
+          # return); timing out in PREPARATION/IN_PROGRESS could still roll back, and dropping blue
+          # connections then would hit a blue that is still primary. In that case discard the pending
+          # marker so it does not linger unactivated for the full marker TTL: we have given up waiting
+          # on this switchover, and no later phase will arrive to activate it.
+          if @latest_status_phase == Phase::POST
+            activate_drop_marker
+          else
+            remove_drop_marker
+            @drop_marker_id = nil
+          end
           update_monitors
           update_status_cache
           log_current_context
@@ -191,7 +208,15 @@ module AwsAdvancedRubyDriverWrapper
           @monitors = { Role::SOURCE => nil, Role::TARGET => nil }
           old_monitors.each_value { |m| m&.stop }
 
-          Plugins::IamAuthPlugin.clear_cache(@storage_service) if iam_enabled?
+          # A completed switchover's marker must survive this reset so connections opened to the old
+          # blue environment are still invalidated on their next use. A rollback leaves blue as primary,
+          # so its marker is discarded instead. Either way a later switchover starts a fresh marker.
+          remove_drop_marker if @rollback
+          @drop_marker_id = nil
+
+          # IamAuthPlugin, FederatedAuthPlugin, and OktaAuthPlugin all cache tokens under the same
+          # shared IAM_TOKEN_CACHE_NAME, so clearing through any one of them clears all of them.
+          Plugins::IamAuthPlugin.clear_cache(@storage_service) if iam_token_cache_in_use?
 
           @rollback = false
           @green_topology_recognized_logged = false
@@ -225,12 +250,15 @@ module AwsAdvancedRubyDriverWrapper
             @switchover_state.update_dns_flags(@bgd_id, role, interim_status)
             @summary_status = @builder.in_progress(@rollback)
             @switchover_state.trigger_in_progress_monitor_reset
+            ensure_drop_marker_pending
           when Phase::POST
             @switchover_state.update_dns_flags(@bgd_id, role, interim_status)
             @summary_status = @builder.post(@rollback)
+            activate_drop_marker
           when Phase::COMPLETED
             @switchover_state.update_dns_flags(@bgd_id, role, interim_status)
             @summary_status = @builder.completed(@rollback)
+            activate_drop_marker
           else
             raise ArgumentError, "[bgd_id: '#{@bgd_id}'] Unknown BG phase '#{@latest_status_phase}'."
           end
@@ -350,14 +378,73 @@ module AwsAdvancedRubyDriverWrapper
           )
         end
 
+        # Creates a pending drop-blue-connections marker the first time a switchover reaches IN_PROGRESS.
+        # The marker is only confirmed (see #activate_drop_marker) once the switchover passes the point
+        # of no return, so nothing is invalidated while it is still in progress and could roll back.
+        def ensure_drop_marker_pending
+          return unless @drop_blue_connections
+          return if @rollback
+          return if @drop_marker_id
+
+          hosts = blue_host_set
+          return if hosts.empty?
+
+          marker = DropBlueConnectionsMarker.new(blue_hosts: hosts, switchover_time_ns: nano_time)
+          @drop_marker_id = marker.id
+          store_drop_marker(marker)
+        end
+
+        # Confirms the switchover's pending marker so stale blue connections begin to be invalidated.
+        # A marker is only created when IN_PROGRESS is observed (see #ensure_drop_marker_pending), since
+        # that is the only point with a reliable switchover time. A provider that first sees the
+        # deployment at POST/COMPLETED cannot tell which of its connections predate the switchover and
+        # would wrongly cover connections opened after switchover, including to green.
+        def activate_drop_marker
+          return unless @drop_blue_connections
+          return if @rollback
+
+          marker = read_drop_marker
+          return if marker.nil? || marker.active?
+
+          store_drop_marker(marker.activated)
+        end
+
+        # The blue host set: downcased SOURCE endpoint names plus their pre-switchover IP addresses.
+        def blue_host_set
+          return Set.new unless @summary_status
+
+          blue = @summary_status.role_by_host.select { |_host, role| role == Role::SOURCE }.keys
+          ips  = blue.filter_map { |host| @host_mapper.host_ip_addresses[host] }
+          (blue + ips).to_set(&:downcase)
+        end
+
+        def store_drop_marker(marker)
+          @storage_service.set(BlueGreenPlugin::DROP_MARKER_NAME, BlueGreenPlugin.drop_marker_key(@bgd_id), marker)
+        end
+
+        def read_drop_marker
+          @storage_service.get(BlueGreenPlugin::DROP_MARKER_NAME, BlueGreenPlugin.drop_marker_key(@bgd_id), register_access: false)
+        end
+
+        def remove_drop_marker
+          @storage_service.remove(BlueGreenPlugin::DROP_MARKER_NAME, BlueGreenPlugin.drop_marker_key(@bgd_id))
+        end
+
+        def nano_time
+          Process.clock_gettime(Process::CLOCK_MONOTONIC, :nanosecond)
+        end
+
         def switchover_finalized?
           completed = (!@rollback && @summary_status.current_phase == Phase::COMPLETED) ||
                       (@rollback && @summary_status.current_phase == Phase::CREATED)
           completed && @event_log.any? { |_, v| v.phase&.active_switchover_or_completed? }
         end
 
-        def iam_enabled?
-          @service_container.plugin_manager&.plugin_in_use?(Plugins::IamAuthPlugin) || false
+        # True when any auth plugin backed by the shared IAM token cache (IAM or AD FS
+        # or Okta) is in use, so a switchover knows to purge that cache regardless
+        # of which one is active.
+        def iam_token_cache_in_use?
+          @service_container.plugin_manager&.iam_token_plugin_in_use? || false
         end
 
         def context_hash

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/custom_endpoint_monitor.rb

@@ -41,7 +41,14 @@ module AwsAdvancedRubyDriverWrapper
           refresh_rate_sec,
           refresh_rate_backoff_factor,
           max_refresh_rate_sec,
-          rds_client_func: ->(_, region) { Aws::RDS::Client.new(region: region) }
+          credentials_provider: nil,
+          rds_client_func: lambda { |_, client_region|
+            # Pass the configured credentials provider when one is set; otherwise omit :credentials
+            # so the SDK resolves from its default chain.
+            options = { region: client_region }
+            options[:credentials] = credentials_provider if credentials_provider
+            Aws::RDS::Client.new(**options)
+          }
         )
           super(termination_timeout_sec: TERMINATION_TIMEOUT_SEC)
 
@@ -58,6 +65,7 @@ module AwsAdvancedRubyDriverWrapper
           @refresh_cv = ConditionVariable.new
           @refresh_required = false
           @connection_issue = false
+          @info_changed_counter = service_container.telemetry_factory.create_counter('custom_endpoint.info_changed.count')
         end
 
         def endpoint_info?
@@ -161,6 +169,7 @@ module AwsAdvancedRubyDriverWrapper
           logger.debug("[#{@endpoint_id}] Custom endpoint info changed: #{endpoint_info}")
           cache_allowed_blocked(endpoint_info)
           cache_info(endpoint_info)
+          @info_changed_counter.inc
           speedup_refresh_rate
 
           elapsed = monotonic_time - start
@@ -241,6 +250,7 @@ module AwsAdvancedRubyDriverWrapper
         end
 
         def cached_info
+          # Explicitly unregistered: a monitor reading its own data must not keep itself alive.
           storage_service.get(ENDPOINT_INFO_CACHE_NAME, @custom_endpoint_host.url, register_access: false)
         end
 

data/lib/aws_advanced_ruby_driver_wrapper/plugins/custom_endpoint/custom_endpoint_plugin.rb

@@ -38,6 +38,7 @@ module AwsAdvancedRubyDriverWrapper
           @props = props
           @should_wait_for_info = PropertyDefinition::WAIT_FOR_CUSTOM_ENDPOINT_INFO.get_bool(props)
           @wait_timeout_sec = PropertyDefinition::WAIT_FOR_CUSTOM_ENDPOINT_INFO_TIMEOUT_SEC.get_float(props)
+          @wait_for_info_counter = service_container.telemetry_factory.create_counter('custom_endpoint.wait_for_info.count')
 
           @custom_endpoint_host = nil
           @endpoint_id = nil
@@ -128,7 +129,8 @@ module AwsAdvancedRubyDriverWrapper
               @region,
               PropertyDefinition::CUSTOM_ENDPOINT_INFO_REFRESH_RATE_SEC.get_float(props),
               PropertyDefinition::CUSTOM_ENDPOINT_INFO_REFRESH_RATE_BACKOFF_FACTOR.get_int(props),
-              PropertyDefinition::CUSTOM_ENDPOINT_INFO_MAX_REFRESH_RATE_SEC.get_float(props)
+              PropertyDefinition::CUSTOM_ENDPOINT_INFO_MAX_REFRESH_RATE_SEC.get_float(props),
+              credentials_provider: PropertyDefinition::AWS_CREDENTIALS_PROVIDER.get(props)
             )
           end
         end
@@ -137,6 +139,7 @@ module AwsAdvancedRubyDriverWrapper
           return if monitor.endpoint_info?
 
           monitor.request_endpoint_info_update
+          @wait_for_info_counter.inc
           logger.debug("CustomEndpointPlugin: waiting up to #{@wait_timeout_sec}s for endpoint info on #{@custom_endpoint_host.url}")
 
           return if monitor.wait_for_info?(@wait_timeout_sec)

data/lib/aws_advanced_ruby_driver_wrapper/plugins/default_plugin.rb

@@ -18,6 +18,7 @@ require 'concurrent'
 require_relative '../errors'
 require_relative '../host/host_info'
 require_relative '../host/host_availability'
+require_relative '../utils/telemetry/telemetry_const'
 
 module AwsAdvancedRubyDriverWrapper
   module Plugins
@@ -44,7 +45,7 @@ module AwsAdvancedRubyDriverWrapper
                              host_info
                            end
 
-        conn = driver_dialect.connect(target_host_info, driver_props)
+        conn = with_driver_span { driver_dialect.connect(target_host_info, driver_props) }
 
         if conn.nil?
           raise Errors::AwsError,
@@ -83,7 +84,7 @@ module AwsAdvancedRubyDriverWrapper
 
       def internal_connect(host_info, driver_props, _wrapper_props, _is_initial_connection, _pipeline_callable)
         driver_dialect = @service_container.dialect_service.driver_dialect
-        driver_dialect.connect(host_info, driver_props)
+        with_driver_span { driver_dialect.connect(host_info, driver_props) }
       end
 
       def execute(target_method_name, target_callable, *args, **, &)
@@ -92,7 +93,7 @@ module AwsAdvancedRubyDriverWrapper
         succeeded = false
 
         begin
-          result = target_callable.call(*args, **, &)
+          result = with_driver_span { target_callable.call(*args, **, &) }
           succeeded = true
           result
         ensure
@@ -103,6 +104,16 @@ module AwsAdvancedRubyDriverWrapper
           end
         end
       end
+
+      private
+
+      # Records the time spent in the driver itself as a span named after the driver.
+      def with_driver_span(&)
+        telemetry_factory = @service_container.telemetry_factory
+        return yield unless telemetry_factory.traces_in_use?
+
+        telemetry_factory.in_context(@service_container.dialect_service.driver_dialect.name, Utils::Telemetry::TraceLevel::NESTED, &)
+      end
     end
   end
 end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/failover_plugin.rb

@@ -20,9 +20,11 @@ require_relative '../host/host_info'
 require_relative '../host/host_role'
 require_relative '../host/host_availability'
 require_relative '../logging'
+require_relative '../services/host_service'
 require_relative '../utils/rds_utils'
 require_relative '../utils/rds_url_type'
 require_relative '../utils/retry_util'
+require_relative '../utils/telemetry/telemetry_const'
 require_relative '../wrapper_property'
 require_relative 'failover_mode'
 
@@ -32,14 +34,25 @@ module AwsAdvancedRubyDriverWrapper
       include Logging
 
       ReaderFailoverResult = Data.define(:connection, :host_info)
+      FailoverCounters = Data.define(:triggered, :succeeded, :failed)
+
+      SUCCESSFUL_FAILOVER_ERRORS = [Errors::FailoverSuccessError, Errors::TransactionStateUnknownError].freeze
+
+      TELEMETRY_WRITER_FAILOVER = 'failover to writer host'
+      TELEMETRY_READER_FAILOVER = 'failover to replica'
 
       def initialize(service_container, props = ::Concurrent::Map.new)
         @service_container = service_container
         @wrapper_props = props
+        @telemetry_factory = service_container.telemetry_factory
+        @telemetry_failover_additional_top_trace = PropertyDefinition::TELEMETRY_FAILOVER_ADDITIONAL_TOP_TRACE.get_bool(props)
+        @writer_failover_counters = failover_counters('writer_failover')
+        @reader_failover_counters = failover_counters('reader_failover')
 
         @retry_util = Utils::RetryUtil.new(service_container)
         @failover_timeout = PropertyDefinition::FAILOVER_TIMEOUT_SEC.get_float(props)
         @reader_selector_strategy = PropertyDefinition::FAILOVER_READER_HOST_SELECTOR_STRATEGY.get(props)
+        Services::HostService.validate_strategy!(@reader_selector_strategy)
         @failover_mode = nil
         @rds_url_type = nil
         @closed_explicitly = false
@@ -235,9 +248,60 @@ module AwsAdvancedRubyDriverWrapper
         end
 
         if @failover_mode == FailoverMode::STRICT_WRITER
-          failover_writer
+          with_failover_telemetry(TELEMETRY_WRITER_FAILOVER) { count_failover(@writer_failover_counters) { failover_writer } }
         else
-          failover_reader
+          with_failover_telemetry(TELEMETRY_READER_FAILOVER) { count_failover(@reader_failover_counters) { failover_reader } }
+        end
+      end
+
+      def failover_counters(prefix)
+        FailoverCounters.new(
+          triggered: @telemetry_factory.create_counter("#{prefix}.triggered.count"),
+          succeeded: @telemetry_factory.create_counter("#{prefix}.completed.success.count"),
+          failed: @telemetry_factory.create_counter("#{prefix}.completed.failed.count")
+        )
+      end
+
+      # Records a failover as a span. Failover always ends in an error: a successful one ends in one
+      # of SUCCESSFUL_FAILOVER_ERRORS, which marks the span as succeeded, and any other error marks it
+      # as failed. With telemetry_failover_additional_top_trace enabled, the span is also submitted as
+      # a trace of its own.
+      def with_failover_telemetry(span_name)
+        telemetry_context = nil
+        @telemetry_factory.in_context(span_name, Utils::Telemetry::TraceLevel::NESTED) do |context|
+          telemetry_context = context
+          yield
+        rescue *SUCCESSFUL_FAILOVER_ERRORS => e
+          context.success = true
+          context.record_exception(e)
+          raise
+        rescue StandardError => e
+          context.success = false
+          context.record_exception(e)
+          raise
+        end
+      ensure
+        submit_failover_copy(telemetry_context) if telemetry_context && @telemetry_failover_additional_top_trace
+      end
+
+      # Runs while the failover's outcome is propagating, so an error here must not replace it.
+      def submit_failover_copy(telemetry_context)
+        @telemetry_factory.post_copy(telemetry_context, Utils::Telemetry::TraceLevel::FORCE_TOP_LEVEL)
+      rescue StandardError => e
+        logger.debug { "Unable to submit the failover trace as a trace of its own: #{e.message}" }
+      end
+
+      # Counts a failover as triggered, and then as succeeded or failed according to how it ends.
+      def count_failover(counters)
+        counters.triggered.inc
+        begin
+          yield
+        rescue *SUCCESSFUL_FAILOVER_ERRORS
+          counters.succeeded.inc
+          raise
+        rescue StandardError
+          counters.failed.inc
+          raise
         end
       end
 
@@ -414,7 +478,8 @@ module AwsAdvancedRubyDriverWrapper
           Host::HostRole::READER,
           @reader_selector_strategy
         )
-      rescue StandardError
+      rescue StandardError => e
+        logger.warn("Failed to select a reader using the '#{@reader_selector_strategy}' host selection strategy: #{e.message}")
         nil
       end
 

data/lib/aws_advanced_ruby_driver_wrapper/plugins/federated_auth/base_saml_auth_plugin.rb

@@ -0,0 +1,266 @@
+# frozen_string_literal: true
+
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+require 'concurrent'
+require_relative '../../errors'
+require_relative '../../property_definition'
+require_relative '../iam_auth_plugin'
+require_relative '../../utils/iam_auth_utils'
+require_relative '../../utils/rds_utils'
+require_relative '../../utils/rds_url_type'
+require_relative '../../utils/saml_utils'
+require_relative '../../utils/telemetry/telemetry_const'
+
+module AwsAdvancedRubyDriverWrapper
+  module Plugins
+    module FederatedAuth
+      # Shared connect flow for the SAML-based authentication plugins (AD FS and
+      # Okta).
+      #
+      # A subclass defines +self.fetch_saml_assertion+, +TELEMETRY_PREFIX+ (the prefix of its metric
+      # names) and +TELEMETRY_FETCH_SAML+ (the name of the span its SAML assertion fetch is recorded as).
+      class BaseSamlAuthPlugin
+        SUBSCRIBED_METHODS = Set['connect', 'internal_connect'].freeze
+        # Reuse the IAM plugin's cache name.
+        IAM_TOKEN_CACHE_NAME = AwsAdvancedRubyDriverWrapper::Plugins::IamAuthPlugin::IAM_TOKEN_CACHE_NAME
+
+        attr_reader :subscribed_methods
+
+        def initialize(service_container, props = ::Concurrent::Map.new)
+          ensure_aws_sdk!
+          @service_container = service_container
+          @wrapper_props = props
+          expiration = PropertyDefinition::IAM_EXPIRATION_SEC.get_float(props)
+          PropertyDefinition::IAM_EXPIRATION_SEC.validate!(expiration)
+          service_container.storage_service.register(IAM_TOKEN_CACHE_NAME, ttl: expiration)
+          @subscribed_methods = SUBSCRIBED_METHODS
+
+          # Each subclass names its metrics after itself, and its SAML assertion span after its IdP.
+          @telemetry_factory = service_container.telemetry_factory
+          @fetch_token_counter = @telemetry_factory.create_counter("#{self.class::TELEMETRY_PREFIX}.fetch_token.count")
+          @telemetry_factory.create_gauge(
+            "#{self.class::TELEMETRY_PREFIX}.token_cache.size",
+            &IamAuthPlugin.token_cache_size(service_container.storage_service)
+          )
+        end
+
+        def connect(host_info, driver_props, _is_initial_connection, pipeline_callable)
+          saml_connect(host_info, driver_props, pipeline_callable)
+        end
+
+        def internal_connect(host_info, driver_props, wrapper_props_override, _is_initial_connection, pipeline_callable)
+          saml_connect(host_info, driver_props, pipeline_callable, wrapper_props_override || @wrapper_props)
+        end
+
+        def self.clear_cache(storage_service)
+          storage_service.clear(IAM_TOKEN_CACHE_NAME)
+        end
+
+        # Exchanges a SAML assertion for temporary AWS credentials via STS
+        # AssumeRoleWithSAML.
+        #
+        # @return [Aws::Credentials]
+        def self.credentials_from_saml(saml_assertion, role_arn:, idp_arn:, region:)
+          require 'aws-sdk-core'
+          # AssumeRoleWithSAML is unsigned, so no AWS credentials are needed. Passing
+          # credentials: nil keeps the SDK from running the default credential chain (which
+          # on a credential-less, non-EC2 host wastes ~1s probing IMDS and logs a misleading
+          # "Error retrieving instance profile credentials").
+          client_opts = { credentials: nil }
+          client_opts[:region] = region if region && !region.to_s.empty?
+          resp = Aws::STS::Client.new(**client_opts).assume_role_with_saml(
+            role_arn:, principal_arn: idp_arn, saml_assertion:
+          )
+          creds = resp.credentials
+          Aws::Credentials.new(creds.access_key_id, creds.secret_access_key, creds.session_token)
+        end
+
+        private
+
+        def saml_connect(host_info, driver_props, pipeline_callable, wrapper_props_override = nil)
+          wrapper_props_override ||= @wrapper_props
+          saml_props = wrapper_props_override
+
+          validate_required!(saml_props)
+
+          # The connection's driver_props :user/:password double as the fallback IdP login when
+          # idp_username / idp_password are not set. We overwrite driver_props with db_user and the
+          # RDS token further down, and that hash is reused across reconnects (failover, retry,
+          # Blue/Green rerouting), so capture the fallback into the wrapper props once, up front.
+          # A later reconnect then still logs in to the IdP with the original credentials rather
+          # than db_user and a stale token.
+          persist_idp_credentials_fallback(saml_props, driver_props)
+          db_user = PropertyDefinition::DB_USER.get(saml_props)
+          idp_username = PropertyDefinition::IDP_USERNAME.get(saml_props)
+          idp_password = PropertyDefinition::IDP_PASSWORD.get(saml_props)
+
+          host = Utils::IamAuthUtils.resolve_host(
+            PropertyDefinition::IAM_HOST.get(wrapper_props_override), host_info
+          )
+          port = Utils::IamAuthUtils.resolve_port(
+            PropertyDefinition::IAM_PORT.get(wrapper_props_override),
+            host_info,
+            @service_container.dialect_service.db_dialect.default_port
+          )
+          rds_type = Utils::RdsUtils.identify_rds_type(host)
+
+          # Only a global writer cluster with no explicit iam_region needs AWS credentials to look
+          # its region up. In that one case the SAML exchange has to happen before the region lookup
+          # (and so before the token-cache check); every other case - including a global cluster
+          # whose iam_region is set - defers the IdP round-trip to the token-cache miss path, so a
+          # pool or failover storm does not trigger one IdP login per connection.
+          explicit_region = PropertyDefinition::IAM_REGION.get(wrapper_props_override)
+          preassumed_credentials = nil
+          if rds_type == Utils::RdsUrlType::RDS_GLOBAL_WRITER_CLUSTER && (explicit_region.nil? || explicit_region.to_s.empty?)
+            preassumed_credentials = obtain_credentials(saml_props, nil, idp_username, idp_password)
+          end
+
+          region = Utils::IamAuthUtils.region_for(
+            host:, props: wrapper_props_override, rds_type:,
+            credentials_provider: preassumed_credentials,
+            rds_client_func: preassumed_credentials ? -> { Aws::RDS::Client.new(credentials: preassumed_credentials) } : nil
+          )
+          unless region
+            raise Errors::SamlAuthError,
+                  "#{self.class.name}: unable to determine connection region. " \
+                  "If you are using a non-standard RDS URL, please set the 'iam_region' property."
+          end
+
+          fingerprint = Utils::SamlUtils.principal_fingerprint(
+            PropertyDefinition::IAM_ROLE_ARN.get(saml_props),
+            PropertyDefinition::IAM_IDP_ARN.get(saml_props)
+          )
+          cache_key = Utils::IamAuthUtils.cache_key(region, host, port, db_user, fingerprint)
+          entry = @service_container.storage_service.get(IAM_TOKEN_CACHE_NAME, cache_key)
+          expiration = PropertyDefinition::IAM_EXPIRATION_SEC.get_float(wrapper_props_override)
+          PropertyDefinition::IAM_EXPIRATION_SEC.validate!(expiration)
+
+          if Utils::IamAuthUtils.valid_entry?(entry)
+            token = entry.token
+            is_cached_token = true
+          else
+            token = fetch_and_cache_token(saml_props, region, host, port, db_user, cache_key, expiration,
+                                          idp_username, idp_password, preassumed_credentials)
+            is_cached_token = false
+          end
+
+          user_key = @service_container.dialect_service.driver_dialect.user_property_key
+          driver_props[user_key] = db_user
+          driver_props[:password] = token
+
+          begin
+            pipeline_callable.call
+          rescue StandardError => e
+            raise unless is_cached_token && @service_container.dialect_service.login_error?(e)
+
+            driver_props[:password] = fetch_and_cache_token(saml_props, region, host, port, db_user, cache_key,
+                                                            expiration, idp_username, idp_password, nil)
+            pipeline_callable.call
+          end
+        end
+
+        # Mints a fresh RDS token signed with SAML-derived credentials and caches it.
+        def fetch_and_cache_token(saml_props, region, host, port, db_user, cache_key, expiration,
+                                  idp_username, idp_password, credentials)
+          credentials ||= obtain_credentials(saml_props, region, idp_username, idp_password)
+          @fetch_token_counter.inc
+          token = Utils::IamAuthUtils.generate_token(
+            region:, hostname: host, port:, user: db_user, credentials_provider: credentials,
+            telemetry_factory: @telemetry_factory
+          )
+          @service_container.storage_service.set(
+            IAM_TOKEN_CACHE_NAME, cache_key, Utils::IamAuthUtils.build_token_entry(token, expiration)
+          )
+          token
+        end
+
+        # Captures the connection's user/password as the IdP login when idp_username /
+        # idp_password are unset, persisting them into the wrapper props so the value survives
+        # the driver_props overwrite later in the flow and any subsequent reconnect. Never stores
+        # a nil (Concurrent::Map rejects nil values), so a connection with no fallback simply
+        # keeps idp_username / idp_password unset.
+        def persist_idp_credentials_fallback(saml_props, driver_props)
+          if PropertyDefinition::IDP_USERNAME.get(saml_props).nil?
+            fallback_user = driver_props[:user] || driver_props[:username]
+            PropertyDefinition::IDP_USERNAME.set(saml_props, fallback_user) unless fallback_user.nil?
+          end
+          return unless PropertyDefinition::IDP_PASSWORD.get(saml_props).nil?
+
+          fallback_password = driver_props[:password]
+          PropertyDefinition::IDP_PASSWORD.set(saml_props, fallback_password) unless fallback_password.nil?
+        end
+
+        # Fetches a SAML assertion from the IdP and exchanges it for AWS credentials.
+        def obtain_credentials(saml_props, region, idp_username, idp_password)
+          assertion = @telemetry_factory.in_context(self.class::TELEMETRY_FETCH_SAML, Utils::Telemetry::TraceLevel::NESTED,
+                                                    record_failure: true) do
+            self.class.fetch_saml_assertion(**saml_fetch_options(saml_props, idp_username, idp_password))
+          end
+          self.class.credentials_from_saml(
+            assertion,
+            role_arn: PropertyDefinition::IAM_ROLE_ARN.get(saml_props),
+            idp_arn: PropertyDefinition::IAM_IDP_ARN.get(saml_props),
+            region:
+          )
+        rescue Errors::SamlAuthError
+          raise
+        rescue StandardError => e
+          # Surface IdP/STS failures (DNS, TLS, STS errors) as a SamlAuthError.
+          raise Errors::SamlAuthError, "#{self.class.name}: failed to obtain credentials via SAML: #{e.message}"
+        end
+
+        def validate_required!(saml_props)
+          missing = required_identity(saml_props).select { |_name, value| value.nil? || value.to_s.empty? }
+          return if missing.empty?
+
+          names = missing.map(&:first)
+          raise Errors::SamlAuthError,
+                "#{self.class.name}: #{names.join(', ')} #{names.length > 1 ? 'are' : 'is'} required"
+        end
+
+        # The properties that must be present for this IdP. Subclasses override to
+        # append their own (e.g. Okta's app_id).
+        def required_identity(saml_props)
+          [
+            ['db_user', PropertyDefinition::DB_USER.get(saml_props)],
+            ['idp_endpoint', PropertyDefinition::IDP_ENDPOINT.get(saml_props)],
+            ['iam_role_arn', PropertyDefinition::IAM_ROLE_ARN.get(saml_props)],
+            ['iam_idp_arn', PropertyDefinition::IAM_IDP_ARN.get(saml_props)]
+          ]
+        end
+
+        # The keyword arguments passed to .fetch_saml_assertion. Subclasses override.
+        def saml_fetch_options(saml_props, idp_username, idp_password)
+          {
+            idp_endpoint: PropertyDefinition::IDP_ENDPOINT.get(saml_props),
+            idp_username:,
+            idp_password:,
+            ssl_insecure: PropertyDefinition::SSL_INSECURE.get_bool(saml_props),
+            socket_timeout_sec: PropertyDefinition::HTTP_CLIENT_SOCKET_TIMEOUT_SEC.get_float(saml_props),
+            connect_timeout_sec: PropertyDefinition::HTTP_CLIENT_CONNECT_TIMEOUT_SEC.get_float(saml_props)
+          }
+        end
+
+        def ensure_aws_sdk!
+          require 'aws-sdk-rds'
+        rescue LoadError
+          raise LoadError,
+                "The SAML auth plugins require 'aws-sdk-rds'. Add it to your Gemfile: gem 'aws-sdk-rds'"
+        end
+      end
+    end
+  end
+end