aws_advanced_ruby_driver_wrapper 1.0.0

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

@@ -0,0 +1,137 @@
+# 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 sliding TTL expiration. Accessing an existing entry
+      # via {#compute_if_absent} renews its expiration, and {#extend_expiration} allows
+      # manual renewal. Provides conditional removal via {#remove_if} and
+      # {#remove_if_expired} for entries that require careful lifecycle management.
+      # Disposal of removed items should be handled by the caller.
+      #
+      # For simple data that does not need expiration renewal or controlled removal,
+      # see {ExpirationCache}.
+      class SlidingExpirationCache
+        DEFAULT_TTL = 900 # 15 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
+
+        # Retrieves the value at the given key. Returns nil if absent or expired.
+        # Extends the expiration if the entry exists and is not 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? || entry.expired?
+
+            entry.extend_expiration(@ttl)
+            entry.value
+          end
+        end
+
+        # Computes and stores a value if the key is absent.
+        # If the key exists, extends its expiration and returns the existing value, even if it is expired.
+        # @param key [Object] the cache key.
+        # @yield [key] block to compute the value if absent.
+        # @return [Object] the current (existing or computed) value.
+        def compute_if_absent(key)
+          @lock.synchronize do
+            entry = @cache[key]
+
+            if entry.nil?
+              value = yield(key)
+              @cache[key] = new_cache_entry(value)
+              return value
+            end
+
+            entry.extend_expiration(@ttl)
+            entry.value
+          end
+        end
+
+        # Extends the expiration of the entry at the given key, if it exists.
+        # @param key [Object] the cache key.
+        def extend_expiration(key)
+          @lock.synchronize do
+            entry = @cache[key]
+            entry&.extend_expiration(@ttl)
+          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 the value at the given key only if the predicate returns true.
+        # @param key [Object] the cache key.
+        # @yield [value] block that returns true if the value should be removed.
+        # @return [Object, nil] the removed value, or nil.
+        def remove_if(key)
+          @lock.synchronize do
+            entry = @cache[key]
+            return nil if entry.nil?
+            return nil unless yield(entry.value)
+
+            @cache.delete(key)
+            entry.value
+          end
+        end
+
+        # Removes the value at the given key only if it is expired.
+        # @param key [Object] the cache key.
+        # @return [Object, nil] the removed value, or nil.
+        def remove_if_expired(key)
+          @lock.synchronize do
+            entry = @cache[key]
+            return nil if entry.nil? || !entry.expired?
+
+            @cache.delete(key)
+            entry.value
+          end
+        end
+
+        # Returns a hash copy of all entries (including expired ones).
+        # @return [Hash]
+        def entries
+          @lock.synchronize do
+            @cache.transform_values(&:value)
+          end
+        end
+
+        private
+
+        def new_cache_entry(value)
+          CacheEntry.new(value, Process.clock_gettime(Process::CLOCK_MONOTONIC) + @ttl)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,172 @@
+# 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 'expiration_cache'
+require_relative '../../logging'
+require_relative '../events/data_access_event'
+
+module AwsAdvancedRubyDriverWrapper
+  module Utils
+    module Storage
+      # A centralized, shared cache registry with per-type TTL and background cleanup.
+      # Each named cache partition stores items independently with its own expiration policy.
+      class StorageService
+        include Logging
+
+        DEFAULT_CLEANUP_INTERVAL = 300 # 5 minutes in seconds
+
+        # @param event_publisher [#publish] the event publisher for data access events.
+        # @param cleanup_interval [Numeric] seconds between cleanup runs.
+        def initialize(event_publisher:, cleanup_interval: DEFAULT_CLEANUP_INTERVAL)
+          @caches = {}
+          @event_publisher = event_publisher
+          @lock = Mutex.new
+          @running = true
+          @cleanup_thread = start_cleanup_thread(cleanup_interval)
+        end
+
+        # Registers a named cache. No-op if already registered.
+        #
+        # @param name [Symbol] cache partition name.
+        # @param ttl [Numeric] time-to-live in seconds.
+        def register(name, ttl:)
+          @lock.synchronize do
+            return if @caches.key?(name)
+
+            @caches[name] = ExpirationCache.new(ttl:)
+          end
+        end
+
+        # Stores an item.
+        # @param name [Symbol] registered cache name.
+        # @param key [Object] item key.
+        # @param value [Object] item to store.
+        def set(name, key, value)
+          fetch_cache!(name).put(key, value)
+        end
+
+        # Retrieves an item. Returns nil if absent or expired.
+        # @param name [Symbol] registered cache name.
+        # @param key [Object] item key.
+        # @param register_access [Boolean] whether to publish a DataAccessEvent.
+        # @return [Object, nil]
+        def get(name, key, register_access: true)
+          value = fetch_cache!(name).get(key)
+          return nil unless value
+
+          @event_publisher.publish(Events::DataAccessEvent.new(data_type: name, key:)) if register_access && @event_publisher
+
+          value
+        end
+
+        # Retrieves an item without raising if the named cache has not been registered.
+        # Returns nil when the cache partition is absent, or when the item is absent or expired.
+        # Use this for reads of optional/foreign caches that may not exist in the current
+        # plugin configuration (e.g. a plugin reading a cache owned by another, optional plugin).
+        # @param name [Symbol] cache name (need not be registered).
+        # @param key [Object] item key.
+        # @param register_access [Boolean] whether to publish a DataAccessEvent.
+        # @return [Object, nil]
+        def get_if_registered(name, key, register_access: true)
+          cache = @caches[name]
+          return nil if cache.nil?
+
+          value = cache.get(key)
+          return nil unless value
+
+          @event_publisher.publish(Events::DataAccessEvent.new(data_type: name, key:)) if register_access && @event_publisher
+
+          value
+        end
+
+        # Returns true if a non-expired item exists at the given name + key.
+        # @param name [Symbol] registered cache name.
+        # @param key [Object] item key.
+        # @return [Boolean]
+        def exists?(name, key)
+          !get(name, key, register_access: false).nil?
+        end
+
+        # Removes an item.
+        # @param name [Symbol] registered cache name.
+        # @param key [Object] item key.
+        def remove(name, key)
+          fetch_cache!(name).remove(key)
+        end
+
+        # Clears all items for a given name.
+        # @param name [Symbol] registered cache name.
+        def clear(name)
+          fetch_cache!(name).clear
+        end
+
+        # Clears all items from all registered caches.
+        def clear_all
+          @lock.synchronize { @caches.keys }.each { |name| clear(name) }
+        end
+
+        # Returns the number of items (including expired) for a given name.
+        # @param name [Symbol] registered cache name.
+        # @return [Integer]
+        def size(name)
+          fetch_cache!(name).size
+        end
+
+        # Stops the cleanup thread.
+        def shutdown
+          @running = false
+          begin
+            @cleanup_thread&.wakeup
+          rescue ThreadError
+            nil
+          end
+          @cleanup_thread&.join(5)
+        end
+
+        private
+
+        def fetch_cache!(name)
+          cache = @caches[name]
+          raise ArgumentError, "Cache not registered: #{name.inspect}" unless cache
+
+          cache
+        end
+
+        def start_cleanup_thread(interval)
+          thread = Thread.new do
+            while @running
+              sleep(interval)
+              remove_expired_items
+            end
+          end
+          thread.name = 'storage-service-cleanup'
+          thread
+        end
+
+        def remove_expired_items
+          @lock.synchronize { @caches.keys }.each do |name|
+            cache = @caches[name]
+            next unless cache
+
+            cache.remove_expired_entries
+          rescue StandardError => e
+            logger.debug("StorageService cleanup failed for #{name}: #{e.message}")
+          end
+        end
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/utils/topology_utils.rb

