feat(orm): extend store with StaleRowError and IndexExists

lib/prism/orm/store.rb

@@ -7,6 +7,10 @@ module Prism
       class RecordNotFound < StandardError; end
       class UnknownTable < StandardError; end
       class TableExists < StandardError; end
+      class StaleRowError < StandardError; end
+      class IndexExists < StandardError; end
+      class UnknownIndex < StandardError; end
+      class UniqueConstraintViolation < StandardError; end
 
       def initialize(path)
         @path = path
@@ -21,11 +25,46 @@ module Prism
         within_transaction { yield }
       end
 
+      def transaction_registry
+        @transaction_registry ||= TransactionRegistry.new
+      end
+
+      # Wraps `transaction` with participant bookkeeping so after_commit/
+      # after_rollback fire exactly once, only for the outermost caller -
+      # this is what Model#transaction and Orm.transaction both call.
+      #
+      # `requires_new: true` opens a real SAVEPOINT instead of flattening
+      # into the surrounding transaction (the default, see TransactionRegistry
+      # docs): a raise from inside rolls back only the writes made since this
+      # call, while the outer transaction keeps whatever it had before and is
+      # still free to commit normally afterwards. Asking for requires_new
+      # with no surrounding `transaction` block open is just a plain
+      # transaction - there is nothing to nest a savepoint inside.
+      def transaction_with_callbacks(requires_new: false, &block)
+        if requires_new && transaction_registry.nested?
+          within_transaction { transaction_registry.within_savepoint(self, &block) }
+        else
+          transaction_registry.within { transaction(&block) }
+        end
+      end
+
+      # Deep-copies the current table data so a SAVEPOINT rollback can
+      # restore exactly this state later. Must run inside an open
+      # transaction (PStore raises otherwise) - callers only reach this via
+      # transaction_with_callbacks, which already guarantees that.
+      def snapshot_tables
+        Marshal.load(Marshal.dump(@pstore[:tables]))
+      end
+
+      def restore_tables(snapshot)
+        @pstore[:tables] = snapshot
+      end
+
       def create_table(name, columns = [])
         table = name.to_s
         within_transaction do
           raise TableExists, "table #{table} already exists" if @pstore[:tables].key?(table)
-          @pstore[:tables][table] = { columns: columns.map(&:to_s), rows: {}, next_id: 1 }
+          @pstore[:tables][table] = { columns: columns.map(&:to_s), rows: {}, next_id: 1, defaults: {}, indexes: {} }
         end
       end
 
@@ -33,21 +72,46 @@ module Prism
         within_transaction { @pstore[:tables].delete(name.to_s) }
       end
 
+      # A non-nil `default:` is remembered on the table definition (not just
+      # backfilled onto rows that already exist) so it also applies to every
+      # row inserted afterwards, the same way a real DB column default
+      # would - without this, a column added with `default: 0` reads back
+      # as nil, not 0, on any record created after the migration ran.
+      #
+      # Any existing index that covers this column is rebuilt from the
+      # backfilled rows. Without this, an index built while the column was
+      # still absent (every row keyed under nil) would stay stuck on that
+      # stale map forever - #lookup_by_index would report zero matches for
+      # the new default value even though every row now actually has it,
+      # which #indexed_candidates trusts as a definitive empty result
+      # (see Relation#rows) instead of falling back to a full scan.
       def add_column(table, column, default: nil)
         within_transaction do
           definition = table_definition(table)
           column = column.to_s
+          definition[:defaults] ||= {}
+          definition[:defaults][column] = default unless default.nil?
           next if definition[:columns].include?(column)
           definition[:columns] << column
           definition[:rows].each_value { |row| row[column] = default }
+          rebuild_indexes_covering(definition, column)
         end
       end
 
+      # Cascades to any index that covers the removed column, the same way a
+      # real database drops an index when the column it is built on is
+      # dropped. Leaving such an index in place would be strictly worse than
+      # having no index at all: its map/sorted_values would keep referencing
+      # a column no longer present on any row, and it would keep matching
+      # future #add_index/#lookup_by_index calls by column name alone.
       def remove_column(table, column)
         within_transaction do
           definition = table_definition(table)
-          definition[:columns].delete(column.to_s)
+          column = column.to_s
-          definition[:rows].each_value { |row| row.delete(column.to_s) }
+          definition[:columns].delete(column)
+          definition[:rows].each_value { |row| row.delete(column) }
+          (definition[:defaults] ||= {}).delete(column)
+          (definition[:indexes] ||= {}).reject! { |_name, index| index[:columns].include?(column) }
         end
       end
 
@@ -60,39 +124,113 @@ module Prism
       end
 
       def insert(table, attributes)
