aws_advanced_ruby_driver_wrapper 1.0.0 → 1.1.0

data/lib/aws_advanced_ruby_driver_wrapper/plugins/gdb/gdb_failover_plugin.rb

@@ -38,6 +38,8 @@ module AwsAdvancedRubyDriverWrapper
       # when that endpoint carries a region. When +accessible_regions+ is set, hosts outside those
       # regions are never selected.
       class GdbFailoverPlugin < FailoverPlugin
+        TELEMETRY_FAILOVER = 'global database failover'
+
         def initialize(service_container, props = ::Concurrent::Map.new)
           super
 
@@ -176,6 +178,10 @@ module AwsAdvancedRubyDriverWrapper
             return
           end
 
+          with_failover_telemetry(TELEMETRY_FAILOVER) { failover_with_plan }
+        end
+
+        def failover_with_plan
           failover_start = Time.now
           failover_deadline = failover_start + @failover_timeout
           @regionless_hosts.clear
@@ -189,6 +195,25 @@ module AwsAdvancedRubyDriverWrapper
               "failover_timeout_sec=#{@failover_timeout}"
           end
 
+          writer_candidate, writer_region, mode = failover_plan
+          log_failover_plan(writer_candidate, writer_region, mode)
+          if mode == GdbFailoverMode::STRICT_WRITER
+            count_failover(@writer_failover_counters) { failover_to_writer(writer_candidate, writer_region, failover_deadline) }
+          else
+            count_failover(@reader_failover_counters) { failover_to_allowed_host(mode, failover_deadline) }
+          end
+        ensure
+          duration_ms = ((Time.now - failover_start) * 1000).round if failover_start
+          logger.debug { "Failover duration: #{duration_ms}ms" } if duration_ms
+        end
+
+        # Finds the writer the topology settled on, and the failover mode that applies to its region.
+        #
+        # When this fails, failover cannot start, and so it is not known whether it would have been a
+        # writer or a reader failover. It is counted as a failed writer failover.
+        #
+        # @return [Array(Host::HostInfo, String, Symbol)] the writer, its region, and the mode
+        def failover_plan
           # This is expected to return once the topology has stabilized, i.e. once the cluster
           # control plane has already chosen a new writer.
           unless host_service.force_refresh_host_list?(verify_writer: true, timeout_sec: @failover_timeout)
@@ -210,15 +235,11 @@ module AwsAdvancedRubyDriverWrapper
             mode = current_failover_mode(writer_region)
           end
 
-          log_failover_plan(writer_candidate, writer_region, mode)
-          if mode == GdbFailoverMode::STRICT_WRITER
-            failover_to_writer(writer_candidate, writer_region, failover_deadline)
-          else
-            failover_to_allowed_host(mode, failover_deadline)
-          end
-        ensure
-          duration_ms = ((Time.now - failover_start) * 1000).round if failover_start
-          logger.debug { "Failover duration: #{duration_ms}ms" } if duration_ms
+          [writer_candidate, writer_region, mode]
+        rescue StandardError
+          @writer_failover_counters.triggered.inc
+          @writer_failover_counters.failed.inc
+          raise
         end
 
         # Explains which mode was chosen, and what the plugin will do with it.

data/lib/aws_advanced_ruby_driver_wrapper/plugins/iam_auth_plugin.rb

@@ -41,6 +41,17 @@ module AwsAdvancedRubyDriverWrapper
         PropertyDefinition::IAM_EXPIRATION_SEC.validate!(expiration)
         service_container.storage_service.register(IAM_TOKEN_CACHE_NAME, ttl: expiration)
         @subscribed_methods = SUBSCRIBED_METHODS
+
+        @telemetry_factory = service_container.telemetry_factory
+        @fetch_token_counter = @telemetry_factory.create_counter('iam.fetch_token.count')
+        @telemetry_factory.create_gauge('iam.token_cache.size', &self.class.token_cache_size(service_container.storage_service))
+      end
+
+      # The callback of a token cache size gauge. It is built here rather than in an instance, because
+      # a gauge is kept for the whole process and a block built in an instance would keep that
+      # instance, and the connection it belongs to, alive with it.
+      def self.token_cache_size(storage_service)
+        -> { storage_service.size(IAM_TOKEN_CACHE_NAME) }
       end
 
       def connect(host_info, driver_props, _is_initial_connection, pipeline_callable)
@@ -128,8 +139,11 @@ module AwsAdvancedRubyDriverWrapper
       # and, when there is a cache key, caches it. Without a snapshot the provider is passed as is,
       # and signing fails for lack of credentials.
       def fetch_and_cache_token(region, host, port, user, credentials, cache_key, expiration)
