aws_advanced_ruby_driver_wrapper 1.0.1 → 1.1.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/sql/trigger_pg.sql

@@ -0,0 +1,54 @@
+-- pgcrypto provides hmac(). If it is already installed in a schema other than public,
+-- replace public in public.hmac below with that schema.
+CREATE EXTENSION IF NOT EXISTS pgcrypto WITH SCHEMA public;
+
+-- Replace 'encrypt' below if encryption_metadata_schema is set to something else.
+-- A trigger runs with the search_path of the session that fires it. Pinning it to pg_catalog
+-- stops a session from putting its own functions or operators, such as hmac() or <>, ahead of
+-- the built-in ones to get a plaintext value past this check.
+CREATE OR REPLACE FUNCTION enforce_encrypted_column() RETURNS trigger
+SET search_path = pg_catalog, pg_temp
+AS $$
+DECLARE
+  col_name  text := TG_ARGV[0];
+  col_value bytea;
+  hmac_key  bytea;
+  unchanged boolean;
+BEGIN
+  -- An UPDATE that leaves the encrypted value as it was is not re-checked; see
+  -- "Updates after a key rotation".
+  IF TG_OP = 'UPDATE' THEN
+    EXECUTE format('SELECT ($1).%I IS NOT DISTINCT FROM ($2).%I', col_name, col_name)
+      INTO unchanged USING NEW, OLD;
+    IF unchanged THEN
+      RETURN NEW;
+    END IF;
+  END IF;
+
+  EXECUTE format('SELECT ($1).%I', col_name) INTO col_value USING NEW;
+  IF col_value IS NULL THEN
+    RETURN NEW;
+  END IF;
+
+  -- Read the column's current HMAC key from key_storage.
+  SELECT ks.hmac_key INTO hmac_key
+    FROM encrypt.encryption_metadata em
+    JOIN encrypt.key_storage ks ON em.key_id = ks.id
+   WHERE em.table_name = TG_TABLE_NAME AND em.column_name = col_name;
+
+  IF hmac_key IS NULL THEN
+    RAISE EXCEPTION 'No HMAC key is configured for %.%', TG_TABLE_NAME, col_name;
+  END IF;
+
+  IF length(col_value) < 65
+     OR substring(col_value from 1 for 32) <> public.hmac(substring(col_value from 33), hmac_key, 'sha256'::text) THEN
+    RAISE EXCEPTION 'Column %.% does not carry a valid HMAC tag (plaintext or tampered value)', TG_TABLE_NAME, col_name;
+  END IF;
+
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Records which version of this SQL is installed. Check it with:
+--   SELECT obj_description('enforce_encrypted_column'::regproc, 'pg_proc');
+COMMENT ON FUNCTION enforce_encrypted_column() IS 'aws-advanced-ruby-driver-wrapper kms_encryption trigger v2';

data/lib/aws_advanced_ruby_driver_wrapper/plugins/secrets_manager_plugin.rb

@@ -19,6 +19,7 @@ require_relative '../errors'
 require_relative '../logging'
 require_relative '../property_definition'
 require_relative '../utils/aws_credentials_utils'
+require_relative '../utils/telemetry/telemetry_const'
 
 module AwsAdvancedRubyDriverWrapper
   module Plugins
@@ -27,6 +28,7 @@ module AwsAdvancedRubyDriverWrapper
 
       SUBSCRIBED_METHODS = Set['connect', 'internal_connect'].freeze
       SECRETS_MANAGER_CACHE_NAME = :secrets_manager
+      TELEMETRY_FETCH_CREDENTIALS = 'fetch credentials'
       MIN_EXPIRATION_SEC = 300
       SYNC_FETCH_TIMEOUT_SEC = 60
       # The partition segment of an ARN varies across AWS partitions, so match any of them.
@@ -126,6 +128,8 @@ module AwsAdvancedRubyDriverWrapper
         )
 
         @subscribed_methods = SUBSCRIBED_METHODS
+        @telemetry_factory = service_container.telemetry_factory
+        @fetch_credentials_counter = @telemetry_factory.create_counter('secrets_manager.fetch_credentials.count')
       end
 
       def connect(_host_info, driver_props, _is_initial_connection, pipeline_callable)
