pg_ha_migrations 1.6.0 → 1.8.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,100 @@
1
+ module PgHaMigrations
2
+ class LockMode
3
+ include Comparable
4
+
5
+ MODE_CONFLICTS = ActiveSupport::OrderedHash.new
6
+
7
+ MODE_CONFLICTS[:access_share] = %i[
8
+ access_exclusive
9
+ ]
10
+
11
+ MODE_CONFLICTS[:row_share] = %i[
12
+ exclusive
13
+ access_exclusive
14
+ ]
15
+
16
+ MODE_CONFLICTS[:row_exclusive] = %i[
17
+ share
18
+ share_row_exclusive
19
+ exclusive
20
+ access_exclusive
21
+ ]
22
+
23
+ MODE_CONFLICTS[:share_update_exclusive] = %i[
24
+ share_update_exclusive
25
+ share
26
+ share_row_exclusive
27
+ exclusive
28
+ access_exclusive
29
+ ]
30
+
31
+ MODE_CONFLICTS[:share] = %i[
32
+ row_exclusive
33
+ share_update_exclusive
34
+ share_row_exclusive
35
+ exclusive
36
+ access_exclusive
37
+ ]
38
+
39
+ MODE_CONFLICTS[:share_row_exclusive] = %i[
40
+ row_exclusive
41
+ share_update_exclusive
42
+ share
43
+ share_row_exclusive
44
+ exclusive
45
+ access_exclusive
46
+ ]
47
+
48
+ MODE_CONFLICTS[:exclusive] = %i[
49
+ row_share
50
+ row_exclusive
51
+ share_update_exclusive
52
+ share
53
+ share_row_exclusive
54
+ exclusive
55
+ access_exclusive
56
+ ]
57
+
58
+ MODE_CONFLICTS[:access_exclusive] = %i[
59
+ access_share
60
+ row_share
61
+ row_exclusive
62
+ share_update_exclusive
63
+ share
64
+ share_row_exclusive
65
+ exclusive
66
+ access_exclusive
67
+ ]
68
+
69
+ attr_reader :mode
70
+
71
+ delegate :to_s, to: :mode
72
+
73
+ def initialize(mode)
74
+ @mode = mode
75
+ .to_s
76
+ .underscore
77
+ .delete_suffix("_lock")
78
+ .to_sym
79
+
80
+ if !MODE_CONFLICTS.keys.include?(@mode)
81
+ raise ArgumentError, "Unrecognized lock mode #{@mode.inspect}. Valid modes: #{MODE_CONFLICTS.keys}"
82
+ end
83
+ end
84
+
85
+ def to_sql
86
+ mode
87
+ .to_s
88
+ .upcase
89
+ .gsub("_", " ")
90
+ end
91
+
92
+ def <=>(other)
93
+ MODE_CONFLICTS.keys.index(mode) <=> MODE_CONFLICTS.keys.index(other.mode)
94
+ end
95
+
96
+ def conflicts_with?(other)
97
+ MODE_CONFLICTS[mode].include?(other.mode)
98
+ end
99
+ end
100
+ end
@@ -0,0 +1,11 @@
1
+ # This is an internal class that is not meant to be used directly
2
+ class PgHaMigrations::PartmanConfig < ActiveRecord::Base
3
+ self.primary_key = :parent_table
4
+
5
+ # This method is called by unsafe_partman_update_config to set the fully
6
+ # qualified table name, as partman is often installed in a schema that
7
+ # is not included the application's search path
8
+ def self.schema=(schema)
9
+ self.table_name = "#{schema}.part_config"
10
+ end
11
+ end
@@ -0,0 +1,155 @@
1
+ module PgHaMigrations
2
+ Relation = Struct.new(:name, :schema, :mode) do
3
+ def self.connection
4
+ ActiveRecord::Base.connection
5
+ end
6
+
7
+ delegate :inspect, to: :name
8
+ delegate :connection, to: :class
9
+
10
+ def initialize(name, schema, mode=nil)
11
+ super(name, schema)
12
+
13
+ self.mode = LockMode.new(mode) if mode.present?
14
+ end
15
+
16
+ def conflicts_with?(other)
17
+ self == other && (
18
+ mode.nil? || other.mode.nil? || mode.conflicts_with?(other.mode)
19
+ )
20
+ end
21
+
22
+ def fully_qualified_name
23
+ @fully_qualified_name ||= [
24
+ PG::Connection.quote_ident(schema),
25
+ PG::Connection.quote_ident(name),
26
+ ].join(".")
27
+ end
28
+
29
+ def present?
30
+ name.present? && schema.present?
31
+ end
32
+
33
+ def ==(other)
34
+ other.is_a?(Relation) && name == other.name && schema == other.schema
35
+ end
36
+ end
37
+
38
+ class Table < Relation
39
+ def self.from_table_name(table, mode=nil)
40
+ pg_name = ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(table.to_s)
41
+
42
+ schema_conditional = if pg_name.schema
43
+ "#{connection.quote(pg_name.schema)}"
44
+ else
45
+ "ANY (current_schemas(false))"
46
+ end
47
+
48
+ schema = connection.select_value(<<~SQL)
49
+ SELECT schemaname
50
+ FROM pg_tables
51
+ WHERE tablename = #{connection.quote(pg_name.identifier)} AND schemaname = #{schema_conditional}
52
+ ORDER BY array_position(current_schemas(false), schemaname)
53
+ LIMIT 1
54
+ SQL
55
+
56
+ raise UndefinedTableError, "Table #{pg_name.quoted} does not exist#{" in search path" unless pg_name.schema}" unless schema.present?
57
+
58
+ new(pg_name.identifier, schema, mode)
59
+ end
60
+
61
+ def natively_partitioned?
62
+ return @natively_partitioned if defined?(@natively_partitioned)
63
+
64
+ @natively_partitioned = !!connection.select_value(<<~SQL)
65
+ SELECT true
66
+ FROM pg_partitioned_table, pg_class, pg_namespace
67
+ WHERE pg_class.oid = pg_partitioned_table.partrelid
68
+ AND pg_class.relnamespace = pg_namespace.oid
69
+ AND pg_class.relname = #{connection.quote(name)}
70
+ AND pg_namespace.nspname = #{connection.quote(schema)}
71
+ SQL
72
+ end
73
+
74
+ def partitions(include_sub_partitions: false, include_self: false)
75
+ tables = connection.structs_from_sql(self.class, <<~SQL)
76
+ SELECT child.relname AS name, child_ns.nspname AS schema, NULLIF('#{mode}', '') AS mode
77
+ FROM pg_inherits
78
+ JOIN pg_class parent ON pg_inherits.inhparent = parent.oid
79
+ JOIN pg_class child ON pg_inherits.inhrelid = child.oid
80
+ JOIN pg_namespace parent_ns ON parent.relnamespace = parent_ns.oid
81
+ JOIN pg_namespace child_ns ON child.relnamespace = child_ns.oid
82
+ WHERE parent.relname = #{connection.quote(name)}
83
+ AND parent_ns.nspname = #{connection.quote(schema)}
84
+ ORDER BY child.oid -- Ensure consistent ordering for tests
85
+ SQL
86
+
87
+ if include_sub_partitions
88
+ sub_partitions = tables.each_with_object([]) do |table, arr|
89
+ arr.concat(table.partitions(include_sub_partitions: true))
90
+ end
91
+
92
+ tables.concat(sub_partitions)
93
+ end
94
+
95
+ tables.prepend(self) if include_self
96
+
97
+ tables
98
+ end
99
+
100
+ def has_rows?
101
+ connection.select_value("SELECT EXISTS (SELECT 1 FROM #{fully_qualified_name} LIMIT 1)")
102
+ end
103
+
104
+ def total_bytes
105
+ connection.select_value(<<~SQL)
106
+ SELECT pg_total_relation_size(pg_class.oid)
107
+ FROM pg_class, pg_namespace
108
+ WHERE pg_class.relname = #{connection.quote(name)}
109
+ AND pg_namespace.nspname = #{connection.quote(schema)}
110
+ SQL
111
+ end
112
+ end
113
+
114
+ class Index < Relation
115
+ MAX_NAME_SIZE = 63 # bytes
116
+
117
+ def self.from_table_and_columns(table, columns)
118
+ name = connection.index_name(table.name, columns)
119
+
120
+ # modified from https://github.com/rails/rails/pull/47753
121
+ if name.bytesize > MAX_NAME_SIZE
122
+ hashed_identifier = "_#{OpenSSL::Digest::SHA256.hexdigest(name).first(10)}"
123
+ description = name.sub("index_#{table.name}_on", "idx_on")
124
+
125
+ short_limit = MAX_NAME_SIZE - hashed_identifier.bytesize
126
+ short_description = description.mb_chars.limit(short_limit).to_s
127
+
128
+ name = "#{short_description}#{hashed_identifier}"
129
+ end
130
+
131
+ new(name, table)
132
+ end
133
+
134
+ attr_accessor :table
135
+
136
+ def initialize(name, table)
137
+ super(name, table.schema)
138
+
139
+ self.table = table
140
+
141
+ connection.send(:validate_index_length!, table.name, name)
142
+ end
143
+
144
+ def valid?
145
+ !!connection.select_value(<<~SQL)
146
+ SELECT pg_index.indisvalid
147
+ FROM pg_index, pg_class, pg_namespace
148
+ WHERE pg_class.oid = pg_index.indexrelid
149
+ AND pg_class.relnamespace = pg_namespace.oid
150
+ AND pg_namespace.nspname = #{connection.quote(schema)}
151
+ AND pg_class.relname = #{connection.quote(name)}
152
+ SQL
153
+ end
154
+ end
155
+ end
@@ -144,6 +144,22 @@ module PgHaMigrations::SafeStatements
144
144
  end