@@ -0,0 +1,127 @@
+# 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 'logger'
+require_relative '../host/host_info'
+require_relative '../host/host_role'
+require_relative '../host/host_availability'
+require_relative '../logging'
+
+module AwsAdvancedRubyDriverWrapper
+  module Utils
+    # A mixin providing shared utility methods for retrieving and processing single-cluster topology information.
+    # Classes that include this module must implement:
+    #   - #build_hosts(conn, results, initial_host_info, instance_template) => Array<HostInfo> or nil
+    module TopologyUtils
+      include Logging
+
+      attr_reader :dialect
+
+      # Builds a HostInfo from the given topology information.
+      #
+      # @param instance_id [String, nil] the database instance identifier, e.g. "mydb-instance-1" (Aurora) or "123456789" (Multi-AZ).
+      # @param instance_name [String, nil] the database instance name, e.g. "mydb-instance-1" (Aurora and Multi-AZ).
+      # @param is_writer [Boolean] true if this is a writer instance.
+      # @param weight [Integer] the instance weight for load balancing.
+      # @param last_update_time [Time] the timestamp of the last update to this instance's information.
+      # @param initial_host_info [AwsAdvancedRubyDriverWrapper::Host::HostInfo] the original host info used for connecting.
+      # @param instance_template [AwsAdvancedRubyDriverWrapper::Host::HostInfo] the template used to construct the new HostInfo.
+      # @return [AwsAdvancedRubyDriverWrapper::Host::HostInfo] a HostInfo representing the given information.
+      def build_host(instance_id, instance_name, is_writer, weight, last_update_time, initial_host_info, instance_template)
+        instance_name = '?' if instance_name.nil?
+        endpoint = instance_template.host.sub('?', instance_name)
+        port = resolve_port(instance_template, initial_host_info)
+        role = is_writer ? Host::HostRole::WRITER : Host::HostRole::READER
+
+        Host::HostInfo.new(
+          host: endpoint,
+          port:,
+          role:,
+          availability: Host::HostAvailability::AVAILABLE,
+          weight:,
+          id: instance_id,
+          last_update_time:
+        )
+      end
+
+      # Retrieves a value from a row hash, trying both string and symbol keys.
+      def row_value(row, key)
+        row[key] || row[key.to_sym]
+      end
+
+      # Query the database for information for each instance in the database topology.
+      #
+      # @param conn [Object] the connection to use to query the database.
+      # @param initial_host_info [AwsAdvancedRubyDriverWrapper::Host::HostInfo] the HostInfo used to initially connect.
+      # @param instance_template [AwsAdvancedRubyDriverWrapper::Host::HostInfo] the template HostInfo to use when
+      #   constructing new HostInfo objects from the data returned by the topology query.
+      # @return [Array<AwsAdvancedRubyDriverWrapper::Host::HostInfo>, nil] a list of HostInfo objects representing
+      #   the results of the topology query, or nil if the query returned unexpected results.
+      def query_topology(conn, initial_host_info, instance_template)
+        results = @dialect.execute(conn, @dialect.topology_query)
+        # We expect at least 4 columns. Note that the server may return 0 columns if failover has occurred.
+        if results.fields.empty?
+          logger.debug('The topology query returned a result with 0 columns. ' \
+                       'This may occur if the topology query is executed when the server is failing over.')
+          return nil
+        end
+
+        verify_writer(build_hosts(conn, results, initial_host_info, instance_template))
+      end
+
+      private
+
+      # Verifies the writer in the host list. If multiple writers exist, the one with the most recent
+      # last_update_time is used as the current writer.
+      #
+      # @param all_hosts [Array<AwsAdvancedRubyDriverWrapper::Host::HostInfo>, nil] the list of all hosts.
+      # @return [Array<AwsAdvancedRubyDriverWrapper::Host::HostInfo>, nil] the verified host list, or nil if no writer found.
+      def verify_writer(all_hosts)
+        return nil if all_hosts.nil?
+
+        hosts = []
+        writers = []
+
+        all_hosts.each do |host|
+          if host.role == Host::HostRole::WRITER
+            writers << host
+          else
+            hosts << host
+          end
+        end
+
+        return nil if writers.empty?
+
+        if writers.size == 1
+          hosts << writers.first
+        else
+          # Assume the latest updated writer instance is the current writer.
+          sorted_writers = writers.sort_by { |w| w.last_update_time || Time.at(0) }.reverse
+          hosts << sorted_writers.first
+        end
+
+        hosts
+      end
+
+      def resolve_port(instance_template, initial_host_info)
+        return instance_template.port if instance_template.port_specified?
+        return initial_host_info.port if initial_host_info
+
+        Host::HostInfo::NO_PORT
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/version.rb

