Commit 3e74bbbd authored by pbair's avatar pbair

Update DynamicModelHelpers to require a connection

Update the methods in DynamicModelHelpers to require a connection that
dynamic models will use to talk to the appropriate database. This change
will ensure that usages of dynamic models are compatible with decomposed
databases.
parent fdbca440
......@@ -2,7 +2,6 @@
class MoveContainerRegistryEnabledToProjectFeatures3 < ActiveRecord::Migration[6.0]
include Gitlab::Database::MigrationHelpers
include Gitlab::Database::DynamicModelHelpers
BATCH_SIZE = 21_000
MIGRATION = 'MoveContainerRegistryEnabledToProjectFeature'
......
......@@ -10,7 +10,7 @@ class CleanUpPendingBuildsTable < ActiveRecord::Migration[6.0]
def up
return unless Gitlab.dev_or_test_env? || Gitlab.com?
each_batch_range('ci_pending_builds', of: BATCH_SIZE) do |min, max|
each_batch_range('ci_pending_builds', connection: connection, of: BATCH_SIZE) do |min, max|
execute <<~SQL
DELETE FROM ci_pending_builds
USING ci_builds
......
......@@ -8,7 +8,7 @@ class MigrateProtectedAttributeToPendingBuilds < ActiveRecord::Migration[6.1]
def up
return unless Gitlab.dev_or_test_env? || Gitlab.com?
each_batch_range('ci_pending_builds', of: 1000) do |min, max|
each_batch_range('ci_pending_builds', connection: connection, of: 1000) do |min, max|
execute <<~SQL
UPDATE ci_pending_builds
SET protected = true
......
......@@ -9,7 +9,7 @@ module Gitlab
include Gitlab::Database::DynamicModelHelpers
def perform(start_id, stop_id, batch_table, batch_column, sub_batch_size, pause_ms)
parent_batch_relation = define_batchable_model(batch_table)
parent_batch_relation = define_batchable_model(batch_table, connection: connection)
.where(batch_column => start_id..stop_id)
parent_batch_relation.each_batch(column: batch_column, of: sub_batch_size) do |sub_batch|
......
......@@ -29,7 +29,7 @@ module Gitlab
end
def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
define_batchable_model(source_table)
define_batchable_model(source_table, connection: connection)
.where(source_key_column => start_id..stop_id)
.where(type: nil)
end
......
......@@ -19,7 +19,7 @@ module Gitlab
# batch_size - The size of the next batch
# job_arguments - The migration job arguments
def next_batch(table_name, column_name, batch_min_value:, batch_size:, job_arguments:)
model_class = define_batchable_model(table_name)
model_class = define_batchable_model(table_name, connection: ActiveRecord::Base.connection)
quoted_column_name = model_class.connection.quote_column_name(column_name)
relation = model_class.where("#{quoted_column_name} >= ?", batch_min_value)
......
......@@ -57,7 +57,7 @@ module Gitlab
end
def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
define_batchable_model(source_table).where(source_key_column => start_id..stop_id)
define_batchable_model(source_table, connection: connection).where(source_key_column => start_id..stop_id)
end
def column_assignment_clauses(copy_from, copy_to)
......
......@@ -15,7 +15,7 @@ module Gitlab
end
def orphaned_deployments
define_batchable_model('deployments')
define_batchable_model('deployments', connection: ActiveRecord::Base.connection)
.where('NOT EXISTS (SELECT 1 FROM environments WHERE deployments.environment_id = environments.id)')
end
......
......@@ -10,7 +10,9 @@ module Gitlab
include Gitlab::Database::DynamicModelHelpers
def perform(start_id, stop_id)
define_batchable_model('vulnerability_finding_links').where(id: start_id..stop_id).delete_all
define_batchable_model('vulnerability_finding_links', connection: ActiveRecord::Base.connection)
.where(id: start_id..stop_id)
.delete_all
end
end
end
......
......@@ -9,7 +9,9 @@ module Gitlab
BATCH_SIZE = 100
def perform(start_id, stop_id)
define_batchable_model('timelogs').where(spent_at: nil, id: start_id..stop_id).each_batch(of: 100) do |subbatch|
define_batchable_model('timelogs', connection: connection)
.where(spent_at: nil, id: start_id..stop_id)
.each_batch(of: 100) do |subbatch|
batch_start, batch_end = subbatch.pluck('min(id), max(id)').first
update_timelogs(batch_start, batch_end)
......@@ -25,9 +27,12 @@ module Gitlab
SQL
end
def execute(sql)
def connection
@connection ||= ::ActiveRecord::Base.connection
@connection.execute(sql)
end
def execute(sql)
connection.execute(sql)
end
end
end
......
......@@ -5,16 +5,19 @@ module Gitlab
module DynamicModelHelpers
BATCH_SIZE = 1_000
def define_batchable_model(table_name)
Class.new(ActiveRecord::Base) do
def define_batchable_model(table_name, connection:)
klass = Class.new(ActiveRecord::Base) do
include EachBatch
self.table_name = table_name
self.inheritance_column = :_type_disabled
end
klass.connection = connection
klass
end
def each_batch(table_name, scope: ->(table) { table.all }, of: BATCH_SIZE)
def each_batch(table_name, connection:, scope: ->(table) { table.all }, of: BATCH_SIZE)
if transaction_open?
raise <<~MSG.squish
each_batch should not run inside a transaction, you can disable
......@@ -23,12 +26,12 @@ module Gitlab
MSG
end
scope.call(define_batchable_model(table_name))
scope.call(define_batchable_model(table_name, connection: connection))
.each_batch(of: of) { |batch| yield batch }
end
def each_batch_range(table_name, scope: ->(table) { table.all }, of: BATCH_SIZE)
each_batch(table_name, scope: scope, of: of) do |batch|
def each_batch_range(table_name, connection:, scope: ->(table) { table.all }, of: BATCH_SIZE)
each_batch(table_name, connection: connection, scope: scope, of: of) do |batch|
yield batch.pluck('MIN(id), MAX(id)').first
end
end
......
......@@ -9,6 +9,18 @@ module Gitlab
include RenameTableHelpers
include AsyncIndexes::MigrationHelpers
def define_batchable_model(table_name, connection: self.connection)
super(table_name, connection: connection)
end
def each_batch(table_name, connection: self.connection, **kwargs)
super(table_name, connection: connection, **kwargs)
end
def each_batch_range(table_name, connection: self.connection, **kwargs)
super(table_name, connection: connection, **kwargs)
end
# https://www.postgresql.org/docs/current/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
MAX_IDENTIFIER_NAME_LENGTH = 63
DEFAULT_TIMESTAMP_COLUMNS = %i[created_at updated_at].freeze
......
......@@ -49,7 +49,7 @@ module Gitlab
end
def relation_scoped_to_range(source_table, source_key_column, start_id, stop_id)
define_batchable_model(source_table)
define_batchable_model(source_table, connection: connection)
.where(source_key_column => start_id..stop_id)
end
......
......@@ -5,7 +5,6 @@ module Gitlab
module PartitioningMigrationHelpers
module TableManagementHelpers
include ::Gitlab::Database::SchemaHelpers
include ::Gitlab::Database::DynamicModelHelpers
include ::Gitlab::Database::MigrationHelpers
include ::Gitlab::Database::Migrations::BackgroundMigrationHelpers
......
......@@ -4,10 +4,11 @@ require 'spec_helper'
RSpec.describe Gitlab::Database::DynamicModelHelpers do
let(:including_class) { Class.new.include(described_class) }
let(:table_name) { 'projects' }
let(:table_name) { Project.table_name }
let(:connection) { Project.connection }
describe '#define_batchable_model' do
subject { including_class.new.define_batchable_model(table_name) }
subject { including_class.new.define_batchable_model(table_name, connection: connection) }
it 'is an ActiveRecord model' do
expect(subject.ancestors).to include(ActiveRecord::Base)
......@@ -40,7 +41,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
it 'iterates table in batches' do
each_batch_size = ->(&block) do
subject.each_batch(table_name, of: 1) do |batch|
subject.each_batch(table_name, connection: connection, of: 1) do |batch|
block.call(batch.size)
end
end
......@@ -56,7 +57,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
end
it 'raises an error' do
expect { subject.each_batch(table_name, of: 1) { |batch| batch.size } }
expect { subject.each_batch(table_name, connection: connection, of: 1) { |batch| batch.size } }
.to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
end
end
......@@ -74,7 +75,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
end
it 'iterates table in batch ranges' do
expect { |b| subject.each_batch_range(table_name, of: 1, &b) }
expect { |b| subject.each_batch_range(table_name, connection: connection, of: 1, &b) }
.to yield_successive_args(
[first_project.id, first_project.id],
[second_project.id, second_project.id]
......@@ -82,13 +83,13 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
end
it 'yields only one batch if bigger than the table size' do
expect { |b| subject.each_batch_range(table_name, of: 2, &b) }
expect { |b| subject.each_batch_range(table_name, connection: connection, of: 2, &b) }
.to yield_successive_args([first_project.id, second_project.id])
end
it 'makes it possible to apply a scope' do
each_batch_limited = ->(&b) do
subject.each_batch_range(table_name, scope: ->(table) { table.limit(1) }, of: 1, &b)
subject.each_batch_range(table_name, connection: connection, scope: ->(table) { table.limit(1) }, of: 1, &b)
end
expect { |b| each_batch_limited.call(&b) }
......@@ -102,7 +103,7 @@ RSpec.describe Gitlab::Database::DynamicModelHelpers do
end
it 'raises an error' do
expect { subject.each_batch_range(table_name, of: 1) { 1 } }
expect { subject.each_batch_range(table_name, connection: connection, of: 1) { 1 } }
.to raise_error(RuntimeError, /each_batch should not run inside a transaction/)
end
end
......
......@@ -14,6 +14,54 @@ RSpec.describe Gitlab::Database::MigrationHelpers do
allow(model).to receive(:puts)
end
describe 'overridden dynamic model helpers' do
let(:test_table) { '__test_batching_table' }
before do
model.connection.execute(<<~SQL)
CREATE TABLE #{test_table} (
id integer NOT NULL PRIMARY KEY,
name text NOT NULL
);
INSERT INTO #{test_table} (id, name)
VALUES (1, 'bob'), (2, 'mary'), (3, 'amy');
SQL
end
describe '#define_batchable_model' do
it 'defines a batchable model with the migration connection' do
expect(model.define_batchable_model(test_table).count).to eq(3)
end
end
describe '#each_batch' do
before do
allow(model).to receive(:transaction_open?).and_return(false)
end
it 'calls each_batch with the migration connection' do
each_batch_name = ->(&block) do
model.each_batch(test_table, of: 2) do |batch|
block.call(batch.pluck(:name))
end
end
expect { |b| each_batch_name.call(&b) }.to yield_successive_args(%w[bob mary], %w[amy])
end
end
describe '#each_batch_range' do
before do
allow(model).to receive(:transaction_open?).and_return(false)
end
it 'calls each_batch with the migration connection' do
expect { |b| model.each_batch_range(test_table, of: 2, &b) }.to yield_successive_args([1, 2], [3, 3])
end
end
end
describe '#remove_timestamps' do
it 'can remove the default timestamps' do
Gitlab::Database::MigrationHelpers::DEFAULT_TIMESTAMP_COLUMNS.each do |column_name|
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment