Class: ActiveRecord::Duplicator::Session

Inherits:
Object
  • Object
show all
Defined in:
lib/active_record/duplicator/session.rb,
sig/generated/active_record/duplicator/session.rbs

Overview

One-shot execution context for a duplication run.

A Session owns the record id map (old_id -> new_id, per class) and exposes the API a handler needs: mark_skip / store_new_id / fetch_new_id / bulk_insert / attributes_for.

Users normally do not instantiate Session directly; a fresh Session is created by Duplicator#duplicate for every run, so state does not leak across runs. Duplicator#duplicate yields the Session to its block so callers can mark_skip / store_new_id before the run starts.

Composite primary keys (Rails 7.1+ self.primary_key = [:a, :b]) are supported. An id is a scalar for a single-column primary key and an Array for a composite one; the shape simply mirrors what record.class.primary_key reports.

Constant Summary collapse

DEFAULT_BATCH_SIZE =

Returns:

  • (::Integer) 
1000

Instance Method Summary collapse

Constructor Details

#initialize(handlers: {}, skipped_classes: Set.new) ⇒ Session

: (?handlers: Hash[Class, Proc], ?skipped_classes: Set) -> void

Parameters:

  • handlers: (Hash[Class, Proc]) (defaults to: {})
  • skipped_classes: (Set[Class]) (defaults to: Set.new) 


30
31
32
33
34
35
36
37
 
# File 'lib/active_record/duplicator/session.rb', line 30

def initialize(handlers: {}, skipped_classes: Set.new)
  @record_id_map = Hash.new { |hash, klass| hash[klass] = {} }
  @handlers = handlers
  @user_handlers = {}
  # dup so that Session#mark_skip_class additions never leak back to the
  # Duplicator-level set shared across every run.
  @skipped_classes = skipped_classes.dup
end

Instance Method Details

#assign_foreign_key(attrs, reflection, old_record) ⇒ void

This method returns an undefined value.

: (Hash[String, untyped], untyped, ActiveRecord::Base) -> void

Parameters:

  • (Hash[String, untyped])
  • (Object)
  • (ActiveRecord::Base) 


261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
 
# File 'lib/active_record/duplicator/session.rb', line 261

def assign_foreign_key(attrs, reflection, old_record)
  fk = reflection.foreign_key

  if fk.is_a?(Array)
    old_values = fk.map { |c| old_record[c] }
    if old_values.all?(&:nil?)
      fk.each { |c| attrs[c] = nil }
    else
      new_values = fetch_new_id(reflection.klass, old_id: old_values)
      fk.zip(new_values).each { |c, v| attrs[c] = v }
    end
  else
    old_value = old_record[fk]
    attrs[fk] = old_value.nil? ? nil : fetch_new_id(reflection.klass, old_id: old_value)
  end
end

#attributes_for(klass, old_record) ⇒ Hash[String, untyped]

Build the attribute Hash for a single record, ready to be handed to insert_all!. Timestamps (as reported by klass.all_timestamp_attributes_in_model, which covers created_at / updated_at / created_on / updated_on), auto-generated primary key columns, and every belongs_to foreign key column are stripped; each belongs_to foreign key is then re-populated from the id map so it points to the new parent. Non-auto columns of a composite primary key (e.g. tenant_id) are carried over from the source record unchanged. : (Class, ActiveRecord::Base) -> Hash[String, untyped]

Parameters:

  • (Class)
  • (ActiveRecord::Base)

Returns:

  • (Hash[String, untyped]) 


164
165
166
167
168
169
170
171
172
173
 
# File 'lib/active_record/duplicator/session.rb', line 164

def attributes_for(klass, old_record)
  belongs_to = klass.reflect_on_all_associations(:belongs_to)
  fk_columns = belongs_to.flat_map { |r| Array.wrap(r.foreign_key) }
  ignored = auto_generated_pk_columns(klass) + klass.all_timestamp_attributes_in_model + fk_columns
  attrs = old_record.attributes.except(*ignored)

  belongs_to.each { |reflection| assign_foreign_key(attrs, reflection, old_record) }

  attrs
