Handling primary key type in generators (based on Devise implementation)

Author: whois-marvin-42

lib/generators/devise_token_auth/install_generator.rb

@@ -4,6 +4,8 @@ module DeviseTokenAuth
   class InstallGenerator < Rails::Generators::Base
     include Rails::Generators::Migration
 
+    class_option :primary_key_type, type: :string, desc: 'The type for primary key'
+
     source_root File.expand_path('templates', __dir__)
 
     argument :user_class, type: :string, default: 'User'
@@ -156,5 +158,18 @@ def database_name
     def database_version
       ActiveRecord::Base.connection.select_value('SELECT VERSION()')
     end
+
+    def rails5?
+      Rails.version.start_with? '5'
+    end
+
+    def primary_key_type
+      primary_key_string if rails5?
+    end
+
+    def primary_key_string
+      key_string = options[:primary_key_type]
+      ", id: :#{key_string}" if key_string
+    end
   end
 end

lib/generators/devise_token_auth/templates/devise_token_auth_create_users.rb.erb

@@ -1,7 +1,7 @@
 class DeviseTokenAuthCreate<%= user_class.pluralize.gsub("::","") %> < ActiveRecord::Migration<%= "[#{Rails::VERSION::STRING[0..2]}]" if Rails::VERSION::MAJOR > 4 %>
   def change
     <% table_name = @user_class.pluralize.gsub("::","").underscore %>
-    create_table(:<%= table_name %>) do |t|
+    create_table(:<%= table_name %><%= primary_key_type %>) do |t|
       ## Required
       t.string :provider, :null => false, :default => "email"
       t.string :uid, :null => false, :default => ""

test/lib/generators/devise_token_auth/install_generator_test.rb

@@ -41,6 +41,15 @@ class InstallGeneratorTest < Rails::Generators::TestCase
       test 'subsequent runs raise no errors' do
         run_generator
       end
+
+      test 'add primary key type with rails 5 when specified in rails generator' do
+        run_generator %w[--primary_key_type=uuid --force]
+        if Rails::VERSION::MAJOR >= 5
+          assert_migration 'db/migrate/devise_token_auth_create_users.rb', /create_table\(:users, id: :uuid\) do/
+        else
+          assert_migration 'db/migrate/devise_token_auth_create_users.rb', /create_table\(:users\) do/
+        end
+      end
     end
 
     describe 'existing user model' do

test/lib/generators/devise_token_auth/install_generator_with_namespace_test.rb

@@ -46,6 +46,15 @@ class InstallGeneratorTest < Rails::Generators::TestCase
       test 'subsequent runs raise no errors' do
         run_generator %W[#{user_class} auth]
       end
+
+      test 'add primary key type with rails 5 when specified in rails generator' do
+        run_generator %W[#{user_class} auth --primary_key_type=uuid --force]
+        if Rails::VERSION::MAJOR >= 5
+          assert_migration "db/migrate/devise_token_auth_create_#{table_name}.rb", /create_table\(:#{table_name}, id: :uuid\) do/
+        else
+          assert_migration "db/migrate/devise_token_auth_create_#{table_name}.rb", /create_table\(:#{table_name}\) do/
+        end
+      end
     end
 
     describe 'existing user model' do