console_kit 1.4.0 → 1.5.1

data/lib/console_kit/connections/sql_connection_handler.rb

@@ -1,69 +1,105 @@
 # frozen_string_literal: true
 
 require_relative 'base_connection_handler'
+require_relative 'sql_strategy'
 
 module ConsoleKit
   module Connections
-    # Handles SQL connections
     class SqlConnectionHandler < BaseConnectionHandler
+      backend :sql,
+              display_name: 'SQL',
+              context_attribute: :tenant_shard,
+              constants_key: :shard,
+              detail_label: 'Shard'
+
+      DEFAULT_BASE_CLASS = 'ApplicationRecord'
+      MISSING_BASE_CLASS = 'the configured sql_base_class %<name>s could not be resolved. Check the class name and ' \
+                           'that the class is loaded'
+
       class << self
+        def target_error(value) = identifier_error(value)
+
         def sql_version(conn)
           conn.select_value('SELECT version()')
-        rescue StandardError
+        rescue StandardError => e
+          raise e if ConsoleKit.programming_error?(e)
+
           nil
         end
 
         def base_class_name = ConsoleKit.configuration.sql_base_class
       end
 
-      def connect
-        shard = context_attribute(:tenant_shard).presence&.to_sym
-        Output.print_info("#{connection_message(shard)} via #{base_class}")
-        disconnect_existing_pool
-        shard ? base_class.establish_connection(shard) : base_class.establish_connection
+      def available? = resolved_base_class.present?
+
+      def unavailable_reason
+        name = self.class.base_class_name
+        return nil if name.to_s == DEFAULT_BASE_CLASS || resolved_base_class.present?
+
+        format(MISSING_BASE_CLASS, name: name.inspect)
+      end
+
+      def prepare(target)
+        validate_target!(target)
+        unless strategy.switchable?
+          raise UnsupportedBackendError, "#{display_name} base class #{base_class} cannot switch connections."
+        end
+        return if strategy.resolvable?(normalize(target))
+
+        raise ConfigurationError,
+              "ConsoleKit: SQL shard #{scrub(target.inspect)} is not a registered shard or database configuration."
       end
 
-      def available? = self.class.base_class_name.to_s.safe_constantize.present?
+      def snapshot = strategy.snapshot
 
-      def diagnostics
-        return unavailable_diagnostics('SQL') unless available?
+      def connect!(target)
+        shard = normalize(target)
+        Output.print_info("#{connection_message(shard)} via #{base_class}")
+        strategy.apply(shard)
+      end
+
+      def verify!(target)
+        expected, actual = strategy.identity(normalize(target))
+        return true if expected.to_s == actual.to_s
 
-        perform_diagnostics
-      rescue StandardError => e
-        error_diagnostics('SQL', e)
+        raise verification_error(expected, actual)
       end
 
+      def restore(state) = strategy.restore(state)
+
+      def diagnostic_identity = strategy.pool_details
+
       private
 
-      def perform_diagnostics
-        conn = base_class.connection
-        latency = measure_latency { conn.execute('SELECT 1') }
-        build_sql_diagnostics(conn, latency)
+      def basic_diagnostics
+        details = strategy.pool_details
+        { name: display_name, status: details.empty? ? :unknown : :connected, latency_ms: nil, details: details }
       end
 
-      def disconnect_existing_pool
-        pool = base_class.try(:connection_pool)
-        pool&.disconnect!
+      def full_diagnostics
+        conn = base_class.connection
+        latency = measure_latency { conn.execute('SELECT 1') }
+        { name: display_name, status: :connected, latency_ms: latency, details: full_details(conn) }
       end
 
-      def build_sql_diagnostics(conn, latency)
+      def full_details(conn)
         {
-          name: 'SQL',
-          status: :connected,
-          latency_ms: latency,
-          details: {
-            adapter: conn.adapter_name,
-            pool_size: base_class.connection_pool.size,
-            version: self.class.sql_version(conn).to_s.truncate(50)
-          }
+          adapter: conn.adapter_name,
+          pool_size: base_class.connection_pool.size,
+          version: self.class.sql_version(conn).to_s.truncate(50)
         }
       end
 
+      def strategy = @strategy ||= SqlStrategy.new(base_class)
+      def normalize(target) = target.presence&.to_sym
+
+      def resolved_base_class = self.class.base_class_name.to_s.safe_constantize
+
       def base_class
         @base_class ||= begin
           name = self.class.base_class_name
           klass = name.to_s.safe_constantize
