allow to do a change_table to change partition key

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} DROP CONSTRAINT IF EXISTS #{table_name}_pkey"
           execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(id, \"#{options[:partition_key]}\")"
         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

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