legion-data 1.10.8 → 1.10.17

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 869153710ece6717285edf2de7483f4c9f5fe5b493837681b80d05b77b014971
-  data.tar.gz: 2813fd707f0f7a80858802046b7ba2a2f23611f94a1515c66e5c7a5db1771dd6
+  metadata.gz: cc2433eff09e371c9b406ba8790fe20da105a2920fb5490947ae590158eee403
+  data.tar.gz: 6d4ec2038079bbf0f01be66b2093313af264651e69ff7d38bf0a89a1f8656b43
 SHA512:
-  metadata.gz: 41aff7bbce675a95d8a2ec6683eb0350decde2cf66854971b0c43fed1d4a09bdc50563e922bfe29d5a01660c522d579be7f2719495150d136ba262f433ba4ac4
-  data.tar.gz: c766be4f0a13cf3ac3dc91e305d275742d5c49206ec90819631d05739137fb1c9eadbc550feeef1d6e635961ec8bf439b9b5fc19dfc334d1798859573d87902a
+  metadata.gz: fdf39d3c2f11766069af4432713e353e87196478585d1c3fe782f1db0aa298f4c0520b579dadb8a343e2b74afdfba485619c29d23188340b45eab641ea726a73
+  data.tar.gz: 4a92b1cd8ab39a28fb9b144286602a851a99d64f468678cc704da1fdb79001cfd60354778da7ea41e45c6f300241ae025fe537743c507763b3d2c0967dd8522b

data/CHANGELOG.md

@@ -1,5 +1,55 @@
 # Legion::Data Changelog
 