@@ -211,7 +215,17 @@ module AwsAdvancedRubyDriverWrapper
       end
 
       # Returns true if credentials were freshly fetched from the service (suppresses login-error retry).
+      #
+      # Each lookup is counted and recorded as a span, whether it is answered from the cache or from
+      # the service.
       def fetch_secret_and_report_if_fresh?(force:)
+        @fetch_credentials_counter.inc
+        @telemetry_factory.in_context(TELEMETRY_FETCH_CREDENTIALS, Utils::Telemetry::TraceLevel::NESTED, record_failure: true) do
+          load_secret_and_report_if_fresh?(force:)
+        end
+      end
+
+      def load_secret_and_report_if_fresh?(force:)
         credentials = Utils::AwsCredentialsUtils.snapshot(@credentials_provider)
         cache_key = secret_cache_key(credentials)
         entry = cache_key && @service_container.storage_service.get(SECRETS_MANAGER_CACHE_NAME, cache_key)

data/lib/aws_advanced_ruby_driver_wrapper/postgresql.rb

@@ -20,6 +20,7 @@ require_relative '../aws_advanced_ruby_driver_wrapper'
 require_relative 'utils/connection_config_parser'
 require_relative 'utils/sql_encoding'
 require_relative 'services/service_utility'
+require_relative 'utils/telemetry/telemetry_factory'
 require_relative 'ruby_method'
 require_relative 'errors'
 
@@ -45,17 +46,20 @@ module AwsAdvancedRubyDriverWrapper
     def initialize(*, **)
       ensure_pg!
       config = Utils::ConnectionConfigParser.parse(:postgresql, *, **)
-      @service_container = Services::ServiceUtility.create_standard_container(config)
-      @service_container.host_service.refresh_host_list
-      @prepared_on = {}
-      @prepared_sql = {}
-      @async_conn = nil
-      @async_sql = nil
-      @copy_conn = nil
-      @copy_sql = nil
-      @lo_conn = nil
-      conn_service = @service_container.connection_service
-      @service_container.plugin_manager.connect(conn_service.initial_host_info, conn_service.driver_props, true)
+      telemetry_factory = Utils::Telemetry::TelemetryFactory.new(config.wrapper_props)
+      telemetry_factory.in_context("#{self.class.name}.new", Utils::Telemetry::TraceLevel::TOP_LEVEL, record_result: true) do
+        @service_container = Services::ServiceUtility.create_standard_container(config, telemetry_factory:)
+        @service_container.host_service.refresh_host_list
+        @prepared_on = {}
+        @prepared_sql = {}
+        @async_conn = nil
+        @async_sql = nil
+        @copy_conn = nil
+        @copy_sql = nil
+        @lo_conn = nil
+        conn_service = @service_container.connection_service
+        @service_container.plugin_manager.connect(conn_service.initial_host_info, conn_service.driver_props, true)
+      end
     end
 
     # Every canonical pg operation that talks to the server, and everything this class has to know about one: the
@@ -173,6 +177,10 @@ module AwsAdvancedRubyDriverWrapper
       lotell: :lo_tell, lotruncate: :lo_truncate
     }.freeze
 
+    # The operations pg yields the result of when given a block, clearing it once the block returns and
+    # returning the block's value in its place.
+    YIELDS_RESULT = %i[exec async_exec exec_params exec_prepared].freeze
+
     # A statement can also be prepared by sending a +PREPARE+ rather than by calling pg's own
     # +prepare+, and the +exec_prepared+ that runs it looks no different either way. The two are read
     # here so that a statement prepared the first way is remembered like one prepared the second, and
@@ -194,24 +202,24 @@ module AwsAdvancedRubyDriverWrapper
 
     # Explicitly define critical methods (bypass method_missing to avoid method_missing overhead).
 
-    def exec(sql, *params)
-      execute_operation(:exec, [sql, *params])
+    def exec(sql, *params, &)
+      execute_operation(:exec, [sql, *params], &)
     end
 
     # pg spells this operation +exec+, +query+, +async_exec+ and +async_query+, all of which run the
     # same libpq call. +query+ is defined here rather than left to method_missing because it is the
     # spelling applications use most after +exec+, and it enters the pipeline as +connection.exec+,
     # since that is the libpq operation being performed.
