diff --git a/lib/activerecord-multi-tenant.rb b/lib/activerecord-multi-tenant.rb index 760126ed..0ac701dc 100644 --- a/lib/activerecord-multi-tenant.rb +++ b/lib/activerecord-multi-tenant.rb @@ -10,3 +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/migrations.rb b/lib/activerecord-multi-tenant/migrations.rb index b6104353..7f358dae 100644 --- a/lib/activerecord-multi-tenant/migrations.rb +++ b/lib/activerecord-multi-tenant/migrations.rb @@ -43,12 +43,49 @@ module ActiveRecord module ConnectionAdapters # :nodoc: module SchemaStatements alias :orig_create_table :create_table + alias :orig_change_table :change_table + + def change_primary_key(table_name, options) + 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 + 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) + + 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) ret = orig_create_table(table_name, options.except(:partition_key), &block) - 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]}\")" + change_primary_key(table_name, options) + ret + end + + def change_table(table_name, options, &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 end 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/schema_dumper_extension.rb b/lib/activerecord-multi-tenant/schema_dumper_extension.rb new file mode 100644 index 00000000..49cf5817 --- /dev/null +++ b/lib/activerecord-multi-tenant/schema_dumper_extension.rb @@ -0,0 +1,58 @@ +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 + +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 new file mode 100644 index 00000000..d9e6c81b --- /dev/null +++ b/lib/activerecord-multi-tenant/tasks/db.rake @@ -0,0 +1,36 @@ +require "active_record" + + +Rake::Task["db:structure:dump"].enhance do + next unless ActiveRecord::SchemaDumper.include_distribute_statements + + if ActiveRecord::VERSION::MAJOR >= 6 + databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml + else + databases = [ActiveRecord::Tasks::DatabaseTasks.current_config] + end + + 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 + 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 + + unless filenames.present? + Rails.application.config.paths['db'].each do |path| + filenames << File.join(path, 'structure.sql') + 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 diff --git a/spec/schema.rb b/spec/schema.rb index 04ab6880..5ae0558f 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,18 @@ 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 + + # 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 class Account < ActiveRecord::Base