Add Returning Attributes for UPDATEs

tbuehlmann · tbuehlmann · commit c2872959920d · 2018-05-21T11:24:27.000+02:00
… and add specs.
diff --git a/Gemfile b/Gemfile
@@ -2,5 +2,3 @@ source 'https://rubygems.org'
 
 git_source(:github) { |repo| "https://github.com/#{repo}.git" }
 gemspec
-
-gem 'pry'
diff --git a/Gemfile.lock b/Gemfile.lock
@@ -2,8 +2,8 @@ PATH
   remote: .
   specs:
     active_record-returning_attributes (0.1.0)
-      activerecord
-      activesupport
+      activerecord (= 5.2.0)
+      pg
 
 GEM
   remote: https://rubygems.org/
@@ -27,9 +27,13 @@ GEM
       concurrent-ruby (~> 1.0)
     method_source (0.9.0)
     minitest (5.11.3)
+    pg (1.0.0)
     pry (0.11.3)
       coderay (~> 1.1.0)
       method_source (~> 0.9.0)
+    pry-doc (0.13.4)
+      pry (~> 0.11)
+      yard (~> 0.9.11)
     rake (10.5.0)
     rspec (3.7.0)
       rspec-core (~> 3.7.0)
@@ -47,6 +51,7 @@ GEM
     thread_safe (0.3.6)
     tzinfo (1.2.5)
       thread_safe (~> 0.1)
+    yard (0.9.12)
 
 PLATFORMS
   ruby
@@ -55,6 +60,7 @@ DEPENDENCIES
   active_record-returning_attributes!
   bundler (~> 1.16)
   pry
+  pry-doc
   rake (~> 10.0)
   rspec (~> 3.0)
 
diff --git a/activerecord-returning_attributes.gemspec b/activerecord-returning_attributes.gemspec
@@ -18,11 +18,12 @@ Gem::Specification.new do |spec|
   spec.executables   = spec.files.grep(%r{^exe/}) { |f| File.basename(f) }
   spec.require_paths = ['lib']
 
-  spec.add_runtime_dependency 'activesupport'
-  spec.add_runtime_dependency 'activerecord'
+  spec.add_runtime_dependency 'activerecord', '5.2.0'
+  spec.add_runtime_dependency 'pg'
 
   spec.add_development_dependency 'bundler', '~> 1.16'
   spec.add_development_dependency 'rake', '~> 10.0'
   spec.add_development_dependency 'rspec', '~> 3.0'
   spec.add_development_dependency 'pry'
+  spec.add_development_dependency 'pry-doc'
 end
diff --git a/lib/active_record/returning_attributes.rb b/lib/active_record/returning_attributes.rb
@@ -1,84 +1,43 @@
-require_relative 'returning_attributes/accessor'
+require_relative 'returning_attributes/patching'
 require_relative 'returning_attributes/version'
-require_relative 'returning_attributes/with_returning'
 
 require 'active_support/lazy_load_hooks'
 
 module ActiveRecord
   module ReturningAttributes
-    def self.monkey_patch_persistence!
-      Persistence.module_eval do
-        def _create_record(attribute_names = self.attribute_names)
-          attribute_names &= self.class.column_names
-          attributes_values = attributes_with_values_for_create(attribute_names)
-
-          new_id, result = self.class.connection.with_returning(returning) do
-            self.class._insert_record(attributes_values)
-          end
-
-          self.id ||= new_id if self.class.primary_key
-
-          self.class.returning.each do |column|
-            write_attribute(column, result.to_hash.first[column.to_s])
-          end
-
-          @new_record = false
-
-          yield(self) if block_given?
-
-          id
-        end
+    def self.postgresql_only
+      if postgresql?
+        yield
+      else
+        warn 'ActiveRecord::ReturningAttributes only works with PostgreSQL, skipping patching ActiveRecord.'
       end
     end
 
-    def self.monkey_patch_database_statements!
-      ConnectionAdapters::DatabaseStatements.module_eval do
-        def insert(arel, name = nil, pk = nil, id_value = nil, sequence_name = nil, binds = [])
-          sql, binds = to_sql_and_binds(arel, binds)
-          @_insert_result = exec_insert(sql, name, binds, pk, sequence_name)
-          id_value || last_inserted_id(@_insert_result)
-        end
+    def self.postgresql?
+      begin
+        require 'pg'
+        true
+      rescue LoadError
+        false
       end
     end
 