end

#auto_generated_pk_columns(klass) ⇒ Array[String]

Primary key columns whose values are filled in by the database (e.g. a bigserial id). These must be stripped from the attributes before calling insert_all! so the DB can assign fresh values. Composite pk columns without a default_function (e.g. tenant_id) are kept. : (Class) -> Array

Parameters:

  • (Class)

Returns:

  • (Array[String]) 


283
284
285
286
287
288
 
# File 'lib/active_record/duplicator/session.rb', line 283

def auto_generated_pk_columns(klass)
  Array.wrap(klass.primary_key).select do |key|
    column = klass.column_for_attribute(key)
    column.default_function.present?
  end
end

#bulk_insert(klass, records, batch_size: DEFAULT_BATCH_SIZE, cursor: nil) ⇒ void

This method returns an undefined value.

Bulk-insert copies of the given records into klass, skipping any that are already present in the id map. Handlers can reuse this via HandlerApi to defer a subset of records to the default duplication path.

When klass has been passed to mark_skip_class the call is a no-op that never touches the DB: it is equivalent to having pre-registered every incoming record via mark_skip, only without the per-record enumeration. : (Class, untyped, ?batch_size: Integer, ?cursor: untyped) -> void

Parameters:

  • (Class)
  • (Object)
  • batch_size: (Integer) (defaults to: DEFAULT_BATCH_SIZE)
  • cursor: (Object) (defaults to: nil) 


134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
 
# File 'lib/active_record/duplicator/session.rb', line 134

def bulk_insert(klass, records, batch_size: DEFAULT_BATCH_SIZE, cursor: nil)
  return if @skipped_classes.include?(klass)

  batches = enumerate_batches(records, batch_size:, cursor:)

  batches.each do |batch|
    old_records = batch.to_a
    old_records.reject! { |r| @record_id_map[klass].key?(current_id(r)) }
    next if old_records.empty?

    rows = old_records.map { |r| attributes_for(klass, r) }
    inserted = klass.insert_all!(rows)

    inserted.to_a.each.with_index do |returned_row, i|
      new_id = extract_id_from_row(klass, returned_row)
      old_id = current_id(old_records[i])
      store_new_id(klass, old_id:, new_id:)
    end
  end
end

#current_id(record) ⇒ Object

Return the id of the given record shaped like the class's primary_key: a scalar for a single-column pk, an Array for a composite pk. : (ActiveRecord::Base) -> untyped

Parameters:

  • (ActiveRecord::Base)

Returns:

  • (Object) 


239
240
241
242
243
244
245
246
 
# File 'lib/active_record/duplicator/session.rb', line 239

def current_id(record)
  pk = record.class.primary_key
  if pk.is_a?(Array)
    pk.map { |c| record[c] }
  else
    record[pk]
  end
end

#duplicate_records(klass, relation_or_records) ⇒ void

This method returns an undefined value.

Dispatch to a handler (or the default bulk_insert path) for klass. All records in relation_or_records are assumed to be of klass; the caller is responsible for splitting mixed inputs (e.g. STI subclasses) into per-class Relations and passing the matching class in.

mark_skip_class does NOT short-circuit here: handlers must still fire (mirroring the per-record mark_skip case), and the class-wide skip is applied inside bulk_insert instead. : (Class, untyped) -> void

Parameters:

  • (Class)
  • (Object) 


208
209
210
211
212
213
214
215
216
 
# File 'lib/active_record/duplicator/session.rb', line 208

def duplicate_records(klass, relation_or_records)
  handler = resolve_handler(klass)

  if handler
    handler.call(HandlerApi.new(self), klass, relation_or_records)
  else
    bulk_insert(klass, relation_or_records)
  end
end