-    def query(sql, *params)
-      execute_operation(:exec, [sql, *params], spelling: :query)
+    def query(sql, *params, &)
+      execute_operation(:exec, [sql, *params], spelling: :query, &)
     end
 
-    def exec_params(sql, params, result_format = 0, type_map = nil)
-      execute_operation(:exec_params, [sql, params, result_format, type_map])
+    def exec_params(sql, params, result_format = 0, type_map = nil, &)
+      execute_operation(:exec_params, [sql, params, result_format, type_map], &)
     end
 
-    def async_exec(sql, *params)
-      execute_operation(:async_exec, [sql, *params])
+    def async_exec(sql, *params, &)
+      execute_operation(:async_exec, [sql, *params], &)
     end
 
     def transaction(&)
@@ -250,8 +258,8 @@ module AwsAdvancedRubyDriverWrapper
       execute_operation(:send_prepare, [stmt_name, sql, param_types])
     end
 
-    def exec_prepared(stmt_name, params = [], result_format = 0, type_map = nil)
-      execute_operation(:exec_prepared, [stmt_name, params, result_format, type_map])
+    def exec_prepared(stmt_name, params = [], result_format = 0, type_map = nil, &)
+      execute_operation(:exec_prepared, [stmt_name, params, result_format, type_map], &)
     end
 
     def describe_prepared(stmt_name)
@@ -312,7 +320,7 @@ module AwsAdvancedRubyDriverWrapper
 
     def method_missing(method_name, *args, **kwargs, &)
       conn = current_conn
-      raise NoMethodError, 'Connection not initialized' if conn.nil?
+      raise PG::ConnectionBad, 'connection is not initialized' if conn.nil?
       raise NoMethodError, "undefined method `#{method_name}' for #{self.class}" unless conn.respond_to?(method_name)
 
       operation = operation_for(method_name)
@@ -344,10 +352,14 @@ module AwsAdvancedRubyDriverWrapper
     # made, since an `after` step may be what forgets it, and it is handed to the result as well, so that a
     # plugin which has to inspect the statement still sees it when the rows are read.
     def execute_operation(operation, args = [], kwargs = {}, spelling: operation, &)
+      if block_given? && YIELDS_RESULT.include?(operation)
+        return yield_result(execute_operation(operation, args, kwargs, spelling: spelling), &)
+      end
+
       spec = OPERATIONS[operation] || { method: "connection.#{operation}" }
       conn = current_conn
       # Guard against a missing connection. Fail loudly instead.
-      raise NoMethodError, 'Connection not initialized' if conn.nil?
+      raise PG::ConnectionBad, 'connection is not initialized' if conn.nil?
 
       # Only forward keyword arguments when there are any.
       sql = sql_for(spec, args)
@@ -518,6 +530,12 @@ module AwsAdvancedRubyDriverWrapper
       @network_bound_methods ||= driver_dialect.network_bound_methods
     end
 
+    def yield_result(result)
+      yield result
+    ensure
+      result.clear
+    end
+
     def wrap_pg_result(result, sql = nil)
       return result unless result.is_a?(PG::Result)
 

data/lib/aws_advanced_ruby_driver_wrapper/property_definition.rb

@@ -54,8 +54,41 @@ module AwsAdvancedRubyDriverWrapper
                                                                                             type: String)
     DIALECT = WrapperProperty.new(:wrapper_dialect, 'The database dialect identifier for the database in use.', type: String)
     AWS_CREDENTIALS_PROVIDER = WrapperProperty.new(:aws_credentials_provider,
-                                                   'AWS credentials provider for IAM token generation or Secrets Manager',
+                                                   'AWS credentials provider for the plugins that call AWS services',
                                                    default_value: nil)