-        token = Aws::RDS::AuthTokenGenerator.new(credentials: credentials || @credentials_provider)
-                                            .auth_token(region:, endpoint: "#{host}:#{port}", user_name: user)
+        @fetch_token_counter.inc
+        token = Utils::IamAuthUtils.generate_token(
+          region:, hostname: host, port:, user:, credentials_provider: credentials || @credentials_provider,
+          telemetry_factory: @telemetry_factory
+        )
         if cache_key
           @service_container.storage_service.set(
             IAM_TOKEN_CACHE_NAME, cache_key, Utils::IamAuthUtils.build_token_entry(token, expiration)

data/lib/aws_advanced_ruby_driver_wrapper/plugins/initial_connection_strategy_plugin.rb

@@ -20,6 +20,7 @@ require_relative '../host/host_role'
 require_relative '../host/host_availability'
 require_relative '../logging'
 require_relative '../property_definition'
+require_relative '../services/host_service'
 require_relative '../utils/accessible_regions'
 require_relative '../utils/rds_utils'
 require_relative '../utils/rds_url_type'
@@ -53,6 +54,7 @@ module AwsAdvancedRubyDriverWrapper
         @retry_interval_sec = PropertyDefinition::INITIAL_CONNECTION_RETRY_INTERVAL_SEC.get_float(props)
         @wait_for_topology_sec = PropertyDefinition::INITIAL_CONNECTION_WAIT_FOR_TOPOLOGY_SEC.get_float(props)
         @host_selector_strategy = PropertyDefinition::INITIAL_CONNECTION_HOST_SELECTOR_STRATEGY.get(props)
+        Services::HostService.validate_strategy!(@host_selector_strategy)
         @accessible_regions = Utils::AccessibleRegions.parse(props)
         @subscribed_methods = SUBSCRIBED_METHODS
 
@@ -323,7 +325,8 @@ module AwsAdvancedRubyDriverWrapper
         end
 
         host_service.select_host(available_hosts, target_role, @host_selector_strategy)
-      rescue StandardError
+      rescue StandardError => e
+        logger.warn("Failed to select a host using the '#{@host_selector_strategy}' host selection strategy: #{e.message}")
         nil
       end
 

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/key_management_utility.rb

@@ -14,7 +14,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-require_relative '../../driver_dialects/mysql_driver_dialect'
+require_relative '../../driver_dialects/mysql2_driver_dialect'
 require_relative '../../driver_dialects/pg_driver_dialect'
 require_relative '../../logging'
 require_relative '../../property_definition'
@@ -324,7 +324,10 @@ module AwsAdvancedRubyDriverWrapper
         def dialect_for_connection(connection)
           name = connection.class.name.to_s
           return DriverDialects::PgDriverDialect.new if name.start_with?('PG::') || name.end_with?('WrapperPgConnection')
-          return DriverDialects::MysqlDriverDialect.new if name.start_with?('Mysql2::') || name.end_with?('WrapperMysql2Client')
+          return DriverDialects::Mysql2DriverDialect.new if name.start_with?('Mysql2::') || name.end_with?('WrapperMysql2Client')
+          if name == 'Trilogy' || name.end_with?('WrapperTrilogyClient')
+            raise ArgumentError, 'KMS encryption does not support the trilogy driver'
+          end
 
           raise ArgumentError,
                 "Cannot infer the driver dialect from #{name.empty? ? connection.class : name}; " \
@@ -334,7 +337,7 @@ module AwsAdvancedRubyDriverWrapper
         def dialect_for_driver(driver)
           case driver.to_sym
           when :postgresql, :pg then DriverDialects::PgDriverDialect.new
-          when :mysql2, :mysql then DriverDialects::MysqlDriverDialect.new
+          when :mysql2, :mysql then DriverDialects::Mysql2DriverDialect.new
           else raise ArgumentError, "Unknown driver #{driver.inspect}; use :postgresql or :mysql2"
           end
         end

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/key_manager.rb

@@ -14,7 +14,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-require 'base64'
 require 'digest'
 require 'securerandom'
 require_relative '../../logging'
@@ -105,7 +104,7 @@ module AwsAdvancedRubyDriverWrapper
 
           plaintext = with_retry('DECRYPT_DATA_KEY') do
             response = @kms_client.decrypt(
-              ciphertext_blob: Base64.strict_decode64(encrypted_data_key.to_s.delete("\n")),
+              ciphertext_blob: strict_decode64(encrypted_data_key.to_s.delete("\n")),
               **(master_key_arn ? { key_id: master_key_arn } : {})
             )
             response.plaintext.dup.b
@@ -130,7 +129,7 @@ module AwsAdvancedRubyDriverWrapper
             response = @kms_client.generate_data_key(key_id: master_key_arn, key_spec: KMS_KEY_SPEC)
             GeneratedDataKey.new(
               plaintext: response.plaintext.dup.b,
-              encrypted_data_key: Base64.strict_encode64(response.ciphertext_blob),
+              encrypted_data_key: strict_encode64(response.ciphertext_blob),
               hmac_key: SecureRandom.bytes(HMAC_KEY_LENGTH)
             )
           end
@@ -272,11 +271,21 @@ module AwsAdvancedRubyDriverWrapper
         # @param encrypted_data_key [String]
         # @return [String]
         def data_key_cache_key(encrypted_data_key)
-          "#{DATA_KEY_CACHE_PREFIX}#{Base64.strict_encode64(Digest::SHA256.digest(encrypted_data_key.to_s))}"
+          "#{DATA_KEY_CACHE_PREFIX}#{strict_encode64(Digest::SHA256.digest(encrypted_data_key.to_s))}"
         end
 
         private
 
+        # Strict base64 (RFC 4648: no line breaks, padding required), using the encoding built into
+        # Array#pack rather than the base64 gem. Decoding raises ArgumentError on invalid input.
+        def strict_encode64(bytes)
+          [bytes].pack('m0')
+        end
+
+        def strict_decode64(encoded)
+          encoded.unpack1('m0')
+        end
+
         def authorize_master_key!(master_key_arn)
           return if @config.master_key_allowed?(master_key_arn)
 

data/lib/aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/kms_encryption_plugin.rb

@@ -15,6 +15,7 @@
 #  limitations under the License.
 
 require 'concurrent'
+require_relative '../../driver_dialects/trilogy_driver_dialect'
 require_relative '../../logging'
 require_relative '../../ruby_method'
 require_relative '../../utils/parser/encryption_annotation_parser'
@@ -205,6 +206,12 @@ module AwsAdvancedRubyDriverWrapper
       # @param encryption_utility [Encryption::KmsEncryptionUtility, nil] a utility to use instead
       #   of building one
       def initialize(service_container, props = ::Concurrent::Map.new, encryption_utility: nil)
+        # Trilogy has no prepared statements, so a parameter could only reach the server written into the
+        # SQL, in plaintext, before it could be encrypted.
+        if service_container.dialect_service.driver_dialect.is_a?(DriverDialects::TrilogyDriverDialect)
+          raise Errors::EncryptionPluginError, 'The kms_encryption plugin does not support the trilogy driver'
+        end
+
         @service_container = service_container
         @encryption_utility = encryption_utility || Encryption::KmsEncryptionUtility.new(service_container, props)
         # Built up front so that a parser dependency the application has not installed (pg_query, on
@@ -586,9 +593,10 @@ module AwsAdvancedRubyDriverWrapper
       # @raise [Errors::MetadataError] when the statement writes a column the plugin has confirmed is
       #   encrypted with a value it cannot encrypt (see {#check_write})
       def parameter_columns(sql)
-        annotations = Utils::Parser::EncryptionAnnotationParser.parse_annotations(sql)
+        version = executable_comment_version
+        annotations = Utils::Parser::EncryptionAnnotationParser.parse_annotations(sql, server_version: version)
         stripped = stripped_sql(sql)
-        analysis = analysis_of(stripped)
+        analysis = analysis_of(stripped, version)
         write = write_statement?(analysis, stripped)
         inferred = analysis.parameter_column_names
         return {} if !write && annotations.empty? && inferred.empty?
@@ -607,9 +615,10 @@ module AwsAdvancedRubyDriverWrapper
       # that writes nothing, so it is reported as nothing having been read rather than as nothing
       # being there to read.
       #
+      # @param version [Integer, nil] see {#executable_comment_version}
       # @return [Utils::Parser::SqlParser::SqlAnalysisResult]
-      def analysis_of(sql)
-        sql_parser.analyze_sql(sql)
+      def analysis_of(sql, version)
+        sql_parser.analyze_sql(sql, server_version: version)
       rescue StandardError => e
         logger.warn("The statement could not be analysed: #{e.message}")
         Utils::Parser::SqlParser::SqlAnalysisResult.new(
@@ -693,8 +702,9 @@ module AwsAdvancedRubyDriverWrapper
       def column_configs(sql)
         return {} unless ready_for_statement?(sql)
 
-        annotations = Utils::Parser::EncryptionAnnotationParser.parse_annotations(sql)
-        tables = statement_tables(analysis_of(stripped_sql(sql)), annotations)
+        version = executable_comment_version
+        annotations = Utils::Parser::EncryptionAnnotationParser.parse_annotations(sql, server_version: version)
+        tables = statement_tables(analysis_of(stripped_sql(sql), version), annotations)
         return {} if tables.empty?
 
         # When two of the statement's tables encrypt a column of the same name, the first of them
@@ -837,6 +847,19 @@ module AwsAdvancedRubyDriverWrapper
         Utils::Parser::EncryptionAnnotationParser.strip_annotations(sql)
       end
 
+      # The version of the MySQL server the statement is sent to, which decides whether a versioned
+      # /*!80023 ... */ comment is part of the statement, and so which column each placeholder after it
+      # belongs to (see DriverDialects::DriverDialect#executable_comment_version). It is read from the
+      # connection every time rather than kept, since a failover or a blue/green switchover can move
+      # the connection to a server of another version.
+      #
+      # @return [Integer, nil] nil on PostgreSQL, which has no executable comments, and when the
+      #   connection cannot say
+      def executable_comment_version
+        @service_container.dialect_service.driver_dialect
+                          .executable_comment_version(@service_container.connection_service.current_connection)
+      end
+
       # -- Lazily built collaborators --
 
       # Builds the parts of the plugin that need a database connection, the first time a statement

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

@@ -0,0 +1,108 @@
+DELIMITER $$
+
+-- HMAC-SHA256 built from SHA2(), which MySQL has no native HMAC for.
+DROP FUNCTION IF EXISTS hmac_sha256$$
+CREATE FUNCTION hmac_sha256(
+    message VARBINARY(65535),
+    key_data VARBINARY(64)
+)
+RETURNS VARBINARY(32)
+DETERMINISTIC
+NO SQL
+BEGIN
+    DECLARE block_size INT DEFAULT 64;  -- SHA-256 block size
+    DECLARE key_padded VARBINARY(64);
+    DECLARE ipad VARBINARY(64);
+    DECLARE opad VARBINARY(64);
+    DECLARE i INT DEFAULT 1;
+    DECLARE key_byte INT;
+
+    -- Reduce an over-long key to its hash, then zero-pad the key to the block size.
+    IF LENGTH(key_data) > block_size THEN
+        SET key_padded = CONCAT(UNHEX(SHA2(key_data, 256)), REPEAT(X'00', block_size - 32));
+    ELSE
+        SET key_padded = CONCAT(key_data, REPEAT(X'00', block_size - LENGTH(key_data)));
+    END IF;
+
+    -- ipad = key XOR 0x36 per byte, opad = key XOR 0x5C per byte, one raw byte at a time.
+    SET ipad = X'';
+    SET opad = X'';
+    WHILE i <= block_size DO
+        SET key_byte = ASCII(SUBSTRING(key_padded, i, 1));
+        SET ipad = CONCAT(ipad, UNHEX(LPAD(HEX(key_byte ^ 0x36), 2, '0')));
+        SET opad = CONCAT(opad, UNHEX(LPAD(HEX(key_byte ^ 0x5C), 2, '0')));
+        SET i = i + 1;
+    END WHILE;
+
+    RETURN UNHEX(SHA2(CONCAT(opad, UNHEX(SHA2(CONCAT(ipad, message), 256))), 256));
+END$$
+
+-- Verifies that a stored value carries a valid HMAC for the given HMAC key.
+DROP FUNCTION IF EXISTS verify_encrypted_data_hmac$$
+CREATE FUNCTION verify_encrypted_data_hmac(
+    data VARBINARY(65535),
+    hmac_key VARBINARY(32)
+)
+RETURNS BOOLEAN
+DETERMINISTIC
+NO SQL
+BEGIN
+    -- Minimum payload: 32 (HMAC) + 4 (key id) + 1 (type) + 12 (IV) + 0 (ciphertext) + 16 (GCM tag) = 65.
+    IF data IS NULL OR LENGTH(data) < 65 THEN
+        RETURN FALSE;
+    END IF;
+
+    RETURN SUBSTRING(data, 1, 32) = hmac_sha256(SUBSTRING(data, 33), hmac_key);
+END$$
+
+-- Called from a BEFORE INSERT trigger (and by the BEFORE UPDATE procedure) to reject a value that is
+-- not a valid kms_encryption payload. Replace SCHEMA_NAME with encryption_metadata_schema.
+-- The COMMENT records which version of this SQL is installed. Check it with:
+--   SELECT ROUTINE_COMMENT FROM information_schema.ROUTINES
+--    WHERE ROUTINE_SCHEMA = DATABASE() AND ROUTINE_NAME = 'validate_encrypted_data_hmac_before_insert';
+DROP PROCEDURE IF EXISTS validate_encrypted_data_hmac_before_insert$$
+CREATE PROCEDURE validate_encrypted_data_hmac_before_insert(
+    IN p_table_name VARCHAR(64),
+    IN p_column_name VARCHAR(64),
+    IN column_value VARBINARY(65535)
+)
+COMMENT 'aws-advanced-ruby-driver-wrapper kms_encryption trigger v1'
+BEGIN
+    DECLARE v_hmac_key VARBINARY(32);
+
+    IF column_value IS NOT NULL THEN
+        SELECT ks.hmac_key INTO v_hmac_key
+        FROM SCHEMA_NAME.encryption_metadata em
+        JOIN SCHEMA_NAME.key_storage ks ON em.key_id = ks.id
+        WHERE em.table_name = p_table_name
+          AND em.column_name = p_column_name
+        LIMIT 1;
+
+        IF v_hmac_key IS NULL THEN
+            SIGNAL SQLSTATE '45000'
+            SET MESSAGE_TEXT = 'No HMAC key is configured for the encrypted column';
+        END IF;
+
+        IF NOT verify_encrypted_data_hmac(column_value, v_hmac_key) THEN
+            SIGNAL SQLSTATE '45000'
+            SET MESSAGE_TEXT = 'Column does not carry a valid HMAC tag (plaintext or tampered value)';
+        END IF;
+    END IF;
+END$$
+
+-- Called from a BEFORE UPDATE trigger. An UPDATE that leaves the value as it was is not
+-- re-checked; see "Updates after a key rotation".
+DROP PROCEDURE IF EXISTS validate_encrypted_data_hmac_before_update$$
+CREATE PROCEDURE validate_encrypted_data_hmac_before_update(
+    IN p_table_name VARCHAR(64),
+    IN p_column_name VARCHAR(64),
+    IN new_value VARBINARY(65535),
+    IN old_value VARBINARY(65535)
+)
+BEGIN
+    IF NOT (new_value <=> old_value) THEN
+        CALL validate_encrypted_data_hmac_before_insert(p_table_name, p_column_name, new_value);
+    END IF;
+END$$
+
+DELIMITER ;

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.
@@ -79,6 +81,12 @@ module AwsAdvancedRubyDriverWrapper
           storage_service.clear(SECRETS_MANAGER_CACHE_NAME)
         end
 
+        # Forgets fetches inherited by a forked child. A fetch running at fork time lost its thread,
+        # so its future never resolves and would otherwise be returned to every later fetch for its key.
+        def release_pending_refreshes_after_fork
+          @pending_refreshes.clear
+        end
+
         # The secret's value depends on which secret is read, from which region and endpoint, and
         # with which AWS credentials, so all of them are part of the key. Connections that differ in
         # any of them neither share a cached secret nor wait on each other's fetch.
@@ -120,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)
@@ -205,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)
@@ -244,7 +264,9 @@ module AwsAdvancedRubyDriverWrapper
           else
             Concurrent::Promises.future_on(:io) { fetch_and_store_secret(nil, credentials) }
           end