-          klass || raise(Error, "ConsoleKit: sql_base_class '#{name}' could not be found.")
+          klass || raise(ConfigurationError, "ConsoleKit: sql_base_class '#{name}' could not be found.")
         end
       end
 

data/lib/console_kit/connections/sql_strategy.rb

@@ -0,0 +1,248 @@
+# frozen_string_literal: true
+
+require_relative '../errors'
+require_relative '../output'
+require_relative '../instrumentation'
+
+module ConsoleKit
+  module Connections
+    class ShardFrame
+      KEY = :console_kit_sql_connected_to_frame
+      STOLEN = 'ConsoleKit: a `connected_to` block removed the shard frame ConsoleKit had committed, so shard ' \
+               '%<shard>p was not live while that block was open. Re-applying it now.'
+      REASSERTED = 'console_kit.sql_frame_reasserted'
+
+      attr_reader :base_class
+
+      def initialize(base_class) = @base_class = base_class
+
+      def apply(shard)
+        base_class.connecting_to(shard: shard || base_class.default_shard, role: base_class.current_role)
+        place(shard)
+      end
+
+      def live_shard
+        reassert if taken?
+        base_class.current_shard
+      end
+
+      def applied_shard
+        return nil unless on?
+
+        owned[:shard] || base_class.default_shard
+      end
+
+      def forget_unless_on
+        Thread.current.thread_variable_set(KEY, nil) if owned && !on?
+      end
+
+      def on? = !index_in.nil?
+
+      private
+
+      def stack = base_class.try(:connected_to_stack)
+      def taken? = !owned.nil? && !on?
+
+      def reassert
+        shard = owned[:shard]
+        Instrumentation.increment(REASSERTED)
+        Output.print_warning(format(STOLEN, shard: shard || base_class.default_shard))
+        apply(shard)
+      end
+
+      def place(shard)
+        current = stack
+        index = index_in
+        current[index] = current.pop if index
+        remember(current, index || (current.size - 1), shard)
+      end
+
+      def index_in
+        current = stack
+        entry = owned
+        return nil unless entry && current
+
+        index = entry[:index]
+        index if current[index].equal?(entry[:frame])
+      end
+
+      def owned
+        entry = Thread.current.thread_variable_get(KEY)
+        entry if entry && entry[:base].equal?(base_class)
+      end
+
+      def remember(current, index, shard)
+        Thread.current.thread_variable_set(
+          KEY, { base: base_class, frame: current[index], index: index, shard: shard }
+        )
+      end
+    end
+
+    class PoolSlot
+      attr_reader :base_class
+
+      def initialize(base_class) = @base_class = base_class
+
+      def absent? = lookup.nil?
+
+      def remove
+        return if absent?
+        return handler.remove_connection_pool(spec, role: role, shard: shard) if removable_handler?
+
+        base_class.remove_connection if base_class.respond_to?(:remove_connection)
+      end
+
+      private
+
+      def handler = base_class.try(:connection_handler)
+      def spec = base_class.try(:connection_specification_name)
+      def role = base_class.try(:current_role)
+      def shard = base_class.try(:current_shard)
+      def removable_handler? = handler.respond_to?(:remove_connection_pool) && spec
+
+      def lookup
+        return handler.retrieve_connection_pool(spec, role: role, shard: shard) if retrievable_handler?
+
+        base_class.try(:connection_pool)
+      rescue StandardError => e
+        raise e if ConsoleKit.programming_error?(e)
+
+        nil
+      end
+
+      def retrievable_handler? = handler.respond_to?(:retrieve_connection_pool) && spec
+    end
+
+    class SqlStrategy
+      NATIVE_METHODS = %i[connecting_to connected_to_stack default_shard current_shard current_role].freeze
+
+      attr_reader :base_class
+
+      def initialize(base_class) = @base_class = base_class
+
+      def switchable? = base_class.respond_to?(:establish_connection) || native_capable?
+
+      def native?(shard)
+        return false unless native_capable?
+        return !connected_to_stack.to_a.empty? if shard.nil?
+
+        !shard_pool(shard).nil?
+      end
+
+      def resolvable?(shard)
+        return true if !shard || native?(shard)
+
+        configs = env_configs
+        configs.empty? || configs.any? { |cfg| config_name(cfg).to_s == shard.to_s }
+      end
+
+      def snapshot
+        name = current_db_config_name
+        {
+          shard: frame.applied_shard || base_class.try(:current_shard),
+          role: base_class.try(:current_role),
+          stack_depth: connected_to_stack&.size,
+          db_config_name: name,
+          pool_absent: !name && pool.absent?
+        }
+      end
+
+      def apply(shard) = native?(shard) ? apply_native(shard) : apply_fallback(shard)
+
+      def restore(state)
+        return unless state
+
+        unwind_stack(state[:stack_depth])
+        restore_shard(state[:shard])
+        state[:pool_absent] ? pool.remove : reestablish(state[:db_config_name])
+      end
+
+      def identity(shard)
+        return [shard || base_class.default_shard, frame.live_shard] if native?(shard)
+
+        [expected_db_config_name(shard), current_db_config_name]
+      end
+
+      def pool_details
+        describe_pool(base_class.try(:connection_pool))
+      rescue StandardError => e
+        raise e if ConsoleKit.programming_error?(e)
+
+        {}
+      end
+
+      private
+
+      def describe_pool(live_pool)
+        return {} unless live_pool
+
+        { adapter: live_pool.try(:db_config).try(:adapter), pool_size: live_pool.try(:size),
+          config: current_db_config_name, shard: base_class.try(:current_shard) }.compact
+      end
+
+      def native_capable? = NATIVE_METHODS.all? { |method| base_class.respond_to?(method) }
+      def connected_to_stack = base_class.try(:connected_to_stack)
+
+      def frame = @frame ||= ShardFrame.new(base_class)
+      def pool = @pool ||= PoolSlot.new(base_class)
+      def apply_native(shard) = frame.apply(shard)
+
+      def restore_shard(shard)
+        return if !shard || !native_capable? || (frame.applied_shard || base_class.current_shard) == shard
+
+        apply_native(shard)
+      end
+
+      def apply_fallback(shard)
+        desired = expected_db_config_name(shard)
+        return if desired && desired.to_s == current_db_config_name.to_s
+
+        shard ? base_class.establish_connection(shard.to_sym) : base_class.establish_connection
+      end
+
+      def unwind_stack(depth)
+        stack = connected_to_stack
+        return unless stack && depth
+
+        stack.pop while stack.size > depth
+        frame.forget_unless_on
+      end
+
+      def reestablish(name)
+        return if !name || name.to_s == current_db_config_name.to_s
+
+        base_class.establish_connection(name.to_sym)
+      end
+
+      def shard_pool(shard)
+        handler = base_class.try(:connection_handler)
+        spec_name = base_class.try(:connection_specification_name)
+        return nil unless spec_name && handler.respond_to?(:retrieve_connection_pool)
+
+        handler.retrieve_connection_pool(spec_name, role: base_class.current_role, shard: shard.to_sym)
+      end
+
+      def expected_db_config_name(shard) = shard ? shard.to_s : config_name(env_configs.first)
+
+      def env_configs
+        configs = base_class.try(:configurations)
+        env = current_db_config.try(:env_name)
+        return [] unless env && configs.respond_to?(:configs_for)
+
+        configs.configs_for(env_name: env)
+      end
+
+      def current_db_config
+        base_class.try(:connection_pool).try(:db_config)
+      rescue StandardError => e
+        raise e if ConsoleKit.programming_error?(e)
+
+        nil
+      end
+
+      def current_db_config_name = config_name(current_db_config)
+
+      def config_name(config) = config&.name
+    end
+  end
+end