+    # -- Telemetry --
+    ENABLE_TELEMETRY = WrapperProperty.new(
+      :enable_telemetry,
+      'Enables telemetry and observability of the wrapper.',
+      default_value: false,
+      type: :boolean
+    )
+    TELEMETRY_SUBMIT_TOP_LEVEL = WrapperProperty.new(
+      :telemetry_submit_top_level,
+      'Submit the traces of wrapper calls as top-level traces, rather than only as children of a trace already in progress.',
+      default_value: false,
+      type: :boolean
+    )
+    TELEMETRY_TRACES_BACKEND = WrapperProperty.new(
+      :telemetry_traces_backend,
+      'Method to export the telemetry traces of the wrapper. Valid values are otlp and none.',
+      default_value: nil,
+      type: String
+    )
+    TELEMETRY_METRICS_BACKEND = WrapperProperty.new(
+      :telemetry_metrics_backend,
+      'Method to export the telemetry metrics of the wrapper. Valid values are otlp and none.',
+      default_value: nil,
+      type: String
+    )
+    TELEMETRY_FAILOVER_ADDITIONAL_TOP_TRACE = WrapperProperty.new(
+      :telemetry_failover_additional_top_trace,
+      'Also submit the trace of each failover as a top-level trace of its own. Only a failover that was ' \
+      'traced is submitted: one inside a trace in progress, or inside a wrapper call that was submitted ' \
+      'as a top-level trace.',
+      default_value: false,
+      type: :boolean
+    )
     # -- Failover --
     FAILOVER_TIMEOUT_SEC = WrapperProperty.new(
       :failover_timeout_sec,
@@ -159,6 +192,15 @@ module AwsAdvancedRubyDriverWrapper
                                                default_value: 0.1, type: Float, validator: POSITIVE_FLOAT)
     BG_SWITCHOVER_TIMEOUT_SEC = WrapperProperty.new(:bg_switchover_timeout_sec, 'Blue/Green Deployment switchover timeout (in seconds).',
                                                     default_value: 180.0, type: Float, validator: POSITIVE_FLOAT)
+    BG_DROP_BLUE_CONNECTIONS = WrapperProperty.new(
+      :bg_drop_blue_connections,
+      'Whether to invalidate connections that were opened to the old blue environment once a Blue/Green ' \
+      'switchover completes. When enabled, the next use of such a connection closes it: with the failover ' \
+      'plugin enabled it reconnects through Blue/Green routing and raises FailoverSuccessError, and without ' \
+      'failover the driver raises its closed-connection error. Either way the caller gets an explicit error ' \
+      'instead of a stale read from old blue.',
+      default_value: true, type: :boolean
+    )
 
     # -- IAM Authentication --
     IAM_HOST = WrapperProperty.new(:iam_host, 'Overrides the host used to generate the IAM token', default_value: nil, type: String)
@@ -171,6 +213,26 @@ module AwsAdvancedRubyDriverWrapper
     IAM_ACCESS_TOKEN_PROPERTY_NAME = WrapperProperty.new(:iam_access_token_property_name, 'Property name used to pass the IAM token',
                                                          default_value: :password, type: Symbol)
 
+    # -- Federated / Okta Authentication (SAML) --
+    IDP_ENDPOINT = WrapperProperty.new(:idp_endpoint, 'The hosting URL of the Identity Provider', default_value: nil, type: String)
+    IDP_PORT = WrapperProperty.new(:idp_port, 'The hosting port of the Identity Provider', default_value: 443, type: Integer)
+    IDP_USERNAME = WrapperProperty.new(:idp_username, 'The federated user name', default_value: nil, type: String)
+    IDP_PASSWORD = WrapperProperty.new(:idp_password, 'The federated user password', default_value: nil, type: String)
+    IAM_ROLE_ARN = WrapperProperty.new(:iam_role_arn, 'The ARN of the IAM Role that is to be assumed', default_value: nil, type: String)
+    IAM_IDP_ARN = WrapperProperty.new(:iam_idp_arn, 'The ARN of the Identity Provider', default_value: nil, type: String)
+    RP_IDENTIFIER = WrapperProperty.new(:rp_identifier, 'The relying party identifier', default_value: 'urn:amazon:webservices',
+                                                                                        type: String)
+    DB_USER = WrapperProperty.new(:db_user, 'The database user used to access the database', default_value: nil, type: String)
+    SSL_INSECURE = WrapperProperty.new(:ssl_insecure, "Whether the IdP server's TLS certificate chain should not be verified",
+                                       default_value: false, type: :boolean)
+    APP_ID = WrapperProperty.new(:app_id, 'The ID of the AWS application configured on Okta', default_value: nil, type: String)
+    HTTP_CLIENT_SOCKET_TIMEOUT_SEC = WrapperProperty.new(:http_client_socket_timeout_sec,
+                                                         'Socket timeout in seconds for IdP HTTP requests',
+                                                         default_value: 60.0, type: Float, validator: POSITIVE_FLOAT)
+    HTTP_CLIENT_CONNECT_TIMEOUT_SEC = WrapperProperty.new(:http_client_connect_timeout_sec,
+                                                          'Connect timeout in seconds for IdP HTTP requests',
+                                                          default_value: 60.0, type: Float, validator: POSITIVE_FLOAT)
+
     # -- Initial Connection Strategy --
     INITIAL_CONNECTION_SUBSTITUTE_HOST = WrapperProperty.new(
       :initial_connection_substitute_host,

data/lib/aws_advanced_ruby_driver_wrapper/ruby_method.rb

@@ -46,6 +46,9 @@ module AwsAdvancedRubyDriverWrapper
     CONNECTION_ABANDON_RESULTS     = define('connection.abandon_results!', check_bounded_connection: true)
     CONNECTION_SET_SERVER_OPTION   = define('connection.set_server_option', check_bounded_connection: false)
 
+    # -- Connection methods (trilogy-specific) --
+    CONNECTION_CHANGE_DB           = define('connection.change_db', check_bounded_connection: false)
+
     # -- Connection methods (pg-specific) --
     CONNECTION_EXEC                = define('connection.exec', check_bounded_connection: false)
     CONNECTION_ASYNC_EXEC          = define('connection.async_exec', check_bounded_connection: false)

data/lib/aws_advanced_ruby_driver_wrapper/services/dialect_service.rb

@@ -61,7 +61,7 @@ module AwsAdvancedRubyDriverWrapper
       }.freeze
 
       # @param connection_service [ConnectionService]
