aws_advanced_ruby_driver_wrapper 1.0.0

data/lib/aws_advanced_ruby_driver_wrapper/monitoring/global_cluster_topology_monitor.rb

@@ -0,0 +1,72 @@
+# 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 'cluster_topology_monitor'
+
+module AwsAdvancedRubyDriverWrapper
+  module Monitoring
+    # Topology monitor for Aurora Global Databases spanning multiple AWS regions.
+    # Extends ClusterTopologyMonitor by using region-aware instance templates and
+    # a global topology query that returns region information per instance.
+    class GlobalClusterTopologyMonitor < ClusterTopologyMonitor
+      # @param service_container [Services::ServiceContainer]
+      # @param cluster_id [String]
+      # @param instance_template [Host::HostInfo] default template (used for initial connection)
+      # @param instance_templates_by_region [Hash{String => Host::HostInfo}] region -> instance template
+      # @param topology_utils [Utils::GlobalAuroraTopologyUtils]
+      # @param monitoring_driver_props [Hash]
+      # @param monitoring_wrapper_props [Hash]
+      def initialize(
+        service_container:,
+        cluster_id:,
+        instance_template:,
+        instance_templates_by_region:,
+        topology_utils:,
+        monitoring_driver_props:,
+        monitoring_wrapper_props: {}
+      )
+        super(
+          service_container:,
+          cluster_id:,
+          instance_template:,
+          topology_utils:,
+          monitoring_driver_props:,
+          monitoring_wrapper_props:
+        )
+        @instance_templates_by_region = instance_templates_by_region
+      end
+
+      private
+
+      def query_topology(conn)
+        @topology_utils.query_global_topology(conn, initial_host_info, @instance_templates_by_region)
+      end
+
+      def resolve_instance_template(instance_id, conn)
+        region = @topology_utils.query_region(instance_id, conn)
+        return @instance_template if region.nil?
+
+        template = @instance_templates_by_region[region]
+        if template.nil?
+          logger.warn("[#{@cluster_id}] No instance template for region '#{region}'")
+          return @instance_template
+        end
+
+        template
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/monitoring/monitor.rb

@@ -0,0 +1,99 @@
+# 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 'monitor_state'
+require_relative '../logging'
+require 'concurrent'
+
+module AwsAdvancedRubyDriverWrapper
+  module Monitoring
+    # Base class for all monitors. Manages thread lifecycle, state transitions,
+    # and activity tracking. Subclasses implement #monitor with their polling logic.
+    class Monitor
+      include Logging
+
+      attr_reader :last_activity_sec
+
+      def state
+        @state.value
+      end
+
+      def initialize(termination_timeout_sec: 30.0)
+        @termination_timeout_sec = termination_timeout_sec
+        @state = Concurrent::AtomicReference.new(nil)
+        @stop_flag = Concurrent::AtomicBoolean.new(false)
+        update_activity
+        @thread = nil
+        @lock = Mutex.new
+      end
+
+      def start
+        @lock.synchronize do
+          return if @state.value == MonitorState::RUNNING
+
+          @stop_flag.make_false
+          @state.set(MonitorState::RUNNING)
+          @thread = Thread.new { run }
+          @thread.name = "monitor-#{monitor_thread_suffix}"
+          logger.debug("Started monitoring thread: #{@thread.name}")
+        end
+      end
+
+      def stop
+        @stop_flag.make_true
+
+        thread = @thread
+        if thread&.alive? && thread != Thread.current
+          thread.join(@termination_timeout_sec)
+          thread.kill if thread.alive?
+        end
+        @thread = nil
+
+        @state.set(MonitorState::STOPPED)
+        close
+        logger.debug("Stopped monitoring thread: monitor-#{monitor_thread_suffix}")
+      end
+
+      def stopped?
+        @stop_flag.true?
+      end
+
+      def close; end
+
+      private
+
+      def run
+        monitor
+      rescue StandardError => e
+        logger.error("Exception in monitoring thread monitor-#{monitor_thread_suffix}: #{e.message}")
+        @state.set(MonitorState::ERROR)
+      ensure
+        @state.compare_and_set(MonitorState::RUNNING, MonitorState::STOPPED)
+      end
+
+      def monitor_thread_suffix
+        class_name = self.class.name
+        return 'unknown' unless class_name
+
+        class_name.split('::').last.downcase
+      end
+
+      def update_activity
+        @last_activity_sec = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/monitoring/monitor_connection.rb