-    def self.monkey_patch_postgresql_database_statements!
-      require 'active_record/connection_adapters/postgresql/database_statements'
-
-      ConnectionAdapters::PostgreSQL::DatabaseStatements.module_eval do
-        private
-
-        def sql_for_insert(sql, pk, id_value, sequence_name, binds)
-          if pk.nil?
-            # Extract the table from the insert sql. Yuck.
-            table_ref = extract_table_ref_from_insert_sql(sql)
-            pk = primary_key(table_ref) if table_ref
-          end
-
-          if pk = suppress_composite_primary_key(pk)
-            if @_returning
-              sql = "#{sql} RETURNING #{[pk, @_returning].flatten.map { |column| quote_column_name(column) }.join(', ')}"
-            else
-              raise 'This case shouldn’t happen.'
-            end
-          end
-
-          super
-        end
+    def self.active_record_5_2_tested_only
+      unless ActiveRecord.version == Gem::Version.new('5.2.0')
+        warn 'ActiveRecord::ReturningAttributes was only tested on ActiveRecord 5.2.0 and might not work with your version.'
       end
     end
 
-    def self.extend_postgresql_adapter!
-      require 'active_record/connection_adapters/postgresql_adapter'
-      ConnectionAdapters::PostgreSQLAdapter.include(WithReturning)
-    end
-
     ActiveSupport.on_load(:active_record) do
-      include Accessor
-
-      ReturningAttributes.monkey_patch_persistence!
-      ReturningAttributes.monkey_patch_database_statements!
-      ReturningAttributes.monkey_patch_postgresql_database_statements!
-      ReturningAttributes.extend_postgresql_adapter!
+      ReturningAttributes.postgresql_only do
+        ReturningAttributes::Patching.patch_base
+        ReturningAttributes::Patching.patch_persistence
+        ReturningAttributes::Patching.patch_database_statements
+        ReturningAttributes::Patching.patch_postgresql_database_statements
+        ReturningAttributes::Patching.patch_postgresql_adapter
+
+        ReturningAttributes.active_record_5_2_tested_only
+      end
     end
   end
 end
diff --git a/lib/active_record/returning_attributes/accessor.rb b/lib/active_record/returning_attributes/accessor.rb
@@ -6,7 +6,7 @@ module Accessor
       extend ActiveSupport::Concern
 
       included do
-        class_attribute :returning, default: []
+        class_attribute :returning_attributes, default: []
       end
     end
   end
diff --git a/lib/active_record/returning_attributes/patching.rb b/lib/active_record/returning_attributes/patching.rb
@@ -0,0 +1,135 @@
+require_relative '../returning_attributes/accessor'
+require_relative '../returning_attributes/with_returning'
+
+module ActiveRecord
+  module ReturningAttributes
+  	module Patching
+      def self.patch_base
+        Base.include(Accessor)
+      end
+
+      def self.patch_persistence
+        Persistence.module_eval do
+          private
+
+          def _create_record(attribute_names = self.attribute_names)
+            attribute_names &= self.class.column_names
+            attributes_values = attributes_with_values_for_create(attribute_names)
+
+            new_id, returned_attributes = self.class.connection.with_returning_attributes(returning_attributes) do
+              self.class._insert_record(attributes_values)
+            end
+
+            self.id ||= new_id if self.class.primary_key
+
+            returning_attributes.each do |attribute|
+              write_attribute(attribute, returned_attributes[attribute.to_s]) if returned_attributes.key?(attribute.to_s)
+            end
+
+            @new_record = false
+
+            yield(self) if block_given?
+
+            id
+          end
+
+          def _update_record(attribute_names = self.attribute_names)
+            attribute_names &= self.class.column_names
+            attribute_names = attributes_for_update(attribute_names)
+
+            if attribute_names.empty?
+              affected_rows = 0
+              @_trigger_update_callback = true
+            else
+              affected_rows, returned_attributes = self.class.connection.with_returning_attributes(returning_attributes) do
+                _update_row(attribute_names)
+              end
+
+              returning_attributes.each do |attribute|
+                write_attribute(attribute, returned_attributes[attribute.to_s]) if returned_attributes.key?(attribute.to_s)
+              end
+
+              @_trigger_update_callback = affected_rows == 1
+            end
+
+            yield(self) if block_given?
+
+            affected_rows
+          end
+        end
+      end
+
+      def self.patch_database_statements
+        ConnectionAdapters::DatabaseStatements.module_eval do
+          def exec_insert(sql, name = nil, binds = [], pk = nil, sequence_name = nil)
+            sql, binds = sql_for_insert(sql, pk, nil, sequence_name, binds)
+
+            exec_query(sql, name, binds).tap do |result|
+              @_returned_attributes = result.first.to_hash
+            end
+          end
+
+          def update(arel, name = nil, binds = [])
+            sql, binds = to_sql_and_binds(arel, binds)
+
+            if @_returning_attributes && !@_returning_attributes.empty?
+              sql = "#{sql} RETURNING #{@_returning_attributes.map { |column| quote_column_name(column) }.join(', ')}"
+            end
+
+            exec_update(sql, name, binds)
+          end
+        end
+      end
+
+      def self.patch_postgresql_database_statements
+        require 'active_record/connection_adapters/postgresql/database_statements'
+
+        ConnectionAdapters::PostgreSQL::DatabaseStatements.module_eval do
+          private
+
+          def sql_for_insert(sql, pk, id_value, sequence_name, binds)
+            if pk.nil?
+              # Extract the table from the insert sql. Yuck.
+              table_ref = extract_table_ref_from_insert_sql(sql)
+              pk = primary_key(table_ref) if table_ref
+            end
+
+            if pk = suppress_composite_primary_key(pk)
+              sql = "#{sql} RETURNING #{[pk, @_returning_attributes].flatten.map { |column| quote_column_name(column) }.join(', ')}"
+            end
+
+            super
+          end
+        end
+      end
+
+      def self.patch_postgresql_adapter
+        require 'active_record/connection_adapters/postgresql_adapter'
+
+        ConnectionAdapters::PostgreSQLAdapter.class_eval do
+          include WithReturning
+
+          private
+
+          def execute_and_clear(sql, name, binds, prepare: false)
+            if without_prepared_statement?(binds)
+              result = exec_no_cache(sql, name, [])
+            elsif !prepare
+              result = exec_no_cache(sql, name, binds)
+            else
+              result = exec_cache(sql, name, binds)
+            end
+
+            if @_returning_attributes && !@_returning_attributes.empty?
+              @_returned_attributes = result.to_a.first
+            end
+
+            ret = yield result
+            result.clear
+            ret
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/active_record/returning_attributes/with_returning.rb b/lib/active_record/returning_attributes/with_returning.rb
@@ -1,13 +1,13 @@
 module ActiveRecord
   module ReturningAttributes
   	module WithReturning