145
145
  end
146
146
 
147
+ def safe_add_index_on_empty_table(table, columns, options={})
148
+ if options[:algorithm] == :concurrently
149
+ raise ArgumentError, "Cannot call safe_add_index_on_empty_table with :algorithm => :concurrently"
150
+ end
151
+
152
+ # Avoids taking out an unnecessary SHARE lock if the table does have data
153
+ ensure_small_table!(table, empty: true)
154
+
155
+ safely_acquire_lock_for_table(table, mode: :share) do
156
+ # Ensure data wasn't written in the split second after the first check
157
+ ensure_small_table!(table, empty: true)
158
+
159
+ unsafe_add_index(table, columns, **options)
160
+ end
161
+ end
162
+
147
163
  def safe_add_concurrent_index(table, columns, options={})
148
164
  unsafe_add_index(table, columns, **options.merge(:algorithm => :concurrently))
149
165
  end
@@ -160,26 +176,96 @@ module PgHaMigrations::SafeStatements
160
176
  unsafe_remove_index(table, **options.merge(:algorithm => :concurrently))
161
177
  end
162
178
 
163
- def safe_set_maintenance_work_mem_gb(gigabytes)
164
- unsafe_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
165
- end
179
+ def safe_add_concurrent_partitioned_index(
180
+ table,
181
+ columns,
182
+ name: nil,
183
+ if_not_exists: nil,
184
+ using: nil,
185
+ unique: nil,
186
+ where: nil,
187
+ comment: nil
188
+ )
189
+
190
+ if ActiveRecord::Base.connection.postgresql_version < 11_00_00
191
+ raise PgHaMigrations::InvalidMigrationError, "Concurrent partitioned index creation not supported on Postgres databases before version 11"
192
+ end
166
193
 