data/lib/console_kit/connections/table_renderer.rb

@@ -1,43 +1,53 @@
 # frozen_string_literal: true
 
-require_relative 'table_formatter'
-
 module ConsoleKit
   module Connections
-    # Renders diagnostic data into a Unicode box-drawing table
     module TableRenderer
+      HEADERS = %w[Service Status Latency Details].freeze
+      STATUS = {
+        connected: "\u2713 Connected",
+        error: "\u2717 Error",
+        unavailable: "\u2014 N/A"
+      }.freeze
+
       class << self
         def render(rows)
-          headers = %w[Service Status Latency Details]
-          table_rows = rows.map { |row| TableFormatter.format_row(row) }
-          widths = calculate_widths(headers, table_rows)
+          table_rows = rows.map { |row| format_row(row) }
+          widths = ([HEADERS] + table_rows).transpose.map { |column| column.map(&:length).max }
 
-          build_table(headers, table_rows, widths)
+          build_table(table_rows, widths)
         end
 
         private
 
-        def calculate_widths(headers, rows)
-          all_rows = [headers] + rows
-          headers.each_index.map do |index|
-            column_max_width(all_rows, index)
-          end
+        def format_row(diag)
+          latency = diag[:latency_ms]
+          [
+            diag[:name],
+            STATUS.fetch(diag[:status], '? Unknown'),
+            latency ? "#{latency}ms" : "\u2014",
+            format_details(diag[:details])
+          ]
         end
 
