aws_advanced_ruby_driver_wrapper 1.0.0 → 1.1.0

data/lib/aws_advanced_ruby_driver_wrapper/active_record/aws_postgresql_adapter.rb

@@ -17,40 +17,82 @@
 require 'active_record/connection_adapters/postgresql_adapter'
 require_relative '../postgresql'
 require_relative '../errors'
-require_relative '../utils/ar_constants'
+require_relative 'type_adapter_alias'
+require_relative 'aws_connection_handler'
+require_relative 'aws_action_cable_postgresql_support'
 
 module ActiveRecord
   module ConnectionAdapters
     class AwsPostgreSQLAdapter < PostgreSQLAdapter
-      ADAPTER_NAME = 'AwsPostgreSQL'
-
-      # ActiveRecord-only keys that should not be passed to the wrapper or native driver.
-      AR_ONLY_KEYS = (AwsAdvancedRubyDriverWrapper::Utils::AR_COMMON_KEYS + %i[advisory_locks schema_search_path]).freeze
-
       # ActiveRecord uses :username and :database, but the PG gem expects :user and :dbname.
       # The parent adapter translates these in @connection_parameters, but also strips
       # non-PG keys (including wrapper properties) via slice!. Since we rebuild from @config
       # to preserve wrapper properties, we need to redo this translation ourselves.
       AR_TO_PG_KEY_MAP = { username: :user, database: :dbname }.freeze
 
+      # The connection options PG accepts, the same set the parent adapter forwards.
+      PG_CONNECTION_KEYS = (PG::Connection.conndefaults_hash.keys + [:requiressl]).to_set.freeze
+
       def initialize(...)
         # Capture the full config before the parent's initialize strips it
         # down to only PG-recognized keys in @connection_parameters.
         super
         @connection_broken = false
-        # Store non-AR parameters in @wrapper_config.
-        wrapper_config = @config.compact.except(*AR_ONLY_KEYS)
+        wrapper_config = @config.compact
         # Remap ActiveRecord key names to PG gem key names.
         AR_TO_PG_KEY_MAP.each { |ar_key, pg_key| wrapper_config[pg_key] = wrapper_config.delete(ar_key) if wrapper_config.key?(ar_key) }
-        @wrapper_config = wrapper_config
+        # Keep only what PG or the wrapper understands. Every other key is an ActiveRecord setting,
+        # and PG rejects any connection option it does not recognize.
+        @wrapper_config = wrapper_config.select { |key, _| self.class.forwarded_key?(key) }
+      end
+
+      # @return [Boolean] whether the config key is passed on to the wrapper connection
+      def self.forwarded_key?(key)
+        PG_CONNECTION_KEYS.include?(key) ||
+          AwsAdvancedRubyDriverWrapper::PropertyDefinition.wrapper_property?(key) ||
+          AwsAdvancedRubyDriverWrapper::PropertyDefinition::KNOWN_PREFIXES.any? { |prefix| key.to_s.start_with?(prefix) }
+      end
+
+      # Resolve native database types from the wrapped PostgreSQLAdapter, overriding only :datetime.
+      # PostgreSQLAdapter builds its type hash from two things: the shared NATIVE_DATABASE_TYPES
+      # constant and the per-class datetime_type class_attribute. Delegating to the parent keeps a
+      # single source of truth, but the parent reads *its own* datetime_type, so this adapter's
+      # datetime_type (e.g. :timestamptz) must be re-applied here or it is silently ignored and
+      # t.datetime columns are created as plain timestamp.
+      def self.native_database_types
+        superclass.native_database_types.merge(datetime: NATIVE_DATABASE_TYPES[datetime_type])
       end
 
-      def adapter_name = ADAPTER_NAME
+      # SQLSTATEs and FATAL-message signals for an authentication/authorization rejection (invalid
+      # password, PAM failure, pg_hba.conf rejection). A failed *connect* usually carries no result, so
+      # the message text is the reliable signal and the SQLSTATE is only a bonus when one is present.
+      AUTH_REJECTION_SQLSTATES = %w[28P01 28000].freeze
+      AUTH_REJECTION_MESSAGE = /password authentication failed|PAM authentication failed|pg_hba\.conf/
 
       # Note that this config includes wrapper properties.
