legion-data 1.10.9 → 1.10.28

data/lib/legion/data/retention.rb

@@ -12,24 +12,36 @@ module Legion
       class << self
         include Legion::Logging::Helper
 
-        def archive_old_records(table:, date_column: nil, archive_after_days: DEFAULT_ARCHIVE_AFTER_DAYS)
+        def archive_old_records(table:, date_column: nil, archive_after_days: DEFAULT_ARCHIVE_AFTER_DAYS, batch_size: 1000)
           db = Legion::Data.connection
           return { archived: 0, table: table } unless db
+          raise ArgumentError, 'batch_size must be a positive Integer' unless batch_size.is_a?(Integer) && batch_size.positive?
 
           date_column = resolve_date_column(table, date_column)
           cutoff = Time.now - (archive_after_days * 86_400)
           archive_table = archive_table_name(table)
+          primary_keys = db.schema(table).filter_map { |name, info| name if info[:primary_key] }
+          raise Sequel::Error, "#{table} needs a primary key for bounded retention" if primary_keys.empty?
 
           ensure_archive_table!(db, table, archive_table)
 
           count = 0
-          db.transaction do
-            records = db[table].where(Sequel.identifier(date_column) < cutoff)
-            count = records.count
-            if count.positive?
-              db[archive_table].multi_insert(records.all)
-              records.delete
+          loop do
+            archived = db.transaction(savepoint: true, **(db.database_type == :sqlite ? { mode: :immediate } : {})) do
+              records = db[table].where(Sequel.identifier(date_column) < cutoff).order(*primary_keys).limit(batch_size)
+              records = records.for_update unless db.database_type == :sqlite
+              rows = records.all
+              next 0 if rows.empty?
+
+              db[archive_table].multi_insert(rows)
+              deleted = delete_archived_rows(db: db, table: table, rows: rows, primary_keys: primary_keys)
+              raise Sequel::Error, "Retention copied #{rows.length} rows but deleted #{deleted} from #{table}" unless deleted == rows.length
+
+              deleted
             end
+            break if archived.zero?
+
+            count += archived
           end
 
           log.info "Archived #{count} row(s) from #{table}" if count.positive?
@@ -113,18 +125,37 @@ module Legion
           end
         end
 
+        def delete_archived_rows(db:, table:, rows:, primary_keys:)
+          if primary_keys.one?
+            key = primary_keys.first
+            return db[table].where(key => rows.map { |row| row.fetch(key) }).delete
+          end
+
+          # SQLite expands composite-key predicates into OR expressions with a depth limit.
+          identity_batch_size = db.database_type == :sqlite ? 100 : rows.length
+          rows.each_slice(identity_batch_size).sum do |batch|
+            db[table].where(primary_keys => batch.map { |row| row.values_at(*primary_keys) }).delete
+          end
+        end
+
         def ensure_archive_table!(db, source_table, archive_table)
           return if db.table_exists?(archive_table)
 
           source_schema = db.schema(source_table).to_h
 
           log.info "Creating archive table #{archive_table} from #{source_table}"
-          db.create_table(archive_table) do
-            source_schema.each do |col_name, col_info|
-              column col_name, col_info[:db_type]
+          db.transaction(savepoint: true) do
+            db.create_table?(archive_table) do
+              source_schema.each do |col_name, col_info|
+                column col_name, col_info[:db_type]
+              end
+              DateTime :archived_at, default: Sequel::CURRENT_TIMESTAMP
             end
-            DateTime :archived_at, default: Sequel::CURRENT_TIMESTAMP
           end
+        rescue Sequel::UniqueConstraintViolation => e
+          raise unless db.table_exists?(archive_table)
+
+          handle_exception(e, level: :debug, handled: true, operation: :concurrent_archive_creation, table: archive_table)
         end
       end
     end

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/spool.rb

