active_postgres 0.9.4 → 0.9.6

data/lib/active_postgres/components/repmgr.rb

@@ -114,6 +114,7 @@ module ActivePostgres
         if replication_user == repmgr_user && replication_password != repmgr_password
           raise Error, 'replication_user matches repmgr user but passwords differ. Use a distinct replication_user or the same password.'
         end
+
         effective_replication_password = replication_user == repmgr_user ? repmgr_password : replication_password
 
         # Variables used in ERB templates via binding
@@ -174,13 +175,13 @@ module ActivePostgres
           info 'Creating repmgr database and user...'
           repmgr_sql = repmgr_component.send(:build_repmgr_setup_sql, repmgr_user, repmgr_db, repmgr_password)
           executor.run_sql_on_backend(self, repmgr_sql, postgres_user: 'postgres', port: 5432, tuples_only: false,
-                                           capture: false)
+                                                        capture: false)
 
           if replication_user != repmgr_user
             info 'Ensuring replication user exists...'
             repl_sql = repmgr_component.send(:build_replication_user_sql, replication_user, effective_replication_password)
             executor.run_sql_on_backend(self, repl_sql, postgres_user: 'postgres', port: 5432, tuples_only: false,
-                                             capture: false)
+                                                        capture: false)
           end
 
           info 'Reloading PostgreSQL configuration to apply pg_hba.conf changes...'
@@ -202,14 +203,14 @@ module ActivePostgres
           cluster_show = nil
           5.times do |attempt|
             cluster_show = capture(:sudo, '-u', 'postgres', 'bash', '-lc',
-                                   "repmgr cluster show -f /etc/repmgr.conf 2>&1", raise_on_non_zero_exit: false).to_s
+                                   'repmgr cluster show -f /etc/repmgr.conf 2>&1', raise_on_non_zero_exit: false).to_s
 
             break if cluster_show.match?(/primary/i)
 
             sleep 2 if attempt < 4
           end
 
-          unless cluster_show && cluster_show.match?(/primary/i)
+          unless cluster_show&.match?(/primary/i)
             # Fallback: verify via repmgr metadata in the repmgr database
             db_check = executor.run_sql_on_backend(self,
                                                    'SELECT type FROM repmgr.nodes WHERE node_id = 1;',
@@ -255,13 +256,14 @@ module ActivePostgres
         if replication_user == repmgr_user && replication_password != repmgr_password
           raise Error, 'replication_user matches repmgr user but passwords differ. Use a distinct replication_user or the same password.'
         end
+
         effective_replication_password = replication_user == repmgr_user ? repmgr_password : replication_password
 
         ensure_primary_registered
         ensure_primary_replication_ready(repmgr_password, effective_replication_password)
 
         setup_pgpass_file(standby_host, repmgr_password, replication_password: effective_replication_password,
-                                                      primary_ip: primary_replication_host)
+                                                         primary_ip: primary_replication_host)
 
         if standby_already_configured?(standby_host)
           puts '  Standby already configured, updating configs...'
@@ -418,9 +420,10 @@ module ActivePostgres
 
       def setup_inter_node_ssh
         dns_config = dns_failover_config
-        key_path = dns_config && dns_config[:ssh_key_path] || '/var/lib/postgresql/.ssh/active_postgres_dns'
+        key_path = (dns_config && dns_config[:ssh_key_path]) || '/var/lib/postgresql/.ssh/active_postgres_dns'
         postgres_user = config.postgres_user
         all_hosts = config.all_hosts
+        replication_hosts = all_hosts.to_h { |host| [host, config.replication_host_for(host)] }
         pub_keys = {}
 
         all_hosts.each do |host|
@@ -430,18 +433,35 @@ module ActivePostgres
         all_hosts.each do |host|
           ssh_executor.execute_on_host(host) do
             other_keys = pub_keys.reject { |h, _| h == host }.values
