From 6ddb39c0be69007698dba7dc6ee7f183be68d05b Mon Sep 17 00:00:00 2001 From: Dawid Szymczak Date: Mon, 11 Mar 2024 21:00:17 +0100 Subject: [PATCH] fix: refactor ActiveRecord monkey patches to modules and use Module prepend and include to apply them as traits --- lib/activerecord-multi-tenant.rb | 14 +- lib/activerecord-multi-tenant/migrations.rb | 123 ------------------ .../migrations_extension.rb | 64 +++++++++ .../schema_dumper.rb | 35 +++++ .../schema_statements.rb | 15 +++ 5 files changed, 124 insertions(+), 127 deletions(-) delete mode 100644 lib/activerecord-multi-tenant/migrations.rb create mode 100644 lib/activerecord-multi-tenant/migrations_extension.rb create mode 100644 lib/activerecord-multi-tenant/schema_dumper.rb create mode 100644 lib/activerecord-multi-tenant/schema_statements.rb diff --git a/lib/activerecord-multi-tenant.rb b/lib/activerecord-multi-tenant.rb index 2a9c591c..9bc985bc 100644 --- a/lib/activerecord-multi-tenant.rb +++ b/lib/activerecord-multi-tenant.rb @@ -1,9 +1,6 @@ -if Object.const_defined?(:ActionController) - require_relative 'activerecord-multi-tenant/controller_extensions' -end +require_relative 'activerecord-multi-tenant/controller_extensions' if Object.const_defined?(:ActionController) require_relative 'activerecord-multi-tenant/copy_from_client' require_relative 'activerecord-multi-tenant/fast_truncate' -require_relative 'activerecord-multi-tenant/migrations' require_relative 'activerecord-multi-tenant/model_extensions' require_relative 'activerecord-multi-tenant/multi_tenant' require_relative 'activerecord-multi-tenant/query_rewriter' @@ -11,3 +8,12 @@ require_relative 'activerecord-multi-tenant/version' require_relative 'activerecord-multi-tenant/with_lock' require_relative 'activerecord-multi-tenant/delete_operations' + +ActiveSupport.on_load(:active_record) do + require_relative 'activerecord-multi-tenant/migrations_extension' + require_relative 'activerecord-multi-tenant/schema_dumper' + require_relative 'activerecord-multi-tenant/schema_statements' + + ActiveRecord::Migration.include MultiTenant::MigrationExtensions + ActiveRecord::SchemaDumper.prepend(MultiTenant::SchemaDumper) +end diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb deleted file mode 100644 index cda58112..00000000 --- a/lib/activerecord-multi-tenant/migrations.rb +++ /dev/null @@ -1,123 +0,0 @@ -module MultiTenant - module MigrationExtensions - def create_distributed_table(table_name, partition_key) - return unless citus_version.present? - - reversible do |dir| - dir.up do - execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)" - end - dir.down do - undistribute_table(table_name) - end - end - end - - def create_reference_table(table_name) - return unless citus_version.present? - - reversible do |dir| - dir.up do - execute "SELECT create_reference_table($$#{table_name}$$)" - end - dir.down do - undistribute_table(table_name) - end - end - end - - def undistribute_table(table_name) - return unless citus_version.present? - - execute "SELECT undistribute_table($$#{table_name}$$))" - end - - def rebalance_table_shards - return unless citus_version.present? - - execute 'SELECT rebalance_table_shards()' - end - - def execute_on_all_nodes(sql) - execute sql - - case citus_version - when '6.0' - execute "SELECT citus_run_on_all_workers($$#{sql}$$)" # initial citus_tools.sql with different names - when nil - # Do nothing, this is regular Postgres - else # 6.1 and newer - execute "SELECT run_command_on_workers($$#{sql}$$)" - end - end - - def enable_extension_on_all_nodes(extension) - execute_on_all_nodes "CREATE EXTENSION IF NOT EXISTS \"#{extension}\"" - end - - def citus_version - execute("SELECT extversion FROM pg_extension WHERE extname = 'citus'").getvalue(0, 0).try(:split, '-').try(:first) - rescue ArgumentError => e - raise unless e.message == 'invalid tuple number 0' - end - end -end - -ActiveRecord::Migration.include MultiTenant::MigrationExtensions if defined?(ActiveRecord::Migration) - -module ActiveRecord - module ConnectionAdapters # :nodoc: - module SchemaStatements - alias orig_create_table create_table - def create_table(table_name, options = {}, &block) - ret = orig_create_table(table_name, **options.except(:partition_key), &block) - if options[:id] != false && options[:partition_key] && options[:partition_key].to_s != 'id' - execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{table_name}_pkey" - execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)" - end - ret - end - end - end -end - -module ActiveRecord - class SchemaDumper - private - - alias initialize_without_citus initialize - def initialize(connection, options = {}) - initialize_without_citus(connection, options) - - citus_version = begin - ActiveRecord::Migration.citus_version - rescue StandardError - # Handle the case where this gem is used with MySQL https://github.com/citusdata/activerecord-multi-tenant/issues/166 - nil - end - @distribution_columns = - if citus_version.present? - @connection.execute('SELECT logicalrelid::regclass AS table_name, column_to_column_name(logicalrelid, partkey) AS dist_col_name FROM pg_dist_partition').to_h do |v| - [v['table_name'], v['dist_col_name']] - end - else - {} - end - end - - # Support for create_distributed_table & create_reference_table - alias table_without_citus table - def table(table, stream) - table_without_citus(table, stream) - table_name = remove_prefix_and_suffix(table) - distribution_column = @distribution_columns[table_name] - if distribution_column - stream.puts " create_distributed_table(#{table_name.inspect}, #{distribution_column.inspect})" - stream.puts - elsif @distribution_columns.key?(table_name) - stream.puts " create_reference_table(#{table_name.inspect})" - stream.puts - end - end - end -end diff --git a/lib/activerecord-multi-tenant/migrations_extension.rb b/lib/activerecord-multi-tenant/migrations_extension.rb new file mode 100644 index 00000000..b49d947f --- /dev/null +++ b/lib/activerecord-multi-tenant/migrations_extension.rb @@ -0,0 +1,64 @@ +module MultiTenant + module MigrationExtensions + def create_distributed_table(table_name, partition_key) + return unless citus_version.present? + + reversible do |dir| + dir.up do + execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)" + end + dir.down do + undistribute_table(table_name) + end + end + end + + def create_reference_table(table_name) + return unless citus_version.present? + + reversible do |dir| + dir.up do + execute "SELECT create_reference_table($$#{table_name}$$)" + end + dir.down do + undistribute_table(table_name) + end + end + end + + def undistribute_table(table_name) + return unless citus_version.present? + + execute "SELECT undistribute_table($$#{table_name}$$))" + end + + def rebalance_table_shards + return unless citus_version.present? + + execute 'SELECT rebalance_table_shards()' + end + + def execute_on_all_nodes(sql) + execute sql + + case citus_version + when '6.0' + execute "SELECT citus_run_on_all_workers($$#{sql}$$)" # initial citus_tools.sql with different names + when nil + # Do nothing, this is regular Postgres + else # 6.1 and newer + execute "SELECT run_command_on_workers($$#{sql}$$)" + end + end + + def enable_extension_on_all_nodes(extension) + execute_on_all_nodes "CREATE EXTENSION IF NOT EXISTS \"#{extension}\"" + end + + def citus_version + execute("SELECT extversion FROM pg_extension WHERE extname = 'citus'").getvalue(0, 0).try(:split, '-').try(:first) + rescue ArgumentError => e + raise unless e.message == 'invalid tuple number 0' + end + end +end diff --git a/lib/activerecord-multi-tenant/schema_dumper.rb b/lib/activerecord-multi-tenant/schema_dumper.rb new file mode 100644 index 00000000..4166a991 --- /dev/null +++ b/lib/activerecord-multi-tenant/schema_dumper.rb @@ -0,0 +1,35 @@ +module MultiTenant + module SchemaDumper + def initialize(connection, options = {}) + super(connection, options) + + citus_version = begin + ActiveRecord::Migration.citus_version + rescue StandardError + # Handle the case where this gem is used with MySQL https://github.com/citusdata/activerecord-multi-tenant/issues/166 + nil + end + @distribution_columns = + if citus_version.present? + @connection.execute('SELECT logicalrelid::regclass AS table_name, column_to_column_name(logicalrelid, partkey) AS dist_col_name FROM pg_dist_partition').to_h do |v| + [v['table_name'], v['dist_col_name']] + end + else + {} + end + end + + def table(table, stream) + super(table, stream) + table_name = remove_prefix_and_suffix(table) + distribution_column = @distribution_columns[table_name] + if distribution_column + stream.puts " create_distributed_table(#{table_name.inspect}, #{distribution_column.inspect})" + stream.puts + elsif @distribution_columns.key?(table_name) + stream.puts " create_reference_table(#{table_name.inspect})" + stream.puts + end + end + end +end diff --git a/lib/activerecord-multi-tenant/schema_statements.rb b/lib/activerecord-multi-tenant/schema_statements.rb new file mode 100644 index 00000000..6103e42e --- /dev/null +++ b/lib/activerecord-multi-tenant/schema_statements.rb @@ -0,0 +1,15 @@ +module ActiveRecord + module ConnectionAdapters # :nodoc: + module SchemaStatements + alias orig_create_table create_table + def create_table(table_name, options = {}, &block) + ret = orig_create_table(table_name, **options.except(:partition_key), &block) + if options[:id] != false && options[:partition_key] && options[:partition_key].to_s != 'id' + execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{table_name}_pkey" + execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)" + end + ret + end + end + end +end