-        def column_max_width(rows, index)
-          rows.map { |row| row[index].length }.max
+        def format_details(details)
+          return '' unless details&.any?
+
+          details.compact.map { |key, value| "#{key}: #{value}" }.join(', ')
         end
 
-        def build_table(headers, rows, widths)
-          lines = [table_top(widths), table_line(headers, widths), table_mid(widths)]
+        def build_table(rows, widths)
+          lines = [rule(widths, "\u250C\u252C\u2510"), table_line(HEADERS, widths),
+                   rule(widths, "\u251C\u253C\u2524")]
           rows.each { |row| lines << table_line(row, widths) }
-          lines << table_bottom(widths)
+          lines << rule(widths, "\u2514\u2534\u2518")
           lines.join("\n")
         end
 
-        def table_top(widths) = "\u250C#{widths.map { |width| "\u2500" * (width + 2) }.join("\u252C")}\u2510"
-        def table_mid(widths) = "\u251C#{widths.map { |width| "\u2500" * (width + 2) }.join("\u253C")}\u2524"
-        def table_bottom(widths) = "\u2514#{widths.map { |width| "\u2500" * (width + 2) }.join("\u2534")}\u2518"
+        def rule(widths, corners)
+          left, join, right = corners.chars
+          "#{left}#{widths.map { |width| "\u2500" * (width + 2) }.join(join)}#{right}"
+        end
 
         def table_line(cells, widths)
           content = cells.each_with_index.map { |cell, index| " #{cell.ljust(widths[index])} " }.join("\u2502")

data/lib/console_kit/console_helpers.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 module ConsoleKit
-  # Helper methods available in the Rails console
   module ConsoleHelpers
     def switch_tenant
       ConsoleKit.reset_current_tenant
@@ -9,56 +8,43 @@ module ConsoleKit
     end
 
     def tenant_info
-      tenant = ConsoleKit::Setup.current_tenant
-      return no_tenant_warning unless tenant
+      tenant = ConsoleKit::StateStore.tenant_key
+      return ConsoleKit::Output.print_warning('No tenant is currently configured.') unless tenant
 
-      display_tenant_info(tenant)
+      constants = ConsoleKit.configuration.tenants[tenant]&.[](:constants) || {}
+      ConsoleHelpers.print_tenant_details(tenant, constants)
       nil
     end
 
-    def dashboard
-      ConsoleKit::Connections::Dashboard.display
+    def dashboard(level: :basic)
+      ConsoleKit::Connections::Dashboard.display(level: level)
       self
     end
 
     def tenants
       names = ConsoleKit.configuration.tenants&.keys || []
-      print_available_tenants(names)
-      names
-    end
-
-    DETAIL_LABELS = {
-      'Partner' => :partner_code, 'Shard' => :shard, 'Mongo DB' => :mongo_db,
-      'Redis DB' => :redis_db, 'ES Prefix' => :elasticsearch_prefix, 'Environment' => :environment
-    }.freeze
-
-    private
-
-    def no_tenant_warning
-      ConsoleKit::Output.print_warning('No tenant is currently configured.')
-      self
-    end
-
-    def display_tenant_info(tenant)
-      constants = ConsoleKit.configuration.tenants[tenant]&.[](:constants) || {}
-      ConsoleHelpers.print_tenant_details(tenant, constants)
-      self
-    end
-
-    def print_available_tenants(names)
       ConsoleKit::Output.print_list(names, header: 'Available Tenants')
-      self
+      names
     end
 
     class << self
       def print_tenant_details(tenant, constants)
         ConsoleKit::Output.print_header("Tenant: #{tenant}")
-        DETAIL_LABELS.each do |label, key|
+        detail_labels.each do |label, key|
           next unless constants.key?(key)
 
           ConsoleKit::Output.print_info("  #{label.ljust(13)}#{constants[key]}")
         end
       end
+
+      private
+
+      def detail_labels
+        backend_labels = ConsoleKit::Connections::BaseConnectionHandler.registry.to_h do |handler|
+          [handler.detail_label, handler.constants_key]
+        end
+        { 'Partner' => :partner_code }.merge(backend_labels).merge('Environment' => :environment)
+      end
     end
   end
 end