167
- def safe_add_unvalidated_check_constraint(table, expression, name:)
168
- unsafe_add_check_constraint(table, expression, name: name, validate: false)
169
- end
194
+ parent_table = PgHaMigrations::Table.from_table_name(table)
170
195
 
171
- def unsafe_add_check_constraint(table, expression, name:, validate: true)
172
- raise ArgumentError, "Expected <name> to be present" unless name.present?
196
+ raise PgHaMigrations::InvalidMigrationError, "Table #{parent_table.inspect} is not a partitioned table" unless parent_table.natively_partitioned?
173
197
 
174
- quoted_table_name = connection.quote_table_name(table)
175
- quoted_constraint_name = connection.quote_table_name(name)
176
- sql = "ALTER TABLE #{quoted_table_name} ADD CONSTRAINT #{quoted_constraint_name} CHECK (#{expression}) #{validate ? "" : "NOT VALID"}"
198
+ parent_index = if name.present?
199
+ PgHaMigrations::Index.new(name, parent_table)
200
+ else
201
+ PgHaMigrations::Index.from_table_and_columns(parent_table, columns)
202
+ end
177
203
 
178
- safely_acquire_lock_for_table(table) do
179
- say_with_time "add_check_constraint(#{table.inspect}, #{expression.inspect}, name: #{name.inspect}, validate: #{validate.inspect})" do
180
- connection.execute(sql)
204
+ # Short-circuit when if_not_exists: true and index already valid
205
+ return if if_not_exists && parent_index.valid?
206
+
207
+ child_indexes = parent_table.partitions.map do |child_table|
208
+ PgHaMigrations::Index.from_table_and_columns(child_table, columns)
209
+ end
210
+
211
+ # TODO: take out ShareLock after issue #39 is implemented
212
+ safely_acquire_lock_for_table(parent_table.fully_qualified_name) do
213
+ # CREATE INDEX ON ONLY parent_table
214
+ unsafe_add_index(
215
+ parent_table.fully_qualified_name,
216
+ columns,
217
+ name: parent_index.name,
218
+ if_not_exists: if_not_exists,
219
+ using: using,
220
+ unique: unique,
221
+ where: where,
222
+ comment: comment,
223
+ algorithm: :only, # see lib/pg_ha_migrations/hacks/add_index_on_only.rb
224
+ )
225
+ end
226
+
227
+ child_indexes.each do |child_index|
228
+ add_index_method = if child_index.table.natively_partitioned?
229
+ :safe_add_concurrent_partitioned_index
230
+ else
231
+ :safe_add_concurrent_index
232
+ end
233
+
234
+ send(
235
+ add_index_method,
236
+ child_index.table.fully_qualified_name,
237
+ columns,
238
+ name: child_index.name,
239
+ if_not_exists: if_not_exists,
240
+ using: using,
241
+ unique: unique,
242
+ where: where,
243
+ )
244
+ end
245
+
246
+ # Avoid taking out an unnecessary lock if there are no child tables to attach
247
+ if child_indexes.present?
248
+ safely_acquire_lock_for_table(parent_table.fully_qualified_name) do
249
+ child_indexes.each do |child_index|
250
+ say_with_time "Attaching index #{child_index.inspect} to #{parent_index.inspect}" do
251
+ connection.execute(<<~SQL)
252
+ ALTER INDEX #{parent_index.fully_qualified_name}
253
+ ATTACH PARTITION #{child_index.fully_qualified_name}
254
+ SQL
255
+ end
256
+ end
181
257
  end