@@ -0,0 +1,57 @@
+# 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.
+
+module AwsAdvancedRubyDriverWrapper
+  module Monitoring
+    # Thread-safe connection wrapper with compare-and-set semantics.
+    # Automatically closes the old connection when replaced.
+    class MonitorConnection
+      # @param driver_dialect [DriverDialects::DriverDialect] the driver dialect used to safely close connections.
+      def initialize(driver_dialect)
+        @connection = Concurrent::AtomicReference.new(nil)
+        @driver_dialect = driver_dialect
+      end
+
+      # Returns the current connection, or nil.
+      def get
+        @connection.value
+      end
+
+      # Replaces the current connection. Closes the old one unless close_old is false.
+      # @param new_conn [Object, nil] the new connection.
+      # @param close_old [Boolean] whether to close the previous connection.
+      def set(new_conn, close_old: true)
+        old = @connection.get_and_set(new_conn)
+        return unless close_old && old && !old.equal?(new_conn)
+
+        @driver_dialect.close_connection(old)
+      end
+
+      # Atomically sets the connection only if the current value is `expected` (identity check).
+      # @param expected [Object, nil] the expected current connection.
+      # @param new_conn [Object] the new connection to set.
+      # @return [Boolean] true if the swap succeeded.
+      def compare_and_set(expected, new_conn)
+        @connection.compare_and_set(expected, new_conn)
+      end
+
+      # Closes and nils the connection.
+      def close
+        set(nil)
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/monitoring/monitor_state.rb

@@ -0,0 +1,25 @@
+# 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.
+
+module AwsAdvancedRubyDriverWrapper
+  module Monitoring
+    module MonitorState
+      RUNNING = :running
+      STOPPED = :stopped
+      ERROR = :error
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/mysql.rb

