aws_advanced_ruby_driver_wrapper 1.0.0

data/lib/aws_advanced_ruby_driver_wrapper/services/dialect_service.rb

@@ -0,0 +1,267 @@
+# 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 '../db_dialects/aurora_mysql_dialect'
+require_relative '../db_dialects/aurora_pg_dialect'
+require_relative '../db_dialects/dialect_codes'
+require_relative '../db_dialects/global_mysql_dialect'
+require_relative '../db_dialects/global_pg_dialect'
+require_relative '../db_dialects/multi_az_cluster_mysql_dialect'
+require_relative '../db_dialects/multi_az_cluster_pg_dialect'
+require_relative '../db_dialects/mysql_dialect'
+require_relative '../db_dialects/pg_dialect'
+require_relative '../db_dialects/rds_mysql_dialect'
+require_relative '../db_dialects/rds_pg_dialect'
+require_relative '../db_dialects/unknown_dialect'
+require_relative '../driver_dialects/driver_dialect_manager'
+require_relative '../errors'
+require_relative '../property_definition'
+require_relative '../utils/rds_url_type'
+require_relative '../utils/rds_utils'
+require_relative '../utils/storage/expiration_cache'
+
+module AwsAdvancedRubyDriverWrapper
+  module Services
+    class DialectService
+      attr_reader :driver_dialect, :db_dialect
+
+      ENDPOINT_CACHE_EXPIRATION = 86_400 # 24 hours in seconds
+
+      @known_endpoint_dialects = Utils::Storage::ExpirationCache.new(ttl: ENDPOINT_CACHE_EXPIRATION)
+
+      class << self
+        attr_reader :known_endpoint_dialects
+      end
+
+      KNOWN_DIALECT_CLASSES = {
+        DialectCodes::MYSQL => DbDialects::MysqlDialect,
+        DialectCodes::RDS_MYSQL => DbDialects::RdsMysqlDialect,
+        DialectCodes::AURORA_MYSQL => DbDialects::AuroraMysqlDialect,
+        DialectCodes::GLOBAL_AURORA_MYSQL => DbDialects::GlobalMysqlDialect,
+        DialectCodes::MULTI_AZ_MYSQL_CLUSTER => DbDialects::MultiAzClusterMysqlDialect,
+        DialectCodes::PG => DbDialects::PgDialect,
+        DialectCodes::RDS_PG => DbDialects::RdsPgDialect,
+        DialectCodes::AURORA_PG => DbDialects::AuroraPgDialect,
+        DialectCodes::GLOBAL_AURORA_PG => DbDialects::GlobalPgDialect,
+        DialectCodes::MULTI_AZ_PG_CLUSTER => DbDialects::MultiAzClusterPgDialect,
+        DialectCodes::UNKNOWN => DbDialects::UnknownDialect
+      }.freeze
+
+      # @param connection_service [ConnectionService]
+      # @param driver_name [Symbol] :mysql2 or :postgresql
+      def initialize(connection_service, driver_name)
+        @connection_service = connection_service
+        @dialect_cache = {}
+        @can_update = false
+        @driver_dialect = DriverDialects::DriverDialectManager.get_dialect(driver_name)
+        @error_handler = AwsAdvancedRubyDriverWrapper.config.custom_error_handler ||
+                         DriverDialects::DriverDialectManager.get_error_handler(driver_name)
+        @db_dialect = init_dialect
+      end
+
+      # @return [Boolean] whether the dialect is known for certain and will not be swapped. This
+      #   is the true when the dialect can no longer be updated, either because it was set
+      #   explicitly / resolved from an unambiguous URL (such as a global writer cluster endpoint)
+      #   or because it has already been confirmed via live connection queries. When true, topology
+      #   fetches may proceed.
+      def dialect_final?
+        !@can_update
+      end
+
+      # Lazily instantiates and caches a dialect by code.
+      #
+      # @param code [String] the dialect code
+      # @return [Object, nil] the dialect instance or nil if unknown
+      def dialect_for_code(code)
+        @dialect_cache[code] ||= begin
+          known_class = KNOWN_DIALECT_CLASSES[code]
+          known_class&.new(@driver_dialect)
+        end
+      end
+
+      # Refines the dialect after a connection is established by querying the server
+      # (e.g. checking for Aurora-specific functions/tables).
+      #
+      # @param connection [Object] the live database connection
+      # @return [Object] the updated database dialect
+      def update_dialect(connection)
+        original_dialect_code = @dialect_code
+
+        if @can_update
+          host = @connection_service.initial_host_info&.host
+          host_url = @connection_service.initial_host_info&.url
+
+          candidates = @db_dialect.dialect_update_candidates
+          candidates&.each do |candidate_code|
+            candidate = dialect_for_code(candidate_code)
+            raise Errors::AwsError, "Unknown dialect code: #{candidate_code}" unless candidate
+
+            next unless candidate.dialect?(connection)
+
+            @can_update = false
+            @dialect_code = candidate_code
+            @db_dialect = candidate
+
+            self.class.known_endpoint_dialects.put(host, candidate_code) if host
+            self.class.known_endpoint_dialects.put(host_url, candidate_code) if host_url
+
+            break
+          end
+
+          if @can_update
+            # No candidate matched
+            raise Errors::AwsError, 'Unable to determine dialect' if @dialect_code == DialectCodes::UNKNOWN
+
+            @can_update = false
+            self.class.known_endpoint_dialects.put(host, @dialect_code) if host
+            self.class.known_endpoint_dialects.put(host_url, @dialect_code) if host_url
+          end
+        end
+
+        swap_host_list_provider if @dialect_code != original_dialect_code
+        @service_container.host_service.refresh_host_list
+        @db_dialect
+      end
+
+      # Creates the initial host list provider from the URL-guessed dialect.
+      # Called after the service container is fully assembled.
+      #
+      # @param service_container [ServiceContainer]
+      def setup_initial_provider(service_container)
+        @service_container = service_container
+        provider = @db_dialect.create_host_list_provider(service_container)
+        service_container.host_service.host_list_provider = provider if provider
+      end
+
+      # @param error [Exception]
+      # @return [Boolean]
+      def network_error?(error)
+        @error_handler.network_error?(error)
+      end
+
+      # @param error [Exception]
+      # @return [Boolean]
+      def login_error?(error)
+        @error_handler.login_error?(error)
+      end
+
+      # @param error [Exception]
+      # @return [Boolean]
+      def read_only_error?(error)
+        @error_handler.read_only_error?(error)
+      end
+
+      private
+
+      # Resolves the initial database dialect from the initial connection info.
+      # Uses RdsUtils to classify the host (Aurora cluster, RDS instance, etc.)
+      # and selects the appropriate dialect.
+      #
+      # @return [Object] the resolved database dialect
+      def init_dialect
+        @can_update = false
+        @db_dialect = nil
+
+        # Check for a custom dialect object injected via config
+        custom = AwsAdvancedRubyDriverWrapper.config.custom_dialect
+        if custom
+          @dialect_code = :custom
+          @db_dialect = custom
+          @can_update = false
+          return @db_dialect
+        end
+
+        user_dialect_setting = PropertyDefinition::DIALECT.get(@connection_service.wrapper_props)&.to_s
+        user_dialect_setting = user_dialect_setting&.strip&.downcase
+        host = @connection_service.initial_host_info&.host
+
+        dialect_code = if user_dialect_setting.nil? || user_dialect_setting.empty?
+                         self.class.known_endpoint_dialects.get(host) unless host.nil?
+                       else
+                         user_dialect_setting
+                       end
+
+        if dialect_code
+          dialect = dialect_for_code(dialect_code)
+          raise Errors::AwsError, "Unknown dialect code: #{dialect_code}" unless dialect
+
+          @dialect_code = dialect_code
+          @db_dialect = dialect
+          return @db_dialect
+        end
+
+        rds_type = Utils::RdsUtils.identify_rds_type(host)
+
+        @dialect_code = if @driver_dialect == DriverDialects::DriverDialectManager::MYSQL_DIALECT
+                          resolve_mysql_dialect(rds_type)
+                        elsif @driver_dialect == DriverDialects::DriverDialectManager::PG_DIALECT
+                          resolve_pg_dialect(rds_type)
+                        else
+                          DialectCodes::UNKNOWN
+                        end
+
+        @db_dialect = dialect_for_code(@dialect_code)
+        @db_dialect
+      end
+
+      def resolve_mysql_dialect(rds_type)
+        if rds_type == Utils::RdsUrlType::RDS_GLOBAL_WRITER_CLUSTER
+          @can_update = false
+          DialectCodes::GLOBAL_AURORA_MYSQL
+        elsif rds_type.rds_cluster?
+          @can_update = true
+          DialectCodes::AURORA_MYSQL
+        elsif rds_type.rds?
+          @can_update = true
+          DialectCodes::RDS_MYSQL
+        else
+          @can_update = true
+          DialectCodes::MYSQL
+        end
+      end
+
+      def resolve_pg_dialect(rds_type)
+        if rds_type == Utils::RdsUrlType::RDS_GLOBAL_WRITER_CLUSTER
+          @can_update = false
+          DialectCodes::GLOBAL_AURORA_PG
+        elsif rds_type == Utils::RdsUrlType::RDS_AURORA_LIMITLESS_DB_SHARD_GROUP
+          @can_update = false
+          DialectCodes::AURORA_PG
+        elsif rds_type.rds_cluster?
+          @can_update = true
+          DialectCodes::AURORA_PG
+        elsif rds_type.rds?
+          @can_update = true
+          DialectCodes::RDS_PG
+        else
+          @can_update = true
+          DialectCodes::PG
+        end
+      end
+
+      def swap_host_list_provider
+        return unless @service_container
+
+        host_service = @service_container.host_service
+        old_provider = host_service.host_list_provider
+        old_provider&.stop_monitor
+
+        new_provider = @db_dialect.create_host_list_provider(@service_container)
+        host_service.host_list_provider = new_provider
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/services/host_service.rb