#enumerate_batches(records, batch_size:, cursor:) ⇒ Object

: (untyped, batch_size: Integer, cursor: untyped) -> untyped

Parameters:

  • (Object)
  • batch_size: (Integer)
  • cursor: (Object)

Returns:

  • (Object) 


226
227
228
229
230
231
232
233
234
 
# File 'lib/active_record/duplicator/session.rb', line 226

def enumerate_batches(records, batch_size:, cursor:)
  if records.respond_to?(:in_batches) && cursor
    records.in_batches(of: batch_size, cursor:)
  elsif records.respond_to?(:find_in_batches)
    records.find_in_batches(batch_size:)
  else
    records.each_slice(batch_size)
  end
end

#extract_id_from_row(klass, row) ⇒ Object

Extract the primary key value from a row returned by insert_all!. Scalar for single-column pk, Array for composite pk. : (Class, Hash[String, untyped]) -> untyped

Parameters:

  • (Class)
  • (Hash[String, untyped])

Returns:

  • (Object) 


251
252
253
254
255
256
257
258
 
# File 'lib/active_record/duplicator/session.rb', line 251

def extract_id_from_row(klass, row)
  pk = klass.primary_key
  if pk.is_a?(Array)
    pk.map { |c| row[c.to_s] }
  else
    row[pk.to_s]
  end
end

#fetch_new_id(klass, old_id:) ⇒ Object

Look up a new_id by (klass, old_id). Returns a scalar for a single-column pk and an Array for a composite pk.

Raises MissingNewIdError when no mapping exists. This usually means a belongs_to parent has not been duplicated yet; place it earlier in the associations tree or register/skip it explicitly. : (Class, old_id: untyped) -> untyped

Parameters:

  • (Class)
  • old_id: (Object)

Returns:

  • (Object) 


118
119
120
121
122
123
124
 
# File 'lib/active_record/duplicator/session.rb', line 118

def fetch_new_id(klass, old_id:)
  return old_id if @skipped_classes.include?(klass)

  @record_id_map[klass].fetch(old_id) do
    raise MissingNewIdError, "no new_id registered for #{klass}##{old_id.inspect}"
  end
end

#mark_skip(*records) ⇒ void

This method returns an undefined value.

Declare that the given records must not be duplicated: their new_id is set equal to their old_id, so any record duplicated later that references them will keep pointing at the same rows. : (*ActiveRecord::Base) -> void

Parameters:

  • (ActiveRecord::Base) 


57
58
59
60
61
62
 
# File 'lib/active_record/duplicator/session.rb', line 57

def mark_skip(*records)
  records.each do |record|
    id = current_id(record)
    store_new_id(record.class, old_id: id, new_id: id)
  end
end

#mark_skip_class(*klasses) ⇒ void

This method returns an undefined value.

Class-wide equivalent of mark_skip: it is exactly as if every record of each given class had been passed to mark_skip. Handlers still fire (matching the mark_skip case where a handler for the class is invoked with a Relation that happens to contain pre-skipped records); the default bulk_insert path becomes a no-op, and fetch_new_id echoes the queried old_id back unchanged.

Use when the target database is expected to already hold every referenced row (shared lookup tables, tenant-wide catalogs, etc.), so per-record mark_skip would only differ from this in cost, not effect. : (*Class) -> void

Parameters:

  • (Class) 


85
86
87
 
# File 'lib/active_record/duplicator/session.rb', line 85

def mark_skip_class(*klasses)
  @skipped_classes.merge(klasses)
end

#mark_skip_id(klass, *ids) ⇒ void

This method returns an undefined value.

Declare that the given ids of the given class must not be duplicated, without loading the records. Equivalent to calling mark_skip on each record but avoids the database round-trip when the ids are already known. Each id should have the same shape as the class's primary key (scalar for a single-column pk, Array for a composite pk). : (Class, *untyped) -> void

Parameters:

  • (Class)
  • (Object) 


70
71
72
 