+        instrument('INSERT', table) do
           within_transaction do
             definition = table_definition(table)
             id = definition[:next_id]
-          definition[:rows][id] = attributes.transform_keys(&:to_s).merge('id' => id)
+            row = attributes.transform_keys(&:to_s)
+            (definition[:defaults] || {}).each { |column, default| row[column] = default unless row.key?(column) }
+            row['id'] = id
+            validate_unique_indexes(definition, id, row)
+            index_add(definition, id, row)
+            definition[:rows][id] = row
             definition[:next_id] = id + 1
             id
           end
         end
+      end
 
-      def update(table, id, attributes)
+      # `lock` (a two-element [column, expected_version] array, or nil to skip
+      # the check entirely) implements optimistic locking: the version check
+      # and the increment happen inside the same transaction as the merge, so
+      # a concurrent writer racing on this exact row can never both succeed
+      # against the same lock_version - the loser gets StaleRowError. It is a
+      # trailing positional argument rather than keyword arguments on purpose:
+      # every existing call site passes `attributes` as a bare trailing hash
+      # (`store.update(:users, id, role: 'owner')`), and Ruby 3 would swallow
+      # that hash as keyword arguments instead of the `attributes` positional
+      # the moment this method declared any keywords of its own, breaking
+      # every one of them with a bogus arity error.
+      def update(table, id, attributes, lock = nil)
+        instrument('UPDATE', table, id: id) do
           within_transaction do
-          row = table_definition(table)[:rows][id.to_i]
+            definition = table_definition(table)
+            row = definition[:rows][id.to_i]
             raise RecordNotFound, "no #{table} with id #{id}" unless row
-          row.merge!(attributes.transform_keys(&:to_s))
+            attributes = attributes.transform_keys(&:to_s)
+            if lock
+              lock_column, lock_version = lock
+              actual = row[lock_column].to_i
+              unless actual == lock_version.to_i
+                raise StaleRowError, "expected #{table}##{id} #{lock_column}=#{lock_version}, found #{actual}"
+              end
+              attributes = attributes.merge(lock_column => actual + 1)
+            end
+            new_row = row.merge(attributes)
+            validate_unique_indexes(definition, id.to_i, new_row)
+            index_remove(definition, id.to_i, row)
+            index_add(definition, id.to_i, new_row)
+            row.merge!(attributes)
             row.dup
           end
         end
+      end
 
-      def delete(table, id)
+      # Atomic read-then-write increment/decrement of a single numeric
-        within_transaction { table_definition(table)[:rows].delete(id.to_i) }
+      # column, used by belongs_to counter_cache instead of a full
+      # Model#update - no timestamps bumped, no callbacks/validations run,
+      # matching the raw-SQL UPDATE Rails issues for the same feature. Runs
+      # inside the surrounding transaction (reentrant via within_transaction)
+      # so a counter bump from a child's after_create/after_destroy rolls
+      # back together with the child write that triggered it.
+      def increment_counter(table, id, column, by = 1)
+        instrument('UPDATE', table, id: id) do
+          within_transaction do
+            row = table_definition(table)[:rows][id.to_i]
+            raise RecordNotFound, "no #{table} with id #{id}" unless row
+            column = column.to_s
+            row[column] = row[column].to_i + by
+            row.dup
+          end
+        end
+      end
+
+      def delete(table, id, lock = nil)
+        instrument('DELETE', table, id: id) do
+          within_transaction do
+            definition = table_definition(table)
+            row = definition[:rows][id.to_i]
+            if lock
+              lock_column, lock_version = lock
+              raise RecordNotFound, "no #{table} with id #{id}" unless row
+              actual = row[lock_column].to_i
+              unless actual == lock_version.to_i
+                raise StaleRowError, "expected #{table}##{id} #{lock_column}=#{lock_version}, found #{actual}"
+              end
+            end
+            index_remove(definition, id.to_i, row) if row
+            definition[:rows].delete(id.to_i)
+          end
+        end
       end
 
       def find(table, id)
+        instrument('SELECT', table, id: id) do
           within_transaction(read_only: true) { table_definition(table)[:rows][id.to_i]&.dup }
         end
+      end
 
       def all(table)
+        instrument('SELECT', table) do
           within_transaction(read_only: true) { table_definition(table)[:rows].values.map(&:dup) }
         end
+      end
 
       def count(table)
+        instrument('SELECT COUNT', table) do
           within_transaction(read_only: true) { table_definition(table)[:rows].size }
         end
+      end
 
       def applied_migrations
         within_transaction(read_only: true) { (@pstore[:schema_migrations] || []).dup }