-      # @param driver_name [Symbol] :mysql2 or :postgresql
+      # @param driver_name [Symbol] :mysql2, :trilogy or :postgresql
       def initialize(connection_service, driver_name)
         @connection_service = connection_service
         @dialect_cache = {}
@@ -164,6 +164,21 @@ module AwsAdvancedRubyDriverWrapper
         @error_handler.read_only_error?(error)
       end
 
+      # Whether the driver reported the error as a database error (it carries a SQLSTATE, or for
+      # trilogy a MySQL error code). Not a strict server-vs-client split; see ErrorHandler#database_error?.
+      # @param error [Exception]
+      # @return [Boolean]
+      def database_error?(error)
+        @error_handler.database_error?(error)
+      end
+
+      # Whether the error is a schema/syntax error (SQLSTATE class '42' or its per-driver equivalent).
+      # @param error [Exception]
+      # @return [Boolean]
+      def schema_error?(error)
+        @error_handler.schema_error?(error)
+      end
+
       private
 
       # Resolves the initial database dialect from the initial connection info.
@@ -205,12 +220,10 @@ module AwsAdvancedRubyDriverWrapper
 
         rds_type = Utils::RdsUtils.identify_rds_type(host)
 
-        @dialect_code = if @driver_dialect == DriverDialects::DriverDialectManager::MYSQL_DIALECT
-                          resolve_mysql_dialect(rds_type)
-                        elsif @driver_dialect == DriverDialects::DriverDialectManager::PG_DIALECT
-                          resolve_pg_dialect(rds_type)
-                        else
-                          DialectCodes::UNKNOWN
+        @dialect_code = case @driver_dialect.database_engine
+                        when :mysql then resolve_mysql_dialect(rds_type)
+                        when :postgresql then resolve_pg_dialect(rds_type)
+                        else DialectCodes::UNKNOWN
                         end
 
         @db_dialect = dialect_for_code(@dialect_code)

data/lib/aws_advanced_ruby_driver_wrapper/services/host_service.rb

@@ -15,6 +15,7 @@
 #  limitations under the License.
 
 require 'concurrent/map'
+require_relative '../errors'
 require_relative '../host/random_host_selector'
 
 module AwsAdvancedRubyDriverWrapper
@@ -47,18 +48,45 @@ module AwsAdvancedRubyDriverWrapper
         # Register a non-default host selector. The selector is shared by every HostService in the process,
         # so it must be safe to call from multiple threads.
         #
+        # Strategy names are case-insensitive: they are normalized here and in #host_selector, so a
+        # selector registered as 'LeastLag' is found when a connection names 'leastlag'.
+        #
         # @param name [String] strategy name
         # @param selector [#select_host] any object responding to select_host(hosts, role, props)
         def register_host_selector(name, selector)