+      #
+      # Mirrors the parent PostgreSQLAdapter.new_client so ActiveRecord can tell a missing database
+      # (which db:prepare creates) or rejected credentials from other failures.
       def self.new_client(config)
         AwsAdvancedRubyDriverWrapper::WrapperPgConnection.new(**config)
+      rescue ::PG::Error => e
+        dbname, user, host = config.values_at(:dbname, :user, :host)
+        # The postgres maintenance database always exists, so a failure there is not a missing database.
+        raise ActiveRecord::ConnectionNotEstablished, e.message if dbname == 'postgres'
+        raise ActiveRecord::NoDatabaseError.db_error(dbname) if dbname && e.message.include?(dbname) && !auth_rejection?(e)
+        raise ActiveRecord::DatabaseConnectionError.username_error(user) if user && e.message.include?(user)
+        raise ActiveRecord::DatabaseConnectionError.hostname_error(host) if host && e.message.include?(host)
+
+        raise ActiveRecord::ConnectionNotEstablished, e.message
+      end
+
+      # @return [Boolean] whether the PG error is an authentication/authorization rejection, by SQLSTATE
+      #   when the error carries a result and otherwise by its canonical FATAL text.
+      def self.auth_rejection?(error)
+        sqlstate = error.result&.error_field(PG::PG_DIAG_SQLSTATE)
+        AUTH_REJECTION_SQLSTATES.include?(sqlstate) || error.message.match?(AUTH_REJECTION_MESSAGE)
       end
+      private_class_method :auth_rejection?
 
       def connect
         @connection_broken = false
@@ -71,10 +113,23 @@ module ActiveRecord
         super
       end
 
+      # Used by db:drop, db:reset, db:test:prepare and db:purge. ActiveRecord disconnects its own
+      # connections first, but the topology monitor shared by this cluster's connections, and the Blue/Green
+      # status providers, keep their own connections to the database, and PostgreSQL refuses to drop a
+      # database other sessions are using.
+      def drop_database(name)
+        raw_connection.stop_topology_monitor
+        AwsAdvancedRubyDriverWrapper::Plugins::BlueGreen::BlueGreenPlugin.clean_up_providers
+        super
+      end
+
       def translate_exception(exception, message:, sql:, binds:)
         return super unless exception.is_a?(AwsAdvancedRubyDriverWrapper::Errors::AwsError)
 
         if exception.needs_reconfiguration
+          # The wrapper has reconnected this connection object to a new physical connection, where the
+          # statements prepared on the old one do not exist. Forget them so they are prepared again.
+          clear_cache!(new_connection: true)
           configure_connection
           exception
         elsif exception.is_a?(AwsAdvancedRubyDriverWrapper::Errors::FailoverFailedError)

data/lib/aws_advanced_ruby_driver_wrapper/active_record/aws_trilogy_adapter.rb

@@ -0,0 +1,80 @@
+# 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 'active_record/connection_adapters/trilogy_adapter'
+require_relative '../trilogy'
+require_relative '../errors'
+require_relative 'type_adapter_alias'
+require_relative 'aws_connection_handler'
+
+module ActiveRecord
+  module ConnectionAdapters
+    class AwsTrilogyAdapter < TrilogyAdapter
+      # Connects as the parent adapter does, through the wrapper instead of trilogy itself: the ssl_mode
+      # name is turned into the constant trilogy reads, and the whole config is passed on, since the
+      # wrapper takes its own properties out and trilogy ignores the ActiveRecord-only keys.
+      #
+      # Translates connection errors the same way the parent adapter does, so ActiveRecord can tell a
+      # missing database (which db:prepare creates) or rejected credentials from other failures.
+      def self.new_client(config)
+        config[:ssl_mode] = parse_ssl_mode(config[:ssl_mode]) if config[:ssl_mode]
+        AwsAdvancedRubyDriverWrapper::WrapperTrilogyClient.new(config)
+      rescue ::Trilogy::Error => e
+        raise translate_connect_error(config, e)
+      end
+
+      def connect
+        @connection_broken = false
+        super
+      end
+
+      def reconnect
+        @connection_broken = false
+        super
+      end
+
+      def active?
+        return false if @connection_broken
+
+        super
+      end
+
+      def translate_exception(exception, message:, sql:, binds:)
+        return super unless exception.is_a?(AwsAdvancedRubyDriverWrapper::Errors::AwsError)
+
+        if exception.needs_reconfiguration
+          # The wrapper has reconnected this connection object to a new physical connection, where the
+          # session ActiveRecord configured on the old one does not exist. Forget what was cached about
+          # that connection and configure the new one.
+          clear_cache!(new_connection: true)
+          configure_connection
+          exception
+        elsif exception.is_a?(AwsAdvancedRubyDriverWrapper::Errors::FailoverFailedError)
+          @connection_broken = true
+          ActiveRecord::ConnectionFailed.new(message, sql:, binds:, connection_pool: @pool)
+        else
+          exception
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters.register(
+  'aws_trilogy',
+  'ActiveRecord::ConnectionAdapters::AwsTrilogyAdapter',
+  'aws_advanced_ruby_driver_wrapper/active_record/aws_trilogy_adapter'
+)