-            other_keys.each do |key|
-              next if key.to_s.empty?
-
-              upload! StringIO.new("#{key}\n"), '/tmp/pg_peer_key.pub'
-              execute :sudo, '-u', postgres_user, 'bash', '-c',
-                      "grep -qxF -f /tmp/pg_peer_key.pub /var/lib/postgresql/.ssh/authorized_keys 2>/dev/null || cat /tmp/pg_peer_key.pub >> /var/lib/postgresql/.ssh/authorized_keys"
-              execute :rm, '-f', '/tmp/pg_peer_key.pub'
+            authorized_keys_path = '/var/lib/postgresql/.ssh/authorized_keys'
+            existing_keys = if test(:sudo, 'test', '-f', authorized_keys_path)
+                              capture(:sudo, 'cat', authorized_keys_path).lines.map(&:strip)
+                            else
+                              []
+                            end
+            authorized_keys = (existing_keys + other_keys.map(&:to_s).map(&:strip)).reject(&:empty?).uniq.join("\n")
+            unless authorized_keys.empty?
+              upload! StringIO.new("#{authorized_keys}\n"), '/tmp/pg_authorized_keys'
+              execute :sudo, 'mv', '/tmp/pg_authorized_keys', authorized_keys_path
+              execute :sudo, 'chown', "#{postgres_user}:#{postgres_user}", authorized_keys_path
+              execute :sudo, 'chmod', '600', authorized_keys_path
             end
 
-            peer_ips = all_hosts.reject { |h| h == host }.map { |h| config.replication_host_for(h) }
-            scan_cmd = "ssh-keyscan #{peer_ips.join(' ')} >> /var/lib/postgresql/.ssh/known_hosts 2>/dev/null || true"
-            execute :sudo, '-u', postgres_user, 'bash', '-c', scan_cmd
+            peer_ips = all_hosts.reject { |h| h == host }.map { |h| replication_hosts[h] }
+            known_hosts_path = '/var/lib/postgresql/.ssh/known_hosts'
+            existing_hosts = if test(:sudo, 'test', '-f', known_hosts_path)
+                               capture(:sudo, 'cat', known_hosts_path).lines.map(&:strip)
+                             else
+                               []
+                             end
+            scanned_hosts = capture(:sudo, '-u', postgres_user, 'ssh-keyscan', *peer_ips, raise_on_non_zero_exit: false)
+            known_hosts = (existing_hosts + scanned_hosts.lines.map(&:strip)).reject(&:empty?).uniq.join("\n")
+            unless known_hosts.empty?
+              upload! StringIO.new("#{known_hosts}\n"), '/tmp/pg_known_hosts'
+              execute :sudo, 'mv', '/tmp/pg_known_hosts', known_hosts_path
+              execute :sudo, 'chown', "#{postgres_user}:#{postgres_user}", known_hosts_path
+              execute :sudo, 'chmod', '600', known_hosts_path
+            end
           end
         end
       end
@@ -518,10 +538,21 @@ module ActivePostgres
           public_key = capture(:sudo, '-u', postgres_user, 'cat', "#{key_path}.pub").strip
 
           unless dns_servers.empty?
-            execute :sudo, '-u', postgres_user, 'touch', '/var/lib/postgresql/.ssh/known_hosts'
-            scan_cmd = "ssh-keyscan -H #{dns_servers.join(' ')} >> /var/lib/postgresql/.ssh/known_hosts 2>/dev/null || true"
-            execute :sudo, '-u', postgres_user, 'bash', '-c', scan_cmd
-            execute :sudo, '-u', postgres_user, 'chmod', '600', '/var/lib/postgresql/.ssh/known_hosts'
+            known_hosts_path = '/var/lib/postgresql/.ssh/known_hosts'
+            existing_hosts = if test(:sudo, 'test', '-f', known_hosts_path)
+                               capture(:sudo, 'cat', known_hosts_path).lines.map(&:strip)
+                             else
+                               []
+                             end
+            scanned_hosts = capture(:sudo, '-u', postgres_user, 'ssh-keyscan', '-H', *dns_servers,
+                                    raise_on_non_zero_exit: false)
+            known_hosts = (existing_hosts + scanned_hosts.lines.map(&:strip)).reject(&:empty?).uniq.join("\n")
+            unless known_hosts.empty?
+              upload! StringIO.new("#{known_hosts}\n"), '/tmp/pg_known_hosts'
+              execute :sudo, 'mv', '/tmp/pg_known_hosts', known_hosts_path
+              execute :sudo, 'chown', "#{postgres_user}:#{postgres_user}", known_hosts_path
+              execute :sudo, 'chmod', '600', known_hosts_path
+            end
           end
         end
 
