From 419f234475bf85d9d94e594adbf2f5da23da5bb6 Mon Sep 17 00:00:00 2001 From: Louise Grandjonc Date: Tue, 17 Sep 2019 10:32:22 -0700 Subject: [PATCH] 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