182
258
  end
259
+
260
+ raise PgHaMigrations::InvalidMigrationError, "Unexpected state. Parent index #{parent_index.inspect} is invalid" unless parent_index.valid?
261
+ end
262
+
263
+ def safe_set_maintenance_work_mem_gb(gigabytes)
264
+ unsafe_execute("SET maintenance_work_mem = '#{PG::Connection.escape_string(gigabytes.to_s)} GB'")
265
+ end
266
+
267
+ def safe_add_unvalidated_check_constraint(table, expression, name:)
268
+ unsafe_add_check_constraint(table, expression, name: name, validate: false)
183
269
  end
184
270
 
185
271
  def safe_validate_check_constraint(table, name:)
@@ -224,6 +310,197 @@ module PgHaMigrations::SafeStatements
224
310
  end
225
311
  end
226
312
 
313
+ def safe_create_partitioned_table(table, partition_key:, type:, infer_primary_key: nil, **options, &block)
314
+ raise ArgumentError, "Expected <partition_key> to be present" unless partition_key.present?
315
+
316
+ unless PgHaMigrations::PARTITION_TYPES.include?(type)
317
+ raise ArgumentError, "Expected <type> to be symbol in #{PgHaMigrations::PARTITION_TYPES} but received #{type.inspect}"
318
+ end
319
+
320
+ if ActiveRecord::Base.connection.postgresql_version < 10_00_00
321
+ raise PgHaMigrations::InvalidMigrationError, "Native partitioning not supported on Postgres databases before version 10"
322
+ end
323
+
324
+ if type == :hash && ActiveRecord::Base.connection.postgresql_version < 11_00_00
325
+ raise PgHaMigrations::InvalidMigrationError, "Hash partitioning not supported on Postgres databases before version 11"
326
+ end
327
+
328
+ if infer_primary_key.nil?
329
+ infer_primary_key = PgHaMigrations.config.infer_primary_key_on_partitioned_tables
330
+ end
331
+
332
+ # Newer versions of Rails will set the primary key column to the type :primary_key.
333
+ # This performs some extra logic that we can't easily undo which causes problems when
334
+ # trying to inject the partition key into the PK. Now, it would be nice to lookup the
335
+ # default primary key type instead of simply using :bigserial, but it doesn't appear
336
+ # that we have access to the Rails configuration from within our migrations.
337
+ if options[:id].nil? || options[:id] == :primary_key
338
+ options[:id] = :bigserial
339
+ end
340
+
341
+ quoted_partition_key = if partition_key.is_a?(Proc)
342
+ # Lambda syntax, like in other migration methods, implies an expression that
343
+ # cannot be easily sanitized.
344
+ #
345
+ # e.g ->{ "(created_at::date)" }
346
+ partition_key.call.to_s
347
+ else
348
+ # Otherwise, assume key is a column name or array of column names
349
+ Array.wrap(partition_key).map { |col| connection.quote_column_name(col) }.join(",")
350
+ end
351
+
352
+ options[:options] = "PARTITION BY #{type.upcase} (#{quoted_partition_key})"
353
+
354
+ safe_create_table(table, options) do |td|
355
+ block.call(td) if block
356
+
357
+ next unless options[:id]
358
+
359
+ pk_columns = td.columns.each_with_object([]) do |col, arr|
360
+ next unless col.options[:primary_key]
361
+
362
+ col.options[:primary_key] = false
363
+
364
+ arr << col.name
365
+ end
366
+
367
+ if infer_primary_key && !partition_key.is_a?(Proc) && ActiveRecord::Base.connection.postgresql_version >= 11_00_00
368
+ td.primary_keys(pk_columns.concat(Array.wrap(partition_key)).map(&:to_s).uniq)
369
+ end
370
+ end
371
+ end
372
+
373
+ def safe_partman_create_parent(table, **options)
374
+ if options[:retention].present? || options[:retention_keep_table] == false
375
+ raise PgHaMigrations::UnsafeMigrationError.new(":retention and/or :retention_keep_table => false can potentially result in data loss if misconfigured. Please use unsafe_partman_create_parent if you want to set these options")
376
+ end
377
+
378
+ unsafe_partman_create_parent(table, **options)
379
+ end
380
+
381
+ def unsafe_partman_create_parent(
382
+ table,
383
+ partition_key:,
384
+ interval:,
385
+ infinite_time_partitions: true,
386
+ inherit_privileges: true,
387
+ premake: nil,
388
+ start_partition: nil,
389
+ template_table: nil,
390
+ retention: nil,
391
+ retention_keep_table: nil
392
+ )
393
+ raise ArgumentError, "Expected <partition_key> to be present" unless partition_key.present?
394
+ raise ArgumentError, "Expected <interval> to be present" unless interval.present?
395
+
396
+ if ActiveRecord::Base.connection.postgresql_version < 11_00_00
397
+ raise PgHaMigrations::InvalidMigrationError, "Native partitioning with partman not supported on Postgres databases before version 11"
398
+ end
399
+
400
+ formatted_start_partition = nil
401
+
402
+ if start_partition.present?
403
+ if !start_partition.is_a?(Date) && !start_partition.is_a?(Time) && !start_partition.is_a?(DateTime)
404
+ raise PgHaMigrations::InvalidMigrationError, "Expected <start_partition> to be Date, Time, or DateTime object but received #{start_partition.class}"
405
+ end
406
+
407
+ formatted_start_partition = if start_partition.respond_to?(:to_fs)
408
+ start_partition.to_fs(:db)
409
+ else
410
+ start_partition.to_s(:db)
411
+ end
412
+ end
413
+
414
+ create_parent_options = {
415
+ parent_table: _fully_qualified_table_name_for_partman(table),
416
+ template_table: template_table ? _fully_qualified_table_name_for_partman(template_table) : nil,
417
+ control: partition_key,
418
+ type: "native",
419
+ interval: interval,
420
+ premake: premake,
421
+ start_partition: formatted_start_partition,
422
+ }.compact
423
+
424
+ create_parent_sql = create_parent_options.map { |k, v| "p_#{k} := #{connection.quote(v)}" }.join(", ")
425
+
426
+ log_message = "partman_create_parent(#{table.inspect}, " \
427
+ "partition_key: #{partition_key.inspect}, " \
428
+ "interval: #{interval.inspect}, " \
429
+ "premake: #{premake.inspect}, " \
430
+ "start_partition: #{start_partition.inspect}, " \
431
+ "template_table: #{template_table.inspect})"
432
+
433
+ say_with_time(log_message) do
434
+ connection.execute("SELECT #{_quoted_partman_schema}.create_parent(#{create_parent_sql})")
435
+ end
436
+
437
+ update_config_options = {
438
+ infinite_time_partitions: infinite_time_partitions,
439
+ inherit_privileges: inherit_privileges,
440
+ retention: retention,
441
+ retention_keep_table: retention_keep_table,
442
+ }.compact
443
+
444
+ unsafe_partman_update_config(table, **update_config_options)
445
+ end
446
+
447
+ def safe_partman_update_config(table, **options)
448
+ if options[:retention].present? || options[:retention_keep_table] == false
449
+ raise PgHaMigrations::UnsafeMigrationError.new(":retention and/or :retention_keep_table => false can potentially result in data loss if misconfigured. Please use unsafe_partman_update_config if you want to set these options")
450
+ end
451
+
452
+ unsafe_partman_update_config(table, **options)
453
+ end
454
+
455
+ def unsafe_partman_update_config(table, **options)
456
+ invalid_options = options.keys - PgHaMigrations::PARTMAN_UPDATE_CONFIG_OPTIONS
457
+
458
+ raise ArgumentError, "Unrecognized argument(s): #{invalid_options}" unless invalid_options.empty?
459
+
460
+ PgHaMigrations::PartmanConfig.schema = _quoted_partman_schema
461
+
462
+ config = PgHaMigrations::PartmanConfig.find(_fully_qualified_table_name_for_partman(table))
463
+
464
+ config.assign_attributes(**options)
465
+
466
+ inherit_privileges_changed = config.inherit_privileges_changed?
467
+
468
+ say_with_time "partman_update_config(#{table.inspect}, #{options.map { |k,v| "#{k}: #{v.inspect}" }.join(", ")})" do
469
+ config.save!
470
+ end
471
+
472
+ safe_partman_reapply_privileges(table) if inherit_privileges_changed
473
+ end
474
+
475
+ def safe_partman_reapply_privileges(table)
476
+ say_with_time "partman_reapply_privileges(#{table.inspect})" do
477
+ connection.execute("SELECT #{_quoted_partman_schema}.reapply_privileges('#{_fully_qualified_table_name_for_partman(table)}')")
478
+ end
479
+ end
480
+
481
+ def _quoted_partman_schema
482
+ schema = connection.select_value(<<~SQL)
483
+ SELECT nspname
484
+ FROM pg_namespace JOIN pg_extension
485
+ ON pg_namespace.oid = pg_extension.extnamespace
486
+ WHERE pg_extension.extname = 'pg_partman'
487
+ SQL
488
+
489
+ raise PgHaMigrations::InvalidMigrationError, "The pg_partman extension is not installed" unless schema.present?
490
+
491
+ connection.quote_schema_name(schema)
492
+ end
493
+
494
+ def _fully_qualified_table_name_for_partman(table)
495
+ table = PgHaMigrations::Table.from_table_name(table)
496
+
497
+ [table.schema, table.name].each do |identifier|
498
+ if identifier.to_s !~ /^[a-z_][a-z_\d]*$/
499
+ raise PgHaMigrations::InvalidMigrationError, "Partman requires schema / table names to be lowercase with underscores"
500
+ end
501
+ end.join(".")
502
+ end
503
+
227
504
  def _per_migration_caller
