aws_advanced_ruby_driver_wrapper 1.0.0 → 1.1.0

data/lib/aws_advanced_ruby_driver_wrapper/trilogy.rb

@@ -0,0 +1,167 @@
+# 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 '../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
+  # A drop-in for a Trilogy connection. Trilogy has no prepared statements and no asynchronous or
+  # streaming reads, and it reads a result off the wire in full before returning it, so the only calls
+  # that reach the server are made on the connection itself, and results are handed back unwrapped.
+  class WrapperTrilogyClient
+    # Trilogy.new takes its options as a single hash, so this does too, and keyword arguments as well.
+    def self.new(options = {}, **)
+      instance = allocate
+      instance.send(:initialize, options, **)
+      instance
+    end
+
+    def initialize(options = {}, **)
+      ensure_trilogy!
+      config = Utils::ConnectionConfigParser.parse(:trilogy, **options, **)
+      @service_container = Services::ServiceUtility.create_standard_container(config)
+      @service_container.host_service.refresh_host_list
+      @query_conn = 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)
+
+    # The connection a statement was sent on is remembered, and so is its SQL, because the further
+    # result sets a statement leaves are read by calls of their own (next_result, abandon_results!),
+    # which carry neither and can only be made on that connection.
+    #
+    # What is remembered is the copy of the SQL that plugins inspect. The driver is sent the SQL as
+    # the application wrote it.
+    def query(sql)
+      ensure_connected!
+      inspected = Utils::SqlEncoding.inspectable(sql)
+      result = pm.execute(RubyMethod::CONNECTION_QUERY, current_conn, ->(*a) { current_conn.query(*a) }, sql,
+                          sql: inspected)
+      @query_conn = current_conn
+      @last_sql = inspected
+      result
+    end
+
+    def ping
+      ensure_connected!
+      pm.execute(RubyMethod::CONNECTION_PING, current_conn, -> { current_conn.ping })
+    end
+
+    def close
+      return if current_conn.nil?
+
+      pm.execute(RubyMethod::CONNECTION_CLOSE, current_conn, -> { current_conn.close })
+    end
+
+    def next_result
+      pm.execute(RubyMethod::CONNECTION_NEXT_RESULT, current_conn, -> { current_conn.next_result },
+                 bounded_conn: @query_conn, sql: @last_sql)
+    end
+
+    # Trilogy's own query_with_flags calls query on the driver connection directly, so it is defined
+    # here to send the statement through the plugin pipeline instead.
+    def query_with_flags(sql, flags)
+      ensure_connected!
+      old_flags = current_conn.query_flags
+      current_conn.query_flags = flags
+      query(sql)
+    ensure
+      current_conn.query_flags = old_flags if old_flags
+    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.
+    DYNAMIC_METHODS = {
+      abandon_results!: RubyMethod::CONNECTION_ABANDON_RESULTS,
+      change_db: RubyMethod::CONNECTION_CHANGE_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_QUERY = Set[:abandon_results!].freeze
+
+    def method_missing(method_name, *, **, &)
+      conn = ensure_connected!
+      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, *, **, &) unless network_bound_methods.include?(method_key)
+
+      bounded_conn = BOUNDED_TO_QUERY.include?(method_name) ? @query_conn : nil
+      pm.execute(
+        DYNAMIC_METHODS[method_name] || method_key, current_conn,
+        ->(*a, **opts, &b) { current_conn.send(method_name, *a, **opts, &b) },
+        *, **, bounded_conn: bounded_conn, &
+      )
+    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
+
+    def ensure_trilogy!
+      require 'trilogy'
+    rescue LoadError
+      raise LoadError, "WrapperTrilogyClient requires 'trilogy'. Add it to your Gemfile: gem 'trilogy'"
+    end
+
+    def current_conn
+      @service_container.connection_service.current_connection
+    end
+
+    # Guards the entry points that reach for the current connection directly (the query paths and
+    # method_missing). Without a live connection AR expects the driver's own "not connected" error,
+    # which it translates to a connection error; a bare NoMethodError on nil would not be recognized.
+    def ensure_connected!
+      conn = current_conn
+      raise Trilogy::BaseConnectionError, 'MySQL client is not connected' if conn.nil?
+
+      conn
+    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
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/utils/connection_config.rb

