Support Composite Primary keys in ActiveRecordColumns compiler

lib/tapioca/dsl/compilers/active_record_columns.rb

@@ -110,8 +110,8 @@ def decorate
 
           root.create_path(constant) do |model|
             model.create_module(AttributeMethodsModuleName) do |mod|
-              constant.attribute_names.each do |column_name|
-                add_methods_for_attribute(mod, column_name)
+              (constant.attribute_names + ["id"]).uniq.each do |attribute_name|
+                add_methods_for_attribute(mod, attribute_name)
               end
 
               constant.attribute_aliases.each do |attribute_name, column_name|
@@ -127,7 +127,7 @@ def decorate
                 old_method_names = patterns.map { |m| m.method_name(column_name) }
                 methods_to_add = new_method_names - old_method_names
 
-                add_methods_for_attribute(mod, column_name, attribute_name, methods_to_add)
+                add_methods_for_attribute(mod, attribute_name, column_name, methods_to_add)
               end
             end
 
@@ -166,13 +166,15 @@ def add_method(klass, name, methods_to_add, return_type: "void", parameters: [])
         sig do
           params(
             klass: RBI::Scope,
-            column_name: String,
             attribute_name: String,
+            column_name: String,
             methods_to_add: T.nilable(T::Array[String]),
           ).void
         end
-        def add_methods_for_attribute(klass, column_name, attribute_name = column_name, methods_to_add = nil)
-          getter_type, setter_type = Helpers::ActiveRecordColumnTypeHelper.new(constant).type_for(column_name)
+        def add_methods_for_attribute(klass, attribute_name, column_name = attribute_name, methods_to_add = nil)
+          getter_type, setter_type = Helpers::ActiveRecordColumnTypeHelper
+            .new(constant)
+            .type_for(attribute_name, column_name)
 
           # Added by ActiveRecord::AttributeMethods::Read
           #

lib/tapioca/dsl/helpers/active_record_column_type_helper.rb

@@ -13,8 +13,26 @@ def initialize(constant)
           @constant = constant
         end
 
+        sig { params(attribute_name: String, column_name: String).returns([String, String]) }
+        def type_for(attribute_name, column_name = attribute_name)
+          return id_type if attribute_name == "id"
+
+          column_type_for(column_name)
+        end
+
+        private
+
+        sig { returns([String, String]) }
+        def id_type
+          if @constant.respond_to?(:composite_primary_key?) && T.unsafe(@constant).composite_primary_key?
+            @constant.primary_key.map(&method(:column_type_for)).map { |tuple| "[#{tuple.join(", ")}]" }
+          else
+            column_type_for(@constant.primary_key)
+          end
+        end
+
         sig { params(column_name: String).returns([String, String]) }
-        def type_for(column_name)
+        def column_type_for(column_name)
           return ["T.untyped", "T.untyped"] if do_not_generate_strong_types?(@constant)
 
           column = @constant.columns_hash[column_name]
@@ -33,7 +51,7 @@ def type_for(column_name)
             return [getter_type, as_nilable_type(setter_type)]
           end
 
-          if column_name == @constant.primary_key ||
+          if Array(@constant.primary_key).include?(column_name) ||
               column_name == "created_at" ||
               column_name == "updated_at"
             getter_type = as_nilable_type(getter_type)
@@ -42,8 +60,6 @@ def type_for(column_name)
           [getter_type, setter_type]
         end
 
-        private
-
         sig { params(column_type: T.untyped).returns(String) }
         def type_for_activerecord_value(column_type)
           case column_type

lib/tapioca/helpers/test/template.rb

@@ -24,12 +24,12 @@ def rails_version(selector)
           ::Gem::Requirement.new(selector).satisfied_by?(ActiveSupport.gem_version)
         end
 
-        sig { params(src: String).returns(String) }
-        def template(src)
+        sig { params(src: String, trim_mode: String).returns(String) }
+        def template(src, trim_mode: ">")
           erb = if ERB_SUPPORTS_KVARGS
-            ::ERB.new(src, trim_mode: ">")
+            ::ERB.new(src, trim_mode: trim_mode)
           else
-            ::ERB.new(src, nil, ">")
+            ::ERB.new(src, nil, trim_mode)
           end
 
           erb.result(binding)