-      def with_returning(returning)
+      def with_returning_attributes(attributes)
         begin
-          @_returning = returning
-          [yield, @_insert_result]
+          @_returning_attributes = attributes
+          [yield, @_returned_attributes || {}]
         ensure
-          @_returning = nil
-          @_insert_result = nil
+          @_returning_attributes = nil
+          @_returned_attributes = nil
         end
       end
     end
diff --git a/spec/active_record/returning_attributes_spec.rb b/spec/active_record/returning_attributes_spec.rb
@@ -1,9 +1,41 @@
 RSpec.describe ActiveRecord::ReturningAttributes do
-  it 'has a version number' do
-    expect(ActiveRecord::ReturningAttributes::VERSION).not_to be nil
+  let(:project) { Project.create(name: 'Freshly Created') }
+
+  context 'without returning attributes' do
+    before do
+      Project.returning_attributes = []
+    end
+
+    context 'on insert' do
+      it 'doesn’t return database-backed default values' do
+        expect(project.uuid).to be_nil
+      end
+    end
+
+    context 'on update' do
+      it 'doesn’t return values set via database triggers' do
+        project.update(name: 'Just Updated')
+        expect(project.update_count).to eq(0)
+      end
+    end
   end
 
-  it 'does something useful' do
-    expect(false).to eq(true)
+  context 'with returning attributes' do
+    before do
+      Project.returning_attributes = [:uuid, :update_count]
+    end
+
+    context 'on insert' do
+      it 'returns database-backed default values' do
+        expect(project.uuid).to be_present
+      end
+    end
+
+    context 'on update' do
+      it 'returns values set via database triggers' do
+        project.update(name: 'Just Updated')
+        expect(project.update_count).to eq(1)
+      end
+    end
   end
 end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
@@ -1,5 +1,9 @@
 require 'bundler/setup'
+require 'active_record'
 require 'active_record/returning_attributes'
+require 'pry'
+
+require_relative 'support/schema_and_models'
 
 RSpec.configure do |config|
   # Enable flags like --only-failures and --next-failure
@@ -11,4 +15,6 @@
   config.expect_with :rspec do |c|
     c.syntax = :expect
   end
+
+  config.filter_run_when_matching :focus
 end
diff --git a/spec/support/schema_and_models.rb b/spec/support/schema_and_models.rb