@@ -0,0 +1,199 @@
+# 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 'concurrent/map'
+require_relative '../host/random_host_selector'
+
+module AwsAdvancedRubyDriverWrapper
+  module Services
+    class HostService
+      DEFAULT_HOST_SELECTORS = {
+        Host::RandomHostSelector::STRATEGY_NAME => Host::RandomHostSelector.new
+      }.freeze
+
+      attr_accessor :host_list_provider
+
+      @host_id_cache = Concurrent::Map.new
+      @strategies = Concurrent::Map.new
+      DEFAULT_HOST_SELECTORS.each { |name, selector| @strategies[name] = selector }
+
+      def initialize(service_container)
+        @service_container = service_container
+        @all_hosts = []
+        @availability_cache = Utils::Storage::ExpirationCache.new
+        @host_list_provider = nil
+      end
+
+      class << self
+        attr_reader :host_id_cache
+
+        def clear_id_cache
+          @host_id_cache.clear
+        end
+
+        # Register a non-default host selector. The selector is shared by every HostService in the process,
+        # so it must be safe to call from multiple threads.
+        #
+        # @param name [String] strategy name
+        # @param selector [#select_host] any object responding to select_host(hosts, role, props)
+        def register_host_selector(name, selector)
+          raise Errors::AwsError, "Cannot override default host selection strategy: '#{name}'" if DEFAULT_HOST_SELECTORS.key?(name)
+
+          @strategies[name] = selector
+        end
+
+        # @param name [String] strategy name
+        # @return [#select_host, nil] the registered selector, or nil if the name is unknown
+        def host_selector(name)
+          @strategies[name]
+        end
+
+        # Removes every non-default host selector. For testing only.
+        # @api private
+        def reset_host_selectors
+          # Snapshot the names first rather than iterating the map while deleting from it.
+          custom_names = @strategies.keys.reject { |name| DEFAULT_HOST_SELECTORS.key?(name) }
+          custom_names.each { |name| @strategies.delete(name) }
+        end
+      end
+
+      # @param hosts [Array<Host::HostInfo>]
+      # @param role [Symbol, nil]
+      # @param strategy [String]
+      # @param props [Hash, nil]
+      # @return [Host::HostInfo]
+      def select_host(hosts, role, strategy, props = nil)
+        normalized = strategy.to_s.strip.downcase
+        selector = self.class.host_selector(normalized)
+        raise Errors::AwsError, "Unsupported host selection strategy: '#{strategy}'" if selector.nil?
+
+        selector.select_host(hosts, role, props)
+      end
+
+      # @return [Array<Host::HostInfo>] all hosts in the topology, including blocked/unavailable
+      attr_reader :all_hosts
+
+      # @return [Array<Host::HostInfo>] hosts filtered by allowed/blocked rules from the custom endpoint plugin
+      def hosts
+        rules = @service_container.storage_service.get_if_registered(
+          :custom_endpoint_allowed_blocked,
+          @service_container.connection_service.initial_host_info&.url,
+          register_access: false
+        )
+        return @all_hosts if rules.nil?
+
+        allowed = rules[:allowed]
+        blocked = rules[:blocked]
+        required_role = rules[:required_role]
+
+        hosts = @all_hosts
+        hosts = hosts.select { |h| allowed.include?(h.id) } if allowed
+        hosts = hosts.reject { |h| h.id && blocked.include?(h.id) } if blocked
+        hosts = hosts.select { |h| h.role == required_role } if required_role
+        hosts
+      end
+
+      # Updates the availability of a host in the internal host list.
+      #
+      # @param host_info [HostInfo] the host whose availability has been determined
+      # @param availability [Symbol] the new availability status, e.g. :available or :unavailable
+      def set_availability(host_info, availability)
+        host = @all_hosts.find { |h| h.id == host_info.id || h.host.casecmp?(host_info.host) }
+        return if host.nil?
+
+        host.availability = availability
+        @availability_cache.put(host_info.url, availability)
+      end
+
+      # Refresh the host list from the host list provider.
+      def refresh_host_list
+        updated_hosts = @host_list_provider&.refresh
+        return if updated_hosts.nil? || updated_hosts == @all_hosts
+
+        apply_cached_availability(updated_hosts)
+        @all_hosts = updated_hosts
+      end
+
+      # Force a refresh of the host list, bypassing any caching.
+      #
+      # @param verify_writer [Boolean]
+      # @param timeout_sec [Float]
+      # @return [Boolean] whether the refresh was successful
+      def force_refresh_host_list?(verify_writer: false, timeout_sec: 5.0)
+        updated_hosts = @host_list_provider&.force_refresh(verify_writer, timeout_sec)
+        return false if updated_hosts.nil?
+
+        if updated_hosts != @all_hosts
+          apply_cached_availability(updated_hosts)
+          @all_hosts = updated_hosts
+        end
+
+        true
+      end
+
+      # Identify which host in the topology a given connection belongs to.
+      #
+      # @param connection [Object]
+      # @param connection_host_info [Host::HostInfo, nil] the host info used to establish the connection
+      # @return [Host::HostInfo, nil]
+      def identify_host(connection, connection_host_info = nil)
+        return find_host(*query_id_and_name(connection)) if connection_host_info.nil?
+
+        url_type = Utils::RdsUtils.identify_rds_type(connection_host_info&.host)
+        case url_type
+        when Utils::RdsUrlType::RDS_INSTANCE
+          connection_host_info
+        when Utils::RdsUrlType::IP_ADDRESS, Utils::RdsUrlType::OTHER
+          get_cached_host_info(connection, connection_host_info)
+        else
+          find_host(*query_id_and_name(connection))
+        end
+      end
+
+      private
+
+      def apply_cached_availability(hosts)
+        hosts.each do |host|
+          availability = @availability_cache.get(host.url)
+          next if availability.nil?
+
+          host.availability = availability
+        end
+      end
+
+      def get_cached_host_info(connection, connection_host_info)
+        host = connection_host_info.host
+        instance_id, instance_name = self.class.host_id_cache.compute_if_absent(host) do
+          query_id_and_name(connection)
+        end
+        find_host(instance_id, instance_name)
+      end
+
+      def query_id_and_name(connection)
+        @service_container.dialect_service.db_dialect.instance_identity(connection)
+      rescue StandardError
+        [nil, nil]
+      end
+
+      def find_host(instance_id, instance_name)
+        topology = @host_list_provider&.refresh
+        return nil if topology.nil? || topology.empty?
+
+        topology.find { |h| h.id == instance_id || h.host == instance_name }
+      end
+    end
+  end
+end

