aws_advanced_ruby_driver_wrapper 1.0.0

data/lib/aws_advanced_ruby_driver_wrapper/utils/retry_util.rb

@@ -0,0 +1,153 @@
+# 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 '../host/host_role'
+require_relative '../host/host_availability'
+require_relative '../logging'
+
+module AwsAdvancedRubyDriverWrapper
+  module Utils
+    class RetryUtil
+      include Logging
+
+      SHORT_DELAY_SEC = 0.1
+      DEFAULT_STRATEGY = 'random'
+
+      Result = Data.define(:connection, :host_info)
+
+      def initialize(service_container)
+        @host_service = service_container.host_service
+        @dialect_service = service_container.dialect_service
+        @connection_service = service_container.connection_service
+      end
+
+      # Repeatedly refreshes the topology and attempts to open a connection to the current writer,
+      # until one is established or the deadline passes.
+      #
+      # @param plugin_to_skip [Object] the plugin that should be skipped in the connect pipeline
+      # @param plugin_manager [Services::PluginManager]
+      # @param deadline [Time] the point in time at which to give up
+      # @return [Result] the new connection and the host it was opened to
+      # @raise [Timeout::Error] if no connection could be established before the deadline
+      def connect_to_writer(plugin_to_skip, plugin_manager, deadline:)
+        connect_to_allowed_host(plugin_to_skip, plugin_manager, verify_role: Host::HostRole::WRITER, deadline: deadline) do |allowed_hosts|
+          writer_candidate = @host_service.all_hosts.find { |h| h.role == Host::HostRole::WRITER }
+
+          if writer_candidate.nil?
+            logger.debug { 'No writer host found in topology' }
+            nil
+          elsif allowed_hosts.none? { |h| h.host_and_port == writer_candidate.host_and_port }
+            logger.debug { "New writer not in allowed hosts: #{writer_candidate.url}" }
+            nil
+          else
+            [writer_candidate]
+          end
+        end
+      end
+
+      # Repeatedly refreshes the topology and attempts to open a connection to one of the hosts
+      # selected by the given block, until one is established or the deadline passes.
+      #
+      # Hosts that cannot be connected to, or whose role does not match +verify_role+, are dropped
+      # from the candidate list. Once every candidate has been ruled out, the topology is refreshed
+      # and the block is consulted again.
+      #
+      # @param plugin_to_skip [Object] the plugin that should be skipped in the connect pipeline
+      # @param plugin_manager [Services::PluginManager]
+      # @param deadline [Time] the point in time at which to give up
+      # @param verify_role [Symbol, nil] the role the new connection must report, or nil to accept any host
+      # @param strategy [String] the host selection strategy used to order the candidates
+      # @yieldparam allowed_hosts [Array<Host::HostInfo>] the current allowed hosts
+      # @yieldreturn [Array<Host::HostInfo>, nil] the hosts to attempt, or nil/empty to refresh and retry
+      # @return [Result] the new connection and the host it was opened to
+      # @raise [Timeout::Error] if no connection could be established before the deadline
+      def connect_to_allowed_host(plugin_to_skip, plugin_manager, deadline:, verify_role: nil, strategy: DEFAULT_STRATEGY)
+        strategy = DEFAULT_STRATEGY if strategy.nil? || strategy.to_s.empty?
+
+        loop do
+          break if Time.now >= deadline
+
+          # The roles in this list might not be accurate, depending on whether the new topology has become available yet.
+          @host_service.refresh_host_list
+          candidates = yield(@host_service.hosts)
+
+          if candidates.nil? || candidates.empty?
+            sleep(SHORT_DELAY_SEC)
+            next
+          end
+
+          # Copy the candidates and mark them available so that the selection strategy considers all of them.
+          remaining = candidates.map do |host|
+            host.deep_dup.tap { |dup| dup.availability = Host::HostAvailability::AVAILABLE }
+          end
+
+          while !remaining.empty? && Time.now < deadline
+            candidate = select_candidate(remaining, verify_role, strategy)
+            if candidate.nil?
+              logger.debug { "Unable to find #{verify_role || 'a host'} in the updated host list: #{remaining.map(&:url)}" }
+              sleep(SHORT_DELAY_SEC)
+              break # Give up on this candidate list and refresh the topology.
+            end
+
+            result = attempt_connection(candidate, verify_role, plugin_to_skip, plugin_manager)
+            return result if result
+
+            remaining.delete(candidate)
+          end
+        end
+
+        raise Timeout::Error, 'Not able to establish a connection before timing out'
+      end
+
+      private
+
+      # @return [Result, nil] the result, or nil if the connection failed or reported the wrong role
+      def attempt_connection(candidate, verify_role, plugin_to_skip, plugin_manager)
+        conn = plugin_manager.connect(candidate, @connection_service.driver_props, false, plugin_to_skip: plugin_to_skip)
+
+        # Since the roles in the host list might not be accurate, we execute a query to check the instance's role.
+        role = verify_role.nil? ? nil : @dialect_service.db_dialect.host_role(conn)
+        if verify_role.nil?
+          return Result.new(conn, candidate)
+        elsif verify_role == role
+          return Result.new(conn, candidate.deep_dup(role: role))
+        end
+
+        # The role is not the one that was asked for, so the connection is not valid.
+        close_quietly(conn)
+        nil
+      rescue StandardError => e
+        logger.debug { "Exception connecting to #{candidate.host}: #{e.message}" }
+        close_quietly(conn)
+        nil
+      end
+
+      def select_candidate(hosts, role, strategy)
+        @host_service.select_host(hosts, role, strategy)
+      rescue StandardError
+        nil
+      end
+
+      def close_quietly(conn)
+        return if conn.nil?
+
+        @dialect_service.driver_dialect.close_connection(conn)
+      rescue StandardError
+        # ignore
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/utils/sql_encoding.rb