-        future.value!(SYNC_FETCH_TIMEOUT_SEC)
+        # value! returns nil instead of raising when the timeout elapses; a completed fetch always
+        # returns an entry or raises.
+        future.value!(SYNC_FETCH_TIMEOUT_SEC) || raise(Timeout::Error)
       rescue Concurrent::CancelledOperationError, Timeout::Error
         raise Errors::SecretsManagerAuthError,
               "Timed out fetching secret after #{SYNC_FETCH_TIMEOUT_SEC}s"

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(&)
@@ -224,6 +232,14 @@ module AwsAdvancedRubyDriverWrapper
 
     alias finish close
 
+    # Stops the background topology monitor of this connection's cluster, which closes the connections it
+    # holds to the database it was started for. The next connection that needs topology starts a new one.
+    # PostgreSQL refuses to drop a database that other sessions are connected to, so this has to happen
+    # before dropping the database the monitor was started for.
+    def stop_topology_monitor
+      @service_container.host_service.host_list_provider&.stop_monitor
+    end
+
     # Resets the connection through the pipeline and returns this wrapper, so the reset connection stays
     # usable through it. The driver's own reset tears down and re-establishes the underlying socket, which
     # clears any server-side session state, so the tracked session state is reset to match.
@@ -242,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)
@@ -304,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)
@@ -336,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)
@@ -510,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)