From 2b431b650fac1dd26091bb792f4370fc15f8a77f Mon Sep 17 00:00:00 2001
From: Peter Solnica <peter@solnica.online>
Date: Fri, 19 Jan 2024 10:22:13 +0000
Subject: [PATCH] Support UUID as PKs in associations

---
 lib/rom/factory/builder/persistable.rb | 18 ++++++-
 spec/integration/rom/factory_spec.rb   | 72 ++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/lib/rom/factory/builder/persistable.rb b/lib/rom/factory/builder/persistable.rb
index 56a850f..c104b07 100644
--- a/lib/rom/factory/builder/persistable.rb
+++ b/lib/rom/factory/builder/persistable.rb
@@ -42,13 +42,29 @@ def create(*traits, **attrs)
 
         # @api private
         def persist(attrs)
-          relation.with(auto_struct: !tuple_evaluator.has_associations?).command(:create).call(attrs)
+          result = relation
+            .with(auto_struct: !tuple_evaluator.has_associations?)
+            .command(:create)
+            .call(attrs)
+
+          # Handle PK values generated by the factory
+          if pk? && (pks = attrs.values_at(*primary_key_names)).compact.size == primary_key_names.size
+            relation.by_pk(*pks).one!
+          elsif result
+            result
+          else
+            relation.where(attrs).one!
+          end
         end
 
         # @api private
         def primary_key_names
           relation.schema.primary_key.map(&:name)
         end
+
+        def pk?
+          primary_key_names.any?
+        end
       end
     end
   end
diff --git a/spec/integration/rom/factory_spec.rb b/spec/integration/rom/factory_spec.rb
index eb57120..14e55be 100644
--- a/spec/integration/rom/factory_spec.rb
+++ b/spec/integration/rom/factory_spec.rb
@@ -1436,4 +1436,76 @@ class User < ROM::Struct
       expect(user.age).to be_between(0, 150)
     end
   end
+
+  describe "using UUID as PKs" do
+    context "many-to-one" do
+      let(:rom) do
+        ROM.container(:sql, conn) do |conf|
+          conf.default.create_table(:workers) do
+            column :id, :uuid, default: Sequel.function(:uuid_generate_v4), primary_key: true
+            column :name, :string
+          end
+
+          conf.default.create_table(:jobs) do
+            column :id, :uuid, default: Sequel.function(:uuid_generate_v4), primary_key: true
+            foreign_key :worker_id, :workers, type: :uuid
+            column :name, :string
+          end
+
+          conf.relation(:workers) do
+            schema(infer: true) do
+              associations do
+                has_many :jobs
+              end
+            end
+          end
+
+          conf.relation(:jobs) do
+            schema(infer: true) do
+              attribute :worker_id, ROM::SQL::Types.ForeignKey(:workers, ROM::SQL::Types::String)
+
+              associations do
+                belongs_to :worker
+              end
+            end
+          end
+        end
+      end
+
+      before do
+        factories.define(:worker) do |f|
+          f.id { fake(:internet, :uuid) }
+          f.name "Test Worker"
+
+          f.association(:jobs, count: 0)
+
+          f.trait(:with_jobs) do |t|
+            t.association(:jobs, count: 2)
+          end
+        end
+
+        factories.define(:job) do |f|
+          f.id { fake(:internet, :uuid) }
+          f.name "Test Job"
+
+          f.association :worker
+        end
+      end
+
+      it "persists a parent struct" do
+        worker = factories[:worker]
+
+        expect(worker.id).to_not be(nil)
+        expect(worker.name).to eql("Test Worker")
+      end
+
+      it "persists a parent struct with its children" do
+        worker = factories[:worker, :with_jobs]
+
+        expect(worker.id).to_not be(nil)
+        expect(worker.name).to eql("Test Worker")
+        expect(worker.jobs.size).to be(2)
+      end
+    end
+  end
 end