228
505
  @_per_migration_caller ||= Kernel.caller
229
506
  end
@@ -251,17 +528,39 @@ module PgHaMigrations::SafeStatements
251
528
  super(conn, direction)
252
529
  end
253
530
 
254
- def safely_acquire_lock_for_table(table, &block)
531
+ def safely_acquire_lock_for_table(table, mode: :access_exclusive, &block)
532
+ nested_target_table = Thread.current[__method__]
533
+
255
534
  _check_postgres_adapter!
256
- table = table.to_s
257
- quoted_table_name = connection.quote_table_name(table)
535
+
536
+ target_table = PgHaMigrations::Table.from_table_name(table, mode)
537
+
538
+ if nested_target_table
539
+ if nested_target_table != target_table
540
+ raise PgHaMigrations::InvalidMigrationError, "Nested lock detected! Cannot acquire lock on #{target_table.fully_qualified_name} while #{nested_target_table.fully_qualified_name} is locked."
541
+ elsif nested_target_table.mode < target_table.mode
542
+ raise PgHaMigrations::InvalidMigrationError, "Lock escalation detected! Cannot change lock level from :#{nested_target_table.mode} to :#{target_table.mode} for #{target_table.fully_qualified_name}."
543
+ end
544
+ else
545
+ Thread.current[__method__] = target_table
546
+ end
547
+
548
+ # Locking a partitioned table will also lock child tables (including sub-partitions),
549
+ # so we need to check for blocking queries on those tables as well
550
+ target_tables = target_table.partitions(include_sub_partitions: true, include_self: true)
258
551
 