@@ -0,0 +1,19 @@
+# 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
+  VERSION = '1.0.0'
+end

data/lib/aws_advanced_ruby_driver_wrapper/wrapper_property.rb

@@ -0,0 +1,64 @@
+# 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
+  class WrapperProperty
+    attr_reader :name, :default_value, :description, :type
+
+    def initialize(name, description, default_value: nil, type: nil, validator: nil)
+      @name = name.to_sym
+      @description = description
+      @default_value = default_value
+      @type = type
+      @validator = validator
+    end
+
+    # @return [String, Boolean, Integer, nil] the value from props, or the property's default
+    def get(props, override = nil)
+      override&.key?(@name) ? override.fetch(@name, @default_value) : props.fetch(@name, @default_value)
+    end
+
+    def get_bool(props, override = nil)
+      val = get(props, override)
+      return val if val.is_a?(TrueClass) || val.is_a?(FalseClass)
+
+      val.to_s.downcase == 'true'
+    end
+
+    def get_int(props, override = nil)
+      val = get(props, override)
+      val.is_a?(Integer) ? val : val.to_i
+    end
+
+    def get_float(props, override = nil)
+      val = get(props, override)
+      val.is_a?(Float) ? val : val.to_f
+    end
+
+    def get_string(props, override = nil)
+      val = get(props, override)
+      val&.to_s
+    end
+
+    def validate!(value)
+      @validator&.call(value, @name)
+    end
+
+    def set(props, value)
+      props.put(@name, value)
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper.rb