-          raise Errors::AwsError, "Cannot override default host selection strategy: '#{name}'" if DEFAULT_HOST_SELECTORS.key?(name)
+          normalized = normalize_strategy(name)
+          raise Errors::AwsError, "Cannot override default host selection strategy: '#{name}'" if DEFAULT_HOST_SELECTORS.key?(normalized)
 
-          @strategies[name] = selector
+          @strategies[normalized] = selector
         end
 
         # @param name [String] strategy name
         # @return [#select_host, nil] the registered selector, or nil if the name is unknown
         def host_selector(name)
-          @strategies[name]
+          @strategies[normalize_strategy(name)]
+        end
+
+        # Raises when no selector is registered under the given strategy name. Plugins call this when
+        # they are constructed, because they treat a selection failure during connect or failover as
+        # "no eligible host", which would otherwise hide a misspelled strategy.
+        #
+        # @param name [String] strategy name
+        def validate_strategy!(name)
+          raise unsupported_strategy_error(name) if host_selector(name).nil?
+        end
+
+        def normalize_strategy(name)
+          name.to_s.strip.downcase
+        end
+
+        # @param name [String] the strategy name as configured, reported without normalization
+        # @return [Errors::AwsError]
+        def unsupported_strategy_error(name)
+          Errors::AwsError.new(
+            "Unsupported host selection strategy: '#{name}'. Registered strategies: #{@strategies.keys.sort.join(', ')}. " \
+            'Register a custom strategy with AwsAdvancedRubyDriverWrapper::Services::HostService.register_host_selector ' \
+            'before opening a connection that uses it.'
+          )
         end
 
         # Removes every non-default host selector. For testing only.
@@ -76,9 +104,8 @@ module AwsAdvancedRubyDriverWrapper
       # @param props [Hash, nil]
       # @return [Host::HostInfo]
       def select_host(hosts, role, strategy, props = nil)
-        normalized = strategy.to_s.strip.downcase
-        selector = self.class.host_selector(normalized)
-        raise Errors::AwsError, "Unsupported host selection strategy: '#{strategy}'" if selector.nil?
+        selector = self.class.host_selector(strategy)
+        raise self.class.unsupported_strategy_error(strategy) if selector.nil?
 
         selector.select_host(hosts, role, props)
       end
@@ -90,8 +117,7 @@ module AwsAdvancedRubyDriverWrapper
       def hosts
         rules = @service_container.storage_service.get_if_registered(
           :custom_endpoint_allowed_blocked,
-          @service_container.connection_service.initial_host_info&.url,
-          register_access: false
+          @service_container.connection_service.initial_host_info&.url
         )
         return @all_hosts if rules.nil?
 

data/lib/aws_advanced_ruby_driver_wrapper/services/monitor_service.rb

@@ -144,7 +144,7 @@ module AwsAdvancedRubyDriverWrapper
       # Processes events from the event publisher.
       # @param event [Event] the event to process.
       def process_event(event)
-        return unless event == Utils::Events::DataAccessEvent
+        return unless event.is_a?(Utils::Events::DataAccessEvent)
 
         handle_data_access_event(event)
       end

data/lib/aws_advanced_ruby_driver_wrapper/services/plugin_manager.rb

@@ -19,10 +19,13 @@ require_relative '../errors'
 require_relative '../property_definition'
 require_relative '../ruby_method'
 require_relative '../utils/sql_encoding'
+require_relative '../utils/telemetry/telemetry_const'
 require_relative '../plugins/default_plugin'
 require_relative '../plugins/failover_plugin'
 require_relative '../plugins/gdb/gdb_failover_plugin'
 require_relative '../plugins/iam_auth_plugin'
+require_relative '../plugins/federated_auth/federated_auth_plugin'
+require_relative '../plugins/federated_auth/okta_auth_plugin'
 require_relative '../plugins/initial_connection_strategy_plugin'
 require_relative '../plugins/kms_encryption/kms_encryption_plugin'
 require_relative '../plugins/secrets_manager_plugin'
@@ -44,6 +47,8 @@ module AwsAdvancedRubyDriverWrapper
         'failover' => Plugins::FailoverPlugin,
         'gdb_failover' => Plugins::Gdb::GdbFailoverPlugin,
         'iam' => Plugins::IamAuthPlugin,