@@ -542,8 +573,8 @@ module ActivePostgres
 
             upload! StringIO.new("#{key}\n"), '/tmp/active_postgres_dns_key.pub'
             execute :bash, '-c',
-                    "grep -qxF -f /tmp/active_postgres_dns_key.pub ~/.ssh/authorized_keys || " \
-                    "cat /tmp/active_postgres_dns_key.pub >> ~/.ssh/authorized_keys"
+                    'grep -qxF -f /tmp/active_postgres_dns_key.pub ~/.ssh/authorized_keys || ' \
+                    'cat /tmp/active_postgres_dns_key.pub >> ~/.ssh/authorized_keys'
             execute :rm, '-f', '/tmp/active_postgres_dns_key.pub'
           end
         end
@@ -560,11 +591,11 @@ module ActivePostgres
 
         domains = normalize_dns_domains(dns_config)
         primary_records = normalize_dns_records(dns_config[:primary_records] || dns_config[:primary_record],
-                                               default_prefix: 'db-primary',
-                                               domains: domains)
+                                                default_prefix: 'db-primary',
+                                                domains: domains)
         replica_records = normalize_dns_records(dns_config[:replica_records] || dns_config[:replica_record],
-                                               default_prefix: 'db-replica',
-                                               domains: domains)
+                                                default_prefix: 'db-replica',
+                                                domains: domains)
 
         _ = primary_records
         _ = replica_records
@@ -732,10 +763,10 @@ module ActivePostgres
             # Check replication status
             begin
               rep_status = executor.run_sql_on_backend(self,
-                                                      'SELECT pg_is_in_recovery();',
-                                                      postgres_user: postgres_user,
-                                                      tuples_only: true,
-                                                      capture: true).to_s
+                                                       'SELECT pg_is_in_recovery();',
+                                                       postgres_user: postgres_user,
+                                                       tuples_only: true,
+                                                       capture: true).to_s
               if rep_status.include?('t')
                 info '✓ Standby is in recovery mode (receiving replication)'
               else
@@ -750,10 +781,10 @@ module ActivePostgres
             # Check lag
             begin
               lag_result = executor.run_sql_on_backend(self,
-                                                      'SELECT EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))::int AS lag;',
-                                                      postgres_user: postgres_user,
-                                                      tuples_only: true,
-                                                      capture: true).to_s.strip
+                                                       'SELECT EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))::int AS lag;',
+                                                       postgres_user: postgres_user,
+                                                       tuples_only: true,
+                                                       capture: true).to_s.strip
               info "Replication lag: #{lag_result}"
             rescue StandardError => e
               error "Failed to get replication lag: #{e.message}"
@@ -785,7 +816,7 @@ module ActivePostgres
         # Create .pgpass file for postgres user to avoid password exposure in logs
         # Format: hostname:port:database:username:password
         pgpass_content = build_pgpass_content(host, repmgr_password, replication_password: replication_password,
-                                                                   primary_ip: primary_ip)
+                                                                     primary_ip: primary_ip)
 
         ssh_executor.execute_on_host(host) do
           # Create .pgpass in postgres user's home directory
@@ -854,8 +885,13 @@ module ActivePostgres
       def escape_pgpass_value(value)
         return '' if value.nil?
 
-        # Must escape backslashes first, then colons (order matters!)
-        value.to_s.gsub('\\', '\\\\\\\\').gsub(':', '\\:')
+        value.to_s.each_char.map do |char|
+          case char
+          when '\\' then '\\\\'
+          when ':' then '\\:'
+          else char
+          end
+        end.join
       end
 
       def build_primary_conninfo(standby_label)
