Allow to do a change_table to change partition key

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)

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

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

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

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

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

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