@@ -49,7 +49,7 @@ module AwsAdvancedRubyDriverWrapper
       end
 
       def cluster_id
-        PropertyDefinition::CLUSTER_ID.get(wrapper_props)
+        PropertyDefinition::CLUSTER_ID.get_string(wrapper_props)
       end
 
       # Redact sensitive property values (passwords, IAM tokens, Secrets Manager

data/lib/aws_advanced_ruby_driver_wrapper/utils/connection_config_parser.rb

@@ -32,9 +32,10 @@ module AwsAdvancedRubyDriverWrapper
       module_function
 
       # Main entry point. The driver_name is provided by the wrapper class
-      # (e.g. :postgresql from WrapperPgConnection, :mysql2 from WrapperMysql2Client).
+      # (e.g. :postgresql from WrapperPgConnection, :mysql2 from WrapperMysql2Client, :trilogy from
+      # WrapperTrilogyClient).
       #
-      # @param driver_name [Symbol] :postgresql or :mysql2
+      # @param driver_name [Symbol] :postgresql, :mysql2 or :trilogy
       def parse(driver_name, *args, **kwargs)
         if driver_name == :mysql2 && !args.empty? && !args.first.is_a?(Hash)
           raise ArgumentError,
@@ -192,7 +193,9 @@ module AwsAdvancedRubyDriverWrapper
       # Validates that values in a config hash match the expected type of their corresponding WrapperProperty.
       # Raises TypeError if a value is present but incompatible. Strings that represent valid integers
       # are accepted for Integer-typed properties. Float-typed properties also accept Integers and
-      # strings that represent a valid integer or decimal number.
+      # strings that represent a valid integer or decimal number. String-typed properties also accept
+      # numbers and Symbol-typed properties also accept strings, since YAML (e.g. database.yml) parses
+      # unquoted values like `cluster_id: 1` as numbers and never produces symbols.
       def validate_props!(config)
         config.each do |key, value|
           prop = PropertyDefinition::KNOWN_PROPERTIES[key]
@@ -200,23 +203,30 @@ module AwsAdvancedRubyDriverWrapper
           next if value.nil?
           next if prop.type.nil?
 
-          valid = if prop.type == :boolean
-                    value.is_a?(TrueClass) || value.is_a?(FalseClass) ||
-                      (value.is_a?(String) && BOOLEAN_STRINGS.include?(value.downcase))
-                  elsif prop.type == Float
-                    value.is_a?(Float) || value.is_a?(Integer) ||
-                      (value.is_a?(String) && value.match?(/\A-?\d+(\.\d+)?\z/))
-                  else
-                    value.is_a?(prop.type) ||
-                      (prop.type == Integer && value.is_a?(String) && value.match?(/\A-?\d+\z/))
-                  end
-
-          raise TypeError, "#{key}: expected #{prop.type}, got #{value.class}" unless valid
+          raise TypeError, "#{key}: expected #{prop.type}, got #{value.class}" unless valid_type?(prop.type, value)
 
           prop.validate!(value)
         end
       end
 
+      def valid_type?(type, value)
+        if type == :boolean
+          value.is_a?(TrueClass) || value.is_a?(FalseClass) ||
+            (value.is_a?(String) && BOOLEAN_STRINGS.include?(value.downcase))
+        elsif type == Float
+          value.is_a?(Float) || value.is_a?(Integer) ||
+            (value.is_a?(String) && value.match?(/\A-?\d+(\.\d+)?\z/))
+        elsif type == Integer
+          value.is_a?(Integer) || (value.is_a?(String) && value.match?(/\A-?\d+\z/))
+        elsif type == String
+          value.is_a?(String) || value.is_a?(Numeric)
+        elsif type == Symbol
+          value.is_a?(Symbol) || value.is_a?(String)
+        else
+          value.is_a?(type)
+        end
+      end
+
       # Forms a HostInfo object from a URI host section string.
       # Each host entry may or may not have a port; missing ports are represented as -1.
       # The resulting port is a comma-delimited string of per-host ports.