@@ -935,12 +971,12 @@ module ActivePostgres
         ssh_executor.execute_on_host(host) do
           repmgr_sql = repmgr_component.send(:build_repmgr_setup_sql, repmgr_user, repmgr_db, repmgr_password)
           executor.run_sql_on_backend(self, repmgr_sql, postgres_user: 'postgres', port: 5432, tuples_only: false,
-                                           capture: false)
+                                                        capture: false)
 
           if replication_user != repmgr_user
             repl_sql = repmgr_component.send(:build_replication_user_sql, replication_user, effective_replication_password)
             executor.run_sql_on_backend(self, repl_sql, postgres_user: 'postgres', port: 5432, tuples_only: false,
-                                             capture: false)
+                                                        capture: false)
           end
 
           info '✓ Primary replication user is ready'

data/lib/active_postgres/configuration.rb

@@ -39,6 +39,35 @@ module ActivePostgres
       [primary_host] + standby_hosts
     end
 
+    def pgbouncer_app_hosts
+      hosts = component_config(:pgbouncer)[:app_hosts] || component_config(:pgbouncer)['app_hosts'] || []
+      hosts = [hosts] unless hosts.is_a?(Array)
+
+      hosts.filter_map do |entry|
+        case entry
+        when String
+          entry.strip.empty? ? nil : { 'host' => entry.strip }
+        when Hash
+          host = entry['host'] || entry[:host] || entry['ip'] || entry[:ip]
+          next if host.to_s.strip.empty?
+
+          entry.merge('host' => host.to_s.strip)
+        end
+      end
+    end
+
+    def pgbouncer_primary_record
+      pgbouncer = component_config(:pgbouncer)
+      explicit = pgbouncer[:primary_record] || pgbouncer['primary_record']
+      return explicit.to_s.strip unless explicit.to_s.strip.empty?
+
+      dns_failover = component_config(:repmgr)[:dns_failover] || component_config(:repmgr)['dns_failover'] || {}
+      record = dns_failover[:primary_record] || dns_failover['primary_record']
+      return record.to_s.strip unless record.to_s.strip.empty?
+
+      primary_replication_host
+    end
+
     def primary_host
       @primary['host']
     end
@@ -96,11 +125,10 @@ module ActivePostgres
       # Validate required secrets if components are enabled
       raise Error, 'Missing replication_password secret' if component_enabled?(:repmgr) && !secrets_config['replication_password']
       raise Error, 'Missing monitoring_password secret' if component_enabled?(:monitoring) && !secrets_config['monitoring_password']
+
       if component_enabled?(:monitoring)
         grafana_config = component_config(:monitoring)[:grafana] || {}
-        if grafana_config[:enabled] && !secrets_config['grafana_admin_password']
-          raise Error, 'Missing grafana_admin_password secret'
-        end
+        raise Error, 'Missing grafana_admin_password secret' if grafana_config[:enabled] && !secrets_config['grafana_admin_password']
         if grafana_config[:enabled] && grafana_config[:host].to_s.strip.empty?
           raise Error, 'monitoring.grafana.host is required when grafana is enabled'
         end
@@ -114,27 +142,34 @@ module ActivePostgres
         end
       end
 