data/lib/aws_advanced_ruby_driver_wrapper/active_record/type_adapter_alias.rb

@@ -0,0 +1,61 @@
+# 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 'active_record'
+
+module AwsAdvancedRubyDriverWrapper
+  # ActiveRecord's adapter-specific type registry (ActiveRecord::Type) keys every type
+  # registration and lookup by the *config adapter name*. The wrapped adapters register nothing of
+  # their own: they rely entirely on the vanilla PostgreSQL/Mysql2 adapters, which register their
+  # OID / native types under adapter: :postgresql and adapter: :mysql2. A lookup for an
+  # :aws_postgresql / :aws_mysql2 connection would therefore find none of them and raise
+  # "Unknown type" for things like :interval, :point, array/range columns, or MySQL's
+  # :unsigned_integer.
+  #
+  # Rather than copy the parent's registrations (which misses types registered after the adapter
+  # loads, needs a duplicate guard, and reaches into the registry's private internals), alias the
+  # wrapper adapter names to their parents inside the registry's own register/add_modifier/lookup.
+  # Every registration and every lookup — including the ones ActiveRecord routes through
+  # Type.adapter_name_from — passes through the registry, so normalizing the adapter there makes the
+  # wrapped adapters resolve the exact same types as the vanilla ones, for all current and future
+  # types, across Rails 7.2–8.1.
+  #
+  # Side effect: in an app that uses BOTH a plain postgresql and an aws_postgresql connection, a
+  # type explicitly registered for :aws_postgresql is normalized to :postgresql and so also applies
+  # to the plain postgresql connections (and vice versa). This is documented in the CHANGELOG.
+  module TypeAdapterAlias
+    ADAPTER_ALIASES = { aws_postgresql: :postgresql, aws_mysql2: :mysql2, aws_trilogy: :trilogy }.freeze
+
+    # Prepended onto ActiveRecord::Type::AdapterSpecificRegistry. Normalizing the adapter in lookup
+    # covers both the framework path and a direct Type.lookup(..., adapter: :aws_postgresql); any
+    # type registered explicitly for a wrapper adapter lands under the parent's key.
+    module RegistryAlias
+      def register(type_name, klass = nil, adapter: nil, **, &)
+        super(type_name, klass, adapter: ADAPTER_ALIASES.fetch(adapter, adapter), **, &)
+      end
+
+      def add_modifier(options, klass, adapter: nil, **args)
+        super(options, klass, adapter: ADAPTER_ALIASES.fetch(adapter, adapter), **args)
+      end
+
+      def lookup(symbol, *, adapter: nil, **)
+        super(symbol, *, adapter: ADAPTER_ALIASES.fetch(adapter, adapter), **)
+      end
+    end
+  end
+end
+
+ActiveRecord::Type::AdapterSpecificRegistry.prepend(AwsAdvancedRubyDriverWrapper::TypeAdapterAlias::RegistryAlias)

data/lib/aws_advanced_ruby_driver_wrapper/driver_dialects/driver_dialect.rb

@@ -33,6 +33,17 @@ module AwsAdvancedRubyDriverWrapper
 
       DEFAULT_MONITORING_TIMEOUT_SEC = 5
 
+      # @return [String] the name of the Ruby driver gem this dialect drives
+      def name
+        raise NotImplementedError
+      end
+
+      # The database engine the driver talks to, which decides the database dialects that apply.
+      # @return [Symbol] +:mysql+ or +:postgresql+
+      def database_engine
+        raise NotImplementedError
+      end
+
       def connect(host_info, config)
         raise NotImplementedError
       end
