feat(orm): extend store with add_index and remove_index

lib/prism/orm/store.rb

@@ -247,6 +247,210 @@ module Prism
         within_transaction { (@pstore[:schema_migrations] || []).delete(version.to_s) }
       end
 
+      # Builds a secondary index over `columns` (one column, or several for a
+      # composite index) so `Relation#where` equality lookups on those exact
+      # columns can fetch matching rows directly instead of scanning the
+      # whole table (see Relation#rows/#indexed_candidates). Backfills from
+      # every existing row, the same way a real database builds an index at
+      # CREATE INDEX time - not lazily on first query.
+      def add_index(table, columns, unique: false, name: nil)
+        within_transaction do
+          definition = table_definition(table)
+          columns = Array(columns).map(&:to_s)
+          index_name = (name || "index_#{table}_on_#{columns.join('_and_')}").to_s
+          definition[:indexes] ||= {}
+          raise IndexExists, "index #{index_name} already exists on #{table}" if definition[:indexes].key?(index_name)
+
+          definition[:indexes][index_name] =
+            build_index_data(definition, columns, unique: unique, name: index_name)
+          index_name
+        end
+      end
+
+      def remove_index(table, name_or_columns)
+        within_transaction do
+          definition = table_definition(table)
+          definition[:indexes] ||= {}
+          index_name = resolve_index_name(definition, name_or_columns)
+          raise UnknownIndex, "no index #{name_or_columns.inspect} on #{table}" unless index_name
+          definition[:indexes].delete(index_name)
+        end
+      end
+
+      def indexes(table)
+        within_transaction(read_only: true) do
+          (table_definition(table)[:indexes] || {}).values.map do |index|
+            { name: index[:name], columns: index[:columns].dup, unique: index[:unique] }
+          end
+        end
+      end
+
+      # Returns matching row dups for an exact-match lookup on the best
+      # available index covering every key in `equalities` (all columns of
+      # some index must be present), or nil if no such index exists - the
+      # signal for the caller to fall back to a full table scan.
+      def lookup_by_index(table, equalities)
+        return nil if equalities.empty?
+
+        # `id` is always "indexed" for free - `definition[:rows]` is already
+        # a Hash keyed by id, exactly the same reasoning #lookup_by_index_in
+        # already applies to an IN-list of ids. Before this, only the
+        # IN-list path got that fast lookup; a plain single-id equality
+        # (`find_by(id: x)`, `where(id: x).first`, and therefore every
+        # Authenticatable#current_user call on every authenticated request)
+        # fell all the way through to a registered-index search that could
+        # never find one (no framework code ever calls `add_index(table,
+        # :id)`) and landed on a full table scan instead - silently, since
+        # nothing about that path looked like an error. Any other equality
+        # keys alongside `id` are still safe to ignore here: at most one row
+        # can ever match on id anyway, and Relation#rows always re-applies
+        # the full condition afterward.
+        if equalities.key?('id')
+          return instrument('SELECT INDEX', table, columns: ['id']) do
+            rows_by_id(table, [equalities['id']])
+          end
+        end
+
+        available = within_transaction(read_only: true) do
+          (table_definition(table)[:indexes] || {}).values
+        end
+        index = available.select { |candidate| candidate[:columns].all? { |c| equalities.key?(c) } }
+                          .max_by { |candidate| candidate[:columns].size }
+        return nil unless index
+
+        instrument('SELECT INDEX', table, columns: index[:columns]) do
+          within_transaction(read_only: true) do
+            definition = table_definition(table)
+            key = index[:columns].map { |column| equalities[column] }
+            ids = (table_definition(table)[:indexes][index[:name]][:map][key] || [])
+            ids.map { |id| definition[:rows][id]&.dup }.compact
+          end
+        end
+      end
+
+      # Returns matching row dups for a gt/gte/lt/lte range on the best
+      # available range-capable index covering `column` (single-column
+      # indexes only), or nil if none exists - the same "nil means scan"
+      # contract as #lookup_by_index. `operators` is a Hash of Symbol
+      # operator => bound, e.g. { gte: 18, lt: 65 }. Finds the boundary with
+      # a binary search over the index's ascending sorted_values (O(log n)),
+      # then only touches the ids for values inside those bounds - never a
+      # scan of the full distinct-value list, let alone the whole table.
+      def lookup_by_index_range(table, column, operators)
+        column = column.to_s
+        index = within_transaction(read_only: true) do
+          (table_definition(table)[:indexes] || {}).values
+            .find { |candidate| candidate[:range_capable] && candidate[:columns] == [column] }
+        end
+        return nil unless index
+
+        instrument('SELECT INDEX RANGE', table, columns: index[:columns]) do
+          within_transaction(read_only: true) do
+            definition = table_definition(table)
+            sorted = definition[:indexes][index[:name]][:sorted_values]
+            lo, hi = range_bounds(sorted, operators)
+            ids = sorted[lo...hi].flat_map { |value| definition[:indexes][index[:name]][:map][[value]] || [] }
+            ids.uniq.map { |id| definition[:rows][id]&.dup }.compact
+          end
+        end
+      end
+
+      # Returns matching row dups for an equality-prefix-plus-range lookup
+      # (`where(scope: val, position: { lt: x })`) on the best available
+      # composite index whose own columns are `[*some subset of
+      # equalities.keys in the index's declared order, range_column]`, or
+      # nil if none exists - same "nil means scan" contract as every other
+      # lookup_by_index* method. Unlike #lookup_by_index_range (single
+      # column only), this walks a composite index's *last* column as the
+      # range dimension while every other column is pinned to an exact
+      # value - the same thing a real database does with a composite
+      # B-tree index `(scope, position)` answering `WHERE scope = ? AND
+      # position < ?` by seeking straight into the `scope` partition
+      # instead of scanning the whole table (or, without this method, every
+      # row in every OTHER scope too - see #lookup_by_index_range's
+      # single-column limitation).
+      #
+      # `equalities` is allowed to carry MORE columns than the index's own
+      # prefix needs (e.g. a paranoid model's `deleted_at: nil` riding
+      # alongside `featured: true` when only `featured` is actually
+      # indexed) - only the index's own declared prefix columns are read
+      # out of it to build the lookup key, exactly like #lookup_by_index's
+      # "index columns must be a subset of equalities.keys" contract (not
+      # the other way around). Any equality condition the index prefix
+      # doesn't cover is still safe to drop here: Relation#rows re-checks
+      # the full condition against every candidate row afterward. When
+      # several composite indexes could answer the same call, the one
+      # covering the most columns wins, same tie-break as #lookup_by_index.
+      def lookup_by_index_equality_range(table, equalities, range_column, operators)
+        equalities = equalities.transform_keys(&:to_s)
+        range_column = range_column.to_s
+
+        index = within_transaction(read_only: true) do
+          (table_definition(table)[:indexes] || {}).values
+            .select do |candidate|
+              candidate[:composite_range_capable] &&
+                candidate[:columns].last == range_column &&
+                candidate[:columns][0..-2].all? { |column| equalities.key?(column) }
+            end
+            .max_by { |candidate| candidate[:columns].size }
+        end
+        return nil unless index
+
+        instrument('SELECT INDEX RANGE', table, columns: index[:columns]) do
+          within_transaction(read_only: true) do
+            definition = table_definition(table)
+            prefix_key = index[:columns][0..-2].map { |column| equalities[column] }
+            sorted = definition[:indexes][index[:name]][:sorted_by_prefix][prefix_key] || []
+            lo, hi = range_bounds(sorted, operators)
+            ids = sorted[lo...hi].flat_map do |value|
+              definition[:indexes][index[:name]][:map][prefix_key + [value]] || []
+            end
+            ids.uniq.map { |id| definition[:rows][id]&.dup }.compact
+          end
+        end
+      end
+
+      # Returns matching row dups for an IN-list lookup (`where(column:
+      # [v1, v2, ...])`), the same "nil means scan" contract as
+      # #lookup_by_index/#lookup_by_index_range. This is what every
+      # eager-loading batch query goes through
+      # (Relation#preload_belongs_to's `where(id: owner_ids)`,
+      # #preload_has_one/#preload_has_many/#preload_has_many_through's
+      # `where(foreign_key: owner_ids)`) - without it, `includes(...)` on
+      # any association was always a full table scan even though it only
+      # ever needed a handful of specific ids, no matter how many other
+      # secondary indexes existed elsewhere.
+      #
+      # `id` is always treated as indexed, with zero add_index calls
+      # needed: #rows is already a Hash keyed by id, so this is just direct
+      # O(1)-per-value hash lookups, same as #find. Any other column needs
+      # a real single-column index (composite indexes can't answer an
+      # IN-list on just one of their columns - which specific OTHER column
+      # value would pair with each id in the list is ambiguous).
+      def lookup_by_index_in(table, column, values)
+        column = column.to_s
+        values = values.uniq
+
+        if column == 'id'
+          instrument('SELECT INDEX IN', table, columns: ['id']) do
+            rows_by_id(table, values)
+          end
+        else
+          index = within_transaction(read_only: true) do
+            (table_definition(table)[:indexes] || {}).values.find { |candidate| candidate[:columns] == [column] }
+          end
+          return nil unless index
+
+          instrument('SELECT INDEX IN', table, columns: index[:columns]) do
+            within_transaction(read_only: true) do
+              definition = table_definition(table)
+              ids = values.flat_map { |value| definition[:indexes][index[:name]][:map][[value]] || [] }
+              ids.uniq.map { |id| definition[:rows][id]&.dup }.compact
+            end
+          end
+        end
+      end
+
       private
 
       def within_transaction(read_only: false)