-        if component_enabled?(:repmgr)
-          dns_failover = component_config(:repmgr)[:dns_failover]
-          if dns_failover && dns_failover[:enabled]
-            domains = Array(dns_failover[:domains] || dns_failover[:domain]).map(&:to_s).map(&:strip).reject(&:empty?)
-            servers = Array(dns_failover[:dns_servers])
-            provider = (dns_failover[:provider] || 'dnsmasq').to_s.strip
-
-            raise Error, 'dns_failover.domain or dns_failover.domains is required when enabled' if domains.empty?
-            raise Error, 'dns_failover.dns_servers is required when enabled' if servers.empty?
-            raise Error, "Unsupported dns_failover provider '#{provider}'" unless provider == 'dnsmasq'
-
-            servers.each do |server|
-              next unless server.is_a?(Hash)
-
-              ssh_host = server['ssh_host'] || server[:ssh_host] || server['host'] || server[:host]
-              private_ip = server['private_ip'] || server[:private_ip] || server['ip'] || server[:ip]
-              raise Error, 'dns_failover.dns_servers entries must include host/ssh_host or private_ip' if
-                (ssh_host.nil? || ssh_host.to_s.strip.empty?) && (private_ip.nil? || private_ip.to_s.strip.empty?)
-            end
+      if component_enabled?(:pgbouncer)
+        pgbouncer_app_hosts.each do |app_host|
+          host = app_host['host'] || app_host[:host]
+          raise Error, 'pgbouncer.app_hosts entries must include host' if host.to_s.strip.empty?
+        end
+      end
+
+      if component_enabled?(:repmgr)
+        dns_failover = component_config(:repmgr)[:dns_failover]
+        if dns_failover && dns_failover[:enabled]
+          domains = Array(dns_failover[:domains] || dns_failover[:domain]).map(&:to_s).map(&:strip).reject(&:empty?)
+          servers = Array(dns_failover[:dns_servers])
+          provider = (dns_failover[:provider] || 'dnsmasq').to_s.strip
+
+          raise Error, 'dns_failover.domain or dns_failover.domains is required when enabled' if domains.empty?
+          raise Error, 'dns_failover.dns_servers is required when enabled' if servers.empty?
+          raise Error, "Unsupported dns_failover provider '#{provider}'" unless provider == 'dnsmasq'
+
+          servers.each do |server|
+            next unless server.is_a?(Hash)
+
+            ssh_host = server['ssh_host'] || server[:ssh_host] || server['host'] || server[:host]
+            private_ip = server['private_ip'] || server[:private_ip] || server['ip'] || server[:ip]
+            raise Error, 'dns_failover.dns_servers entries must include host/ssh_host or private_ip' if
+              (ssh_host.nil? || ssh_host.to_s.strip.empty?) && (private_ip.nil? || private_ip.to_s.strip.empty?)
           end
         end
+      end
 
       true
     end

data/lib/active_postgres/direct_executor.rb

@@ -62,19 +62,13 @@ module ActivePostgres
       return nil if value.nil?
 
       # Handle Rails credentials
-      if value.start_with?('rails_credentials:')
-        return resolve_rails_credentials(value)
-      end
+      return resolve_rails_credentials(value) if value.start_with?('rails_credentials:')
 
       # Handle environment variables
-      if value.start_with?('$')
-        return ENV[value[1..]]
-      end
+      return ENV.fetch(value[1..], nil) if value.start_with?('$')
 
       # Handle shell commands
-      if value.start_with?('$(') && value.end_with?(')')
-        return `#{value[2..-2]}`.strip
-      end
+      return `#{value[2..-2]}`.strip if value.start_with?('$(') && value.end_with?(')')
 
       value
     end

data/lib/active_postgres/error_handler.rb

@@ -127,24 +127,24 @@ module ActivePostgres
 
       # Identify error type from error message
       def identify_error_type(message)
-        case message.downcase
-        when /ssh|connection refused|network unreachable/
-          :ssh_connection
-        when /private network|vpn|wireguard network/
-          :private_network_connectivity
-        when /postgresql.*not.*start|cluster.*not.*running/
-          :postgresql_not_starting
-        when /repmgr.*clone|data directory/
-          :repmgr_clone_failed
-        when /repmgr.*register|unable to connect to.*primary/
-          :repmgr_register_failed
-        when /ssl|certificate|tls/
-          :ssl_certificate_error
-        when /no space|disk full/
-          :disk_space_error
-        when /authentication|password|pg_hba/
-          :authentication_failed
-        end
+        normalized = message.to_s.downcase
+
+        return :ssh_connection if contains_any?(normalized, 'ssh', 'connection refused', 'network unreachable')
+        return :private_network_connectivity if contains_any?(normalized, 'private network', 'vpn', 'wireguard network')
+        return :postgresql_not_starting if normalized.include?('postgresql') && normalized.include?('not') && normalized.include?('start')
+        return :postgresql_not_starting if normalized.include?('cluster') && normalized.include?('not') && normalized.include?('running')
+        return :repmgr_clone_failed if normalized.include?('repmgr') && normalized.include?('clone')
+        return :repmgr_clone_failed if normalized.include?('data directory')
+        return :repmgr_register_failed if normalized.include?('repmgr') && normalized.include?('register')
+        return :repmgr_register_failed if normalized.include?('unable to connect to') && normalized.include?('primary')
+        return :ssl_certificate_error if contains_any?(normalized, 'ssl', 'certificate', 'tls')
+        return :disk_space_error if contains_any?(normalized, 'no space', 'disk full')
+
+        :authentication_failed if contains_any?(normalized, 'authentication', 'password', 'pg_hba')
+      end
+
+      def contains_any?(value, *needles)
+        needles.any? { |needle| value.include?(needle) }
       end
 
       # Show troubleshooting guide for a specific error type