@@ -115,6 +126,11 @@ module AwsAdvancedRubyDriverWrapper
         raise NotImplementedError
       end
 
+      # Releases a connection inherited from the parent process after a fork, without telling the
+      # server: the parent still owns the session. Garbage collection or process exit in the child
+      # must not end it.
+      def abandon_connection(connection); end
+
       def sql_state(_exception)
         nil
       end
@@ -133,6 +149,29 @@ module AwsAdvancedRubyDriverWrapper
         :user
       end
 
+      # Returns the property key the underlying driver expects for the TLS CA certificate file.
+      # Used by Blue/Green substitute routing for drivers that overload the host for both the socket
+      # address and the TLS server name (mysql2, trilogy): it tells whether TLS is active before
+      # deciding to connect by hostname rather than raw IP, avoiding a certificate-name mismatch.
+      # Drivers that can pin the socket separately (see #props_pinned_to_ip) do not use this.
+      # mysql2 names it +:sslca+ (the default); override where it differs (trilogy uses +:ssl_ca+).
+      def tls_ca_property_key
+        :sslca
+      end
+
+      # Blue/Green substitute routing must connect to a specific instance IP — so a DNS lookup during
+      # the switchover window cannot resolve to the old instance — while still presenting the original
+      # hostname for TLS certificate verification and IAM token generation. A driver that can pin the
+      # socket address separately from the TLS/auth server name returns a copy of +driver_props+ that
+      # connects to +ip+ but keeps the hostname for verification; a driver that overloads the host for
+      # both returns +nil+, and the routing falls back to the #tls_ca_property_key hostname check.
+      # @param driver_props [Hash]
+      # @param ip [String] the instance IP to pin the socket to
+      # @return [Hash, nil] props pinned to +ip+, or +nil+ when the driver cannot split socket from TLS name
+      def props_pinned_to_ip(_driver_props, _ip)
+        nil
+      end
+
       # Applies default socket/connect timeouts to monitoring connection driver props.
       # These ensure that a query or close on a dead connection raises a timeout error
       # rather than hanging indefinitely or segfaulting.
@@ -149,6 +188,15 @@ module AwsAdvancedRubyDriverWrapper
       def reported_in_transaction(_connection)
         nil
       end
+
+      # The server version that decides which versioned /*!80023 ... */ executable comments the server
+      # runs, as MySQL numbers it (8.0.39 is 80039). Only MySQL has executable comments.
+      # @param connection [Object]
+      # @return [Integer, nil] nil when the driver's SQL has no executable comments, or the version
+      #   cannot be read
+      def executable_comment_version(_connection)
+        nil
+      end
     end
   end
 end

data/lib/aws_advanced_ruby_driver_wrapper/driver_dialects/driver_dialect_manager.rb

@@ -14,26 +14,33 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-require_relative 'mysql_driver_dialect'
+require_relative 'mysql2_driver_dialect'
 require_relative 'pg_driver_dialect'
+require_relative 'trilogy_driver_dialect'
 require_relative '../errors'
 require_relative '../errors/mysql_error_handler'
 require_relative '../errors/pg_error_handler'
+require_relative '../errors/trilogy_error_handler'
 
 module AwsAdvancedRubyDriverWrapper
   module DriverDialects
     module DriverDialectManager
-      MYSQL_DIALECT = MysqlDriverDialect.new.freeze
+      MYSQL2_DIALECT = Mysql2DriverDialect.new.freeze
       PG_DIALECT = PgDriverDialect.new.freeze
+      TRILOGY_DIALECT = TrilogyDriverDialect.new.freeze
 
       REGISTRY = {
         mysql2: {
-          driver_dialect: MYSQL_DIALECT,
-          error_handler: Errors::MysqlErrorHandler.new(MYSQL_DIALECT)
+          driver_dialect: MYSQL2_DIALECT,
+          error_handler: Errors::MysqlErrorHandler.new(MYSQL2_DIALECT)
         },
         postgresql: {
           driver_dialect: PG_DIALECT,
           error_handler: Errors::PgErrorHandler.new(PG_DIALECT)
+        },
+        trilogy: {
+          driver_dialect: TRILOGY_DIALECT,
+          error_handler: Errors::TrilogyErrorHandler.new(TRILOGY_DIALECT)
         }
       }.freeze
 