@@ -0,0 +1,56 @@
+# 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 Utils
+    module SqlEncoding
+      module_function
+
+      # A copy of the SQL as valid UTF-8, which is what the patterns and parsers that inspect statements
+      # can read. A statement in an encoding that is not ASCII-compatible, such as UTF-16, cannot be
+      # matched against a UTF-8 pattern at all, and neither can one with bytes that are invalid in its
+      # encoding.
+      #
+      # Only the copy is converted. The SQL a driver is sent is the caller's own, since the driver
+      # converts it to the connection's encoding itself.
+      #
+      # A character with no UTF-8 equivalent, or a byte that is invalid, is replaced rather than raised on,
+      # so that inspecting a statement never fails a call the driver would have made.
+      #
+      # SQL that is binary, or in one of the few encodings Ruby has no converter to UTF-8 for, has no
+      # text to convert. The driver sends such SQL as the bytes it is, so its bytes are read as UTF-8,
+      # which is how the server reads them on a UTF-8 connection. Replacing or skipping them instead
+      # would leave the checks reading a different statement from the one the server runs.
+      #
+      # @param sql [Object] the SQL to inspect
+      # @return [String, Object] the SQL as valid UTF-8, or anything that is not a String as it is
+      def inspectable(sql)
+        return sql unless sql.is_a?(String)
+        return sql.valid_encoding? ? sql : sql.scrub if sql.encoding == Encoding::UTF_8
+        return sql if sql.ascii_only? && sql.encoding.ascii_compatible?
+        return bytes_as_utf8(sql) if sql.encoding == Encoding::BINARY
+
+        sql.encode(Encoding::UTF_8, invalid: :replace, undef: :replace)
+      rescue Encoding::ConverterNotFoundError
+        bytes_as_utf8(sql)
+      end
+
+      def bytes_as_utf8(sql)
+        sql.b.force_encoding(Encoding::UTF_8).scrub
+      end
+      private_class_method :bytes_as_utf8
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/utils/sql_method_analyzer.rb