+## [1.10.17] - 2026-10-29
+
+### Fixed
+- Cache the default encryption key under `nil`, separately from every tenant ID, including the literal `__default__`. Key selection no longer depends on which scope was requested first (#74).
+- Scope database spec cleanup to initialized resources so early setup errors are preserved and original state, including nil connection snapshots, is restored (#83).
+
+## [1.10.16] - 2026-10-29
+
+### Fixed
+- Generate each archive batch UUID before upload and use it in S3, Azure, and local object names and the manifest. Separate runs no longer overwrite earlier archives in the same month. Direct `upload_batch` callers now supply `batch_id:` instead of `batch_n:` (#73).
+- Load the archive compressor's StringIO dependency explicitly.
+
+## [1.10.15] - 2026-10-29
+
+### Fixed
+- Isolate query-log write and close I/O failures from database operations. Full disks no longer turn committed writes into apparent failures or replace the original SQL exception; logging failures remain reported through Legion::Logging (#72).
+
+## [1.10.14] - 2026-10-29
+
+### Fixed
+- Propagate fatal connection configuration, shared migration, model load, and local migration failures. Close failed setup connections, clear readiness, and allow setup to retry rather than continuing with partially initialized databases (#71).
+
+## [1.10.13] - 2026-10-29
+
+### Fixed
+- Build PostgreSQL replica pools before connecting, decode replica URLs into valid Sequel server options, and preserve inherited credentials and session timeouts. Clear replica routing on shutdown and reconfiguration (#70).
+
+## [1.10.12] - 2026-10-29
+
+### Fixed
+- Keep PostgreSQL tenant setup, scoped queries, and restoration on one checked-out connection. Preserve nested tenant scopes and clean up after exceptions without leaking tenant state to another borrower (#69).
+
+## [1.10.11] - 2026-10-29
+
+### Fixed
+- Apply PostgreSQL statement and lock timeouts before any connection is opened, including initial, preconnected, and replacement sessions (#68).
+
+## [1.10.10] - 2026-10-29
+
+### Fixed
+- Release Ruby threads while SQLite waits for write locks, with bounded per-connection busy handlers. Apply local SQLite tuning to every pooled connection and honor explicit `data.timeout` values (#67).
+
+## [1.10.9] - 2026-07-31
+### Fixed
+- **Postgres session timeouts**: Apply `statement_timeout` (default 5000ms) and `lock_timeout` (default 2000ms) via `after_connect` hook on every new Postgres connection. Prevents unbounded `SELECT` or lock-wait from hanging the process indefinitely — the root cause of ~1h50m connection wedges.
+- **TCP keepalives on Postgres**: Wire libpq `keepalives`, `keepalives_idle` (10s), `keepalives_interval` (5s), `keepalives_count` (3), and `tcp_user_timeout` (15000ms) through the Sequel connection hash. Detects half-open TCP sockets within ~25s instead of relying on OS defaults (often 2+ hours).
+- **Force-disconnect on shutdown**: `shutdown` now waits up to `shutdown_timeout` (default 5s) for checked-out connections to return, then force-disconnects all remaining connections. Previously, `disconnect` only closed idle pool connections — a connection stuck in `select()` survived forever, hanging process exit.
+- **Adapter timeout defaults tightened**: Postgres `connect_timeout` reduced from 20s to 5s; MySQL `connect_timeout` reduced from 120s to 5s, `read_timeout` and `write_timeout` default to 5s. The old 20s/120s values allowed connection attempts to block worker threads far too long under network partition.
+- **MySQL read/write timeout wiring preserved**: `read_timeout` and `write_timeout` remain in `ADAPTER_KEYS[:mysql2]` where Sequel's mysql2 adapter honors them. For Postgres, the equivalent is `statement_timeout` (query bound) plus TCP keepalives (socket bound) — `read_timeout` is not a libpq/Sequel-postgres concept, so it is correctly absent from the Postgres key list.
+
 ## [1.10.8] - 2026-07-23
 ### Fixed
 - Detect auth failures (`role does not exist`, `password authentication failed`) on Sequel pool connections and trigger immediate Vault lease reissue via LeaseManager instead of retrying dead credentials forever. 30-second cooldown prevents reissue storms during bulk credential rotation.

data/README.md

@@ -2,7 +2,7 @@
 
 Persistent database storage for the [LegionIO](https://github.com/LegionIO/LegionIO) async job engine and AI coding assistant platform. Provides database connectivity via the [Sequel ORM](https://sequel.jeremyevans.net/), automatic schema migrations (97 numbered migrations), Sequel models for the full LegionIO control plane, and a parallel local SQLite database for on-node agentic cognitive state.
 
-**Version**: 1.8.0 | **Ruby**: >= 3.4 | **License**: Apache-2.0
+**Version**: 1.10.17 | **Ruby**: >= 3.4 | **License**: Apache-2.0
 
 ---
 
@@ -23,6 +23,8 @@ Core responsibilities:
 
 The schema is portable by default across SQLite, MySQL, and PostgreSQL. PostgreSQL-only behavior is isolated to features that need PostgreSQL, such as Apollo vector columns.
 
+`Legion::Data.setup` raises on fatal connection configuration, migration, or model-loading failures rather than announcing successful startup. Failed setup closes its database connections and clears readiness; fix the underlying error before retrying. Local migrations registered after startup also propagate failures and leave the local database disconnected until setup succeeds.
+
 ## Supported Databases
 
 | Database | Adapter | Gem | Default |
@@ -239,6 +241,8 @@ Legion::Data::Rls.with_tenant(tenant_id) do
 end
 ```
 
+The scope holds one pooled connection until the block finishes and its prior tenant state has been restored, including when the block raises. Nested scopes restore the outer tenant. Run scoped database work in the calling thread; newly spawned threads do not inherit its tenant scope.
+
 ### Permission Checks
 
 ```ruby
@@ -266,6 +270,8 @@ All settings live under the `data` key. The adapter controls which options apply
 }
 ```
 
+SQLite lock waits use a Ruby busy handler that releases the GVL between retries, so other threads can finish transactions while a writer waits. `data.timeout` sets the wait limit in milliseconds for both databases; when unset, shared SQLite uses 5,000 ms and local SQLite uses 30,000 ms. Local connection-specific tuning is applied to every pooled connection. `PRAGMA busy_timeout` does not report the Ruby handler's deadline.
+
 ### MySQL
 
 ```json
@@ -300,6 +306,8 @@ All settings live under the `data` key. The adapter controls which options apply
 }
 ```
 
+`data.statement_timeout` (default 5,000 ms) and `data.lock_timeout` (default 2,000 ms) are applied before every PostgreSQL connection enters the pool, including startup and replacement connections. Set a value to `null` to leave that server setting unchanged, or `0` to explicitly disable the corresponding PostgreSQL timeout.
+
 PostgreSQL with `pgvector` is required for Apollo models:
 
 ```sql
@@ -390,6 +398,10 @@ Two independent caching tiers, both disabled by default:
 Legion::Data.reload_static_cache
 ```
 
+### Query File Logging
+
+Enable `data.query_log` (or `data.local.query_log`) to write SQL to a dedicated file. Write and close I/O failures, including a full disk, are reported through Legion::Logging without changing the database operation's result or replacing its exception. Opening the configured log file must still succeed during setup.
+
 ### Read Replicas (PostgreSQL)
 
 ```json
@@ -401,6 +413,16 @@ Legion::Data.reload_static_cache
 }
 ```
 
+Replica URLs initialize a sharded pool at startup. URL-specific options override primary connection options; omitted credentials and session timeouts are inherited. Duplicate URLs are registered once as `read_0`, `read_1`, and so on. Reads remain on the primary unless explicitly routed, for example `dataset.server(Legion::Data::Connection.read_server)`. With no replicas configured, `read_server` returns `:default`.
+
+### Archive Objects
+
+Each archive batch receives a UUID before upload. S3, Azure, and local object paths use `legion-archive/<table>/<year>/<month>/batch_<uuid>.jsonl.gz`; the same UUID is stored in `archive_manifest.batch_id`. Repeated runs do not reuse numbered object names. Existing manifest paths remain valid; no stored objects are renamed. Direct `Archiver.upload_batch` callers must supply `batch_id:` instead of `batch_n:`.
+
+### Encryption Key Cache
+
+The default scope (`tenant_id: nil`) and tenant IDs have separate cache entries, including a tenant literally named `__default__`. Restart processes or clear existing key-provider caches when upgrading. This fix does not rewrite existing ciphertext: if that tenant name was used with the old cache, assess records encrypted with the wrong scope's key and re-encrypt them with the intended key. The reader does not fall back to another scope's key.
+
 ---
 
 ## Common Fields Standard

data/lib/legion/data/archiver.rb

@@ -5,6 +5,7 @@ require 'digest'
 require 'fileutils'
 require 'json'
 require 'securerandom'
+require 'stringio'
 require 'tmpdir'
 require 'zlib'
 
@@ -47,15 +48,15 @@ module Legion
           raise
         end
 
-        def upload_batch(data:, table:, year:, month:, batch_n:, backend:)
-          log.info "Archiver storing batch table=#{table} backend=#{backend || :tmpdir} year=#{year} month=#{month} batch=#{batch_n}"
+        def upload_batch(data:, table:, year:, month:, batch_id:, backend:)
+          log.info "Archiver storing batch table=#{table} backend=#{backend || :tmpdir} year=#{year} month=#{month} batch=#{batch_id}"
           case backend
           when :s3
-            upload_s3(data: data, table: table, year: year, month: month, batch_n: batch_n)
+            upload_s3(data: data, table: table, year: year, month: month, batch_id: batch_id)
           when :azure
-            upload_azure(data: data, table: table, year: year, month: month, batch_n: batch_n)
+            upload_azure(data: data, table: table, year: year, month: month, batch_id: batch_id)
           else
-            upload_tmpdir(data: data, table: table, year: year, month: month, batch_n: batch_n)
+            upload_tmpdir(data: data, table: table, year: year, month: month, batch_id: batch_id)
           end
         end
 
@@ -102,7 +103,6 @@ module Legion
             table:           table,
             cutoff:          cutoff,
             batch_size:      batch_size,
-            batch_n:         batches + 1,
             now:             now,
             storage_backend: storage_backend
           ))
@@ -114,18 +114,19 @@ module Legion
           { batches: batches, total_rows: total_rows, paths: paths }
         end
 
-        def archive_batch(conn:, table:, cutoff:, batch_size:, batch_n:, now:, storage_backend:)
+        def archive_batch(conn:, table:, cutoff:, batch_size:, now:, storage_backend:)
           rows = conn[table].where { created_at < cutoff }.limit(batch_size).all
           return if rows.empty?
 
+          batch_id = SecureRandom.uuid
           compressed = gzip_compress(serialize_rows(rows))
           path = upload_batch(
-            data:    compressed,
-            table:   table.to_s,
-            year:    now.year,
-            month:   now.month,
-            batch_n: batch_n,
-            backend: storage_backend
+            data:     compressed,
+            table:    table.to_s,
+            year:     now.year,
+            month:    now.month,
+            batch_id: batch_id,
+            backend:  storage_backend
           )
 
           record_archived_batch(
@@ -134,16 +135,17 @@ module Legion
             rows:       rows,
             compressed: compressed,
             path:       path,
+            batch_id:   batch_id,
             now:        now
           )
 
           { row_count: rows.size, path: path }
         end
 
-        def record_archived_batch(conn:, table:, rows:, compressed:, path:, now:)
+        def record_archived_batch(conn:, table:, rows:, compressed:, path:, batch_id:, now:)
           conn.transaction do
             conn[:archive_manifest].insert(
-              batch_id:     SecureRandom.uuid,
+              batch_id:     batch_id,
               source_table: table.to_s,
               row_count:    rows.size,
               checksum:     Digest::SHA256.hexdigest(compressed),
@@ -167,47 +169,47 @@ module Legion
           output.string
         end
 
-        def upload_s3(data:, table:, year:, month:, batch_n:)
+        def upload_s3(data:, table:, year:, month:, batch_id:)
           raise UploadError, 'S3 backend not available: Legion::Extensions::S3::Runners::Put not defined' unless defined?(Legion::Extensions::S3::Runners::Put)
 
-          key = "legion-archive/#{table}/#{year}/#{month}/batch_#{batch_n}.jsonl.gz"
+          key = "legion-archive/#{table}/#{year}/#{month}/batch_#{batch_id}.jsonl.gz"
           Legion::Extensions::S3::Runners::Put.run(key: key, body: data)
           log.info "Archiver uploaded batch to s3 key=#{key}"
           "s3://#{key}"
         rescue UploadError => e
-          handle_exception(e, level: :error, handled: false, operation: :upload_s3, table: table, year: year, month: month, batch_n: batch_n)
+          handle_exception(e, level: :error, handled: false, operation: :upload_s3, table: table, year: year, month: month, batch_id: batch_id)
           raise
         rescue StandardError => e
-          handle_exception(e, level: :error, handled: true, operation: :upload_s3, table: table, year: year, month: month, batch_n: batch_n)
+          handle_exception(e, level: :error, handled: true, operation: :upload_s3, table: table, year: year, month: month, batch_id: batch_id)
           raise UploadError, "S3 upload failed: #{e.message}"
         end
 
-        def upload_azure(data:, table:, year:, month:, batch_n:)
+        def upload_azure(data:, table:, year:, month:, batch_id:)
           unless defined?(Legion::Extensions::AzureStorage::Runners::Upload)
             raise UploadError, 'Azure backend not available: Legion::Extensions::AzureStorage::Runners::Upload not defined'
           end
 
-          blob_name = "legion-archive/#{table}/#{year}/#{month}/batch_#{batch_n}.jsonl.gz"
+          blob_name = "legion-archive/#{table}/#{year}/#{month}/batch_#{batch_id}.jsonl.gz"
           Legion::Extensions::AzureStorage::Runners::Upload.run(blob_name: blob_name, data: data)
           log.info "Archiver uploaded batch to azure blob=#{blob_name}"
           "azure://#{blob_name}"
         rescue UploadError => e
-          handle_exception(e, level: :error, handled: false, operation: :upload_azure, table: table, year: year, month: month, batch_n: batch_n)
+          handle_exception(e, level: :error, handled: false, operation: :upload_azure, table: table, year: year, month: month, batch_id: batch_id)
           raise
         rescue StandardError => e
-          handle_exception(e, level: :error, handled: false, operation: :upload_azure, table: table, year: year, month: month, batch_n: batch_n)
+          handle_exception(e, level: :error, handled: false, operation: :upload_azure, table: table, year: year, month: month, batch_id: batch_id)
           raise UploadError, "Azure upload failed: #{e.message}"
         end
 
-        def upload_tmpdir(data:, table:, year:, month:, batch_n:)
+        def upload_tmpdir(data:, table:, year:, month:, batch_id:)
           dir = File.join(Dir.tmpdir, 'legion-archive', table.to_s, year.to_s, month.to_s)
           FileUtils.mkdir_p(dir)
-          path = File.join(dir, "batch_#{batch_n}.jsonl.gz")
+          path = File.join(dir, "batch_#{batch_id}.jsonl.gz")
           File.binwrite(path, data)
           log.info "Archiver stored batch locally path=#{path}"
           "file://#{path}"
         rescue StandardError => e
-          handle_exception(e, level: :error, handled: true, operation: :upload_tmpdir, table: table, year: year, month: month, batch_n: batch_n)
+          handle_exception(e, level: :error, handled: true, operation: :upload_tmpdir, table: table, year: year, month: month, batch_id: batch_id)
           raise UploadError, "Tmpdir upload failed: #{e.message}"
         end
       end

data/lib/legion/data/connection.rb

@@ -5,6 +5,7 @@ require 'legion/data/auth_failure_handler'
 
 require 'fileutils'
 require 'sequel'
+require 'uri'
 
 module Legion
   module Data
@@ -19,14 +20,18 @@ module Legion
 
       ADAPTER_KEYS = {
         sqlite:   %i[timeout readonly disable_dqs],
-        postgres: %i[connect_timeout sslmode sslrootcert search_path],
+        postgres: %i[connect_timeout sslmode sslrootcert search_path
+                     keepalives keepalives_idle keepalives_interval keepalives_count
+                     tcp_user_timeout],
         mysql2:   %i[connect_timeout read_timeout write_timeout encoding sql_mode]
       }.freeze
 
       ADAPTER_DEFAULTS = {
         sqlite:   { timeout: 5000, readonly: false, disable_dqs: true },
-        postgres: { connect_timeout: 20, sslmode: 'disable' },
-        mysql2:   { connect_timeout: 120, encoding: 'utf8mb4' }
+        postgres: { connect_timeout: 5, sslmode: 'disable',
+                    keepalives: 1, keepalives_idle: 10, keepalives_interval: 5,
+                    keepalives_count: 3, tcp_user_timeout: 15_000 },
+        mysql2:   { connect_timeout: 5, read_timeout: 5, write_timeout: 5, encoding: 'utf8mb4' }
       }.freeze
 
       QUERY_LOG_DIR = File.expand_path('~/.legionio/logs').freeze
@@ -138,6 +143,9 @@ module Legion
             @closed = true
             @file.close unless @file.closed?
           end
+        rescue IOError, SystemCallError => e
+          handle_exception(e, level: :warn, handled: true, operation: :query_file_close, path: @path)
+          nil
         end
 
         private
@@ -148,9 +156,7 @@ module Legion
 
             @file.puts "[#{Time.now.strftime('%Y-%m-%d %H:%M:%S.%L')}] #{level} #{message}"
           end
-        rescue IOError => e
-          return nil if @closed || @file.closed?
-
+        rescue IOError, SystemCallError => e
           handle_exception(e, level: :warn, handled: true, operation: :query_file_write, path: @path)
           nil
         end
@@ -166,10 +172,14 @@ module Legion
         end
 
         def setup
+          Legion::Settings[:data][:connected] = false
+          @sequel&.disconnect
+          @sequel = nil
           @adapter = Legion::Settings[:data][:adapter]&.to_sym || :sqlite
           opts = sequel_opts
           log.info("Legion::Data::Connection setup adapter=#{adapter}")
           @fallback_active = false
+          @replica_servers = []
           @sequel = if adapter == :sqlite
                       ::Sequel.connect(opts.merge(adapter: :sqlite, database: sqlite_path))
                     else
@@ -179,7 +189,7 @@ module Legion
                       rescue UnresolvedCredentialError => e
                         handle_exception(e, level: :fatal, handled: false, operation: :data_connect)
                         raise
-                      rescue StandardError => e
+                      rescue ::Sequel::DatabaseConnectionError => e
                         raise unless dev_fallback?
 
                         log.error("Legion::Data FALLING BACK TO SQLITE — #{attempted_adapter} network DB connection failed: #{e.message}")
@@ -192,10 +202,15 @@ module Legion
                         ::Sequel.connect(sqlite_opts.merge(adapter: :sqlite, database: sqlite_path))
                       end
                     end
-          Legion::Settings[:data][:connected] = true
+          @replica_servers = @sequel.opts.fetch(:servers, {}).keys
           log_connection_info
           configure_extensions
-          connect_with_replicas
+          Legion::Settings[:data][:connected] = true
+        rescue StandardError, LoadError => e
+          shutdown
+          @sequel = nil
+          handle_exception(e, level: :fatal, handled: false, operation: :connection_setup)
+          raise
         end
 
         # Returns connection metadata for health checks and diagnostics.
@@ -272,10 +287,14 @@ module Legion
         end
 
         def shutdown
-          @sequel&.disconnect
+          if @sequel
+            timeout = Legion::Settings[:data][:shutdown_timeout]
+            force_disconnect_pool(timeout: timeout)
+          end
           @query_file_logger&.close
           @query_file_logger = nil
           @fallback_active = false
+          @replica_servers = []
           Legion::Settings[:data][:connected] = false
           log.info 'Legion::Data connection closed'
         end
@@ -321,28 +340,6 @@ module Legion
           handle_exception(e, level: :warn, handled: true, operation: :expire_all_pooled_connections)
         end
 
-        def connect_with_replicas
-          return unless adapter == :postgres
-
-          replica_url  = Legion::Settings[:data][:read_replica_url]
-          replica_list = Array(Legion::Settings[:data][:replicas]).dup
-
-          replica_list.prepend(replica_url) if replica_url && !replica_url.empty?
-          replica_list.uniq!
-          replica_list.compact!
-
-          return if replica_list.empty?
-
-          @sequel.extension(:server_block)
-
-          replica_list.each_with_index do |url, idx|
-            @sequel.add_servers("read_#{idx}": url)
-          end
-
-          @replica_servers = replica_list.each_with_index.map { |_, idx| :"read_#{idx}" }
-          log.debug "Registered #{@replica_servers.size} read replica(s)"
-        end
-
         def read_server
           return :default if @replica_servers.nil? || @replica_servers.empty?
 
@@ -379,6 +376,20 @@ module Legion
           creds
         end
 
+        def configure_sqlite_busy_handler(connection:, timeout:)
+          timeout_seconds = timeout.to_f / 1000
+          deadline = nil
+          connection.busy_handler do |attempt|
+            now = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC)
+            deadline = now + timeout_seconds if attempt.zero?
+            remaining = deadline - now
+            next false unless remaining.positive?
+
+            sleep([0.001, remaining].min)
+            true
+          end
+        end
+
         def creds_builder(final_creds = {})
           final_creds.merge! Legion::Data::Settings.creds(adapter)
           final_creds.merge! Legion::Settings[:data][:creds] if Legion::Settings[:data][:creds].is_a? Hash
@@ -430,6 +441,28 @@ module Legion
           data_settings[:dev_mode] == true && data_settings[:dev_fallback] != false
         end
 
+        def force_disconnect_pool(timeout:)
+          pool = @sequel.pool
+          in_use = pool.size - (pool.respond_to?(:available_connections) ? pool.available_connections.size : 0)
+
+          if in_use.positive?
+            log.warn("Legion::Data shutdown: #{in_use} connection(s) still checked out, waiting up to #{timeout}s")
+            deadline = ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) + timeout
+            while ::Process.clock_gettime(::Process::CLOCK_MONOTONIC) < deadline
+              break if pool.size <= (pool.respond_to?(:available_connections) ? pool.available_connections.size : 0)
+
+              sleep 0.1
+            end
+            remaining = pool.size - (pool.respond_to?(:available_connections) ? pool.available_connections.size : 0)
+            log.warn("Legion::Data shutdown: force-disconnecting #{remaining} wedged connection(s)") if remaining.positive?
+          end
+
+          @sequel.disconnect
+        rescue StandardError => e
+          handle_exception(e, level: :warn, handled: true, operation: :force_disconnect_pool)
+          @sequel&.disconnect rescue nil # rubocop:disable Style/RescueModifier
+        end
+
         def sqlite_path
           path = Legion::Settings[:data][:creds][:database] || 'legionio.db'
           return path if File.absolute_path?(path)
@@ -443,9 +476,25 @@ module Legion
           connection_opts = opts.merge(adapter: adapter, **creds_builder)
           validate_no_unresolved_uris!(connection_opts)
           connection_opts[:preconnect] = false if adapter != :sqlite && dev_fallback?
+          if adapter == :postgres
+            servers = postgres_replica_servers
+            connection_opts[:servers] = servers unless servers.empty?
+          end
           connection_opts
         end
 
+        def postgres_replica_servers
+          data = Legion::Settings[:data]
+          urls = [data[:read_replica_url], *Array(data[:replicas])].compact.reject(&:empty?).uniq
+          urls.each_with_index.to_h do |url, index|
+            uri = URI.parse(url)
+            raise ArgumentError, 'Read replicas require PostgreSQL URLs' unless %w[postgres postgresql].include?(uri.scheme)
+
+            options = ::Sequel::Database.send(:options_from_uri, uri).compact
+            [:"read_#{index}", options]
+          end
+        end
+
         def validate_no_unresolved_uris!(connection_opts)
           %i[user username password].each do |key|
             value = connection_opts[key]
@@ -489,10 +538,18 @@ module Legion
             val = data.key?(key) && !data[key].nil? ? data[key] : defaults[key]
             opts[key] = val unless val.nil?
           end
+          opts[:after_connect] = proc { |connection| configure_sqlite_busy_handler(connection: connection, timeout: opts[:timeout]) } if adapter == :sqlite
+          opts[:connect_sqls] = postgres_connect_sqls(data) if adapter == :postgres
 
           opts
         end
 
+        def postgres_connect_sqls(data)
+          %i[statement_timeout lock_timeout].filter_map do |setting|
+            "SET #{setting} = '#{data[setting].to_i}ms'" if data[setting]
+          end
+        end
+
         def tuning_stats(data)
           tuning = {}
 
@@ -513,6 +570,11 @@ module Legion
           tuning[:connection_expiration]          = data[:connection_expiration]
           tuning[:connection_expiration_timeout]  = data[:connection_expiration_timeout]
 
+          # Session timeouts (postgres)
+          tuning[:statement_timeout]   = data[:statement_timeout]
+          tuning[:lock_timeout]        = data[:lock_timeout]
+          tuning[:shutdown_timeout]    = data[:shutdown_timeout]
+
           # Adapter-specific (only keys relevant to current adapter)
           defaults = ADAPTER_DEFAULTS.fetch(adapter, {})
           ADAPTER_KEYS.fetch(adapter, []).each do |key|
@@ -621,6 +683,7 @@ module Legion
           if adapter == :postgres
             Sequel.extension(:pg_array)
             @sequel.extension(:pg_array)
+            @sequel.extension(:server_block) unless @replica_servers.empty?
           end
 
           if data[:connection_validation] != false
@@ -635,13 +698,15 @@ module Legion
 
           install_auth_failure_hook
         rescue StandardError => e
-          handle_exception(e, level: :warn, handled: true, operation: :configure_extensions, adapter: adapter)
+          handle_exception(e, level: :error, handled: false, operation: :configure_extensions, adapter: adapter)
+          raise
         end
 
         def install_auth_failure_hook
           Legion::Data::AuthFailureHandler.install(@sequel)
         rescue StandardError => e
-          handle_exception(e, level: :warn, handled: true, operation: :install_auth_failure_hook)
+          handle_exception(e, level: :error, handled: false, operation: :install_auth_failure_hook)
+          raise
         end
 
         def build_data_logger

data/lib/legion/data/encryption/key_provider.rb

@@ -15,8 +15,7 @@ module Legion
         end
 
         def key_for(tenant_id: nil)
-          cache_key = tenant_id || '__default__'
-          @key_cache[cache_key] ||= derive_key(tenant_id)
+          @key_cache[tenant_id] ||= derive_key(tenant_id)
         end
 
         def clear_cache!

data/lib/legion/data/local.rb

@@ -32,6 +32,11 @@ module Legion
             val = data.key?(key) && !data[key].nil? ? data[key] : sqlite_defaults[key]
             opts[key] = val unless val.nil?
           end
+          opts[:timeout] = data[:timeout] || 30_000
+          opts[:after_connect] = proc do |connection|
+            Legion::Data::Connection.configure_sqlite_busy_handler(connection: connection, timeout: opts[:timeout])
+            connection.execute_batch('PRAGMA synchronous=NORMAL; PRAGMA cache_size=-20000; PRAGMA mmap_size=268435456')
+          end
 
           if local_settings[:query_log]
             log_path = File.join(Legion::Data::Connection::QUERY_LOG_DIR, 'data-local-query.log')
@@ -46,14 +51,11 @@ module Legion
 
           @connection = ::Sequel.connect(opts)
           @connection.run('PRAGMA journal_mode=WAL')
-          @connection.run('PRAGMA busy_timeout=30000')
-          @connection.run('PRAGMA synchronous=NORMAL')
-          @connection.run('PRAGMA cache_size=-20000')
-          @connection.run('PRAGMA mmap_size=268435456')
-          @connected = true
           run_migrations
-          log.info "Legion::Data::Local connected to #{db_file} (WAL mode, 30s busy_timeout, 20MB cache, 256MB mmap)"
-        rescue StandardError => e
+          @connected = true
+          log.info "Legion::Data::Local connected to #{db_file} (WAL mode, #{opts[:timeout]}ms busy timeout, 20MB cache, 256MB mmap)"
+        rescue StandardError, LoadError => e
+          shutdown
           handle_exception(e, level: :error, handled: false, operation: :local_setup, database: db_file)
           raise
         end
@@ -74,6 +76,10 @@ module Legion
           @registered_migrations ||= {}
           @registered_migrations[name] = path
           run_single_migration(name, path) if connected?
+        rescue StandardError, LoadError => e
+          shutdown
+          handle_exception(e, level: :error, handled: false, operation: :register_migrations, name: name, path: path)
+          raise
         end
 
         def registered_migrations
@@ -142,8 +148,9 @@ module Legion
 
           table = :"schema_migrations_#{name}"
           ::Sequel::TimestampMigrator.new(@connection, path, table: table).run
-        rescue StandardError => e
-          handle_exception(e, level: :warn, handled: true, operation: :local_migration, name: name, path: path)
+        rescue StandardError, LoadError => e
+          handle_exception(e, level: :error, handled: false, operation: :local_migration, name: name, path: path)
+          raise
         end
 
         def local_settings

data/lib/legion/data/rls.rb

@@ -47,11 +47,17 @@ module Legion
       end
 
       def with_tenant(tenant_id)
-        previous = current_tenant
-        assign_tenant(tenant_id)
-        yield
-      ensure
-        previous ? assign_tenant(previous) : reset_tenant
+        return yield unless rls_enabled?
+
+        Legion::Data.connection.synchronize do
+          previous = current_tenant
+          begin
+            assign_tenant(tenant_id)
+            yield
+          ensure
+            previous ? assign_tenant(previous) : reset_tenant
+          end
+        end
       end
     end
   end

data/lib/legion/data/settings.rb

@@ -58,6 +58,28 @@ module Legion
           connection_expiration:         true,
           connection_expiration_timeout: 14_400,
 
+          # Postgres session timeouts (milliseconds, applied via SET on each new connection).
+          # statement_timeout bounds any single query; lock_timeout bounds lock acquisition.
+          # These prevent a wedged SELECT or lock wait from hanging the process indefinitely.
+          statement_timeout:             5000,
+          lock_timeout:                  2000,
+
+          # TCP keepalives for Postgres (libpq connection params).
+          # Detect half-open sockets where the peer vanished without FIN/RST.
+          # keepalives_idle: seconds before first keepalive probe after idle
+          # keepalives_interval: seconds between probes
+          # keepalives_count: failed probes before declaring connection dead
+          # tcp_user_timeout: total ms for unacked data before kernel drops connection (Linux)
+          keepalives:                    1,
+          keepalives_idle:               10,
+          keepalives_interval:           5,
+          keepalives_count:              3,
+          tcp_user_timeout:              15_000,
+
+          # Shutdown: max seconds to wait for checked-out connections to return before
+          # force-disconnecting them.
+          shutdown_timeout:              5,
+
           # Adapter-specific (nil = use adapter built-in default)
           connect_timeout:               nil,
           read_timeout:                  nil,

data/lib/legion/data/version.rb

@@ -2,6 +2,6 @@
 
 module Legion
   module Data
-    VERSION = '1.10.8'
+    VERSION = '1.10.17'
   end
 end

data/lib/legion/data.rb

@@ -48,6 +48,11 @@ module Legion
         setup_cache
         setup_local
         log.info 'Legion::Data setup complete'
+      rescue StandardError, LoadError => e
+        shutdown
+        Legion::Data::Connection.sequel = nil
+        handle_exception(e, level: :fatal, handled: false, operation: :setup)
+        raise
       end
 
       def setup_local
@@ -62,10 +67,19 @@ module Legion
 
       def setup_global
         connection_setup
+        Legion::Settings[:data][:connected] = false
+        Legion::Settings[:data][:migrations][:ran] = false
+        Legion::Settings[:data][:models][:loaded] = false
         migrate
         load_models
-      rescue StandardError => e
-        handle_exception(e, level: :fatal, operation: :setup_global)
+        Legion::Settings[:data][:connected] = true
+      rescue StandardError, LoadError => e
+        Legion::Data::Connection.shutdown
+        Legion::Data::Connection.sequel = nil
+        Legion::Settings[:data][:migrations][:ran] = false
+        Legion::Settings[:data][:models][:loaded] = false
+        handle_exception(e, level: :fatal, handled: false, operation: :setup_global)
+        raise
       end
 
       def setup_cache

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: legion-data
 version: !ruby/object:Gem::Version
-  version: 1.10.8
+  version: 1.10.17
 platform: ruby
 authors:
 - Esity