data/lib/console_kit/diagnostics.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+require_relative 'errors'
+require_relative 'tenant_state'
+require_relative 'instrumentation'
+require_relative 'connections/diagnostic_helpers'
+
+module ConsoleKit
+  module Diagnostics
+    LEVELS = %i[basic full].freeze
+    DEFAULT_TIMEOUT = 2
+    CACHE_TTL_SECONDS = 2.0
+    EVENT = 'console_kit.diagnostics'
+    TIMEOUT_COUNTER = 'console_kit.diagnostics_timeout'
+
+    class << self
+      def run(level: :basic, timeout: DEFAULT_TIMEOUT)
+        validate_level!(level)
+        available_handlers.map { |handler| cached(handler, level, timeout) }
+      end
+
+      def clear_cache! = Cache.clear!
+
+      def validate_level!(level)
+        return level if LEVELS.include?(level)
+
+        raise ConfigurationError,
+              "ConsoleKit: unknown diagnostics level #{level.inspect}. Expected one of #{LEVELS.inspect}."
+      end
+
+      private
+
+      def available_handlers
+        Connections::ConnectionManager.available_handlers(ConsoleKit.configuration.context_class)
+      end
+
+      def cached(handler, level, timeout)
+        Cache.fetch_row(handler, level) { handler.safe_diagnostics(timeout: timeout, level: level) }
+      end
+    end
+
+    module Cache
+      STORE_KEY = :console_kit_diagnostics_cache
+      CACHED_LEVEL = :full
+
+      class << self
+        def fetch_row(handler, level)
+          return yield unless level == CACHED_LEVEL
+
+          key = [StateStore.tenant_key, level, handler.backend_key]
+          identity = identity_of(handler)
+          read(key, identity) || write(key, identity, yield)
+        end
+
+        def clear!
+          Thread.current.thread_variable_set(STORE_KEY, nil)
+        end
+
+        private
+
+        def store
+          thread = Thread.current
+          thread.thread_variable_get(STORE_KEY) || thread.thread_variable_set(STORE_KEY, {})
+        end
+
+        def read(key, identity)
+          entry = store[key]
+          return nil unless entry
+
+          unless current?(entry, identity)
+            store.delete(key)
+            return nil
+          end
+
+          entry[:row]
+        end
+
+        def write(key, identity, row)
+          return row if row[:status] == :error
+
+          purge_expired!
+          store[key] = { row: row, expires_at: now + CACHE_TTL_SECONDS, state: state, identity: identity }
+          row
+        end
+
+        def identity_of(handler)
+          handler.diagnostic_identity
+        rescue StandardError => e
+          raise e if ConsoleKit.programming_error?(e)
+
+          Object.new
+        end
+
+        def purge_expired!
+          store.delete_if { |_key, entry| !fresh?(entry) }
+        end
+
+        def fresh?(entry) = entry[:expires_at] > now && entry[:state].equal?(state)
+
+        def current?(entry, identity) = fresh?(entry) && entry[:identity] == identity
+
+        def state = StateStore.stored
+        def now = Connections::DiagnosticHelpers.clock_time
+      end
+    end
+
+    module Runner
+      class << self
+        def call(handler, timeout: DEFAULT_TIMEOUT, level: :basic)
+          Diagnostics.validate_level!(level)
+          started = Connections::DiagnosticHelpers.clock_time
+          outcome = execute(handler, level)
+          report_overrun(started, timeout)
+          resolve(handler, outcome)
+        end
+
+        def execute(handler, level)
+          Instrumentation.instrument(EVENT, backend: handler.backend_key, level: level) do
+            handler.diagnostics(level: level)
+          end
+        rescue StandardError, ScriptError => e
+          e
+        end
+
+        private
+
+        def resolve(handler, outcome) = outcome.is_a?(Hash) ? outcome : failed_row(handler, outcome)
+
+        def report_overrun(started, timeout)
+          return if Connections::DiagnosticHelpers.clock_time - started <= timeout
+
+          Instrumentation.increment(TIMEOUT_COUNTER)
+        end
+
+        def failed_row(handler, error)
+          raise error if ConsoleKit.programming_error?(error)
+
+          Connections::DiagnosticHelpers.error_diagnostics(handler.display_name, error)
+        end
+      end
+    end
+  end
+end