+        'federated_auth' => Plugins::FederatedAuth::FederatedAuthPlugin,
+        'okta' => Plugins::FederatedAuth::OktaAuthPlugin,
         'initial_connection' => Plugins::InitialConnectionStrategyPlugin,
         'kms_encryption' => Plugins::KmsEncryptionPlugin,
         'secrets_manager' => Plugins::SecretsManagerPlugin
@@ -60,9 +65,18 @@ module AwsAdvancedRubyDriverWrapper
         Plugins::Gdb::GdbFailoverPlugin => 500,
         Plugins::IamAuthPlugin => 1800,
         Plugins::SecretsManagerPlugin => 1900,
+        Plugins::FederatedAuth::FederatedAuthPlugin => 2000,
+        Plugins::FederatedAuth::OktaAuthPlugin => 2010,
         Plugins::KmsEncryptionPlugin => 2050
       }
 
+      # Plugins that authenticate by minting an RDS IAM token into the shared IAM token cache.
+      IAM_TOKEN_PLUGIN_CLASSES = [
+        Plugins::IamAuthPlugin,
+        Plugins::FederatedAuth::FederatedAuthPlugin,
+        Plugins::FederatedAuth::OktaAuthPlugin
+      ].freeze
+
       class << self
         def register_plugin(plugin_code, plugin_class, weight: WEIGHT_RELATIVE_TO_PRIOR_PLUGIN)
           @plugin_classes[plugin_code] = plugin_class
@@ -73,30 +87,35 @@ module AwsAdvancedRubyDriverWrapper
       end
 
       def initialize(service_container)
+        @telemetry_factory = service_container.telemetry_factory
         @plugins = load_plugins(service_container)
         @pipeline_cache = {}
       end
 
       def connect(host_info, driver_props, is_initial_connection, plugin_to_skip: nil)
-        execute_with_subscribed_plugins(
-          'connect',
-          lambda do |plugin, next_plugin_callable|
-            plugin.connect(host_info, driver_props, is_initial_connection, next_plugin_callable)
-          end,
-          NOOP_CALLABLE,
-          plugin_to_skip:
-        )
+        @telemetry_factory.in_context('connect', Utils::Telemetry::TraceLevel::NESTED) do
+          execute_with_subscribed_plugins(
+            'connect',
+            lambda do |plugin, next_plugin_callable|
+              plugin.connect(host_info, driver_props, is_initial_connection, next_plugin_callable)
+            end,
+            NOOP_CALLABLE,
+            plugin_to_skip:
+          )
+        end
       end
 
       def internal_connect(host_info, driver_props, wrapper_props, is_initial_connection, plugin_to_skip: nil)
-        execute_with_subscribed_plugins(
-          'internal_connect',
-          lambda do |plugin, next_plugin_callable|
-            plugin.internal_connect(host_info, driver_props, wrapper_props, is_initial_connection, next_plugin_callable)
-          end,
-          NOOP_CALLABLE,
-          plugin_to_skip:
-        )
+        @telemetry_factory.in_context('internal_connect', Utils::Telemetry::TraceLevel::NESTED) do
+          execute_with_subscribed_plugins(
+            'internal_connect',
+            lambda do |plugin, next_plugin_callable|
+              plugin.internal_connect(host_info, driver_props, wrapper_props, is_initial_connection, next_plugin_callable)
+            end,
+            NOOP_CALLABLE,
+            plugin_to_skip:
+          )
+        end
       end
 
       # The context of the call currently being executed, for plugins that need to know more about
@@ -140,15 +159,18 @@ module AwsAdvancedRubyDriverWrapper
         Thread.current[CURRENT_CALL_CONTEXT_KEY] = context
 
         begin
