From 987337aa2d6d1805c17b4dfe43a7fc21806b04f7 Mon Sep 17 00:00:00 2001
From: Louise Grandjonc <louve.grandjonc@gmail.com>
Date: Wed, 11 Sep 2019 14:39:48 -0700
Subject: [PATCH] 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