data/lib/aws_advanced_ruby_driver_wrapper/driver_dialects/mysql2_driver_dialect.rb

@@ -0,0 +1,129 @@
+# 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_relative 'mysql_driver_dialect'
+
+module AwsAdvancedRubyDriverWrapper
+  module DriverDialects
+    # How the mysql2 driver is called. The MySQL grammar and what every MySQL driver shares are
+    # inherited.
+    class Mysql2DriverDialect < MysqlDriverDialect
+      # Every mysql2 call that talks to the server. A call that is not listed here is handed straight
+      # to the driver, bypassing the plugin pipeline.
+      #
+      # Not listed, because libmysql answers them without talking to the server: escape, the row and
+      # column counts, last_id, affected_rows, info, warning_count, thread_id, server_info,
+      # session_track, and the connection's own settings.
+      NETWORK_BOUND_METHODS = (COMMON_NETWORK_BOUND_METHODS | Set[
+        RubyMethod::CONNECTION_QUERY.name,
+        RubyMethod::CONNECTION_ASYNC_RESULT.name,
+        RubyMethod::CONNECTION_SELECT_DB.name,
+        RubyMethod::CONNECTION_MORE_RESULTS.name,
+        RubyMethod::CONNECTION_NEXT_RESULT.name,
+        RubyMethod::CONNECTION_STORE_RESULT.name,
+        RubyMethod::CONNECTION_ABANDON_RESULTS.name,
+        RubyMethod::CONNECTION_SET_SERVER_OPTION.name,
+        RubyMethod::RESULT_EACH.name,
+        RubyMethod::RESULT_FREE.name
+      ]).freeze
+
+      def name
+        'mysql2'
+      end
+
+      def connect(host_info, config)
+        ::Mysql2::Client.new(**prepare_connect_config(host_info, config))
+      end
+
+      def execute(connection, sql)
+        raise Mysql2::Error, 'MySQL client is not connected' if connection.nil? || connection.closed?
+
+        connection.query(sql)
+      end
+
+      def execute_with_params(connection, sql, params)
+        raise Mysql2::Error, 'MySQL client is not connected' if connection.nil? || connection.closed?
+
+        stmt = connection.prepare(sql)
+        result = stmt.execute(*params)
+        # Everything the caller needs must be read off the statement before the ensure closes its
+        # handle. A SELECT result is bound to the handle, so materialize its rows now - otherwise a
+        # later iteration (e.g. SqlRunner calling #to_a) raises "Statement handle already closed".
+        # A non-SELECT statement returns nil, and its affected-row count must come from the statement
+        # too: once the handle is closed the connection-level mysql_affected_rows returns its -1 error
+        # sentinel and mysql2 raises. So return that count for #affected_rows to hand back.
+        result.is_a?(::Mysql2::Result) ? result.to_a : stmt.affected_rows
+      ensure
+        stmt&.close
+      end
+
+      # mysql2 binds a blob parameter as raw binary bytes.
+      def binary_param(bytes)
+        bytes.b
+      end
+
+      # mysql2 hands back a blob column as a string that only needs its encoding forced to binary.
+      def read_binary(value)
+        value.b
+      end
+
+      # A prepared statement (execute_with_params) hands back its own affected-row count, since the
+      # connection-level count is unreadable once the statement handle is closed. A plain
+      # connection.query (execute) leaves nil, so the count is read off the connection instead.
+      def affected_rows(connection, result)
+        result.is_a?(Integer) ? result : connection.affected_rows.to_i
+      end
+
+      # mysql2 has no RETURNING clause, so the generated id is read from the connection afterwards.
+      def insert_returning_id(connection, sql, params, _id_column)
+        execute_with_params(connection, sql, params)
+        id = connection.last_id
+        id&.positive? ? id : nil
+      end
+
+      # With automatic_close off, mysql2 invalidates the socket instead of sending COM_QUIT when
+      # the client is collected, leaving the parent's session open.
+      def abandon_connection(connection)
+        return if connection.nil? || connection.closed?
+
+        connection.automatic_close = false
+      rescue StandardError => e
+        logger.debug("Failed to abandon MySQL connection: #{e.message}")
+      end
+
+      # mysql2 reports the version it learned when it connected, so reading it costs no round trip.
+      def executable_comment_version(connection)
+        return nil if connection.nil?
+
+        info = connection.server_info
+        info.is_a?(Hash) ? info[:id] : nil
+      rescue StandardError => e
+        logger.debug { "The MySQL server version could not be read: #{e.message}" }
+        nil
+      end
+
+      def sql_state(exception)
+        return nil unless exception.respond_to?(:sql_state)
+
+        exception.sql_state
+      end
+
+      def network_bound_methods
+        NETWORK_BOUND_METHODS
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/driver_dialects/mysql_driver_dialect.rb