data/lib/active_postgres/health_checker.rb

@@ -12,14 +12,6 @@ module ActivePostgres
       @executor
     end
 
-    private def create_executor
-      if use_ssh_executor?
-        SSHExecutor.new(config, quiet: true)
-      else
-        DirectExecutor.new(config, quiet: true)
-      end
-    end
-
     def use_ssh_executor?
       mode = ENV['ACTIVE_POSTGRES_STATUS_MODE'].to_s.strip.downcase
       return true if mode == 'ssh'
@@ -131,6 +123,14 @@ module ActivePostgres
 
     private
 
+    def create_executor
+      if use_ssh_executor?
+        SSHExecutor.new(config, quiet: true)
+      else
+        DirectExecutor.new(config, quiet: true)
+      end
+    end
+
     def collect_node_status
       nodes = []
 
@@ -184,9 +184,7 @@ module ActivePostgres
         lag: [3, nodes.map { |n| n[:lag].to_s.length }.max].max
       }
 
-      if config.component_enabled?(:pgbouncer)
-        cols[:pgbouncer] = [9, nodes.map { |n| n[:pgbouncer].to_s.length }.max].max
-      end
+      cols[:pgbouncer] = [9, nodes.map { |n| n[:pgbouncer].to_s.length }.max].max if config.component_enabled?(:pgbouncer)
 
       cols
     end
@@ -194,7 +192,7 @@ module ActivePostgres
     def print_table_header(cols)
       fmt = "%-#{cols[:role]}s  %-#{cols[:host]}s  %-#{cols[:private_ip]}s  " \
             "%-#{cols[:label]}s  %-#{cols[:status]}s  %#{cols[:conn]}s  %#{cols[:lag]}s"
-      headers = %w[Role Host Private\ IP Label Status Conn Lag]
+      headers = ['Role', 'Host', 'Private IP', 'Label', 'Status', 'Conn', 'Lag']
 
       if cols[:pgbouncer]
         fmt += "  %-#{cols[:pgbouncer]}s"

data/lib/active_postgres/overview.rb

@@ -61,14 +61,14 @@ module ActivePostgres
         label = config.node_label_for(host)
         stats = fetch_system_stats(host, paths)
         if stats.nil?