data/lib/aws_advanced_ruby_driver_wrapper/utils/events/batching_event_publisher.rb

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+require 'concurrent'
 require_relative '../../logging'
 
 module AwsAdvancedRubyDriverWrapper
@@ -23,10 +24,11 @@ module AwsAdvancedRubyDriverWrapper
       # Batches deduplicate events via Set semantics (eql?/hash).
       #
       # Public API:
-      #   subscribe(subscriber, event_classes) — register for event types
-      #   unsubscribe(subscriber, event_classes) — deregister
-      #   publish(event) — deliver immediate or queue batched
-      #   release_resources — stop background thread
+      #   subscribe(subscriber, event_classes) - register for event types
+      #   unsubscribe(subscriber, event_classes) - deregister
+      #   publish(event) - deliver immediate or queue batched
+      #   release_resources - stop background thread
+      #   restart_after_fork - restart background thread in a forked child
       class BatchingEventPublisher
         include Logging
 
@@ -37,7 +39,7 @@ module AwsAdvancedRubyDriverWrapper
           @subscribers = {}
           @event_queue = Set.new
           @lock = Mutex.new
-          @running = true
+          @running = Concurrent::AtomicBoolean.new(true)
           @thread = start_publishing_thread
         end
 
@@ -70,7 +72,7 @@ module AwsAdvancedRubyDriverWrapper
         end
 
         def release_resources
-          @lock.synchronize { @running = false }
+          @running.make_false
           begin
             @thread&.wakeup
           rescue ThreadError
@@ -79,11 +81,20 @@ module AwsAdvancedRubyDriverWrapper
           @thread&.join(@message_interval_sec)
         end
 
+        # Restarts the publishing thread in a forked child, where it no longer runs. A thread that is
+        # still alive is kept, so calling this again does not start a second one.
+        def restart_after_fork
+          return if @thread&.alive?
+
+          @running.make_true
+          @thread = start_publishing_thread
+        end
+
         private
 
         def start_publishing_thread
           thread = Thread.new do
-            while @lock.synchronize { @running }
+            while @running.true?
               sleep(@message_interval_sec)
               send_messages
             end

data/lib/aws_advanced_ruby_driver_wrapper/utils/global_aurora_topology_utils.rb

@@ -35,9 +35,9 @@ module AwsAdvancedRubyDriverWrapper
       def query_global_topology(conn, initial_host_info, instance_templates_by_region)
         results = @dialect.execute(conn, @dialect.topology_query)
 
-        if results.fields.empty?
+        if topology_result_empty?(results)
           # We expect at least 4 columns. Note that the server may return 0 columns if failover has occurred.