@@ -19,84 +19,22 @@ require_relative '../host/connection_string_host_list_provider'
 
 module AwsAdvancedRubyDriverWrapper
   module DriverDialects
+    # What every MySQL driver shares: the server's SQL grammar, and the connection options and
+    # connection calls the MySQL drivers name alike. A driver's own dialect subclasses this and adds
+    # how the driver itself is called.
     class MysqlDriverDialect
       include DriverDialect
 
-      # Every mysql2 call that talks to the server. A call that is not listed here is handed straight
-      # to the driver, bypassing the plugin pipeline.
-      #
-      # Not listed, because libmysql answers them without talking to the server: escape, the row and
-      # column counts, last_id, affected_rows, info, warning_count, thread_id, server_info,
-      # session_track, and the connection's own settings.
-      NETWORK_BOUND_METHODS = (COMMON_NETWORK_BOUND_METHODS | Set[
-        RubyMethod::CONNECTION_QUERY.name,
-        RubyMethod::CONNECTION_ASYNC_RESULT.name,
-        RubyMethod::CONNECTION_SELECT_DB.name,
-        RubyMethod::CONNECTION_MORE_RESULTS.name,
-        RubyMethod::CONNECTION_NEXT_RESULT.name,
-        RubyMethod::CONNECTION_STORE_RESULT.name,
-        RubyMethod::CONNECTION_ABANDON_RESULTS.name,
-        RubyMethod::CONNECTION_SET_SERVER_OPTION.name,
-        RubyMethod::RESULT_EACH.name,
-        RubyMethod::RESULT_FREE.name
-      ]).freeze
-
-      def connect(host_info, config)
-        ::Mysql2::Client.new(**prepare_connect_config(host_info, config))
-      end
-
-      def execute(connection, sql)
-        raise Mysql2::Error, 'MySQL client is not connected' if connection.nil? || connection.closed?
-
-        connection.query(sql)
-      end
-
-      def execute_with_params(connection, sql, params)
-        raise Mysql2::Error, 'MySQL client is not connected' if connection.nil? || connection.closed?
-
-        stmt = connection.prepare(sql)
-        result = stmt.execute(*params)
-        # Everything the caller needs must be read off the statement before the ensure closes its
-        # handle. A SELECT result is bound to the handle, so materialize its rows now - otherwise a
-        # later iteration (e.g. SqlRunner calling #to_a) raises "Statement handle already closed".
-        # A non-SELECT statement returns nil, and its affected-row count must come from the statement
-        # too: once the handle is closed the connection-level mysql_affected_rows returns its -1 error
-        # sentinel and mysql2 raises. So return that count for #affected_rows to hand back.
-        result.is_a?(::Mysql2::Result) ? result.to_a : stmt.affected_rows
-      ensure
-        stmt&.close
+      def database_engine
+        :mysql
       end
 
-      # mysql2 uses +?+ placeholders, so the SQL is already in its native form.
+      # MySQL uses +?+ placeholders, so the SQL is already in its native form.
       def translate_placeholders(sql)
         sql
       end
 