@@ -61,8 +61,10 @@ module Legion
 
         def read(sub_namespace)
           sorted_files(sub_namespace).each_with_object([]) do |path, events|
-            event = load_event_file(path, sub_namespace)
-            events << event if event
+            with_locked_event(path) do |file|
+              event = load_event_file(file, path, sub_namespace)
+              events << event if event
+            end
           end
         rescue StandardError => e
           handle_exception(e, level: :error, handled: false, operation: :spool_read, sub_namespace: sub_namespace)
@@ -72,13 +74,16 @@ module Legion
         def flush(sub_namespace)
           count = 0
           path = nil
-          sorted_files(sub_namespace).each do |path|
-            event = load_event_file(path, sub_namespace)
-            next unless event
-
-            yield event
-            File.delete(path)
-            count += 1
+          sorted_files(sub_namespace).each do |entry_path|
+            path = entry_path
+            with_locked_event(path) do |file|
+              event = load_event_file(file, path, sub_namespace)
+              next unless event
+
+              yield event
+              File.delete(path)
+              count += 1
+            end
           end
           log.info "Spool drained #{count} item(s) from #{sub_namespace}" if count.positive?
           count
@@ -95,8 +100,10 @@ module Legion
           dir = sub_dir(sub_namespace)
           return unless Dir.exist?(dir)
 
-          Dir[File.join(dir, '*.json')].each { |f| File.delete(f) }
-          log.info "Spool cleared #{sub_namespace}"
+          sorted_files(sub_namespace).each do |path|
+            with_locked_event(path) { |_file| File.delete(path) }
+          end
+          log.info "Spool cleared unclaimed events from #{sub_namespace}"
         rescue StandardError => e
           handle_exception(e, level: :error, handled: false, operation: :spool_clear, sub_namespace: sub_namespace)
           raise
@@ -115,11 +122,24 @@ module Legion
           Dir.glob(File.join(dir, '*.json'), sort: true)
         end
 
-        def load_event_file(path, sub_namespace)
-          ::JSON.parse(File.binread(path), symbolize_names: true)
-        rescue Errno::ENOENT => e
-          log.debug("spool event file not found: #{path}: #{e.message}")
-          nil
+        def with_locked_event(path)
+          file = begin
+            File.open(path, 'rb')
+          rescue Errno::ENOENT => e
+            log.debug("spool event file not found: #{path}: #{e.message}")
+            nil
+          end
+          return unless file
+          return unless file.flock(File::LOCK_EX | File::LOCK_NB)
+          return unless File.identical?(file, path)
+
+          yield file
+        ensure
+          file&.close
+        end
+
+        def load_event_file(file, path, sub_namespace)
+          ::JSON.parse(file.read, symbolize_names: true)
         rescue ::JSON::ParserError, EOFError, ArgumentError => e
           quarantine_corrupt_file(path, sub_namespace, e)
           nil

data/lib/legion/data/version.rb

@@ -2,6 +2,6 @@
 
 module Legion
   module Data
-    VERSION = '1.10.9'
+    VERSION = '1.10.28'
   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
@@ -119,36 +133,56 @@ module Legion
       def can_write?(table_name)
         return false unless connected?
 
-        adapter = Legion::Settings[:data][:adapter]&.to_s
-        return true if adapter == 'sqlite'
-
-        @write_privileges ||= {}
-        return @write_privileges[table_name] unless @write_privileges[table_name].nil?
-
-        @write_privileges[table_name] = connection
-                                        .fetch("SELECT has_table_privilege(current_user, ?, 'INSERT') AS can", table_name.to_s)
-                                        .first[:can] == true
+        database = connection
+        return false unless database
+        return true if database.database_type == :sqlite
+
+        cache = @write_privileges ||= ObjectSpace::WeakKeyMap.new
+        database.synchronize do |session|
+          privileges = cache[session] ||= {}
+          return privileges[table_name] if privileges.key?(table_name)
+
+          privileges[table_name] = case database.database_type
+                                   when :postgres
+                                     database.fetch("SELECT has_table_privilege(current_user, ?, 'INSERT') AS can", table_name.to_s)
+                                             .server(:default).first[:can] == true
+                                   when :mysql
+                                     database.fetch('EXPLAIN INSERT INTO ? () VALUES ()', Sequel.identifier(table_name)).server(:default).all
+                                     true
+                                   else
+                                     raise Sequel::Error, "Unsupported privilege adapter: #{database.database_type}"
+                                   end
+        end
       rescue StandardError => e
         handle_exception(e, level: :warn, handled: true, operation: :can_write?, table: table_name)