@@ -0,0 +1,429 @@
+# 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.
+
+# Loads the gem's top-level setup (version constant, global configuration) that this client relies on,
+# so requiring this file alone is enough to use the MySQL client.
+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 'ruby_method'
+require_relative 'errors'
+
+module AwsAdvancedRubyDriverWrapper
+  class WrapperMysql2Client
+    def self.new(**)
+      instance = allocate
+      instance.send(:initialize, **)
+      instance
+    end
+
+    def initialize(**)
+      ensure_mysql2!
+      config = Utils::ConnectionConfigParser.parse(:mysql2, **)
+      @service_container = Services::ServiceUtility.create_standard_container(config)
+      @service_container.host_service.refresh_host_list
+      @async_conn = nil
+      @async_sql = nil
+      @last_sql = nil
+      conn_service = @service_container.connection_service
+      @service_container.plugin_manager.connect(conn_service.initial_host_info, conn_service.driver_props, true)
+    end
+
+    # Explicitly define critical methods (bypass method_missing to avoid method_missing overhead)
+
+    # This is how mysql2 sends a statement asynchronously as well: +query(sql, async: true)+ returns
+    # nothing and the result is read afterward by +async_result+. The connection the statement was
+    # sent on is remembered for that read, and so is its SQL, since the read is a call of its own and
+    # carries neither.
+    #
+    # What is remembered is the copy of the SQL that plugins inspect, made once here rather than on
+    # every call that is handed it later. The driver is sent the SQL as the application wrote it.
+    def query(sql, options = {})
+      inspected = Utils::SqlEncoding.inspectable(sql)
+      result = pm.execute(RubyMethod::CONNECTION_QUERY, current_conn, ->(*a) { current_conn.query(*a) }, sql, options,
+                          sql: inspected)
+      @last_sql = inspected
+      if options[:async]
+        @async_conn = current_conn
+        @async_sql = inspected
+      end
+      wrap_mysql_result(result, inspected)
+    end
+
+    def prepare(sql)
+      inspected = Utils::SqlEncoding.inspectable(sql)
+      mysql_stmt = pm.execute(RubyMethod::CONNECTION_PREPARE, current_conn, ->(*a) { current_conn.prepare(*a) }, sql,
+                              sql: inspected)
+      Mysql2WrapperStatement.new(@service_container, current_conn, mysql_stmt, inspected)
+    end
+
+    def escape(string)
+      pm.execute(RubyMethod::CONNECTION_ESCAPE, current_conn, ->(*a) { current_conn.escape(*a) }, string)
+    end
+
+    def ping
+      pm.execute(RubyMethod::CONNECTION_PING, current_conn, -> { current_conn.ping })
+    end
+
+    def close
+      pm.execute(RubyMethod::CONNECTION_CLOSE, current_conn, -> { current_conn.close })
+    end
+
+    # -- Async readers (check bounded to @async_conn) --
+
+    # The result of a statement that was sent with +async: true+. It is read by a call of its own, so
+    # it is handed the SQL of the statement it belongs to, and it is refused on any connection other
+    # than the one that statement was sent on.
+    def async_result
+      result = pm.execute(RubyMethod::CONNECTION_ASYNC_RESULT, current_conn, -> { current_conn.async_result },
+                          bounded_conn: @async_conn, sql: @async_sql)
+      sql = @async_sql
+      @async_conn = nil
+      @async_sql = nil
+      wrap_mysql_result(result, sql)
+    end
+
+    def store_result
+      result = pm.execute(RubyMethod::CONNECTION_STORE_RESULT, current_conn, -> { current_conn.store_result },
+                          bounded_conn: @async_conn, sql: @async_sql)
+      sql = @async_sql
+      @async_conn = nil
+      @async_sql = nil
+      wrap_mysql_result(result, sql)
+    end
+
+    def more_results?
+      pm.execute(RubyMethod::CONNECTION_MORE_RESULTS, current_conn, -> { current_conn.more_results? }, bounded_conn: @async_conn)
+    end
+
+    # A statement that leaves more than one result set is read by moving to each in turn and storing
+    # it. Storing a result forgets the connection the statement was sent on and the SQL that was sent
+    # on it, and every one of those results belongs to that statement, so both are put back for the
+    # read that follows and dropped once there is nothing left to read.
+    def next_result
+      result = pm.execute(RubyMethod::CONNECTION_NEXT_RESULT, current_conn, -> { current_conn.next_result },
+                          bounded_conn: @async_conn)
+      if result
+        @async_conn = current_conn
+        @async_sql = @last_sql
+      else
+        @async_conn = nil
+        @async_sql = nil
+      end
+      result
+    end
+
+    # -- method_missing: covers non-network calls and rarely used network calls --
+
+    # The network calls that are rare enough not to be worth a method of their own. They are entered
+    # into the pipeline under the name the pipeline knows them by rather than as a bare string, so
+    # that the connection each is bound to is checked. mysql2 gives none of its calls a second
+    # spelling, so unlike pg there is nothing here to translate.
+    DYNAMIC_METHODS = {
+      abandon_results!: RubyMethod::CONNECTION_ABANDON_RESULTS,
+      select_db: RubyMethod::CONNECTION_SELECT_DB,
+      set_server_option: RubyMethod::CONNECTION_SET_SERVER_OPTION
+    }.freeze
+
+    # Draining what is left of a statement can only be done on the connection it was sent on.
+    BOUNDED_TO_ASYNC = Set[:abandon_results!].freeze
+
+    def method_missing(method_name, *args, **kwargs, &)
+      conn = current_conn
+      raise NoMethodError, 'Connection not initialized' if conn.nil?
+      raise NoMethodError, "undefined method `#{method_name}' for #{self.class}" unless conn.respond_to?(method_name)
+
+      method_key = "connection.#{method_name}"
+      return conn.send(method_name, *args, **kwargs, &) unless network_bound_methods.include?(method_key)
+
+      execute_dynamic(method_name, method_key, args, kwargs, &)
+    end
+
+    def respond_to_missing?(method, include_private = false)
+      current_conn.respond_to?(method, include_private) || super
+    end
+
+    # A concise representation that never exposes the connection config (which carries
+    # credentials) or the cached SQL text this instance holds. Defined so a default
+    # dump - via logging, interpolation, `pp`, or a backtrace - cannot leak either.
+    def inspect
+      format('#<%<class>s:0x%<addr>016x>', class: self.class.name, addr: object_id << 1)
+    end
+    alias to_s inspect
+
+    def pretty_print(pp)
+      pp.text(inspect)
+    end
+
+    private
+
+    # Runs a call that reached method_missing through the pipeline, telling the plugins the connection
+    # it is bound to, which is known here rather than from the arguments of the call.
+    def execute_dynamic(method_name, method_key, args, kwargs, &)
+      bounded_conn = BOUNDED_TO_ASYNC.include?(method_name) ? @async_conn : nil
+      result = pm.execute(
+        DYNAMIC_METHODS[method_name] || method_key, current_conn,
+        ->(*a, **opts, &b) { current_conn.send(method_name, *a, **opts, &b) },
+        *args, **kwargs, bounded_conn: bounded_conn, &
+      )
+      if method_name == :abandon_results!
+        @async_conn = nil
+        @async_sql = nil
+      end
+      wrap_mysql_result(result)
+    end
+
+    def ensure_mysql2!
+      require 'mysql2'
+    rescue LoadError
+      raise LoadError, "WrapperMysql2Client requires 'mysql2'. Add it to your Gemfile: gem 'mysql2'"
+    end
+
+    def current_conn
+      @service_container.connection_service.current_connection
+    end
+
+    def pm
+      @service_container.plugin_manager
+    end
+
+    def network_bound_methods
+      @network_bound_methods ||= @service_container.dialect_service.driver_dialect.network_bound_methods
+    end
+
+    def wrap_mysql_result(result, sql = nil)
+      return result unless result.is_a?(Mysql2::Result)
+
+      Mysql2WrapperResult.new(result, @service_container, current_conn, sql)
+    end
+  end
+
+  class Mysql2WrapperStatement
+    # @param sql [String, nil] the SQL the statement was prepared with, kept so that plugins
+    #   which inspect statements still see it when the statement is executed
+    def initialize(service_container, connection, mysql_stmt, sql = nil)
+      @service_container = service_container
+      @connection = connection
+      @mysql_stmt = mysql_stmt
+      @sql = sql
+    end
+
+    def execute(*params, **)
+      result = pm.execute(
+        RubyMethod::STATEMENT_EXECUTE, current_conn,
+        ->(*p, **o) { @mysql_stmt.execute(*p, **o) },
+        *params, bounded_conn: @connection, sql: @sql, **
+      )
+      return result unless result.is_a?(Mysql2::Result)
+
+      # The statement is handed to the result so it can read column names from the statement's
+      # metadata rather than from Mysql2::Result#fields, which dereferences the result's field
+      # pointer and segfaults on a prepared statement that returned no rows.
+      Mysql2WrapperResult.new(result, @service_container, @connection, @sql, @mysql_stmt)
+    end
+
+    def close
+      pm.execute(RubyMethod::STATEMENT_CLOSE, current_conn, -> { @mysql_stmt.close })
+    end
+
+    # Delegate non-network methods directly
+    def fields
+      @mysql_stmt.fields
+    end
+
+    def field_count
+      @mysql_stmt.field_count
+    end
+
+    def param_count
+      @mysql_stmt.param_count
+    end
+
+    def affected_rows
+      @mysql_stmt.affected_rows
+    end
+
+    def last_id
+      @mysql_stmt.last_id
+    end
+
+    def closed?
+      @mysql_stmt.closed?
+    end
+
+    def inspect
+      format('#<%<class>s:0x%<addr>016x>', class: self.class.name, addr: object_id << 1)
+    end
+    alias to_s inspect
+
+    def pretty_print(pp)
+      pp.text(inspect)
+    end
+
+    private
+
+    def current_conn
+      @service_container.connection_service.current_connection
+    end
+
+    def pm
+      @service_container.plugin_manager
+    end
+  end
+
+  class Mysql2WrapperResult
+    include Enumerable
+
+    # @param sql [String, nil] the SQL that produced the result, kept so that plugins which
+    #   inspect statements still see it when the rows are read
+    # @param statement [Mysql2::Statement, nil] the prepared statement the result came from, if any,
+    #   used as a safe source of column names (see {#result_field_names})
+    def initialize(result, service_container, connection, sql = nil, statement = nil)
+      @result = result
+      @service_container = service_container
+      @connection = connection
+      @sql = sql
+      @statement = statement
+    end
+
+    def each(*args, &block)
+      pm.execute(RubyMethod::RESULT_EACH, current_conn, ->(&blk) { @result.each(*args, &blk) },
+                 bounded_conn: @connection, sql: @sql, field_names: -> { result_field_names }, &block)
+    end
+
+    def to_a
+      pm.execute(RubyMethod::RESULT_TO_A, current_conn, -> { @result.to_a },
+                 bounded_conn: @connection, sql: @sql, field_names: -> { result_field_names })
+    end
+
+    def [](index)
+      pm.execute(RubyMethod::RESULT_BRACKET, current_conn, ->(*a) { @result[*a] }, index,
+                 bounded_conn: @connection, sql: @sql, field_names: -> { result_field_names })
+    end
+
+    # A buffered result is already in client memory, so letting it go is local. An unbuffered one,
+    # from +query(sql, stream: true)+, still has whatever was not read on the wire, and libmysql has
+    # to drain it before it can free the result. That makes this a call to the server, on the
+    # connection the statement was sent on, so it should go through the pipeline.
+    def free
+      pm.execute(RubyMethod::RESULT_FREE, current_conn, -> { @result.free }, bounded_conn: @connection, sql: @sql)
+    end
+
+    # Mysql2::Statement#execute reads a buffered result to the end before returning it, and when that
+    # read finds no rows mysql2 frees the result's metadata without having cached any column names.
+    # Mysql2::Result#fields then reads the freed metadata - an empty list at best, a segfault at worst
+    # (ActiveRecord calls it on every result, so any finder that matches nothing can crash). The
+    # names for that one case come from the statement's own metadata, symbolized when the result was
+    # asked for symbolize_keys just as Mysql2::Result#fields would. Every other result is read
+    # directly, so a freed streaming result still raises as the driver does.
+    def fields
+      return @result.fields unless empty_buffered_statement_result?
+
+      names = @statement.fields
+      result_query_options[:symbolize_keys] ? names.map(&:to_sym) : names
+    end
+
+    def field_types
+      @result.field_types
+    end
+
+    def count
+      @result.count
+    end
+
+    def size
+      @result.size
+    end
+
+    def server_flags
+      @result.server_flags
+    end
+
+    def query_time
+      @result.query_time
+    end
+
+    def tables
+      @result.tables
+    end
+
+    def dbs
+      @result.dbs
+    end
+
+    def method_missing(method_name, ...)
+      return @result.send(method_name, ...) if @result.respond_to?(method_name)
+
+      super
+    end
+
+    def respond_to_missing?(method_name, include_private = false)
+      @result.respond_to?(method_name, include_private) || super
+    end
+
+    def inspect
+      format('#<%<class>s:0x%<addr>016x>', class: self.class.name, addr: object_id << 1)
+    end
+    alias to_s inspect
+
+    def pretty_print(pp)
+      pp.text(inspect)
+    end
+
+    private
+
+    # Column names for the plugins that inspect a result as it is read, such as one that decrypts
+    # certain columns and has to know which position each holds. This is internal plumbing, not the
+    # driver's own #fields - so unlike #fields it must always yield the real column names, including
+    # for a prepared-statement result that fetched no rows. Mysql2::Result#fields cannot supply them
+    # there: it dereferences the result's field pointer, which mysql2 never populates until a row is
+    # fetched, so it returns an empty list or, on some client libraries, segfaults. In that one case
+    # the names come from the statement's own metadata (mysql_stmt_result_metadata), which carries
+    # them whether or not any rows came back. Once a row has been fetched the field cache is in place
+    # and the result's own fields are safe; results from +query+ carry no statement and read directly.
+    def result_field_names
+      # none? would iterate the result and consume an unbuffered one; count reads the driver's own
+      # row count without touching the rows.
+      return @statement.fields if @statement && @result.count.zero? # rubocop:disable Style/CollectionQuerying
+
+      @result.fields
+    end
+
+    # A streaming result is not read until the caller iterates it, so only a buffered one has had its
+    # metadata freed by the time it is returned. count reads the driver's own row count without
+    # touching the rows.
+    def empty_buffered_statement_result?
+      @statement && !result_query_options[:stream] && @result.count.zero? # rubocop:disable Style/CollectionQuerying
+    end
+
+    # Part of the workaround in #fields for a mysql2 bug that can segfault when an empty
+    # prepared-statement result is asked for its fields. Telling that case apart, and symbolizing the
+    # fallback names, needs the options the result was produced with; mysql2 keeps them in this ivar
+    # and exposes no reader for it.
+    def result_query_options
+      @result.instance_variable_get(:@query_options) || {}
+    end
+
+    def current_conn
+      @service_container.connection_service.current_connection
+    end
+
+    def pm
+      @service_container.plugin_manager
+    end
+  end
+end