Improve support for primary keys in migrations

config/initializers/migrate.rb

@@ -1,50 +1,57 @@
 if defined?(ActiveRecord::ConnectionAdapters::AbstractAdapter)
   module OpenStreetMap
-    module ActiveRecord
-      module PostgreSQLAdapter
-        def quote_column_name(name)
-          Array(name).map { |n| super(n) }.join(", ")
-        end
+    module SchemaStatements
+      def add_primary_key(table_name, column_name, options = {})
+        constraint_name = "#{table_name}_pkey"
 
-        def add_primary_key(table_name, column_name, _options = {})
-          table_name = quote_table_name(table_name)
-          column_name = quote_column_name(column_name)
+        options = options.merge(:name => constraint_name, :unique => true)
 
-          execute "ALTER TABLE #{table_name} ADD PRIMARY KEY (#{column_name})"
-        end
+        add_index(table_name, column_name, **options)
+        set_primary_key table_name, constraint_name
+      end
 
-        def remove_primary_key(table_name)
-          table_name = quote_table_name(table_name)
+      def remove_primary_key(table_name)
+        constraint_name = quote_table_name("#{table_name}_pkey")
+        table_name = quote_table_name(table_name)
 
-          execute "ALTER TABLE #{table_name} DROP PRIMARY KEY"
-        end
+        execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{constraint_name}"
+      end
 
-        def alter_primary_key(table_name, new_columns)
-          constraint_name = quote_table_name("#{table_name}_pkey")
-          table_name = quote_table_name(table_name)
-          new_columns = quote_column_name(new_columns)
+      def alter_primary_key(table_name, column_name, options = {})
+        constraint_name = "#{table_name}_pkey"
+        tmp_constraint_name = "#{table_name}_pkey_tmp"
 
-          execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{constraint_name}"
-          execute "ALTER TABLE #{table_name} ADD PRIMARY KEY (#{new_columns})"
-        end
+        options = options.merge(:name => tmp_constraint_name)
 
-        def create_enumeration(enumeration_name, values)
-          execute "CREATE TYPE #{enumeration_name} AS ENUM ('#{values.join '\',\''}')"
-        end
+        add_index(table_name, column_name, **options)
+        remove_primary_key table_name
+        set_primary_key table_name, tmp_constraint_name
+        rename_index table_name, tmp_constraint_name, constraint_name
+      end
 
-        def drop_enumeration(enumeration_name)
-          execute "DROP TYPE #{enumeration_name}"
-        end
+      def set_primary_key(table_name, constraint_name)
+        constraint_name = quote_table_name(constraint_name)
+        table_name = quote_table_name(table_name)
+
+        execute "ALTER TABLE #{table_name} ADD PRIMARY KEY USING INDEX #{constraint_name}"
+      end
+
+      def create_enumeration(enumeration_name, values)
+        execute "CREATE TYPE #{enumeration_name} AS ENUM ('#{values.join '\',\''}')"
+      end
+
+      def drop_enumeration(enumeration_name)
+        execute "DROP TYPE #{enumeration_name}"
+      end
 
-        def rename_enumeration(old_name, new_name)
-          old_name = quote_table_name(old_name)
-          new_name = quote_table_name(new_name)
+      def rename_enumeration(old_name, new_name)
+        old_name = quote_table_name(old_name)
+        new_name = quote_table_name(new_name)
 
-          execute "ALTER TYPE #{old_name} RENAME TO #{new_name}"
-        end
+        execute "ALTER TYPE #{old_name} RENAME TO #{new_name}"
       end
     end
   end
 
-  ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(OpenStreetMap::ActiveRecord::PostgreSQLAdapter)
+  ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaStatements.prepend(OpenStreetMap::SchemaStatements)
 end