-          puts "  #{host}#{label ? " (#{label})" : ''}: unavailable"
+          puts "  #{host}#{" (#{label})" if label}: unavailable"
           next
         end
 
         mem_used_kb = stats[:mem_total_kb] - stats[:mem_avail_kb]
         mem_pct = stats[:mem_total_kb].positive? ? (mem_used_kb.to_f / stats[:mem_total_kb] * 100).round : 0
 
-        puts "  #{host}#{label ? " (#{label})" : ''}: load #{stats[:loadavg]} | cpu #{stats[:cpu]}% | " \
+        puts "  #{host}#{" (#{label})" if label}: load #{stats[:loadavg]} | cpu #{stats[:cpu]}% | " \
              "mem #{format_kb(mem_used_kb)}/#{format_kb(stats[:mem_total_kb])} (#{mem_pct}%)"
 
         stats[:disks].each do |disk|
@@ -227,7 +227,7 @@ module ActivePostgres
       with_timeout("system stats #{host}") do
         ssh_executor = executor
         safe_paths = paths.map { |p| Shellwords.escape(p) }.join(' ')
-        script = <<~'BASH'
+        script = <<~BASH
           set -e
           loadavg=$(cut -d' ' -f1-3 /proc/loadavg)
 
@@ -256,7 +256,7 @@ module ActivePostgres
 
         ssh_executor.execute_on_host(host) do
           output = capture(:bash, '-lc', "#{script}\n df -kP #{safe_paths} 2>/dev/null | tail -n +2 | " \
-                                        "awk '{print \"disk=\" $6 \"|\" $2 \"|\" $3 \"|\" $4 \"|\" $5}'",
+                                         "awk '{print \"disk=\" $6 \"|\" $2 \"|\" $3 \"|\" $4 \"|\" $5}'",
                            raise_on_non_zero_exit: false).to_s
         end
       end
@@ -299,12 +299,12 @@ module ActivePostgres
       stats
     end
 
-    def format_kb(kb)
-      kb = kb.to_f
-      gb = kb / 1024 / 1024
+    def format_kb(kilobytes)
+      kilobytes = kilobytes.to_f
+      gb = kilobytes / 1024 / 1024
       return format('%.1fG', gb) if gb >= 1
 
-      mb = kb / 1024
+      mb = kilobytes / 1024
       format('%.0fM', mb)
     end
 
@@ -342,8 +342,8 @@ module ActivePostgres
       value.positive? ? value : DEFAULT_TIMEOUT
     end
 
-    def with_timeout(label)
-      Timeout.timeout(overview_timeout) { yield }
+    def with_timeout(label, &)
+      Timeout.timeout(overview_timeout, &)
     rescue Timeout::Error
       raise StandardError, "#{label} timed out after #{overview_timeout}s"
     end

data/lib/active_postgres/secrets.rb

@@ -20,8 +20,8 @@ module ActivePostgres
     end
 
     def resolve_all
-      config.secrets_config.keys.each_with_object({}) do |key, result|
-        result[key] = resolve(key)
+      config.secrets_config.keys.to_h do |key|
+        [key, resolve(key)]
       end
     end
 
@@ -47,8 +47,9 @@ module ActivePostgres
 
       resolve_all.each do |key, value|
         file_path = File.join(directory, key)
-        File.write(file_path, value)
-        File.chmod(0o600, file_path)
+        File.open(file_path, File::WRONLY | File::CREAT | File::TRUNC, 0o600) do |file|
+          file.write(value)
+        end
         puts "Cached #{key} to #{file_path}"
       end
 
@@ -96,12 +97,12 @@ module ActivePostgres
       return if defined?(::Rails) && ::Rails.respond_to?(:application) && ::Rails.application
 
       @rails_boot_attempted = true
+      ENV['RAILS_ENV'] ||= config.environment
+      ENV['BORING_ENVIRONMENT'] ||= config.environment
 
       # Try loading just the Rails application (without full initialization)
       # This makes Rails.application available for credential access
-      if File.exist?('./config/application.rb')
-        require './config/application'
-      end
+      require './config/application' if File.exist?('./config/application.rb')
     rescue StandardError
       # Rails boot failed — CLI running outside a Rails project
       nil

data/lib/active_postgres/ssh_executor.rb

@@ -215,11 +215,11 @@ module ActivePostgres
       execute_on_host(host) do
         backend = self
         result = executor.run_sql_on_backend(backend, sql,
-                                    postgres_user: postgres_user,
-                                    port: port,
-                                    database: database,
-                                    tuples_only: tuples_only,
-                                    capture: capture)
+                                             postgres_user: postgres_user,
+                                             port: port,
+                                             database: database,
+                                             tuples_only: tuples_only,
+                                             capture: capture)
       end
       result
     end

data/lib/active_postgres/version.rb

@@ -1,3 +1,3 @@
 module ActivePostgres
-  VERSION = '0.9.4'.freeze
+  VERSION = '0.9.6'.freeze
 end