# File 'lib/active_record/duplicator/session.rb', line 70

def mark_skip_id(klass, *ids)
  ids.each { |id| store_new_id(klass, old_id: id, new_id: id) }
end

#on(klass) {|arg0, arg1, arg2| ... } ⇒ void

This method returns an undefined value.

Register a handler that overrides Duplicator#on for the current session only. Useful when a single duplication run needs a special copy of a model that the shared Duplicator config should not know about.

Raises DuplicateHandlerError when the same class is registered twice via Session#on in the same session. Overriding a class that also has a Duplicator-level handler is intentional and does not raise. : (Class) { (HandlerApi, Class, untyped) -> void } -> void

Parameters:

  • (Class)

Yields:

Yield Parameters:

Yield Returns:

  • (void) 


47
48
49
50
51
 
# File 'lib/active_record/duplicator/session.rb', line 47

def on(klass, &block)
  raise DuplicateHandlerError, "handler for #{klass} is already overridden in this session" if @user_handlers.key?(klass)

  @user_handlers[klass] = block
end

#resolve_handler(klass) ⇒ Proc?

Session-local handlers (Session#on) take precedence over any Duplicator-level handler registered for the exact same class. : (Class) -> Proc?

Parameters:

  • (Class)

Returns:

  • (Proc, nil) 


221
222
223
 
# File 'lib/active_record/duplicator/session.rb', line 221

def resolve_handler(klass)
  @user_handlers[klass] || @handlers[klass]
end

#run(root_record, associations: [], transaction_options: {}) ⇒ ActiveRecord::Base

Run the duplication inside a transaction and return the freshly loaded root record. Normally invoked by Duplicator#duplicate.

transaction_options is passed through verbatim to root_record.transaction, so callers can request requires_new: true / joinable: false / isolation: :serializable etc. as fits the surrounding call site. It defaults to {}, i.e. plain transaction do ... end behaviour. : (ActiveRecord::Base, ?associations: untyped, ?transaction_options: Hash[Symbol, untyped]) -> ActiveRecord::Base

Parameters:

  • (ActiveRecord::Base)
  • associations: (Object) (defaults to: [])
  • transaction_options: (Hash[Symbol, untyped]) (defaults to: {})

Returns:

  • (ActiveRecord::Base) 


184
185
186
187
188
189
190
191
192
193
194
195
 
# File 'lib/active_record/duplicator/session.rb', line 184

def run(root_record, associations: [], transaction_options: {})
  root_record.transaction(**transaction_options) do
    duplicate_records(root_record.class, [root_record])

    AssociationTraversal.new(root_record, associations).each do |relation|
      duplicate_records(relation.klass, relation)
    end

    new_id = fetch_new_id(root_record.class, old_id: current_id(root_record))
    root_record.class.find(new_id)
  end
end

#store_new_id(klass, old_id:, new_id:) ⇒ void

This method returns an undefined value.

Record a mapping from an old_id to a new_id. Called automatically from bulk_insert; called explicitly by user code (or a handler) when new records are produced outside of the default duplication path.

For a single-column primary key pass a scalar; for a composite primary key pass an Array matching the shape of the class's primary_key.

Raises InvalidRecordIdError when the same old_id is already mapped to a different new_id. : (Class, old_id: untyped, new_id: untyped) -> void

Parameters:

  • (Class)
  • old_id: (Object)
  • new_id: (Object) 


99
100
101
102
103
104
105
106
107
108
109
 
# File 'lib/active_record/duplicator/session.rb', line 99

def store_new_id(klass, old_id:, new_id:)
  map = @record_id_map[klass]

  if !map.key?(old_id)
    map[old_id] = new_id
  elsif map[old_id] != new_id
    raise InvalidRecordIdError,
      "#{klass} old_id=#{old_id.inspect} is already mapped to " \
      "#{map[old_id].inspect}, cannot overwrite with #{new_id.inspect}"
  end
end