-          execute_with_subscribed_plugins(
-            method_name,
-            lambda do |plugin, next_plugin_callable|
-              # Read from the context rather than from args and block, so that a plugin which
-              # replaced either is honoured by the plugins after it and by the target method.
-              plugin.execute(method_name, next_plugin_callable, *context.args, **kwargs, &context.block)
-            end,
-            target_callable
-          )
+          @telemetry_factory.in_context(method_name, Utils::Telemetry::TraceLevel::TOP_LEVEL, record_result: true) do |telemetry_context|
+            telemetry_context.set_attribute(Utils::Telemetry::TelemetryConst::RUBY_CALL_ANNOTATION, method_name)
+            execute_with_subscribed_plugins(
+              method_name,
+              lambda do |plugin, next_plugin_callable|
+                # Read from the context rather than from args and block, so that a plugin which
+                # replaced either is honoured by the plugins after it and by the target method.
+                plugin.execute(method_name, next_plugin_callable, *context.args, **kwargs, &context.block)
+              end,
+              target_callable
+            )
+          end
         ensure
           Thread.current[CURRENT_CALL_CONTEXT_KEY] = previous_context
         end
@@ -164,6 +186,12 @@ module AwsAdvancedRubyDriverWrapper
         @plugins.any?(plugin_class)
       end
 
+      # Whether any plugin that mints RDS IAM tokens into the shared cache is active.
+      # @see IAM_TOKEN_PLUGIN_CLASSES
+      def iam_token_plugin_in_use?
+        IAM_TOKEN_PLUGIN_CLASSES.any? { |plugin_class| plugin_in_use?(plugin_class) }
+      end
+
       private
 
       def load_plugins(service_container)
@@ -185,7 +213,7 @@ module AwsAdvancedRubyDriverWrapper
       end
 
       def ensure_single_auth_plugin(plugin_code_list)
-        auth_plugins_used = plugin_code_list & %w[iam secrets_manager].freeze
+        auth_plugins_used = plugin_code_list & %w[iam secrets_manager federated_auth okta].freeze
 
         return unless auth_plugins_used.length > 1
 
@@ -245,29 +273,52 @@ module AwsAdvancedRubyDriverWrapper
 
       # Builds the plugin pipeline function chain. The pipeline allows plugins to perform logic both before and after
       # the target driver function is called.
+      #
+      # When traces are recorded, each plugin's part of the call is recorded as a span named after the
+      # plugin. That is decided here, once per pipeline, so a pipeline built without traces does no
+      # tracing work at all.
       def make_pipeline(target_method_name)
         subscribed_plugins = @plugins.select do |p|
           p.subscribed_methods.include?('*') || p.subscribed_methods.include?(target_method_name)
         end
         raise Errors::AwsError, 'Plugin pipeline is nil' if subscribed_plugins.empty?
 
+        traced = @telemetry_factory.traces_in_use?
         base_plugin = subscribed_plugins.last
+        base_plugin_name = plugin_name(base_plugin) if traced
         base = lambda do |plugin_callable, target_driver_callable, *, **|
-          plugin_callable.call(base_plugin, target_driver_callable)
+          if traced
+            call_traced(plugin_callable, base_plugin, base_plugin_name, target_driver_callable)
+          else
+            plugin_callable.call(base_plugin, target_driver_callable)
+          end
         end
 
         subscribed_plugins[0...-1].reverse.reduce(base) do |next_plugin_callable, plugin|
+          name = plugin_name(plugin) if traced
           lambda do |plugin_callable, target_driver_callable, method_name, plugin_to_skip|
             if plugin_to_skip == plugin
               next_plugin_callable.call(plugin_callable, target_driver_callable, method_name, plugin_to_skip)
             else
-              plugin_callable.call(plugin, lambda do
+              call_next = lambda do
                 next_plugin_callable.call(plugin_callable, target_driver_callable, method_name, plugin_to_skip)
-              end)
+              end
+              traced ? call_traced(plugin_callable, plugin, name, call_next) : plugin_callable.call(plugin, call_next)
             end
           end
         end
       end
+
+      # A plugin's span is not marked as failed when the call raises: the error passes out through every
+      # plugin's span, and is recorded on the call's top-level span instead.
+      def call_traced(plugin_callable, plugin, name, next_plugin_callable)
+        @telemetry_factory.in_context(name, Utils::Telemetry::TraceLevel::NESTED) { plugin_callable.call(plugin, next_plugin_callable) }
+      end
+
+      # The plugin's code (such as +failover+), or the class name of a plugin that has none.
+      def plugin_name(plugin)
+        self.class.plugin_classes.key(plugin.class) || plugin.class.name&.split('::')&.last || plugin.class.to_s
+      end
     end
   end
 end