@@ -0,0 +1,195 @@
+# 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 '../ruby_method'
+require_relative 'sql_encoding'
+
+module AwsAdvancedRubyDriverWrapper
+  module Utils
+    module SqlMethodAnalyzer
+      EXECUTE_SQL_METHODS = Set[
+        RubyMethod::CONNECTION_EXEC.name,
+        RubyMethod::CONNECTION_ASYNC_EXEC.name,
+        RubyMethod::CONNECTION_EXEC_PARAMS.name,
+        RubyMethod::CONNECTION_EXEC_PREPARED.name,
+        RubyMethod::CONNECTION_QUERY.name,
+        RubyMethod::CONNECTION_SEND_QUERY.name,
+        RubyMethod::CONNECTION_SEND_QUERY_PARAMS.name,
+        RubyMethod::STATEMENT_EXECUTE.name
+      ].freeze
+
+      CLOSE_TRANSACTION_METHODS = Set[
+        RubyMethod::CONNECTION_CLOSE.name
+      ].freeze
+
+      QUOTE_CHARS = ["'", '"', '`'].freeze
+
+      # The combined transaction-state decision for one executed statement. +autocommit_value+ is nil
+      # when the statement does not set autocommit.
+      TransactionEffect = Data.define(:opens_transaction, :closes_transaction, :autocommit_value)
+
+      module_function
+
+      # Computes every transaction-state effect of a statement in a single pass. It runs on every
+      # executed statement, so the SQL is normalized once here rather than once per question asked
+      # of it (open? close? set autocommit?).
+      def transaction_effect(method_name, args, autocommit:, autocommit_before:, mysql_backslash_escapes: false)
+        method_closes = CLOSE_TRANSACTION_METHODS.include?(method_name) ||
+                        method_name == RubyMethod::CONNECTION_TRANSACTION.name
+
+        unless EXECUTE_SQL_METHODS.include?(method_name)
+          return TransactionEffect.new(opens_transaction: false, closes_transaction: method_closes, autocommit_value: nil)
+        end
+
+        stmt = first_statement(args&.first, mysql_backslash_escapes: mysql_backslash_escapes)
+        return TransactionEffect.new(opens_transaction: false, closes_transaction: method_closes, autocommit_value: nil) unless stmt
+
+        sets_autocommit = stmt.start_with?('SET AUTOCOMMIT')
+        autocommit_value = sets_autocommit ? parse_autocommit_value(stmt) : nil
+
+        opens = starts_transaction?(stmt) || (!autocommit && opens_transaction_scope?(stmt))
+        closes = method_closes || ends_transaction?(stmt) ||
+                 (!autocommit_before && sets_autocommit && autocommit_value == true)
+
+        TransactionEffect.new(opens_transaction: opens, closes_transaction: closes, autocommit_value: autocommit_value)
+      end
+
+      # Extracts the boolean autocommit value from an already-normalized statement.
+      def parse_autocommit_value(stmt)
+        sep = stmt.index('=')
+        if sep
+          val_start = sep + 1
+        else
+          to_idx = stmt.index(' TO ')
+          return nil unless to_idx
+
+          val_start = to_idx + 4
+        end
+
+        val = stmt[val_start..].split(';', 2).first.strip
+        case val
+        when 'TRUE', '1', 'ON' then true
+        when 'FALSE', '0', 'OFF' then false
+        end
+      end
+
+      def first_statement(sql, mysql_backslash_escapes: false)
+        sql = SqlEncoding.inspectable(sql)
+        return nil unless sql.is_a?(String) && !sql.strip.empty?
+
+        stmts = strip_comments(sql, mysql_backslash_escapes: mysql_backslash_escapes).split(';')
+        return nil if stmts.empty?
+
+        stmts.first.squeeze(' ').strip.upcase(:ascii)
+      end
+
+      # Removes -- line comments, # line comments (MySQL), and /* */ block comments.
+      # Quoted sections ('', "", ``, $$…$$) are preserved so comment markers inside strings are ignored.
+      # Postgres nested block comments (/* /* */ */) are handled correctly.
+      # Pass mysql_backslash_escapes: true to honour \ as an escape inside quoted strings.
+      # Each comment is replaced by a single space to avoid merging adjacent tokens.
+      def strip_comments(sql, mysql_backslash_escapes: false)
+        result = +''
+        i = 0
+        len = sql.length
+        while i < len
+          c = sql[i]
+          if c == '$' && (i + 1 < len) && (sql[i + 1] == '$' || sql[i + 1].match?(/[A-Za-z_]/)) &&
+             (m = sql[i..].match(/\A(\$[^$]*\$)/n))
+            # Postgres dollar-quoted string: $tag$...$tag$
+            tag = m[1]
+            close = sql.index(tag, i + tag.length)
+            if close.nil?
+              result << sql[i..]
+              i = len
+            else
+              result << sql[i...(close + tag.length)]
+              i = close + tag.length
+            end
+          elsif QUOTE_CHARS.include?(c)
+            j = skip_quoted(sql, i, backslash_escapes: mysql_backslash_escapes)
+            result << sql[i...j]
+            i = j
+          elsif c == '#' || (c == '-' && i + 1 < len && sql[i + 1] == '-')
+            i = skip_to_end_of_line(sql, i)
+            result << ' '
+          elsif c == '/' && i + 1 < len && sql[i + 1] == '*'
+            i = skip_block_comment(sql, i + 2)
+            result << ' '
+          else
+            result << c
+            i += 1
+          end
+        end
+        result
+      end
+
+      def skip_block_comment(sql, start)
+        depth = 1
+        i = start
+        len = sql.length
+        while i < len && depth.positive?
+          if sql[i] == '/' && sql[i + 1] == '*'
+            depth += 1
+            i += 2
+          elsif sql[i] == '*' && sql[i + 1] == '/'
+            depth -= 1
+            i += 2
+          else
+            i += 1
+          end
+        end
+        i
+      end
+
+      def skip_quoted(sql, start, backslash_escapes: false)
+        quote = sql[start]
+        i = start + 1
+        while i < sql.length
+          if backslash_escapes && sql[i] == '\\'
+            i += 2
+          elsif sql[i] == quote
+            return i + 1 if i + 1 >= sql.length || sql[i + 1] != quote
+
+            i += 2
+          else
+            i += 1
+          end
+        end
+        sql.length
+      end
+
+      def skip_to_end_of_line(sql, start)
+        i = start
+        i += 1 while i < sql.length && sql[i] != "\n" && sql[i] != "\r"
+        i
+      end
+
+      def starts_transaction?(stmt)
+        stmt.start_with?('BEGIN', 'START TRANSACTION')
+      end
+
+      def ends_transaction?(stmt)
+        stmt.start_with?('COMMIT', 'ROLLBACK', 'END', 'ABORT')
+      end
+
+      def opens_transaction_scope?(stmt)
+        !starts_transaction?(stmt) && !ends_transaction?(stmt) &&
+          !stmt.start_with?('SET ') && !stmt.start_with?('USE ') && !stmt.start_with?('SHOW ')
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/utils/storage/cache_entry.rb