-      # mysql2 binds a blob parameter as raw binary bytes.
-      def binary_param(bytes)
-        bytes.b
-      end
-
-      # mysql2 hands back a blob column as a string that only needs its encoding forced to binary.
-      def read_binary(value)
-        value.b
-      end
-
-      # A prepared statement (execute_with_params) hands back its own affected-row count, since the
-      # connection-level count is unreadable once the statement handle is closed. A plain
-      # connection.query (execute) leaves nil, so the count is read off the connection instead.
-      def affected_rows(connection, result)
-        result.is_a?(Integer) ? result : connection.affected_rows.to_i
-      end
-
-      # mysql2 has no RETURNING clause, so the generated id is read from the connection afterwards.
-      def insert_returning_id(connection, sql, params, _id_column)
-        execute_with_params(connection, sql, params)
-        id = connection.last_id
-        id&.positive? ? id : nil
-      end
-
-      # mysql2 upserts with ON DUPLICATE KEY UPDATE, reading the incoming row from VALUES().
+      # MySQL upserts with ON DUPLICATE KEY UPDATE, reading the incoming row from VALUES().
       def upsert_clause(_conflict_columns, update_columns)
         assignments = update_columns.map { |column| "#{column} = VALUES(#{column})" }.join(', ')
         "ON DUPLICATE KEY UPDATE #{assignments}"
@@ -131,16 +69,6 @@ module AwsAdvancedRubyDriverWrapper
         logger.error("Failed to close MySQL connection: #{e.message}")
       end
 
-      def sql_state(exception)
-        return nil unless exception.respond_to?(:sql_state)
-
-        exception.sql_state
-      end
-
-      def network_bound_methods
-        NETWORK_BOUND_METHODS
-      end
-
       def prepare_connect_config(host_info, config)
         cfg = {}
         config.each { |k, v| cfg[k] = v }
@@ -153,9 +81,9 @@ module AwsAdvancedRubyDriverWrapper
         :username
       end
 
+      # Read, write and connect timeouts (in seconds) make a query or a close on a dead socket raise a
+      # timeout error rather than hang. With mysql2 they also keep it from segfaulting.
       def apply_monitoring_defaults(driver_props)
-        # mysql2 read_timeout / write_timeout (in seconds) ensure that queries and closes
-        # on a dead socket raise Mysql2::Error::TimeoutError instead of segfaulting.
         driver_props[:read_timeout] ||= DEFAULT_MONITORING_TIMEOUT_SEC
         driver_props[:write_timeout] ||= DEFAULT_MONITORING_TIMEOUT_SEC
         driver_props[:connect_timeout] ||= DEFAULT_MONITORING_TIMEOUT_SEC

data/lib/aws_advanced_ruby_driver_wrapper/driver_dialects/pg_driver_dialect.rb

@@ -85,6 +85,14 @@ module AwsAdvancedRubyDriverWrapper
         RubyMethod::CONNECTION_LO_TRUNCATE.name
       ]).freeze
 
+      def name
+        'pg'
+      end
+
+      def database_engine
+        :postgresql
+      end
+
       def connect(host_info, config)
         ::PG::Connection.new(**prepare_connect_config(host_info, config))
       end
@@ -162,6 +170,16 @@ module AwsAdvancedRubyDriverWrapper
         logger.error("Failed to close PostgreSQL connection: #{e.message}")
       end
 
+      # Points the socket at /dev/null, so the Terminate message libpq sends when the connection
+      # is finished or collected never reaches the server.
+      def abandon_connection(connection)
+        return if connection.nil? || connection.finished?
+
+        connection.socket_io.reopen(IO::NULL)
+      rescue StandardError => e
+        logger.debug("Failed to abandon PostgreSQL connection: #{e.message}")
+      end
+
       def sql_state(exception)
         return nil unless exception.is_a?(::PG::Error) && exception.result
 
@@ -172,6 +190,16 @@ module AwsAdvancedRubyDriverWrapper
         NETWORK_BOUND_METHODS
       end
 
+      # libpq separates the socket target (+hostaddr+) from the TLS/IAM server name (+host+), so a pg
+      # connection can be pinned to a specific instance IP while still verifying the certificate
+      # against - and generating an IAM token for - the original hostname. Blue/Green substitute
+      # routing uses this so a pinned connection reaches the intended instance regardless of what DNS
+      # resolves to mid-switchover, and still passes +verify-full+ whatever the CA source - instead of
+      # connecting to a raw IP (certificate-name mismatch) or falling back to the hostname (DNS risk).
+      def props_pinned_to_ip(driver_props, ip)
+        driver_props.dup.tap { |props| props[:hostaddr] = ip }
+      end
+
       def prepare_connect_config(host_info, config)
         cfg = {}
         config.each { |k, v| cfg[k] = v }