@@ -0,0 +1,116 @@
+# 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/version'
+require_relative 'aws_advanced_ruby_driver_wrapper/custom_configuration'
+require_relative 'aws_advanced_ruby_driver_wrapper/property_definition'
+require_relative 'aws_advanced_ruby_driver_wrapper/logging'
+# Error classes are loaded up front so they can be named before any connection is opened, for
+# example in a Rails `rescue_from` or ActiveJob `retry_on`, which are evaluated at class load.
+require_relative 'aws_advanced_ruby_driver_wrapper/errors'
+require_relative 'aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/errors'
+require_relative 'aws_advanced_ruby_driver_wrapper/utils/rds_utils'
+require_relative 'aws_advanced_ruby_driver_wrapper/utils/connection_config'
+require_relative 'aws_advanced_ruby_driver_wrapper/utils/connection_config_parser'
+require_relative 'aws_advanced_ruby_driver_wrapper/monitoring/monitor_state'
+require_relative 'aws_advanced_ruby_driver_wrapper/monitoring/monitor'
+require_relative 'aws_advanced_ruby_driver_wrapper/plugins/blue_green/blue_green_plugin'
+
+module AwsAdvancedRubyDriverWrapper
+  @config = Configuration.new
+
+  # Direct-driver connection classes are autoloaded so requiring this file alone is enough to use them:
+  # referencing either constant loads its driver-specific file on demand. Loading stays lazy - a client
+  # for one driver never loads the other driver's code, and the underlying 'pg'/'mysql2' gem is only
+  # required when a connection is actually opened.
+  autoload :WrapperPgConnection, 'aws_advanced_ruby_driver_wrapper/postgresql'
+  autoload :WrapperMysql2Client, 'aws_advanced_ruby_driver_wrapper/mysql'
+
+  # Other public entry points are autoloaded for the same reason, so they work from code that runs
+  # before any connection is opened, such as a Rails initializer or a setup script.
+  module Services
+    autoload :HostService, 'aws_advanced_ruby_driver_wrapper/services/host_service'
+    autoload :PluginManager, 'aws_advanced_ruby_driver_wrapper/services/plugin_manager'
+  end
+
+  module Plugins
+    module Encryption
+      autoload :EncryptionConfig, 'aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/encryption_config'
+      autoload :KeyManagementUtility, 'aws_advanced_ruby_driver_wrapper/plugins/kms_encryption/key_management_utility'
+    end
+  end
+
+  class << self
+    attr_reader :config
+  end
+
+  # Gracefully tears down all background resources. Invoked by the at_exit hook (and thus
+  # indirectly by the TERM/INT signal traps, which just exit). Safe to call more than once.
+  def self.shutdown(grace_period_sec: 10)
+    Plugins::BlueGreen::BlueGreenPlugin.clean_up_providers
+    release_resources(grace_period_sec: grace_period_sec)
+  end
+
+  def self.clear_caches
+    require_relative 'aws_advanced_ruby_driver_wrapper/services/service_utility'
+    require_relative 'aws_advanced_ruby_driver_wrapper/services/host_service'
+    Services::CoreServices.storage_service.clear_all
+    Utils::RdsUtils.clear_cache
+    Services::DialectService.known_endpoint_dialects.clear
+    Services::HostService.clear_id_cache
+  end
+
+  def self.release_resources(grace_period_sec: 5)
+    require_relative 'aws_advanced_ruby_driver_wrapper/services/service_utility'
+    Services::CoreServices.monitor_service.shutdown(grace_period: grace_period_sec)
+    Services::CoreServices.event_publisher.release_resources
+    clear_caches
+  end
+end
+
+# Register signal traps and at_exit hook for graceful shutdown. The traps only exit — the real
+# teardown runs in the at_exit hook, since acquiring locks / joining threads inside a signal-trap
+# context is unsafe in Ruby and can deadlock.
+%w[TERM INT].each { |signal| trap(signal) { exit(0) } }
+
+at_exit { AwsAdvancedRubyDriverWrapper.shutdown }
+
+# Register adapters with ActiveRecord via a lazy-load hook so the require order between this file and
+# ActiveRecord does not matter. The block runs the first time ActiveRecord::Base is referenced, whether
+# ActiveRecord is loaded before or after this file. The register call is itself lazy — each adapter file
+# is only loaded when a connection using that adapter is first established, so users of a single adapter
+# never load the other.
+begin
+  require 'active_support/lazy_load_hooks'
+rescue LoadError
+  # ActiveSupport is absent (raw-driver-only usage with no ActiveRecord); there is nothing to register.
+end
+
+if defined?(ActiveSupport) && ActiveSupport.respond_to?(:on_load)
+  ActiveSupport.on_load(:active_record) do
+    ActiveRecord::ConnectionAdapters.register(
+      'aws_postgresql',
+      'ActiveRecord::ConnectionAdapters::AwsPostgreSQLAdapter',
+      'aws_advanced_ruby_driver_wrapper/active_record/aws_postgresql_adapter'
+    )
+
+    ActiveRecord::ConnectionAdapters.register(
+      'aws_mysql2',
+      'ActiveRecord::ConnectionAdapters::AwsMysql2Adapter',
+      'aws_advanced_ruby_driver_wrapper/active_record/aws_mysql2_adapter'
+    )
+  end
+end