259
552
  successfully_acquired_lock = false
260
553
 
261
554
  until successfully_acquired_lock
262
555
  while (
263
556
  blocking_transactions = PgHaMigrations::BlockingDatabaseTransactions.find_blocking_transactions("#{PgHaMigrations::LOCK_TIMEOUT_SECONDS} seconds")
264
- blocking_transactions.any? { |query| query.tables_with_locks.include?(table) }
557
+ blocking_transactions.any? do |query|
558
+ query.tables_with_locks.any? do |locked_table|
559
+ target_tables.any? do |target_table|
560
+ target_table.conflicts_with?(locked_table)
561
+ end
562
+ end
563
+ end
265
564
  )
266
565
  say "Waiting on blocking transactions:"
267
566
  blocking_transactions.each do |blocking_transaction|
@@ -274,13 +573,13 @@ module PgHaMigrations::SafeStatements
274
573
  adjust_timeout_method = connection.postgresql_version >= 9_03_00 ? :adjust_lock_timeout : :adjust_statement_timeout
275
574
  begin
276
575
  method(adjust_timeout_method).call(PgHaMigrations::LOCK_TIMEOUT_SECONDS) do
277
- connection.execute("LOCK #{quoted_table_name};")
576
+ connection.execute("LOCK #{target_table.fully_qualified_name} IN #{target_table.mode.to_sql} MODE;")
278
577
  end