-        @write_privileges[table_name] = false if @write_privileges
         false
       end
 
       def can_read?(table_name)
         return false unless connected?
 
-        adapter = Legion::Settings[:data][:adapter]&.to_s
-        return true if adapter == 'sqlite'
-
-        @read_privileges ||= {}
-        return @read_privileges[table_name] unless @read_privileges[table_name].nil?
-
-        @read_privileges[table_name] = connection
-                                       .fetch("SELECT has_table_privilege(current_user, ?, 'SELECT') AS can", table_name.to_s)
-                                       .first[:can] == true
+        database = connection
+        return false unless database
+        return true if database.database_type == :sqlite
+
+        cache = @read_privileges ||= ObjectSpace::WeakKeyMap.new
+        database.synchronize do |session|
+          privileges = cache[session] ||= {}
+          return privileges[table_name] if privileges.key?(table_name)
+
+          privileges[table_name] = case database.database_type
+                                   when :postgres
+                                     database.fetch("SELECT has_table_privilege(current_user, ?, 'SELECT') AS can", table_name.to_s)
+                                             .server(:default).first[:can] == true
+                                   when :mysql
+                                     database.fetch('SELECT * FROM ? LIMIT 0', Sequel.identifier(table_name)).server(:default).all
+                                     true
+                                   else
+                                     raise Sequel::Error, "Unsupported privilege adapter: #{database.database_type}"
+                                   end
+        end
       rescue StandardError => e
         handle_exception(e, level: :warn, handled: true, operation: :can_read?, table: table_name)
-        @read_privileges[table_name] = false if @read_privileges
         false
       end
 
@@ -191,8 +225,9 @@ module Legion
 
       def shutdown
         Legion::Data::Local.shutdown if defined?(Legion::Data::Local) && Legion::Data::Local.connected?
-        Legion::Data::Connection.shutdown
-        log.info 'Legion::Data shutdown complete'
+        closed = Legion::Data::Connection.shutdown
+        log.info 'Legion::Data shutdown complete' if closed
+        closed
       end
 
       private

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: legion-data
 version: !ruby/object:Gem::Version
-  version: 1.10.9
+  version: 1.10.28
 platform: ruby
 authors:
 - Esity
@@ -126,6 +126,8 @@ files:
 - lib/legion/data/audit_record.rb
 - lib/legion/data/auth_failure_handler.rb
 - lib/legion/data/connection.rb
+- lib/legion/data/connection_pool.rb
+- lib/legion/data/database_lock.rb
 - lib/legion/data/encryption/cipher.rb
 - lib/legion/data/encryption/key_provider.rb
 - lib/legion/data/encryption/sequel_plugin.rb
@@ -284,6 +286,8 @@ files:
 - lib/legion/data/migrations/134_add_route_attempt_columns.rb
 - lib/legion/data/migrations/135_add_llm_context_token_accounting.rb
 - lib/legion/data/migrations/136_drop_llm_messages_seq_unique.rb
+- lib/legion/data/migrations/137_preserve_task_archive_fields.rb
+- lib/legion/data/migrations/138_add_append_claims.rb
 - lib/legion/data/model.rb
 - lib/legion/data/models/apollo/access_log.rb
 - lib/legion/data/models/apollo/entries.rb