From d4ecaf5859af414af95489d40565c1e71a9195d6 Mon Sep 17 00:00:00 2001
From: Matt Larraz <mlarraz@users.noreply.github.com>
Date: Fri, 21 Mar 2025 12:24:49 -0400
Subject: [PATCH] Add YAML loading options to mirror ActiveRecord behavior

---
 lib/dynamoid/config.rb    |  2 ++
 lib/dynamoid/undumping.rb | 28 ++++++++++++++++++----------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/lib/dynamoid/config.rb b/lib/dynamoid/config.rb
index 31ccf4bd..c0180f9c 100644
--- a/lib/dynamoid/config.rb
+++ b/lib/dynamoid/config.rb
@@ -62,6 +62,8 @@ module Config
     option :http_open_timeout, default: nil     #                                                  - default: 15
     option :http_read_timeout, default: nil     #                                                  - default: 60
     option :create_table_on_save, default: true
+    option :use_yaml_unsafe_load, default: false
+    option :yaml_column_permitted_classes, default: [Symbol, Set, Date, Time, DateTime] # classes to allow when using YAML.safe_load
 
     # The default logger for Dynamoid: either the Rails logger or just stdout.
     #
diff --git a/lib/dynamoid/undumping.rb b/lib/dynamoid/undumping.rb
index 094d9d3b..09274911 100644
--- a/lib/dynamoid/undumping.rb
+++ b/lib/dynamoid/undumping.rb
@@ -236,18 +236,13 @@ def process(value)
     end
 
     class SerializedUndumper < Base
-      # We must use YAML.safe_load in Ruby 3.1 to handle serialized Set class
-      minimum_ruby_version = ->(version) { Gem::Version.new(RUBY_VERSION) >= Gem::Version.new(version) }
-
-      # Once we drop support for Rubies older than 2.6 we can remove this condition (with major version bump)!
-      # YAML_SAFE_LOAD = minimum_ruby_version.call("2.6")
-      # But we don't want to change behavior for Ruby <= 3.0 that has been using the gem, without a major version bump
-      YAML_SAFE_LOAD = minimum_ruby_version.call('3.1')
+      # .safe_load is the default behavior after this version
+      YAML_SAFE_LOAD = ::YAML::VERSION >= Gem::Version.new('4.0')
 
       def process(value)
         if @options[:serializer]
           @options[:serializer].load(value)
-        elsif YAML_SAFE_LOAD
+        elsif YAML_SAFE_LOAD && !Dynamoid::Config.use_yaml_unsafe_load
           # The classes listed in permitted classes are added to the default set of "safe loadable" classes.
           # TrueClass
           # FalseClass
@@ -257,9 +252,22 @@ def process(value)
           # String
           # Array
           # Hash
-          YAML.safe_load(value, permitted_classes: [Symbol, Set, Date, Time, DateTime])
+          ::YAML.safe_load(value, permitted_classes: Dynamoid::Config.yaml_column_permitted_classes)
         else
-          YAML.load(value)
+          unsafe_load(value)
+        end
+      end
+
+      private
+
+      # Should behave the same in both cases
+      if ::YAML.respond_to?(:unsafe_load)
+        def unsafe_load(value)
+          ::YAML.unsafe_load(value)
+        end
+      else
+        def unsafe_load(value)
+          ::YAML.load(value)
         end
       end
     end