data/lib/aws_advanced_ruby_driver_wrapper/services/monitor_service.rb

@@ -0,0 +1,186 @@
+# 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 '../monitoring/monitor_state'
+require_relative '../logging'
+require_relative '../utils/storage/sliding_expiration_cache'
+require_relative '../utils/events/data_access_event'
+
+module AwsAdvancedRubyDriverWrapper
+  module Services
+    # Manages monitor lifecycle: registration, deduplication, expiration, and cleanup.
+    # Monitors are grouped by type (symbol) and keyed within each type.
+    # Subscribes to DataAccessEvent to extend monitor TTLs.
+    class MonitorService
+      include Logging
+
+      CLEANUP_INTERVAL_SEC = 60.0
+
+      # Internal container grouping a cache for a monitor type.
+      CacheContainer = Data.define(:cache, :produced_data_type)
+
+      # @param event_publisher [#subscribe] the event publisher to subscribe to.
+      def initialize(event_publisher:)
+        @caches = {}
+        @lock = Mutex.new
+        @running = true
+        @cleanup_thread = start_cleanup_thread
+        event_publisher.subscribe(
+          self,
+          Set[Utils::Events::DataAccessEvent]
+        )
+      end
+
+      # Registers a monitor type. No-op if already registered.
+      # @param monitor_type [Symbol] identifier for the monitor type.
+      # @param expiration_timeout_sec [Numeric] how long an unused monitor lives before expiring.
+      # @param produced_data_type [Symbol, nil] the data type this monitor produces (for DataAccessEvent linking).
+      def register_type(monitor_type, expiration_timeout_sec:, produced_data_type: nil)
+        @lock.synchronize do
+          return if @caches.key?(monitor_type)
+
+          cache = Utils::Storage::SlidingExpirationCache.new(ttl: expiration_timeout_sec)
+          @caches[monitor_type] = CacheContainer.new(cache:, produced_data_type:)
+        end
+      end
+
+      # Returns or creates a monitor for the given type and key.
+      # If the monitor already exists, extends its expiration.
+      # @param monitor_type [Symbol] the registered monitor type.
+      # @param key [Object] unique key for this monitor instance.
+      # @param service_container [Object] passed to the initializer block.
+      # @yield [service_container] block to create the monitor if absent.
+      # @return [Object] the monitor instance.
+      def run_if_absent(monitor_type, key, service_container)
+        container = @lock.synchronize { @caches[monitor_type] }
+        raise ArgumentError, "Monitor type not registered: #{monitor_type}" unless container
+
+        container.cache.compute_if_absent(key) do
+          monitor = yield(service_container)
+          monitor.start
+          monitor
+        end
+      end
+
+      # Retrieves a monitor by type and key. Returns nil if absent or expired.
+      # @param monitor_type [Symbol] the monitor type.
+      # @param key [Object] the monitor key.
+      # @return [Object, nil]
+      def get(monitor_type, key)
+        container = @lock.synchronize { @caches[monitor_type] }
+        container&.cache&.get(key)
+      end
+
+      # Removes a monitor without stopping it.
+      # @param monitor_type [Symbol] the monitor type.
+      # @param key [Object] the monitor key.
+      # @return [Object, nil] the removed monitor.
+      def remove(monitor_type, key)
+        container = @lock.synchronize { @caches[monitor_type] }
+        container&.cache&.remove(key)
+      end
+
+      # Stops and removes a monitor.
+      # @param monitor_type [Symbol] the monitor type.
+      # @param key [Object] the monitor key.
+      def stop_and_remove(monitor_type, key)
+        monitor = remove(monitor_type, key)
+        monitor&.stop
+      end
+
+      # Stops and removes all monitors across all types.
+      def stop_and_remove_all
+        @lock.synchronize { @caches.values }.each do |container|
+          container.cache.entries.each_key do |key|
+            monitor = container.cache.remove(key)
+            monitor&.stop
+          end
+        end
+      end
+
+      def shutdown(grace_period:)
+        @running = false
+        begin
+          @cleanup_thread&.wakeup
+        rescue ThreadError
+          nil
+        end
+        @cleanup_thread&.join([grace_period, 5].min)
+        stop_and_remove_all
+      end
+
+      # Processes events from the event publisher.
+      # @param event [Event] the event to process.
+      def process_event(event)
+        return unless event == Utils::Events::DataAccessEvent
+
+        handle_data_access_event(event)
+      end
+
+      private
+
+      def handle_data_access_event(event)
+        @lock.synchronize { @caches.values }.each do |container|
+          next unless container.produced_data_type == event.data_type
+
+          container.cache.extend_expiration(event.key)
+        end
+      end
+
+      def start_cleanup_thread
+        thread = Thread.new do
+          while @running
+            sleep(CLEANUP_INTERVAL_SEC)
+            run_cleanup
+          end
+        end
+        thread.name = 'monitor-service-cleanup'
+        thread
+      end
+
+      def run_cleanup
+        @lock.synchronize { @caches.values }.each do |container|
+          cleanup_container(container)
+        end
+      rescue StandardError
+        # Cleanup must not crash the thread.
+      end
+
+      def cleanup_container(container)
+        container.cache.entries.each_key do |key|
+          # Remove stopped monitors
+          removed = container.cache.remove_if(key) { |m| m.state == Monitoring::MonitorState::STOPPED }
+          next if removed
+
+          # Stop and remove errored monitors
+          removed = container.cache.remove_if(key) { |m| m.state == Monitoring::MonitorState::ERROR }
+          if removed
+            logger.debug("Removed monitor in error state: #{key}")
+            removed.stop
+            next
+          end
+
+          # Remove expired monitors that can be disposed
+          removed = container.cache.remove_if_expired(key)
+          if removed
+            logger.debug("Removed expired monitor: #{key}")
+            removed.stop
+          end
+        end
+      end
+    end
+  end
+end