Optimize .mark_as_read! and extend to work on relations

.gitignore

@@ -1,4 +1,5 @@
 *.gem
+*.log
 .bundle
 Gemfile.lock
 gemfiles/*.lock

lib/generators/unread/generator_helper.rb

@@ -0,0 +1,13 @@
+module Unread
+  module Generators
+    module GeneratorHelper
+      def next_migration_number(dirname)
+        if ActiveRecord::Base.timestamped_migrations
+          Time.now.utc.strftime("%Y%m%d%H%M%S")
+        else
+          "%.3d" % (current_migration_number(dirname) + 1)
+        end
+      end
+    end
+  end
+end

lib/generators/unread/migration/migration_generator.rb

@@ -1,23 +1,17 @@
 require 'rails/generators'
 require 'rails/generators/migration'
+require 'generators/unread/generator_helper'
 
 module Unread
   class MigrationGenerator < Rails::Generators::Base
     include Rails::Generators::Migration
+    extend Unread::Generators::GeneratorHelper
 
     desc "Generates migration for read_markers"
     source_root File.expand_path('../templates', __FILE__)
 
     def create_migration_file
       migration_template 'migration.rb', 'db/migrate/unread_migration.rb'
     end
-
-    def self.next_migration_number(dirname)
-      if ActiveRecord::Base.timestamped_migrations
-        Time.now.utc.strftime("%Y%m%d%H%M%S")
-      else
-        "%.3d" % (current_migration_number(dirname) + 1)
-      end
-    end
   end
 end

lib/generators/unread/migration/templates/migration.rb

@@ -1,24 +1,31 @@
+require 'generators/unread/migration_helper'
+
 class UnreadMigration < Unread::MIGRATION_BASE_CLASS
+  extend Unread::Generators::MigrationHelper
+
   def self.up
     create_table ReadMark, force: true, options: create_options do |t|
       t.references :readable, polymorphic: { null: false }
       t.references :reader,   polymorphic: { null: false }
       t.datetime :timestamp
     end
 
-    add_index ReadMark, [:reader_id, :reader_type, :readable_type, :readable_id], name: 'read_marks_reader_readable_index', unique: true
+    index_name = 'read_marks_reader_readable_index'
+    add_index ReadMark, [:reader_id, :reader_type, :readable_type, :readable_id], name: index_name, unique: true
+
+    if postgresql_9_5?
+      execute <<-SQL
+        ALTER TABLE #{ReadMark.table_name}
+          ADD CONSTRAINT read_marks_reader_readable_constraint UNIQUE USING INDEX #{index_name}
+      SQL
+    end
   end
 
   def self.down
     drop_table ReadMark
   end
 
   def self.create_options
-    options = ''
-    if defined?(ActiveRecord::ConnectionAdapters::Mysql2Adapter) \
-      && ActiveRecord::Base.connection.instance_of?(ActiveRecord::ConnectionAdapters::Mysql2Adapter)
-      options = 'DEFAULT CHARSET=latin1'
-    end
-    options
+    mysql? ? 'DEFAULT CHARSET=latin1' : ''
   end
 end

lib/generators/unread/migration_helper.rb

@@ -0,0 +1,16 @@
+module Unread
+  module Generators
+    module MigrationHelper
+      def mysql?
+        defined?(ActiveRecord::ConnectionAdapters::Mysql2Adapter) \
+          && ActiveRecord::Base.connection.instance_of?(ActiveRecord::ConnectionAdapters::Mysql2Adapter)
+      end
+
+      def postgresql_9_5?
+        defined?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter) \
+          && ActiveRecord::Base.connection.instance_of?(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter) \
+          && ActiveRecord::Base.connection.send(:postgresql_version) >= 90500
+      end
+    end
+  end
+end

lib/generators/unread/polymorphic_reader_migration/polymorphic_reader_migration_generator.rb

@@ -1,23 +1,17 @@
 require 'rails/generators'
 require 'rails/generators/migration'
+require 'generators/unread/generator_helper'
 
 module Unread
   class PolymorphicReaderMigrationGenerator < Rails::Generators::Base
     include Rails::Generators::Migration
+    extend Unread::Generators::GeneratorHelper
 
     desc "Generates update migration to make reader of read_markers polymorphic"
     source_root File.expand_path('../templates', __FILE__)
 
     def create_migration_file
       migration_template 'unread_polymorphic_reader_migration.rb', 'db/migrate/unread_polymorphic_reader_migration.rb'
     end
-
-    def self.next_migration_number(dirname)
-      if ActiveRecord::Base.timestamped_migrations
-        Time.now.utc.strftime("%Y%m%d%H%M%S")
-      else
-        "%.3d" % (current_migration_number(dirname) + 1)
-      end
-    end
   end
 end

lib/generators/unread/postgresql_unique_constraint_migration/postgresql_unique_constraint_migration_generator.rb

@@ -0,0 +1,17 @@
+require 'rails/generators'
+require 'rails/generators/migration'
+require 'generators/unread/generator_helper'
+
+module Unread
+  class PostgresqlUniqueConstraintMigrationGenerator < Rails::Generators::Base
+    include Rails::Generators::Migration
+    extend Unread::Generators::GeneratorHelper
+
+    desc "Generates update migration to add unique constraint for PostgresSQL database"
+    source_root File.expand_path('../templates', __FILE__)
+
+    def create_migration_file
+      migration_template 'unread_postgresql_unique_constraint_migration.rb', 'db/migrate/unread_postgresql_unique_constraint_migration.rb'
+    end
+  end
+end

lib/generators/unread/postgresql_unique_constraint_migration/templates/unread_postgresql_unique_constraint_migration.rb

@@ -0,0 +1,31 @@
+require 'generators/unread/migration_helper'
+
+class UnreadPostgresqlUniqueConstraintMigration < Unread::MIGRATION_BASE_CLASS
+  extend Unread::Generators::MigrationHelper
+
+  INDEX_NAME = 'read_marks_reader_readable_index'
+  INDEX_COLUMNS = [:reader_id, :reader_type, :readable_type, :readable_id]
+
+  def self.up
+    if postgresql_9_5?
+      unless index_exists?(ReadMark, INDEX_COLUMNS, unique: true)
+        add_index ReadMark, INDEX_COLUMNS, name: INDEX_NAME, unique: true
+      end
+
+      execute <<-SQL
+        ALTER TABLE #{ReadMark.table_name}
+          ADD CONSTRAINT read_marks_reader_readable_constraint UNIQUE USING INDEX #{INDEX_NAME}
+      SQL
+    end
+  end
+
+  def self.down
+    if postgresql_9_5?
+      execute <<-SQL
+        ALTER TABLE #{ReadMark.table_name}
+          DROP CONSTRAINT IF EXISTS read_marks_reader_readable_constraint
+      SQL
+      add_index ReadMark, INDEX_COLUMNS, name: INDEX_NAME, unique: true
+    end
+  end
+end

lib/unread.rb

@@ -1,4 +1,5 @@
 require 'active_record'
+require 'upsert'
 
 require 'unread/base'
 require 'unread/read_mark'
@@ -7,9 +8,11 @@
 require 'unread/readable_scopes'
 require 'unread/reader_scopes'
 require 'unread/garbage_collector'
+require 'unread/active_record_relation'
 require 'unread/version'
 
 ActiveRecord::Base.send :include, Unread
+ActiveRecord::Relation.send :include, ActiveRecordRelation::Unread
 
 Unread::MIGRATION_BASE_CLASS = if ActiveRecord::VERSION::MAJOR >= 5
   ActiveRecord::Migration[5.0]

lib/unread/active_record_relation.rb

@@ -0,0 +1,7 @@
+module ActiveRecordRelation
+  module Unread
+    def mark_as_read!(options)
+      klass.mark_as_read!(self, options)
+    end
+  end
+end

lib/unread/readable.rb

@@ -9,35 +9,46 @@ def mark_as_read!(target, options)
 
         if target == :all
           reset_read_marks_for_user(reader)
-        elsif target.respond_to?(:each)
-          mark_collection_as_read(target, reader)
+        elsif target.is_a?(ActiveRecord::Relation)
+          mark_relation_as_read(target, reader)
         else
-          raise ArgumentError
+          mark_collection_as_read(target, reader)
+        end
+      end
+
+      def mark_relation_as_read(relation, reader)
+        raise ArgumentError unless relation.klass == self
+
+        ReadMark.transaction do
+          global_timestamp = reader.read_mark_global(self).try(:timestamp)
+          on = readable_options[:on]
+
+          if global_timestamp
+            relation = relation.where("#{on} > ?", global_timestamp)
+          end
+
+          Upsert.batch(connection, ReadMark.table_name) do |upsert|
+            relation.pluck(relation.klass.primary_key, on).each do |readable_id, timestamp|
+              upsert.row({ readable_id: readable_id, readable_type: readable_parent.name, reader_id: reader.id, reader_type: reader.class.name }, timestamp: timestamp)
+            end
+          end
+          true
         end
       end
 
       def mark_collection_as_read(collection, reader)
         ReadMark.transaction do
           global_timestamp = reader.read_mark_global(self).try(:timestamp)
 
-          collection.each do |obj|
-            raise ArgumentError unless obj.is_a?(self)
-            timestamp = obj.send(readable_options[:on])
-
-            if global_timestamp && global_timestamp >= timestamp
-              # The object is implicitly marked as read, so there is nothing to do
-            else
-              # This transaction is needed, so that parent transaction won't rollback even there's an error.
-              ReadMark.transaction(requires_new: true) do
-                begin
-                  rm = obj.read_marks.where(reader_id: reader.id, reader_type: reader.class.base_class.name).first || obj.read_marks.build
-                  rm.reader_id   = reader.id
-                  rm.reader_type = reader.class.base_class.name
-                  rm.timestamp   = timestamp
-                  rm.save!
-                rescue ActiveRecord::RecordNotUnique
-                  raise ActiveRecord::Rollback
-                end
+          Upsert.batch(connection, ReadMark.table_name) do |upsert|
+            Array(collection).each do |obj|
+              raise ArgumentError unless obj.is_a?(self)
+              timestamp = obj.send(readable_options[:on])
+
+              if global_timestamp && global_timestamp >= timestamp
+                # The object is implicitly marked as read, so there is nothing to do
+              else
+                upsert.row({ readable_id: obj.id, readable_type: readable_parent.name, reader_id: reader.id, reader_type: reader.class.name }, timestamp: timestamp)
               end
             end
           end

spec/spec_helper.rb

@@ -65,6 +65,8 @@ def setup_db
 
   ActiveRecord::Base.establish_connection(db_name.to_sym)
   ActiveRecord::Base.default_timezone = :utc
+  ActiveRecord::Base.logger = Logger.new(File.join(__dir__, "debug.log"))
+  ActiveRecord::Base.logger.level = ENV["CI"] ? ::Logger::ERROR : ::Logger::DEBUG
   ActiveRecord::Migration.verbose = false
 
   UnreadMigration.up

spec/unread/readable_spec.rb

@@ -269,9 +269,6 @@
     it "should mark the rest as read when the first record is not unique" do
       Email.mark_as_read! [ @email1 ], for: @reader
 
-      allow(@email1).to receive_message_chain("read_marks.build").and_return(@email1.read_marks.build)
-      allow(@email1).to receive_message_chain("read_marks.where").and_return([])
-
       expect do
         Email.mark_as_read! [ @email1, @email2 ], for: @reader
       end.to change(ReadMark, :count).by(1)

unread.gemspec

@@ -21,12 +21,13 @@ Gem::Specification.new do |s|
   s.require_paths = ["lib"]
 
   s.add_dependency 'activerecord', '>= 3'
+  s.add_dependency 'upsert', '~> 2.2.1'
 
   s.add_development_dependency 'rake'
   s.add_development_dependency 'timecop'
   s.add_development_dependency 'sqlite3'
   s.add_development_dependency 'mysql2'
-  s.add_development_dependency 'pg'
+  s.add_development_dependency 'pg', '< 1'
   s.add_development_dependency 'rspec'
   s.add_development_dependency 'simplecov'
   s.add_development_dependency 'term-ansicolor'