-          logger.debug('The topology query returned a result with 0 columns. ' \
+          logger.debug('The topology query returned a result with no rows or rows with no columns. ' \
                        'This may occur if the topology query is executed when the server is failing over.')
           return nil
         end

data/lib/aws_advanced_ruby_driver_wrapper/utils/iam_auth_utils.rb

@@ -19,10 +19,14 @@ require_relative 'rds_utils'
 require_relative 'rds_url_type'
 require_relative '../property_definition'
 require_relative '../logging'
+require_relative 'telemetry/null_telemetry'
+require_relative 'telemetry/telemetry_const'
 
 module AwsAdvancedRubyDriverWrapper
   module Utils
     module IamAuthUtils
+      TELEMETRY_FETCH_TOKEN = 'fetch authentication token'
+
       module_function
 
       TokenEntry = Data.define(:token, :expires_at) do
@@ -83,10 +87,15 @@ module AwsAdvancedRubyDriverWrapper
         arn&.match(/\Aarn:aws(?:-[a-z]+)*:rds:(?<region>[^:]+)/)&.[](:region)
       end
 
-      def generate_token(region:, hostname:, port:, user:, credentials_provider:)
-        Aws::RDS::AuthTokenGenerator
-          .new(credentials: credentials_provider)
-          .auth_token(region:, endpoint: "#{hostname}:#{port}", user_name: user)
+      # @param telemetry_factory [Telemetry::TelemetryFactory, Telemetry::NullTelemetryFactory] records
+      #   the token generation as a span
+      def generate_token(region:, hostname:, port:, user:, credentials_provider:,
+                         telemetry_factory: Telemetry::NullTelemetryFactory::INSTANCE)
+        telemetry_factory.in_context(TELEMETRY_FETCH_TOKEN, Telemetry::TraceLevel::NESTED, record_failure: true) do
+          Aws::RDS::AuthTokenGenerator
+            .new(credentials: credentials_provider)
+            .auth_token(region:, endpoint: "#{hostname}:#{port}", user_name: user)
+        end
       end
 
       def resolve_host(iam_host, host_info)

data/lib/aws_advanced_ruby_driver_wrapper/utils/parser/encryption_annotation_parser.rb

@@ -14,6 +14,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+require_relative 'mysql_text'
+
 module AwsAdvancedRubyDriverWrapper
   module Utils
     module Parser
@@ -21,21 +23,23 @@ module AwsAdvancedRubyDriverWrapper
         # Both placeholder styles are accepted, so that the same annotation works for mysql2's
         # positional +?+ and pg's numbered +$1+.
         ANNOTATION_PATTERN = %r{/\*@encrypt:([\w.]+)\*/\s*(?:\?|\$(\d+))}
-        STRIP_PATTERN = %r{/\*@encrypt:[\w.]+\*/\s*}
-        # The quote characters a literal can be wrapped in. A backtick quotes an identifier rather than
-        # a value, but is included so a +?+ inside a backtick-quoted name is not miscounted either.
-        QUOTES = ["'", '"', '`'].freeze
+        # Whitespace after an annotation is taken with it, but not a line break, which may be what ends
+        # a line comment the annotation sits in.
+        STRIP_PATTERN = %r{/\*@encrypt:[\w.]+\*/[^\S\n]*}
 
         module_function
 
         # Returns a 1-based map of parameter index => "table.column" for each
         # /*@encrypt:table.column*/ ? or /*@encrypt:table.column*/ $n placeholder.
         # @param sql [String]
+        # @param server_version [Integer, nil] the MySQL server's version, which decides whether a
+        #   versioned +/*!+ comment is part of the statement (see MysqlText.each_segment)
         # @return [Hash{Integer => String}]
-        def parse_annotations(sql)
-          return {} unless sql.is_a?(String) && !sql.empty?
+        def parse_annotations(sql, server_version: nil)
+          # Most statements carry no annotation, so they are not scanned for placeholders at all.
+          return {} unless annotations?(sql)
 
-          question_marks = unquoted_question_mark_positions(sql)
+          question_marks = placeholder_positions(sql, server_version)
 
           sql.to_enum(:scan, ANNOTATION_PATTERN)
              .each_with_object({}) do |_, result|
@@ -66,31 +70,11 @@ module AwsAdvancedRubyDriverWrapper
         end
 
         # The character positions of every +?+ that is a bind placeholder, skipping any inside a quoted
-        # literal so a literal question mark does not shift the count of the placeholder an annotation
-        # points at. A doubled quote reads as one closing and one opening quote, and a backslash escapes
-        # the next character except inside a backtick, which matches how the drivers read a literal.
-        def unquoted_question_mark_positions(sql)
+        # literal or a comment so that it does not shift the count of the placeholder an annotation
+        # points at. Only mysql2 binds with +?+, so the text is read by MySQL's rules.
+        def placeholder_positions(sql, server_version)
           positions = []
-          quote = nil
-          index = 0
-
-          while index < sql.length
-            char = sql[index]
-            if quote
-              if char == '\\' && quote != '`'
-                index += 2
-                next
-              end
-
-              quote = nil if char == quote
-            elsif QUOTES.include?(char)
-              quote = char
-            elsif char == '?'
-              positions << index
-            end
-            index += 1
-          end
-
+          MysqlText.each_code_char(sql, server_version: server_version) { |char, index| positions << index if char == '?' }
           positions
         end
       end