From 987337aa2d6d1805c17b4dfe43a7fc21806b04f7 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Wed, 11 Sep 2019 14:39:48 -0700 Subject: [PATCH 01/11] allow to do a change_table to change partition key --- lib/activerecord-multi-tenant/migrations.rb | 20 ++++++++++++++++---- spec/schema.rb | 8 ++++++-- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb index b6104353..c3b8ca73 100644 --- a/lib/activerecord-multi-tenant/migrations.rb +++ b/lib/activerecord-multi-tenant/migrations.rb @@ -43,12 +43,24 @@ 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) + alias :orig_change_table :change_table + + def change_primary_key(table_name, options) if 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(id, \"#{options[:partition_key]}\")" + execute "ALTER TABLE #{table_name} DROP CONSTRAINT IF EXISTS #{table_name}_pkey" + execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)" end + end + + def create_table(table_name, options = {}, &block) + ret = orig_create_table(table_name, options.except(:partition_key), &block) + change_primary_key(table_name, options) + ret + end + + def change_table(table_name, options, &block) + ret = orig_change_table(table_name, options.except(:partition_key), &block) + change_primary_key(table_name, options) ret end end diff --git a/spec/schema.rb b/spec/schema.rb index 04ab6880..e141ced9 100644 --- a/spec/schema.rb +++ b/spec/schema.rb @@ -82,7 +82,7 @@ t.column :name, :string end - create_table :project_categories, force: true, partition_key: :account_id do |t| + create_table :project_categories, force: true do |t| t.column :name, :string t.column :account_id, :integer t.column :project_id, :integer @@ -106,9 +106,13 @@ create_distributed_table :partition_key_not_model_tasks, :non_model_id create_distributed_table :subclass_tasks, :non_model_id create_distributed_table :uuid_records, :organization_id - create_distributed_table :project_categories, :account_id create_distributed_table :allowed_places, :account_id create_reference_table :categories + + change_table :project_categories, partition_key: :account_id do |t| + end + + create_distributed_table :project_categories, :account_id end class Account < ActiveRecord::Base From 99c6fbd66621906013b8b0703cd69af8d777264f Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Wed, 11 Sep 2019 14:55:25 -0700 Subject: [PATCH 02/11] check that the primary key has n column to avoid recreating each time there is a change_table :partition_key --- lib/activerecord-multi-tenant/migrations.rb | 27 ++++++++++++++++++--- spec/schema.rb | 7 ++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb index c3b8ca73..058c4a26 100644 --- a/lib/activerecord-multi-tenant/migrations.rb +++ b/lib/activerecord-multi-tenant/migrations.rb @@ -46,10 +46,31 @@ module SchemaStatements alias :orig_change_table :change_table def change_primary_key(table_name, options) - if options[:partition_key] && options[:partition_key].to_s != 'id' - execute "ALTER TABLE #{table_name} DROP CONSTRAINT IF EXISTS #{table_name}_pkey" - execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)" + if !options[:partition_key] || options[:partition_key].to_s == 'id' + return end + + pkey_columns = ['id', options[:partition_key]] + + # we are here comparing the columns in the primary key on the database and the one in the migration file + columns_result = execute("select kcu.column_name as key_column " \ + "from information_schema.table_constraints tco "\ + "join information_schema.key_column_usage kcu " \ + "ON kcu.constraint_name = tco.constraint_name " \ + "AND kcu.constraint_schema = tco.constraint_schema "\ + "WHERE tco.constraint_type = 'PRIMARY KEY' " \ + "AND tco.constraint_name = '#{table_name}_pkey'") + + if columns_result.present? + columns = [] + columns_result.values.each do |c| columns.push(c[0]) end + + if columns.length != pkey_columns.length + execute "ALTER TABLE #{table_name} DROP CONSTRAINT IF EXISTS #{table_name}_pkey" + execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)" + end + end + end def create_table(table_name, options = {}, &block) diff --git a/spec/schema.rb b/spec/schema.rb index e141ced9..ac37f39d 100644 --- a/spec/schema.rb +++ b/spec/schema.rb @@ -112,6 +112,13 @@ change_table :project_categories, partition_key: :account_id do |t| end + + # primary key shoud not be recreated + change_table :project_categories, partition_key: :account_id do |t| + t.column :subtitle, :string + end + + create_distributed_table :project_categories, :account_id end From f35c0a410cb017392a213390c127aa7ce6dc15f4 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Thu, 12 Sep 2019 13:45:51 -0700 Subject: [PATCH 03/11] clean query --- lib/activerecord-multi-tenant/migrations.rb | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb index 058c4a26..d54f1f3d 100644 --- a/lib/activerecord-multi-tenant/migrations.rb +++ b/lib/activerecord-multi-tenant/migrations.rb @@ -53,17 +53,18 @@ def change_primary_key(table_name, options) pkey_columns = ['id', options[:partition_key]] # we are here comparing the columns in the primary key on the database and the one in the migration file - columns_result = execute("select kcu.column_name as key_column " \ + query = ActiveRecord::Base::sanitize_sql_array(["select kcu.column_name as key_column " \ "from information_schema.table_constraints tco "\ "join information_schema.key_column_usage kcu " \ "ON kcu.constraint_name = tco.constraint_name " \ "AND kcu.constraint_schema = tco.constraint_schema "\ "WHERE tco.constraint_type = 'PRIMARY KEY' " \ - "AND tco.constraint_name = '#{table_name}_pkey'") + "AND tco.constraint_name = '%s_pkey'", table_name]) + columns_result = execute(query) if columns_result.present? columns = [] - columns_result.values.each do |c| columns.push(c[0]) end + columns_result.values.each {|c| columns.push(c[0])} if columns.length != pkey_columns.length execute "ALTER TABLE #{table_name} DROP CONSTRAINT IF EXISTS #{table_name}_pkey" @@ -80,7 +81,10 @@ def create_table(table_name, options = {}, &block) end def change_table(table_name, options, &block) - ret = orig_change_table(table_name, options.except(:partition_key), &block) + ret = nil + if block + ret = orig_change_table(table_name, options.except(:partition_key), &block) + end change_primary_key(table_name, options) ret end From 855a382ceb5f1db03f37e872546c7452adbc92f3 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Thu, 12 Sep 2019 13:47:07 -0700 Subject: [PATCH 04/11] change spec schema.rb to remove useless block --- spec/schema.rb | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/spec/schema.rb b/spec/schema.rb index ac37f39d..5ae0558f 100644 --- a/spec/schema.rb +++ b/spec/schema.rb @@ -109,9 +109,7 @@ create_distributed_table :allowed_places, :account_id create_reference_table :categories - change_table :project_categories, partition_key: :account_id do |t| - end - + change_table :project_categories, partition_key: :account_id # primary key shoud not be recreated change_table :project_categories, partition_key: :account_id do |t| From 8cf51b04eba6f247ec4f576dd767171197763432 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Thu, 12 Sep 2019 13:50:01 -0700 Subject: [PATCH 05/11] clean code --- lib/activerecord-multi-tenant/migrations.rb | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb index d54f1f3d..a6a088d3 100644 --- a/lib/activerecord-multi-tenant/migrations.rb +++ b/lib/activerecord-multi-tenant/migrations.rb @@ -63,8 +63,7 @@ def change_primary_key(table_name, options) columns_result = execute(query) if columns_result.present? - columns = [] - columns_result.values.each {|c| columns.push(c[0])} + columns = columns_result.values.map(&:first) if columns.length != pkey_columns.length execute "ALTER TABLE #{table_name} DROP CONSTRAINT IF EXISTS #{table_name}_pkey" From 0b011790948fa87df558312ca4161f95dbb5433e Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Thu, 12 Sep 2019 13:51:57 -0700 Subject: [PATCH 06/11] start file to have distribute statement in structure.sql --- lib/activerecord-multi-tenant/tasks/db.rake | 53 +++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 lib/activerecord-multi-tenant/tasks/db.rake diff --git a/lib/activerecord-multi-tenant/tasks/db.rake b/lib/activerecord-multi-tenant/tasks/db.rake new file mode 100644 index 00000000..195bbb9a --- /dev/null +++ b/lib/activerecord-multi-tenant/tasks/db.rake @@ -0,0 +1,53 @@ +Rake::Task["db:structure:dump"].enhance do + printf('coucou') + filenames = [] + filenames << ENV['DB_STRUCTURE'] if ENV.key?('DB_STRUCTURE') + + + if ActiveRecord::VERSION::MAJOR >= 6 + # Based on https://github.com/rails/rails/pull/36560/files + databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml + ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |spec_name| + Rails.application.config.paths['db'].each do |path| + filenames << File.join(path, spec_name + '_structure.sql') + end + end + end + + unless filenames.present? + Rails.application.config.paths['db'].each do |path| + filenames << File.join(path, 'structure.sql') + end + end + + + query_distributed = "SELECT logicalrelid, pg_attribute.attname" \ + "FROM pg_dist_partition" \ + "INNER JOIN pg_attribute ON (logicalrelid=attrelid)" \ + "WHERE logicalrelid::varchar(255) = '{}'" \ + "AND partmethod='h'" \ + "AND attnum=substring(partkey from '%:varattno #\"[0-9]+#\"%' for '#')::int" + + query_reference = "SELECT logicalrelid FROM pg_dist_partition WHERE partmethod = 'n';" + + connection = ActiveRecord::Base.establish_connection(database_urls.first) + distributed_tables = connection.execute(query_distributed) + reference_tables = connection.execute(query_reference) + + schema = '' + + distributed_tables.values.each do |distributed_table| + schema << "SELECT create_distributed_table(%s, %s);\n" % [distributed_table[0], distributed_table[1]] + end + + reference_tables.values.each do |reference_table| + schema << "SELECT create_reference_table(%s);\n" % [reference_table[0]] + end + + + filenames.each do |file| + File.open(file, "w") { |f| f.puts schema } + puts "Added distribute statements to #{file}" + end + +end From f5f2f425870b87bdd4468db3a3c037a78c1aeea6 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Thu, 12 Sep 2019 17:35:50 -0700 Subject: [PATCH 07/11] rake task to add distributed tables and reference tables at the end of the structure.sql --- lib/activerecord-multi-tenant.rb | 1 + lib/activerecord-multi-tenant/railtie.rb | 7 ++ lib/activerecord-multi-tenant/tasks/db.rake | 77 +++++++++++---------- 3 files changed, 49 insertions(+), 36 deletions(-) create mode 100644 lib/activerecord-multi-tenant/railtie.rb diff --git a/lib/activerecord-multi-tenant.rb b/lib/activerecord-multi-tenant.rb index 760126ed..014bbc17 100644 --- a/lib/activerecord-multi-tenant.rb +++ b/lib/activerecord-multi-tenant.rb @@ -10,3 +10,4 @@ require_relative 'activerecord-multi-tenant/query_monitor' require_relative 'activerecord-multi-tenant/version' require_relative 'activerecord-multi-tenant/with_lock' +require_relative 'activerecord-multi-tenant/railtie' if defined?(Rails) diff --git a/lib/activerecord-multi-tenant/railtie.rb b/lib/activerecord-multi-tenant/railtie.rb new file mode 100644 index 00000000..241baf88 --- /dev/null +++ b/lib/activerecord-multi-tenant/railtie.rb @@ -0,0 +1,7 @@ +module ActiveRecordDistributeStatementsStructure + class Railtie < Rails::Railtie + rake_tasks do + load 'activerecord-multi-tenant/tasks/db.rake' + end + end +end diff --git a/lib/activerecord-multi-tenant/tasks/db.rake b/lib/activerecord-multi-tenant/tasks/db.rake index 195bbb9a..51752b08 100644 --- a/lib/activerecord-multi-tenant/tasks/db.rake +++ b/lib/activerecord-multi-tenant/tasks/db.rake @@ -1,53 +1,58 @@ -Rake::Task["db:structure:dump"].enhance do - printf('coucou') - filenames = [] - filenames << ENV['DB_STRUCTURE'] if ENV.key?('DB_STRUCTURE') +require "active_record" +Rake::Task["db:structure:dump"].enhance do if ActiveRecord::VERSION::MAJOR >= 6 - # Based on https://github.com/rails/rails/pull/36560/files databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml - ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |spec_name| - Rails.application.config.paths['db'].each do |path| - filenames << File.join(path, spec_name + '_structure.sql') - end - end - end - - unless filenames.present? - Rails.application.config.paths['db'].each do |path| - filenames << File.join(path, 'structure.sql') - end + else + databases = [ActiveRecord::Tasks::DatabaseTasks.current_config] end + query_distributed = 'SELECT logicalrelid, pg_attribute.attname ' \ + 'FROM pg_dist_partition ' \ + 'INNER JOIN pg_attribute ON (logicalrelid=attrelid) ' \ + 'WHERE partmethod=\'h\' ' \ + 'AND attnum=substring(partkey from \'%:varattno #"[0-9]+#"%\' for \'#\')::int;' - query_distributed = "SELECT logicalrelid, pg_attribute.attname" \ - "FROM pg_dist_partition" \ - "INNER JOIN pg_attribute ON (logicalrelid=attrelid)" \ - "WHERE logicalrelid::varchar(255) = '{}'" \ - "AND partmethod='h'" \ - "AND attnum=substring(partkey from '%:varattno #\"[0-9]+#\"%' for '#')::int" query_reference = "SELECT logicalrelid FROM pg_dist_partition WHERE partmethod = 'n';" - connection = ActiveRecord::Base.establish_connection(database_urls.first) - distributed_tables = connection.execute(query_distributed) - reference_tables = connection.execute(query_reference) - - schema = '' + databases.each do |db_config| + # for versions before 6.0, there will only be 1 database in the list + connection = ActiveRecord::Base.establish_connection(db_config) + filenames = [] + if ActiveRecord::VERSION::MAJOR >= 6 + Rails.application.config.paths['db'].each do |path| + filenames << File.join(path, db_config.spec_name + '_structure.sql') + end + end - distributed_tables.values.each do |distributed_table| - schema << "SELECT create_distributed_table(%s, %s);\n" % [distributed_table[0], distributed_table[1]] - end + unless filenames.present? + Rails.application.config.paths['db'].each do |path| + filenames << File.join(path, 'structure.sql') + end + end - reference_tables.values.each do |reference_table| - schema << "SELECT create_reference_table(%s);\n" % [reference_table[0]] - end + schema = '' + begin + distributed_tables = connection.connection.execute(query_distributed) + reference_tables = connection.connection.execute(query_reference) + rescue + # citus is not installed, failing because pg_dist_partition not found + else + distributed_tables.values.each do |distributed_table| + schema << "SELECT create_distributed_table('%s', '%s');\n" % [distributed_table[0], distributed_table[1]] + end + reference_tables.values.each do |reference_table| + schema << "SELECT create_reference_table('%s');\n" % [reference_table[0]] + end + end - filenames.each do |file| - File.open(file, "w") { |f| f.puts schema } - puts "Added distribute statements to #{file}" + filenames.each do |filename| + File.open(filename, "a") { |f| f.puts schema } + end + puts "Added distribute statements to #{filenames}" end end From 876f95c20a4600b0a75dc55608503dc5d9f3b2bf Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Fri, 13 Sep 2019 13:37:32 -0700 Subject: [PATCH 08/11] add an option to add or not distribute statements --- lib/activerecord-multi-tenant.rb | 1 + .../schema_dumper_extension.rb | 11 +++++++++++ lib/activerecord-multi-tenant/tasks/db.rake | 2 ++ 3 files changed, 14 insertions(+) create mode 100644 lib/activerecord-multi-tenant/schema_dumper_extension.rb diff --git a/lib/activerecord-multi-tenant.rb b/lib/activerecord-multi-tenant.rb index 014bbc17..0ac701dc 100644 --- a/lib/activerecord-multi-tenant.rb +++ b/lib/activerecord-multi-tenant.rb @@ -10,4 +10,5 @@ require_relative 'activerecord-multi-tenant/query_monitor' require_relative 'activerecord-multi-tenant/version' require_relative 'activerecord-multi-tenant/with_lock' +require_relative 'activerecord-multi-tenant/schema_dumper_extension' require_relative 'activerecord-multi-tenant/railtie' if defined?(Rails) diff --git a/lib/activerecord-multi-tenant/schema_dumper_extension.rb b/lib/activerecord-multi-tenant/schema_dumper_extension.rb new file mode 100644 index 00000000..9018f542 --- /dev/null +++ b/lib/activerecord-multi-tenant/schema_dumper_extension.rb @@ -0,0 +1,11 @@ +module MultiTenant + module SchemaDumperExtension + cattr_accessor :include_distribute_statements, default: true + + + end +end + +if defined?(ActiveRecord::SchemaDumper) + ActiveRecord::SchemaDumper.extend(MultiTenant::SchemaDumperExtension) +end diff --git a/lib/activerecord-multi-tenant/tasks/db.rake b/lib/activerecord-multi-tenant/tasks/db.rake index 51752b08..f06cd8bb 100644 --- a/lib/activerecord-multi-tenant/tasks/db.rake +++ b/lib/activerecord-multi-tenant/tasks/db.rake @@ -2,6 +2,8 @@ require "active_record" Rake::Task["db:structure:dump"].enhance do + return unless ActiveRecord::SchemaDumper.include_distribute_statements + if ActiveRecord::VERSION::MAJOR >= 6 databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml else From 0e5db1170cbea0857b5c3f850ee0b1d93d68ec75 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Fri, 13 Sep 2019 13:48:48 -0700 Subject: [PATCH 09/11] compat with rails 4 --- lib/activerecord-multi-tenant/migrations.rb | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/lib/activerecord-multi-tenant/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb index a6a088d3..7f358dae 100644 --- a/lib/activerecord-multi-tenant/migrations.rb +++ b/lib/activerecord-multi-tenant/migrations.rb @@ -53,14 +53,15 @@ def change_primary_key(table_name, options) pkey_columns = ['id', options[:partition_key]] # we are here comparing the columns in the primary key on the database and the one in the migration file - query = ActiveRecord::Base::sanitize_sql_array(["select kcu.column_name as key_column " \ - "from information_schema.table_constraints tco "\ - "join information_schema.key_column_usage kcu " \ - "ON kcu.constraint_name = tco.constraint_name " \ - "AND kcu.constraint_schema = tco.constraint_schema "\ - "WHERE tco.constraint_type = 'PRIMARY KEY' " \ - "AND tco.constraint_name = '%s_pkey'", table_name]) - columns_result = execute(query) + query = "select kcu.column_name as key_column " \ + "from information_schema.table_constraints tco "\ + "join information_schema.key_column_usage kcu " \ + "ON kcu.constraint_name = tco.constraint_name " \ + "AND kcu.constraint_schema = tco.constraint_schema "\ + "WHERE tco.constraint_type = 'PRIMARY KEY' " \ + "AND tco.constraint_name = '%s_pkey'" + + columns_result = execute(ActiveRecord::Base.send(:sanitize_sql_array, [query, table_name])) if columns_result.present? columns = columns_result.values.map(&:first) From 6e7d33e17801d503dfd43dae3d5874fea8cd8569 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Fri, 13 Sep 2019 13:51:21 -0700 Subject: [PATCH 10/11] clean code --- lib/activerecord-multi-tenant/schema_dumper_extension.rb | 2 -- 1 file changed, 2 deletions(-) diff --git a/lib/activerecord-multi-tenant/schema_dumper_extension.rb b/lib/activerecord-multi-tenant/schema_dumper_extension.rb index 9018f542..f6ea566a 100644 --- a/lib/activerecord-multi-tenant/schema_dumper_extension.rb +++ b/lib/activerecord-multi-tenant/schema_dumper_extension.rb @@ -1,8 +1,6 @@ module MultiTenant module SchemaDumperExtension cattr_accessor :include_distribute_statements, default: true - - end end From 419f234475bf85d9d94e594adbf2f5da23da5bb6 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Tue, 17 Sep 2019 10:32:22 -0700 Subject: [PATCH 11/11] split task to add unit tests --- .../schema_dumper_extension.rb | 49 ++++++++++ lib/activerecord-multi-tenant/tasks/db.rake | 30 +----- .../schema_dumper_spec.rb | 96 +++++++++++++++++++ 3 files changed, 148 insertions(+), 27 deletions(-) create mode 100644 spec/activerecord-multi-tenant/schema_dumper_spec.rb diff --git a/lib/activerecord-multi-tenant/schema_dumper_extension.rb b/lib/activerecord-multi-tenant/schema_dumper_extension.rb index f6ea566a..49cf5817 100644 --- a/lib/activerecord-multi-tenant/schema_dumper_extension.rb +++ b/lib/activerecord-multi-tenant/schema_dumper_extension.rb @@ -1,6 +1,55 @@ module MultiTenant module SchemaDumperExtension cattr_accessor :include_distribute_statements, default: true + + def get_distributed_tables(connection) + query_distributed = 'SELECT logicalrelid, pg_attribute.attname ' \ + 'FROM pg_dist_partition ' \ + 'INNER JOIN pg_attribute ON (logicalrelid=attrelid) ' \ + 'WHERE partmethod=\'h\' ' \ + 'AND attnum=substring(partkey from \'%:varattno #"[0-9]+#"%\' for \'#\')::int ' \ + 'ORDER BY logicalrelid' + + begin + return connection.execute(query_distributed).values + rescue; end + end + + def get_reference_tables(connection) + query_reference = "SELECT logicalrelid FROM pg_dist_partition WHERE partmethod = 'n' ORDER BY logicalrelid" + begin + return connection.execute(query_reference).values + rescue; end + end + + def get_distribute_statements(connection, reference=false) + if reference + distributed_tables = get_reference_tables(connection) + query = "SELECT create_reference_table('%s');\n" + else + distributed_tables = get_distributed_tables(connection) + query = "SELECT create_distributed_table('%s', '%s');\n" + end + + return unless distributed_tables + + schema = '' + distributed_tables.each do |distributed_table| + attrs = if reference then [distributed_table[0]] else [distributed_table[0], distributed_table[1]] end + schema << query % attrs + end + + schema + end + + def get_full_distribute_statements(connection) + schema = ActiveRecord::SchemaDumper.get_distribute_statements(connection) || '' + schema << (ActiveRecord::SchemaDumper.get_distribute_statements(connection, + reference=true) || '') + + schema + end + end end diff --git a/lib/activerecord-multi-tenant/tasks/db.rake b/lib/activerecord-multi-tenant/tasks/db.rake index f06cd8bb..d9e6c81b 100644 --- a/lib/activerecord-multi-tenant/tasks/db.rake +++ b/lib/activerecord-multi-tenant/tasks/db.rake @@ -2,7 +2,7 @@ require "active_record" Rake::Task["db:structure:dump"].enhance do - return unless ActiveRecord::SchemaDumper.include_distribute_statements + next unless ActiveRecord::SchemaDumper.include_distribute_statements if ActiveRecord::VERSION::MAJOR >= 6 databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml @@ -10,18 +10,9 @@ Rake::Task["db:structure:dump"].enhance do databases = [ActiveRecord::Tasks::DatabaseTasks.current_config] end - query_distributed = 'SELECT logicalrelid, pg_attribute.attname ' \ - 'FROM pg_dist_partition ' \ - 'INNER JOIN pg_attribute ON (logicalrelid=attrelid) ' \ - 'WHERE partmethod=\'h\' ' \ - 'AND attnum=substring(partkey from \'%:varattno #"[0-9]+#"%\' for \'#\')::int;' - - - query_reference = "SELECT logicalrelid FROM pg_dist_partition WHERE partmethod = 'n';" - databases.each do |db_config| # for versions before 6.0, there will only be 1 database in the list - connection = ActiveRecord::Base.establish_connection(db_config) + connection = ActiveRecord::Base.establish_connection(db_config).connection filenames = [] if ActiveRecord::VERSION::MAJOR >= 6 Rails.application.config.paths['db'].each do |path| @@ -35,26 +26,11 @@ Rake::Task["db:structure:dump"].enhance do end end - schema = '' - begin - distributed_tables = connection.connection.execute(query_distributed) - reference_tables = connection.connection.execute(query_reference) - rescue - # citus is not installed, failing because pg_dist_partition not found - else - distributed_tables.values.each do |distributed_table| - schema << "SELECT create_distributed_table('%s', '%s');\n" % [distributed_table[0], distributed_table[1]] - end - - reference_tables.values.each do |reference_table| - schema << "SELECT create_reference_table('%s');\n" % [reference_table[0]] - end - end + schema = ActiveRecord::SchemaDumper.get_full_distribute_statements(connection) filenames.each do |filename| File.open(filename, "a") { |f| f.puts schema } end puts "Added distribute statements to #{filenames}" end - end diff --git a/spec/activerecord-multi-tenant/schema_dumper_spec.rb b/spec/activerecord-multi-tenant/schema_dumper_spec.rb new file mode 100644 index 00000000..262e90fb --- /dev/null +++ b/spec/activerecord-multi-tenant/schema_dumper_spec.rb @@ -0,0 +1,96 @@ +require 'spec_helper' +require 'rake' + + +describe 'Schema Dumper enhancement' do + let(:file_like_object) { double("file like object") } + + it 'should list distributed tables' do + distributed_tables = ActiveRecord::SchemaDumper.get_distributed_tables(ActiveRecord::Base.connection) + + distributed_result = [["accounts", "id"], + ["projects", "account_id"], + ["managers", "account_id"], + ["tasks", "account_id"], + ["sub_tasks", "account_id"], + ["aliased_tasks", "account_id"], + ["custom_partition_key_tasks", "accountID"], + ["comments", "account_id"], + ["partition_key_not_model_tasks", "non_model_id"], + ["subclass_tasks", "non_model_id"], + ["uuid_records", "organization_id"], + ["allowed_places", "account_id"], + ["project_categories", "account_id"]] + + expect(distributed_tables.to_set).to eq(distributed_result.to_set) + end + + it 'should list reference tables' do + reference_tables = ActiveRecord::SchemaDumper.get_reference_tables(ActiveRecord::Base.connection) + reference_result = [["categories"]] + expect(reference_tables.to_set).to eq(reference_result.to_set) + end + + it 'distribute statements' do + distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq("SELECT create_distributed_table('accounts', 'id');\nSELECT create_distributed_table('projects', 'account_id');\nSELECT create_distributed_table('managers', 'account_id');\nSELECT create_distributed_table('tasks', 'account_id');\nSELECT create_distributed_table('sub_tasks', 'account_id');\nSELECT create_distributed_table('aliased_tasks', 'account_id');\nSELECT create_distributed_table('custom_partition_key_tasks', 'accountID');\nSELECT create_distributed_table('comments', 'account_id');\nSELECT create_distributed_table('partition_key_not_model_tasks', 'non_model_id');\nSELECT create_distributed_table('subclass_tasks', 'non_model_id');\nSELECT create_distributed_table('uuid_records', 'organization_id');\nSELECT create_distributed_table('project_categories', 'account_id');\nSELECT create_distributed_table('allowed_places', 'account_id');\n") + end + + it 'reference tables statements' do + distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection, reference=true) + expect(distributed_statements).to eq("SELECT create_reference_table('categories');\n") + end + + it 'no distributed tables' do + ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {[]} + distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq("") + end + + it 'no citus metadata tables' do + ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {nil} + distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq(nil) + end + + it 'no reference tables' do + ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {[]} + distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection, reference=true) + expect(distributed_statements).to eq("") + + end + + it 'no citus metadata tables for reference' do + ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {nil} + distributed_statements = ActiveRecord::SchemaDumper.get_distribute_statements(ActiveRecord::Base.connection, reference=true) + expect(distributed_statements).to eq(nil) + end + + + it 'full statements' do + distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq("SELECT create_distributed_table('accounts', 'id');\nSELECT create_distributed_table('projects', 'account_id');\nSELECT create_distributed_table('managers', 'account_id');\nSELECT create_distributed_table('tasks', 'account_id');\nSELECT create_distributed_table('sub_tasks', 'account_id');\nSELECT create_distributed_table('aliased_tasks', 'account_id');\nSELECT create_distributed_table('custom_partition_key_tasks', 'accountID');\nSELECT create_distributed_table('comments', 'account_id');\nSELECT create_distributed_table('partition_key_not_model_tasks', 'non_model_id');\nSELECT create_distributed_table('subclass_tasks', 'non_model_id');\nSELECT create_distributed_table('uuid_records', 'organization_id');\nSELECT create_distributed_table('project_categories', 'account_id');\nSELECT create_distributed_table('allowed_places', 'account_id');\nSELECT create_reference_table('categories');\n") + + end + + it 'full statements no reference' do + ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {[]} + distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq("SELECT create_distributed_table('accounts', 'id');\nSELECT create_distributed_table('projects', 'account_id');\nSELECT create_distributed_table('managers', 'account_id');\nSELECT create_distributed_table('tasks', 'account_id');\nSELECT create_distributed_table('sub_tasks', 'account_id');\nSELECT create_distributed_table('aliased_tasks', 'account_id');\nSELECT create_distributed_table('custom_partition_key_tasks', 'accountID');\nSELECT create_distributed_table('comments', 'account_id');\nSELECT create_distributed_table('partition_key_not_model_tasks', 'non_model_id');\nSELECT create_distributed_table('subclass_tasks', 'non_model_id');\nSELECT create_distributed_table('uuid_records', 'organization_id');\nSELECT create_distributed_table('project_categories', 'account_id');\nSELECT create_distributed_table('allowed_places', 'account_id');\n") + end + + it 'full statements no distributed' do + ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {nil} + distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq("SELECT create_reference_table('categories');\n") + end + + it 'full statements no citus' do + ActiveRecord::SchemaDumper.stub(:get_distributed_tables).with(anything()) {nil} + ActiveRecord::SchemaDumper.stub(:get_reference_tables).with(anything()) {nil} + + distributed_statements = ActiveRecord::SchemaDumper.get_full_distribute_statements(ActiveRecord::Base.connection) + expect(distributed_statements).to eq("") + + end +end