Merge pull request #2007 from alex-tan/fix-composite-key-types

Fix types for composite primary keys
Shopify · Aug 30, 2024 · e274c5e · e274c5e
2 parents b6db17d + a5c0082
commit e274c5e
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 4 deletions.
diff --git a/lib/tapioca/dsl/helpers/active_record_column_type_helper.rb b/lib/tapioca/dsl/helpers/active_record_column_type_helper.rb
@@ -88,11 +88,18 @@ def type_for(attribute_name, column_name = attribute_name)
         sig { returns([String, String]) }
         def id_type
           if @constant.respond_to?(:composite_primary_key?) && T.unsafe(@constant).composite_primary_key?
-            @constant.primary_key.map do |column|
-              column_type_for(column)
-            end.map do |tuple|
-              "[#{tuple.join(", ")}]"
+            primary_key_columns = @constant.primary_key
+
+            getters = []
+            setters = []
+
+            primary_key_columns.each do |column|
+              getter, setter = column_type_for(column)
+              getters << getter
+              setters << setter
             end
+
+            ["[#{getters.join(", ")}]", "[#{setters.join(", ")}]"]
           else
             column_type_for(@constant.primary_key)
           end

diff --git a/spec/tapioca/dsl/compilers/active_record_columns_spec.rb b/spec/tapioca/dsl/compilers/active_record_columns_spec.rb
@@ -239,6 +239,33 @@ def body?; end
                 assert_includes(rbi_for(:Post), expected)
               end
 
+              it "handles composite primary keys" do
+                if rails_version(">= 7.1")
+                  add_ruby_file("schema.rb", <<~RUBY)
+                    ActiveRecord::Migration.suppress_messages do
+                      ActiveRecord::Schema.define do
+                        create_table :posts, primary_key: [:a, :b] do |t|
+                          t.string :a, null: false
+                          t.integer :b, null: false
+                        end
+                      end
+                    end
+                  RUBY
+
+                  add_ruby_file("post.rb", <<~RUBY)
+                    class Post < ActiveRecord::Base
+                    end
+                  RUBY
+
+                  expected = indented(<<~RBI, 4)
+                    sig { returns([::String, ::Integer]) }
+                    def id; end
+                  RBI
+
+                  assert_includes(rbi_for(:Post), expected)
+                end
+              end
+
               it "uses ActiveModel::Type::Value types when inheriting from EncryptedAttributeType" do
                 add_ruby_file("schema.rb", <<~RUBY)
                   ActiveRecord::Migration.suppress_messages do