@@ -0,0 +1,56 @@
+# 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 Utils
+    module Storage
+      # A container that holds a cached value along with its expiration time.
+      class CacheEntry
+        attr_reader :value
+
+        # @param value [Object] the cached value.
+        # @param expiration_time [Float] monotonic clock time (in seconds) at which this entry expires.
+        def initialize(value, expiration_time)
+          @value = value
+          @expiration_time = expiration_time
+        end
+
+        # @return [Boolean] true if the entry has passed its expiration time.
+        def expired?
+          Process.clock_gettime(Process::CLOCK_MONOTONIC) > @expiration_time
+        end
+
+        # Extends the expiration by the given TTL from the current time.
+        # @param ttl [Numeric] time-to-live in seconds from now.
+        def extend_expiration(ttl)
+          @expiration_time = Process.clock_gettime(Process::CLOCK_MONOTONIC) + ttl
+        end
+
+        # @param other [Object] the object to compare.
+        # @return [Boolean] true if both entries hold the same value.
+        def ==(other)
+          other.is_a?(self.class) && @value == other.value
+        end
+
+        alias eql? ==
+
+        def hash
+          value.hash
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/utils/storage/expiration_cache.rb

@@ -0,0 +1,108 @@
+# 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 'cache_entry'
+
+module AwsAdvancedRubyDriverWrapper
+  module Utils
+    module Storage
+      # A thread-safe cache with fixed TTL expiration. Entries expire after a set duration
+      # and are not renewed on access. Expired entries are removed on read or via
+      # {#remove_expired_entries}. Suitable for simple data that does not require
+      # special cleanup when removed.
+      #
+      # For entries that need sliding expiration or controlled removal,
+      # see {SlidingExpirationCache}.
+      class ExpirationCache
+        DEFAULT_TTL = 300 # 5 minutes in seconds
+
+        # @param ttl [Numeric] time-to-live for cache entries in seconds.
+        def initialize(ttl: DEFAULT_TTL)
+          @cache = {}
+          @ttl = ttl
+          @lock = Mutex.new
+        end
+
+        # Stores the given value at the given key.
+        # @param key [Object] the cache key.
+        # @param value [Object] the value to store.
+        # @return [Object, nil] the previous value, or nil.
+        def put(key, value)
+          @lock.synchronize do
+            previous = @cache[key]
+            @cache[key] = new_cache_entry(value)
+
+            previous&.value
+          end
+        end
+
+        # Retrieves the value at the given key. Returns nil if absent or expired.
+        # @param key [Object] the cache key.
+        # @return [Object, nil] the cached value, or nil.
+        def get(key)
+          @lock.synchronize do
+            entry = @cache[key]
+            return nil if entry.nil?
+
+            if entry.expired?
+              @cache.delete(key)
+              return nil
+            end
+
+            entry.value
+          end
+        end
+
+        # Removes the value at the given key.
+        # @param key [Object] the cache key.
+        # @return [Object, nil] the removed value, or nil.
+        def remove(key)
+          @lock.synchronize do
+            entry = @cache.delete(key)
+            entry&.value
+          end
+        end
+
+        # Removes all entries.
+        def clear
+          @lock.synchronize do
+            @cache.clear
+          end
+        end
+
+        # @return [Integer] the number of entries (including expired ones).
+        def size
+          @lock.synchronize do
+            @cache.size
+          end
+        end
+
+        # Removes expired entries from the cache.
+        def remove_expired_entries
+          @lock.synchronize do
+            @cache.delete_if { |_key, entry| entry.expired? }
+          end
+        end
+
+        private
+
+        def new_cache_entry(value)
+          CacheEntry.new(value, Process.clock_gettime(Process::CLOCK_MONOTONIC) + @ttl)
+        end
+      end
+    end
+  end
+end