279
578
  successfully_acquired_lock = true
280
579
  rescue ActiveRecord::StatementInvalid => e
281
580
  if e.message =~ /PG::LockNotAvailable.+ lock timeout/ || e.message =~ /PG::QueryCanceled.+ statement timeout/
282
581
  sleep_seconds = PgHaMigrations::LOCK_FAILURE_RETRY_DELAY_MULTLIPLIER * PgHaMigrations::LOCK_TIMEOUT_SECONDS
283
- say "Timed out trying to acquire an exclusive lock on the #{quoted_table_name} table."
582
+ say "Timed out trying to acquire #{target_table.mode.to_sql} lock on the #{target_table.fully_qualified_name} table."
284
583
  say "Sleeping for #{sleep_seconds}s to allow potentially queued up queries to finish before continuing."
285
584
  sleep(sleep_seconds)
286
585
 
@@ -295,6 +594,8 @@ module PgHaMigrations::SafeStatements
295
594
  end
296
595
  end
297
596
  end
597
+ ensure
598
+ Thread.current[__method__] = nil unless nested_target_table
298
599
  end
299
600
 
300
601
  def adjust_lock_timeout(timeout_seconds = PgHaMigrations::LOCK_TIMEOUT_SECONDS, &block)
@@ -336,4 +637,16 @@ module PgHaMigrations::SafeStatements
336
637
  end
337
638
  end
338
639
  end
640
+
641
+ def ensure_small_table!(table, empty: false, threshold: PgHaMigrations::SMALL_TABLE_THRESHOLD_BYTES)
642
+ table = PgHaMigrations::Table.from_table_name(table)
643
+
644
+ if empty && table.has_rows?
645
+ raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} has rows"
646
+ end
647
+
648
+ if table.total_bytes > threshold
649
+ raise PgHaMigrations::InvalidMigrationError, "Table #{table.inspect} is larger than #{threshold} bytes"
650
+ end
651
+ end
339
652
  end